[Paddle Tensor No.4] 新增 Tensor.__rmatmul__ (#69406)

LittleHeroZZZX · web-flow · commit b3bd7efac23d · 2024-11-15T22:38:55.000+08:00
* add Tensor.__rmatmul__

* add tests for Tensor.__rmatmul__

* Fix comment typo and rmatmul test tolerance

* Fix code style
diff --git a/paddle/fluid/pybind/eager_math_op_patch.cc b/paddle/fluid/pybind/eager_math_op_patch.cc
@@ -1286,6 +1286,131 @@ static PyObject* tensor__matmul__method(TensorObject* self,
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+static PyObject* tensor__rmatmul__method(TensorObject* self,
+                                         PyObject* args,
+                                         PyObject* kwargs) {
+  phi::RecordEvent pythonc_record_event(
+      "__rmatmul__ pybind_patch_func", phi::TracerEventType::UserDefined, 1);
+  EAGER_TRY
+
+  VLOG(6) << "Running Eager tensor__rmatmul__method";
+
+  // Set Device ID
+  auto place = egr::Controller::Instance().GetExpectedPlace();
+  SetDevice(place);
+
+  paddle::Tensor ret;
+  paddle::Tensor self_tensor = self->tensor;
+
+  PyObject* other_obj = PyTuple_GET_ITEM(args, 0);
+
+  // 1. scalar exists cases
+  // there is no scalar_matmul function for __rmatmul__ now
+  float other_double = 0.0f;
+  bool has_other_double = false;
+  if (PyFloat_Check(other_obj) || PyCheckInteger(other_obj) ||
+      IsNumpyType(other_obj)) {
+    if (PyFloat_Check(other_obj)) {
+      other_double = CastPyArg2Double(other_obj, "__rmatmul__", 0);  // NOLINT
+      has_other_double = true;
+      if (_supported_int_dtype_.find(self_tensor.dtype()) !=
+          _supported_int_dtype_.end()) {
+        eager_gil_scoped_release guard;
+        self_tensor = cast_ad_func(self_tensor, DataType::FLOAT32);
+      }
+    } else if (PyCheckInteger(other_obj) || IsNumpyType(other_obj)) {
+      other_double = CastPyArg2Double(other_obj, "__rmatmul__", 0);  // NOLINT
+      has_other_double = true;
+    }
+  }
+
+  // 2. create or get tensor for other_obj
+  paddle::Tensor other_tensor;
+  if (has_other_double) {
+    eager_gil_scoped_release guard;
+    other_tensor = full_ad_func({1},
+                                phi::Scalar(other_double),
+                                self_tensor.dtype(),
+                                self_tensor.place());
+    const phi::distributed::ProcessMesh* mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+      ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+    }
+  } else if (PyCheckTensor(other_obj)) {
+    auto& self_tensor_ref = self->tensor;
+    auto& other_tensor_ref = CastPyArg2Tensor(other_obj, 0);
+    const phi::distributed::ProcessMesh* mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, self_tensor_ref, other_tensor_ref)) {
+      ConvertAllInputsToDistTensor(mesh, self_tensor_ref, other_tensor_ref);
+    }
+    self_tensor = self_tensor_ref;
+    other_tensor = other_tensor_ref;
+  } else {
+    if (IsNumpyArray(other_obj)) {
+      py::object numpy_value =
+          py::reinterpret_borrow<py::object>(py::handle(other_obj));
+      other_tensor = paddle::empty({}, phi::DataType::FLOAT32, place);
+      InitTensorWithNumpyValue(numpy_value, place, &other_tensor);
+    } else {
+      paddle::experimental::Scalar value =
+          CastPyArg2Scalar(other_obj, "__matmul__", 0);
+      if (PyComplex_Check(other_obj)) {
+        eager_gil_scoped_release guard;
+        other_tensor =
+            full_ad_func({1}, value, DataType::COMPLEX64, self_tensor.place());
+      } else {
+        eager_gil_scoped_release guard;
+        other_tensor =
+            full_ad_func({1}, value, self_tensor.dtype(), self_tensor.place());
+      }
+    }
+    const phi::distributed::ProcessMesh* mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+      ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+    }
+  }
+
+  // 3. promote types or unify right var type to left var
+  phi::DataType lhs_dtype = self_tensor.dtype();
+  phi::DataType rhs_dtype = other_tensor.dtype();
+  if (lhs_dtype != rhs_dtype) {
+    // note: only op_type in _supported_promote_complex_types_ should promote
+    // dtype
+    if (_complex_dtypes.find(lhs_dtype) != _complex_dtypes.end() ||
+        _complex_dtypes.find(rhs_dtype) != _complex_dtypes.end()) {
+      phi::DataType promote_dtype =
+          phi::TransToPhiDataType(framework::PromoteTypesIfComplexExists(
+              framework::TransToProtoVarType(lhs_dtype),
+              framework::TransToProtoVarType(rhs_dtype)));
+      if (lhs_dtype != promote_dtype) {
+        // cast
+        eager_gil_scoped_release guard;
+        self_tensor = cast_ad_func(self_tensor, promote_dtype);
+      }
+      if (rhs_dtype != promote_dtype) {
+        eager_gil_scoped_release guard;
+        other_tensor = cast_ad_func(other_tensor, promote_dtype);
+      }
+    } else {
+      VLOG(6) << "The dtype of left and right Tensor are not the same, left "
+                 "dtype is "
+              << lhs_dtype << ", but right dtype is " << rhs_dtype
+              << ", the right dtype will convert to " << lhs_dtype;
+      eager_gil_scoped_release guard;
+      other_tensor = cast_ad_func(other_tensor, lhs_dtype);
+    }
+  }
+
+  // 4. calculation
+  VLOG(6) << "Calling matmul_ad_func in tensor__rmatmul__method";
+  {
+    eager_gil_scoped_release guard;
+    ret = matmul_ad_func(other_tensor, self_tensor, false, false);
+  }
+  return ToPyObject(ret);
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* tensor__lt__method(TensorObject* self,
                                     PyObject* args,
                                     PyObject* kwargs) {
@@ -2087,6 +2212,10 @@ PyMethodDef math_op_patch_methods[] = {  // NOLINT
      (PyCFunction)(void (*)())tensor__matmul__method,
      METH_VARARGS | METH_KEYWORDS,
      nullptr},
+    {"__rmatmul__",
+     (PyCFunction)(void (*)())tensor__rmatmul__method,
+     METH_VARARGS | METH_KEYWORDS,
+     nullptr},
     {"__gt__",
      (PyCFunction)(void (*)())tensor__gt__method,
      METH_VARARGS | METH_KEYWORDS,
diff --git a/python/paddle/base/dygraph/math_op_patch.py b/python/paddle/base/dygraph/math_op_patch.py
@@ -249,6 +249,7 @@ def _mT_(var: Tensor) -> Tensor:
         '__mod__',
         '__rmod__',
         '__matmul__',
+        '__rmatmul__',
         '__gt__',
         '__ge__',
         '__lt__',
diff --git a/python/paddle/base/layers/math_op_patch.py b/python/paddle/base/layers/math_op_patch.py
@@ -83,6 +83,7 @@
     "__mod__": "A % B",
     "__rmod__": "A %= B",
     "__matmul__": "A @ B",
+    "__rmatmul__": "A @= B",
     "__eq__": "A == B",
     "__ne__": "A != B",
     "__lt__": "A < B",
@@ -891,6 +892,10 @@ def to_dense(var):
             '__matmul__',
             _binary_creator_('__matmul__', "matmul_v2", False, None),
         ),
+        (
+            '__rmatmul__',
+            _binary_creator_('__rmatmul', "matmul_v2", True, None),
+        ),
         #  for logical compare
         ('__eq__', _binary_creator_('__eq__', 'equal', False, None)),
         ('__ne__', _binary_creator_('__ne__', 'not_equal', False, None)),
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -1140,6 +1140,10 @@ def register_hook(self, hook):
             '__matmul__',
             _binary_creator_('__matmul__', paddle.tensor.matmul, False, None),
         ),
