【complex op No.4】support complex for bmm

paddle/phi/kernels/cpu/bmm_grad_kernel.cc

@@ -18,5 +18,11 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/bmm_grad_kernel_impl.h"
 
-PD_REGISTER_KERNEL(
-    bmm_grad, CPU, ALL_LAYOUT, phi::BmmGradKernel, float, double) {}
+PD_REGISTER_KERNEL(bmm_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BmmGradKernel,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/cpu/bmm_kernel.cc

@@ -18,4 +18,11 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/bmm_kernel_impl.h"
 
-PD_REGISTER_KERNEL(bmm, CPU, ALL_LAYOUT, phi::BmmKernel, float, double) {}
+PD_REGISTER_KERNEL(bmm,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BmmKernel,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/bmm_grad_kernel.cu

@@ -25,4 +25,6 @@ PD_REGISTER_KERNEL(bmm_grad,
                    float,
                    double,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/bmm_kernel.cu

@@ -25,4 +25,6 @@ PD_REGISTER_KERNEL(bmm,
                    float,
                    double,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/impl/bmm_grad_kernel_impl.h

@@ -61,11 +61,16 @@ void BmmGradKernel(const Context& dev_ctx,
       &x_help, &y_help, &out_grad_help, false, false);
 
   phi::DDim dx_dims;
+  // for complex
+  DenseTensor x_conj;
+  DenseTensor y_conj;
   if (x_grad) {
     dx_dims = x_grad->dims();
     if (dx_dims != x_help.dims()) {
       x_grad->Resize(x_help.dims());
     }
+
+    y_conj = Conj<T>(dev_ctx, y_help);
   }
 
   phi::DDim dy_dims;
@@ -74,12 +79,14 @@ void BmmGradKernel(const Context& dev_ctx,
     if (dy_dims != y_help.dims()) {
       y_grad->Resize(y_help.dims());
     }
+
+    x_conj = Conj<T>(dev_ctx, x_help);
   }
 
   CalcInputGrad<T, Context>(
-      dev_ctx, out_grad_help, false, y_help, true, x_grad);
+      dev_ctx, out_grad_help, false, y_conj, true, x_grad);
   CalcInputGrad<T, Context>(
-      dev_ctx, x_help, true, out_grad_help, false, y_grad);
+      dev_ctx, x_conj, true, out_grad_help, false, y_grad);
 
   if (x_grad) {
     if (dx_dims != x_help.dims()) {

test/legacy_test/test_bmm_op.py

@@ -157,5 +157,57 @@ def test_api_error(self):
         self.assertRaises(ValueError, paddle.bmm, x_data, y_data_wrong3)
 
 
+class TestBmmOpCaseComplex64(TestBmmOp):
+    def setUp(self):
+        self.op_type = "bmm"
+        self.python_api = paddle.tensor.bmm
+        self.public_python_api = paddle.tensor.bmm
+        self.dtype = np.complex64
+        X = (
+            np.random.uniform(1, 5, (10, 3, 4))
+            + 1j * np.random.uniform(1, 5, (10, 3, 4))
+        ).astype("complex64")
+        Y = (
+            np.random.uniform(1, 5, (10, 4, 2))
+            + 1j * np.random.uniform(1, 5, (10, 4, 2))
+        ).astype("complex64")
+        self.inputs = {'X': X, 'Y': Y}
+        Out = np.matmul(X, Y)
+        self.outputs = {'Out': Out}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_prim=False)
+        # pass
+
+    def test_checkout_grad(self):
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True, check_prim=False)
+
+
+class TestBmmOpCaseComplex128(TestBmmOp):
+    def setUp(self):
+        self.op_type = "bmm"
+        self.python_api = paddle.tensor.bmm
+        self.public_python_api = paddle.tensor.bmm
+        self.dtype = np.complex128
+        X = (
+            np.random.uniform(1, 5, (10, 3, 4))
+            + 1j * np.random.uniform(1, 5, (10, 3, 4))
+        ).astype("complex128")
+        Y = (
+            np.random.uniform(1, 5, (10, 4, 2))
+            + 1j * np.random.uniform(1, 5, (10, 4, 2))
+        ).astype("complex128")
+        self.inputs = {'X': X, 'Y': Y}
+        Out = np.matmul(X, Y)
+        self.outputs = {'Out': Out}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_prim=False)
+        # pass
+
+    def test_checkout_grad(self):
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True, check_prim=False)
+
+
 if __name__ == "__main__":
     unittest.main()