[pir] add __neg__ and clean pir.cc (#60166)

gouzil · web-flow · commit bc7380d7129e · 2023-12-22T11:09:59.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -581,55 +581,6 @@ const phi::DDim &GetValueDims(Value value) {
   }
 }
 
-#define OVERRIDE_OPERATOR(operator, api, other_type)      \
-  value.def(#operator, [](Value self, other_type other) { \
-    return paddle::dialect::api(self, other);             \
-  });
-
-#define OVERRIDE_OPERATOR_WITH_SCALE(operator,            \
-                                     other_type,          \
-                                     scale_value,         \
-                                     bias_value,          \
-                                     bias_after_scale)    \
-  value.def(#operator, [](Value self, other_type other) { \
-    return paddle::dialect::scale(                        \
-        self, scale_value, bias_value, bias_after_scale); \
-  });
-
-#define OVERRIDE_OPERATOR_FOR_EACH(operator,         \
-                                   api,              \
-                                   scale_value,      \
-                                   bias_value,       \
-                                   bias_after_scale) \
-  OVERRIDE_OPERATOR(operator, api, Value)            \
-  OVERRIDE_OPERATOR_WITH_SCALE(operator,             \
-                               int,                  \
-                               scale_value,          \
-                               bias_value,           \
-                               bias_after_scale)     \
-  OVERRIDE_OPERATOR_WITH_SCALE(operator,             \
-                               float,                \
-                               scale_value,          \
-                               bias_value,           \
-                               bias_after_scale)     \
-  OVERRIDE_OPERATOR_WITH_SCALE(operator,             \
-                               double,               \
-                               scale_value,          \
-                               bias_value,           \
-                               bias_after_scale)
-
-#define OVERRIDE_COMPARE_OP_WITH_FULL(operator, api, other_type)         \
-  value.def(#operator, [](Value self, other_type other) {                \
-    auto rhs =                                                           \
-        paddle::dialect::full(/*shape=*/{}, other, GetValueDtype(self)); \
-    return paddle::dialect::api(self, rhs);                              \
-  });
-
-#define OVERRIDE_COMPARE_OP_FOR_EACH(operator, api)   \
-  OVERRIDE_OPERATOR(operator, api, Value)             \
-  OVERRIDE_COMPARE_OP_WITH_FULL(operator, api, int)   \
-  OVERRIDE_COMPARE_OP_WITH_FULL(operator, api, float) \
-  OVERRIDE_COMPARE_OP_WITH_FULL(operator, api, double)
 void BindValue(py::module *m) {
   py::class_<Value> value(*m, "Value", R"DOC(
     Value class represents the SSA value in the IR system. It is a directed edge
@@ -787,23 +738,9 @@ void BindValue(py::module *m) {
              print_stream << ")";
              return print_stream.str();
            })
-      .def("__neg__",
-           [](Value self) {
-             return paddle::dialect::scale(self, -1.0, 0.0, true);
-           })
       .def("is_same", &Value::operator==)
       .def("hash", [](Value self) { return std::hash<pir::Value>{}(self); })
       .def("__repr__", &Value2String);
-  // For basaic operators
-  OVERRIDE_OPERATOR_FOR_EACH(__add__, add, 1.0, other, true);
-  OVERRIDE_OPERATOR_FOR_EACH(__sub__, subtract, 1.0, -1.0 * other, true);
-  OVERRIDE_OPERATOR_FOR_EACH(__mul__, multiply, other, 0.0, false);
-  OVERRIDE_OPERATOR_FOR_EACH(__truediv__, divide, 1.0 / other, 0.0, false);
-  // For compare opeartors
-  OVERRIDE_COMPARE_OP_FOR_EACH(__lt__, less_than);
-  OVERRIDE_COMPARE_OP_FOR_EACH(__le__, less_equal);
-  OVERRIDE_COMPARE_OP_FOR_EACH(__gt__, greater_than);
-  OVERRIDE_COMPARE_OP_FOR_EACH(__ge__, greater_equal);
 }
 
 void BindOpOperand(py::module *m) {
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -281,6 +281,9 @@ def _scalar_mul_(var, value):
     def _scalar_div_(var, value):
         return paddle.scale(var, 1.0 / value, 0.0)
 
+    def _scalar_neg_(var):
+        return paddle.scale(var, -1.0, 0.0)
+
     def _binary_creator_(
         method_name,
         python_api,
@@ -513,6 +516,7 @@ def value_hash(self):
         ('append', append),
         ('set_shape', set_shape),
         ('__hash__', value_hash),
+        # For basic operators
         (
             '__add__',
             _binary_creator_('__add__', paddle.tensor.add, False, _scalar_add_),
@@ -591,7 +595,8 @@ def value_hash(self):
             '__matmul__',
             _binary_creator_('__matmul__', paddle.tensor.matmul, False, None),
         ),
-        # for logical compare
+        ('__neg__', _scalar_neg_),
+        # For compare opeartors
         (
             '__eq__',
             _binary_creator_('__eq__', paddle.tensor.equal, False, None),
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -491,6 +491,25 @@ def test_append(self):
                 with self.assertRaises(TypeError):
                     x.append(array)
 
+    def test_neg(self):
+        x_np = np.random.uniform(-1, 1, [10, 1024]).astype(np.float32)
+        res = -x_np
+        with paddle.pir_utils.IrGuard():
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.static.data(
+                    name='x', shape=[10, 1024], dtype="float32"
+                )
+                a = -x
+                b = x.__neg__()
+                (a_np, b_np) = exe.run(
+                    main_program,
+                    feed={"x": x_np},
+                    fetch_list=[a, b],
+                )
+                np.testing.assert_array_equal(res, a_np)
+                np.testing.assert_array_equal(res, b_np)
+
     def test_math_exists(self):
         with paddle.pir_utils.IrGuard():
             a = paddle.static.data(name='a', shape=[1], dtype='float32')