+        (
+            '__rmatmul__',
+            _binary_creator_('__rmatmul__', paddle.tensor.matmul, True, None),
+        ),
         ('__neg__', _scalar_neg_),
         ('__abs__', _scalar_abs_),
         # For compare operators
diff --git a/python/paddle/tensor/tensor.prototype.pyi b/python/paddle/tensor/tensor.prototype.pyi
@@ -175,6 +175,7 @@ class AbstractTensor:
     def __radd__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rsub__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rmul__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
+    def __rmatmul__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rtruediv__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rmod__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rpow__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -402,6 +402,22 @@ def test_matmul(self):
         )
         np.testing.assert_allclose(a_np @ b_np, c_np, rtol=1e-05)
 
+    @prog_scope()
+    def test_rmatmul(self):
+        a = paddle.static.data(name='a', shape=[2, 3], dtype='float32')
+        b = paddle.static.data(name='b', shape=[3, 5], dtype='float32')
+        c = b.__rmatmul__(a)
+        a_np = np.random.uniform(-1, 1, size=[2, 3]).astype('float32')
+        b_np = np.random.uniform(-1, 1, size=[3, 5]).astype('float32')
+        place = paddle.CPUPlace()
+        exe = paddle.static.Executor(place)
+        (c_np,) = exe.run(
+            paddle.static.default_main_program(),
+            feed={"a": a_np, "b": b_np},
+            fetch_list=[c],
+        )
+        np.testing.assert_allclose(a_np @ b_np, c_np, rtol=1e-05)
+
     @prog_scope()
     def test_builtin_type_conversion(self):
         a = paddle.static.data(name="a", shape=[])
