[Prim][PIR] Surport dynamic shape for matmul_grad (#66835)

zeroRains · web-flow · commit 2fd39872bb21 · 2024-08-05T10:25:05.000+08:00
* surport dynamic shape for matmul

* fix the test case
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1178,17 +1178,22 @@ void matmul_grad(const Tensor& x,
                 reverse_perm[reverse_perm.size() - 2]);
       x_grad_trans = transpose<T>(x_grad_mm, reverse_perm);
     }
-
-    if (x_grad_trans.dims() != x.dims()) {
-      phi::DDim x_reduce_dim = get_reduce_dims_from_out(
-          x_grad_trans.dims(), temp_x_unsqueeze.dims());
-      auto dx_reduce_res = sum<T>(
-          x_grad_trans, common::vectorize(x_reduce_dim), x.dtype(), false);
-      auto x_grad_out = reshape<T>(dx_reduce_res, x.shape());
+    if (has_dynamic_shape(x.shape()) ||
+        has_dynamic_shape(x_grad_trans.shape())) {
+      auto x_grad_out = reduce_as<T>(x_grad_trans, temp_x_unsqueeze);
       set_output<T>(x_grad_out, x_grad);
     } else {
-      auto x_grad_out = x_grad_trans;
-      set_output<T>(x_grad_out, x_grad);
+      if (x_grad_trans.dims() != x.dims()) {
+        phi::DDim x_reduce_dim = get_reduce_dims_from_out(
+            x_grad_trans.dims(), temp_x_unsqueeze.dims());
+        auto dx_reduce_res = sum<T>(
+            x_grad_trans, common::vectorize(x_reduce_dim), x.dtype(), false);
+        auto x_grad_out = reshape<T>(dx_reduce_res, x.shape());
+        set_output<T>(x_grad_out, x_grad);
+      } else {
+        auto x_grad_out = x_grad_trans;
+        set_output<T>(x_grad_out, x_grad);
+      }
     }
   }
 
@@ -1206,17 +1211,22 @@ void matmul_grad(const Tensor& x,
                 reverse_perm[reverse_perm.size() - 2]);
       y_grad_trans = transpose<T>(y_grad_mm, reverse_perm);
     }
-
-    if (y_grad_trans.dims() != y.dims()) {
-      phi::DDim y_reduce_dim = get_reduce_dims_from_out(
-          y_grad_trans.dims(), temp_y_unsqueeze.dims());
-      auto dy_reduce_res = sum<T>(
-          y_grad_trans, common::vectorize(y_reduce_dim), y.dtype(), false);
-      auto y_grad_out = reshape<T>(dy_reduce_res, y.shape());
+    if (has_dynamic_shape(y.shape()) ||
+        has_dynamic_shape(y_grad_trans.shape())) {
+      auto y_grad_out = reduce_as<T>(y_grad_trans, temp_y_unsqueeze);
       set_output<T>(y_grad_out, y_grad);
     } else {
-      auto y_grad_out = y_grad_trans;
-      set_output<T>(y_grad_out, y_grad);
+      if (y_grad_trans.dims() != y.dims()) {
+        phi::DDim y_reduce_dim = get_reduce_dims_from_out(
+            y_grad_trans.dims(), temp_y_unsqueeze.dims());
+        auto dy_reduce_res = sum<T>(
+            y_grad_trans, common::vectorize(y_reduce_dim), y.dtype(), false);
+        auto y_grad_out = reshape<T>(dy_reduce_res, y.shape());
+        set_output<T>(y_grad_out, y_grad);
+      } else {
+        auto y_grad_out = y_grad_trans;
+        set_output<T>(y_grad_out, y_grad);
+      }
     }
   }
 }
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -54,6 +54,7 @@
     "pd_op.pow",
     "pd_op.elementwise_pow",
     "pd_op.softmax",
+    "pd_op.matmul",
 ]
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py
@@ -119,6 +119,10 @@ def softmax_net(x):
     return paddle.nn.functional.softmax(x, axis=-1)
 
 
+def matmul_net(x, y):
+    return paddle.matmul(x, y)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -1370,5 +1374,80 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimMatmulWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 200]
+        self.init_x_shape = [None, None, 200]
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = matmul_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimMatmulWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [1, 30, 40, 200]
+        self.init_x_shape = [None, None, None, 200]
+        self.y_shape = [30, 1, 200, 40]
+        self.init_y_shape = [None, None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = matmul_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimMatmulWithGrad3(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [1, 30, 40, 200]
+        self.init_x_shape = [1, None, None, 200]
+        self.y_shape = [30, 1, 200, 40]
+        self.init_y_shape = [None, 1, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = matmul_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimMatmulWithGrad4(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 1, 40, 200]
+        self.init_x_shape = [None, None, None, 200]
+        self.y_shape = [1, 30, 200, 40]
+        self.init_y_shape = [None, None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = matmul_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimMatmulWithGrad5(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 1, 40, 200]
+        self.init_x_shape = [None, 1, None, 200]
+        self.y_shape = [1, 30, 200, 40]
+        self.init_y_shape = [1, None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = matmul_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,7 @@`
`54`	`54`	`"pd_op.pow",`
`55`	`55`	`"pd_op.elementwise_pow",`
`56`	`56`	`"pd_op.softmax",`
	`57`	`+ "pd_op.matmul",`
`57`	`58`	`]`
`58`	`59`
`59`	`60`