[Prim][PIR] fix prim reshape grad bug

paddle/fluid/primitive/codegen/gen.py

@@ -63,7 +63,6 @@
     'fused_softmax_mask_upper_triangle_grad',
     'matmul_grad',
     'pow_grad',
-    'reshape_grad',
     'rsqrt_grad',
     'slice_grad',
     'transpose_grad',

paddle/fluid/primitive/rule/vjp/details.h

@@ -200,9 +200,14 @@ void gelu_grad(const Tensor& x,
 }
 
 template <typename T>
-void reshape_grad(const Tensor& x, const Tensor& grad_out, Tensor* grad_x) {
+void reshape_grad(const Tensor& xshape,
+                  const Tensor& grad_out,
+                  Tensor* grad_x) {
   if (grad_x) {
-    auto grad_x_tmp = reshape<T>(grad_out, phi::vectorize(x.dims()));
+    // xshape: [0] + x.shape
+    auto xshape_dims = xshape.dims();
+    auto x_dims = phi::slice_ddim(xshape_dims, 1, xshape_dims.size());
+    auto grad_x_tmp = reshape<T>(grad_out, phi::vectorize(x_dims));
     set_output<T>(grad_x_tmp, grad_x);
   }
 }

paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc

@@ -40,5 +40,31 @@ std::vector<std::vector<paddle::Tensor>> add_n_vjp(
   return vjp_res;
 }
 
+std::vector<std::vector<paddle::Tensor>> reshape_vjp(
+    const Tensor& xshape,
+    const Tensor& out_grad,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  std::vector<std::vector<paddle::Tensor>> vjp_res;
+  for (auto arg : stop_gradients) {
+    vjp_res.push_back(std::vector<paddle::Tensor>(arg.size()));
+  }
+  std::string op_name = "reshape_grad";
+  auto need_skip =
+      paddle::prim::StaticCompositeContext::Instance().CheckSkipCompOps(
+          op_name);
+  if (paddle::prim::StaticCompositeContext::Instance().IsBwdPrimEnabled() &&
+      !need_skip) {
+    FLAGS_tensor_operants_mode = "static";
+    paddle::Tensor* x_grad = !stop_gradients[0][0] ? &vjp_res[0][0] : nullptr;
+
+    details::reshape_grad<LazyTensor>(xshape, out_grad, x_grad);
+  } else {
+    auto op_res = backend::reshape_grad<LazyTensor>(xshape, out_grad);
+    vjp_res[0][0] = op_res;
+    vjp_res = ConstructVjpResultByStopGradients(vjp_res, stop_gradients);
+  }
+  return vjp_res;
+}
+
 }  // namespace primitive
 }  // namespace paddle

paddle/fluid/primitive/rule/vjp/manual/manual_vjp.h

@@ -27,5 +27,11 @@ std::vector<std::vector<paddle::Tensor>> add_n_vjp(
     const std::vector<paddle::Tensor>& x,
     const Tensor& out_grad,
     const std::vector<std::vector<bool>>& stop_gradients);
+
+std::vector<std::vector<paddle::Tensor>> reshape_vjp(
+    const Tensor& xshape,
+    const Tensor& out_grad,
+    const std::vector<std::vector<bool>>& stop_gradients);
+
 }  // namespace primitive
 }  // namespace paddle

python/paddle/decomposition/rules.py

@@ -39,7 +39,7 @@ def mean(x, axis, keepdim):
 
 
 @register_decomp('pd_op.gelu')
-def gelu_composite(x, approximate):
+def gelu(x, approximate):
     """define composite rule of op gelu"""
     M_SQRT1_2 = (
         0.70710678118654752440  # /* 1/sqrt(2) */ copy from gelu-kernel.cc
@@ -66,7 +66,7 @@ def gelu_composite(x, approximate):
 
 
 @register_decomp('pd_op.rsqrt')
-def rsqrt_composite(x):
+def rsqrt(x):
     """define composite rule of op rsqrt."""
     # rsqrt(x) = x^(-0.5)
     is_amp = False
@@ -77,7 +77,7 @@ def rsqrt_composite(x):
         is_amp = True
         x = cast(x, "float32")
     y = full(x.shape if len(x.shape) == 0 else [1], -0.5, x.dtype)
-    res = pow(x, y)
+    res = pow_composite(x, y)
     return res if not is_amp else cast(res, dtype)
 
 
@@ -104,7 +104,7 @@ def pow_composite(x, y):
 
 
 @register_decomp('pd_op.layer_norm')
-def layernorm_composite(x, scale, bias, epsilon, begin_norm_axis):
+def layernorm(x, scale, bias, epsilon, begin_norm_axis):
     """
     define composite rule of op layer_norm
     out = (x - mean(x)) / sqrt(var + epsilon))
@@ -146,7 +146,7 @@ def layernorm_composite(x, scale, bias, epsilon, begin_norm_axis):
 
 
 @register_decomp('pd_op.dropout')
-def dropout_composite(x, seed_tensor, p, is_test, mode, seed, fix_seed):
+def dropout(x, seed_tensor, p, is_test, mode, seed, fix_seed):
     """define composite rule of op dropout.
     upscale_in_train:
         train: out = input * mask / ( 1.0 - p )
@@ -207,7 +207,7 @@ def bernoulli(shape, dtype, p, seed=0):
 
 
 @register_decomp('pd_op.add_n')
-def sum_composite(x):
+def add_n(x):
     ans = x[0]
     for xi in x[1:]:
         ans = xi + ans