@@ -536,6 +552,16 @@ def test_dygraph_equal(self):
         np.testing.assert_equal(actual_out, expect_out)
         paddle.enable_static()
 
+    def test_dygraph_rmatmul(self):
+        paddle.disable_static()
+        a_np = np.random.random((2, 3)).astype(np.float32) * 100
+        b_np = np.random.random((3, 5)).astype(np.float32) * 100
+        a = paddle.to_tensor(a_np)
+        b = paddle.to_tensor(b_np)
+        c = b.__rmatmul__(a)
+        np.testing.assert_allclose(a @ b, c.numpy(), rtol=1e-5, atol=1e-5)
+        paddle.enable_static()
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -116,6 +116,7 @@ def test_matmul(self):
         res_np_b = x_np @ y_np  # __matmul__
         res_np_c = paddle.matmul(paddle.to_tensor(x_np), paddle.to_tensor(y_np))
         res_np_d = x_np.__matmul__(y_np)
+        res_np_e = y_np.__rmatmul__(x_np)
         paddle.enable_static()
         with paddle.pir_utils.IrGuard():
             main_program, exe, program_guard = new_program()
@@ -125,14 +126,16 @@ def test_matmul(self):
                 b = x @ y
                 c = x.matmul(y)
                 d = x.__matmul__(y)
-                (b_np, c_np, d_np) = exe.run(
+                e = y.__rmatmul__(x)
+                (b_np, c_np, d_np, e_np) = exe.run(
                     main_program,
                     feed={"x": x_np, "y": y_np},
-                    fetch_list=[b, c, d],
+                    fetch_list=[b, c, d, e],
                 )
                 np.testing.assert_allclose(res_np_b, b_np, atol=1e-05)
                 np.testing.assert_allclose(res_np_c, c_np, atol=1e-05)
                 np.testing.assert_allclose(res_np_d, d_np, atol=1e-05)
+                np.testing.assert_allclose(res_np_e, e_np, atol=1e-05)
 
     def test_floordiv(self):
         paddle.disable_static()
