decomp dot_grad (#68399)

cubehan3 · web-flow · commit 1b27c2478d41 · 2024-09-26T10:04:08.000+08:00
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -89,6 +89,7 @@
     'elementwise_pow_grad',
     'maximum_grad',
     'reduce_as_grad',
+    'dot_grad',
 ]
 
 OTHER_PRIM_VJP_OPS = [
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2808,6 +2808,38 @@ void atan_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
   }
 }
 
+template <typename T>
+void dot_grad(const Tensor& x,
+              const Tensor& y,
+              const Tensor& out_grad,
+              Tensor* x_grad,
+              Tensor* y_grad) {
+  const int64_t out_grad_dim_size = out_grad.dims().size();
+  Tensor out_grad_ = out_grad;
+
+  if (has_dynamic_shape(x.shape()) || has_dynamic_shape(y.shape())) {
+    auto out_grad_shape =
+        get_unsqueeze_dims<T>(shape<T>(out_grad_), {out_grad_dim_size});
+    out_grad_ = backend::reshape<T>(out_grad_, out_grad_shape);
+    out_grad_ = backend::expand_with_tensor<T>(out_grad_, shape<T>(x));
+  } else {
+    std::vector<int64_t> x_dim = common::vectorize<int64_t>(x.dims());
+    auto out_grad_shape = get_unsqueeze_dims(out_grad, {out_grad_dim_size});
+    out_grad_ =
+        expand<T>(reshape<T>(out_grad_, out_grad_shape), IntArray(x_dim));
+  }
+
+  if (x_grad) {
+    Tensor x_grad_tmp = out_grad_ * y;
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+
+  if (y_grad) {
+    Tensor y_grad_tmp = out_grad_ * x;
+    set_output<T>(y_grad_tmp, y_grad);
+  }
+}
+
 template <typename T>
 void logcumsumexp_grad(const Tensor& x,
                        const Tensor& out,
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -39,6 +39,7 @@
     "pd_op.cumprod",
     "pd_op.cumsum",
     "pd_op.divide",
+    "pd_op.dot",
     "pd_op.dropout",
     "pd_op.elementwise_pow",
     "pd_op.erf",
diff --git a/test/legacy_test/test_dot_op.py b/test/legacy_test/test_dot_op.py
@@ -25,7 +25,9 @@
 class DotOp(OpTest):
     def setUp(self):
         self.op_type = "dot"
+        self.prim_op_type = "prim"
         self.python_api = paddle.dot
+        self.public_python_api = paddle.dot
         self.init_dtype()
         self.init_input_output()
 
@@ -40,39 +42,92 @@ def test_check_output(self):
         self.check_output(check_pir=True)
 
     def test_check_grad_normal(self):
-        if core.is_compiled_with_rocm():
-            self.check_grad(
-                ['X', 'Y'],
-                'Out',
-                user_defined_grads=[self.inputs['Y'], self.inputs['X']],
-                check_pir=True,
-            )
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            if core.is_compiled_with_rocm():
+                self.check_grad(
+                    ['X', 'Y'],
+                    'Out',
+                    user_defined_grads=[self.inputs['Y'], self.inputs['X']],
+                    check_pir=True,
+                )
+            else:
+                self.check_grad(['X', 'Y'], 'Out', check_pir=True)
         else:
-            self.check_grad(['X', 'Y'], 'Out', check_pir=True)
+            if core.is_compiled_with_rocm():
+                self.check_grad(
+                    ['X', 'Y'],
+                    'Out',
+                    user_defined_grads=[self.inputs['Y'], self.inputs['X']],
+                    check_pir=True,
+                )
+            else:
+                self.check_grad(
+                    ['X', 'Y'], 'Out', check_pir=True, check_prim_pir=True
+                )
 
     def test_check_grad_ignore_x(self):
-        if core.is_compiled_with_rocm():
-            self.check_grad(
-                ['Y'],
-                'Out',
-                no_grad_set=set("X"),
-                user_defined_grads=[self.inputs['X']],
-                check_pir=True,
-            )
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            if core.is_compiled_with_rocm():
+                self.check_grad(
+                    ['Y'],
+                    'Out',
+                    no_grad_set=set("X"),
+                    user_defined_grads=[self.inputs['X']],
+                    check_pir=True,
+                )
+            else:
+                self.check_grad(
+                    ['Y'], 'Out', no_grad_set=set("X"), check_pir=True
+                )
         else:
-            self.check_grad(['Y'], 'Out', no_grad_set=set("X"), check_pir=True)
+            if core.is_compiled_with_rocm():
+                self.check_grad(
+                    ['Y'],
+                    'Out',
+                    no_grad_set=set("X"),
+                    user_defined_grads=[self.inputs['X']],
+                    check_pir=True,
+                )
+            else:
+                self.check_grad(
+                    ['Y'],
+                    'Out',
+                    no_grad_set=set("X"),
+                    check_pir=True,
+                    check_prim_pir=True,
+                )
 
     def test_check_grad_ignore_y(self):
-        if core.is_compiled_with_rocm():
-            self.check_grad(
-                ['X'],
-                'Out',
-                no_grad_set=set('Y'),
-                user_defined_grads=[self.inputs['Y']],
-                check_pir=True,
-            )
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            if core.is_compiled_with_rocm():
+                self.check_grad(
+                    ['X'],
+                    'Out',
+                    no_grad_set=set('Y'),
+                    user_defined_grads=[self.inputs['Y']],
+                    check_pir=True,
+                )
+            else:
+                self.check_grad(
+                    ['X'], 'Out', no_grad_set=set('Y'), check_pir=True
+                )
         else:
-            self.check_grad(['X'], 'Out', no_grad_set=set('Y'), check_pir=True)
+            if core.is_compiled_with_rocm():
+                self.check_grad(
+                    ['X'],
+                    'Out',
+                    no_grad_set=set('Y'),
+                    user_defined_grads=[self.inputs['Y']],
+                    check_pir=True,
+                )
+            else:
+                self.check_grad(
+                    ['X'],
+                    'Out',
+                    no_grad_set=set('Y'),
+                    check_pir=True,
+                    check_prim_pir=True,
+                )
 
     def init_input_output(self):
         self.x = np.random.uniform(0.1, 1, [121]).astype(self.dtype)
@@ -96,13 +151,25 @@ def init_input_output(self):
         self.out = np.sum(self.x * self.y, axis=1)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X', 'Y'], 'Out', check_pir=True)
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True, check_prim_pir=True)
 
     def test_check_grad_ignore_x(self):
-        self.check_grad(['Y'], 'Out', no_grad_set=set("X"), check_pir=True)
+        self.check_grad(
+            ['Y'],
+            'Out',
+            no_grad_set=set("X"),
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad_ignore_y(self):
-        self.check_grad(['X'], 'Out', no_grad_set=set('Y'), check_pir=True)
+        self.check_grad(
+            ['X'],
+            'Out',
+            no_grad_set=set('Y'),
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestDotOpError(unittest.TestCase):
diff --git a/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py
@@ -87,6 +87,10 @@ def divide_net(x, y):
     return x / y
 
 
+def dot_net(x, y):
+    return paddle.dot(x, y)
+
+
 def dropout_net1(x):
     return paddle.nn.functional.dropout(
         x, training=False, mode='downscale_in_infer'
@@ -737,6 +741,38 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimDotWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.dot_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200]
+        self.init_x_shape = [None, None]
+        self.y_shape = [30, 200]
+        self.init_y_shape = [None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = dot_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimDotWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.dot_grad"
+        self.dtype = "float32"
+        self.x_shape = [200]
+        self.init_x_shape = [None]
+        self.y_shape = [200]
+        self.init_y_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = dot_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
 class TestPrimDropoutWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2023)

Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,7 @@`
`89`	`89`	`'elementwise_pow_grad',`
`90`	`90`	`'maximum_grad',`
`91`	`91`	`'reduce_as_grad',`
	`92`	`+ 'dot_grad',`
`92`	`93`	`]`
`93`	`94`
`94`	`95`	`OTHER_PRIM_VJP_OPS = [`