PaddlePaddle
diff --git a/‎paddle/fluid/operators/generator/parse_utils.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/operators/generator/parse_utils.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/operators/ops_signature/reshape_sig.cc
Copy file name to clipboard
+11-8Lines changed: 11 additions & 8 deletions b/‎paddle/fluid/operators/ops_signature/reshape_sig.cc
Copy file name to clipboard
+11-8Lines changed: 11 additions & 8 deletions
diff --git a/‎paddle/fluid/operators/reshape_op.cc
Copy file name to clipboard
+16-15Lines changed: 16 additions & 15 deletions b/‎paddle/fluid/operators/reshape_op.cc
Copy file name to clipboard
+16-15Lines changed: 16 additions & 15 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
Copy file name to clipboard
-4Lines changed: 0 additions & 4 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
Copy file name to clipboard
-4Lines changed: 0 additions & 4 deletions
diff --git a/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+3-5Lines changed: 3 additions & 5 deletions b/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+3-5Lines changed: 3 additions & 5 deletions
diff --git a/‎paddle/phi/api/generator/api_base.py
Copy file name to clipboard
+41-9Lines changed: 41 additions & 9 deletions b/‎paddle/phi/api/generator/api_base.py
Copy file name to clipboard
+41-9Lines changed: 41 additions & 9 deletions
diff --git a/‎paddle/phi/api/generator/dist_api_gen.py
Copy file name to clipboard
+4-6Lines changed: 4 additions & 6 deletions b/‎paddle/phi/api/generator/dist_api_gen.py
Copy file name to clipboard
+4-6Lines changed: 4 additions & 6 deletions
diff --git a/‎paddle/phi/infermeta/backward.cc
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions b/‎paddle/phi/infermeta/backward.cc
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions
diff --git a/‎paddle/phi/infermeta/backward.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎paddle/phi/infermeta/backward.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/reshape.cc
Copy file name to clipboard
+14-19Lines changed: 14 additions & 19 deletions b/‎paddle/phi/infermeta/spmd_rules/reshape.cc
Copy file name to clipboard
+14-19Lines changed: 14 additions & 19 deletions
@@ -695,7 +695,7 @@ def cross_validate(ops):
 
                 assert len(fw_call["outputs"]) == len(
                     fw_op["outputs"]
-                ), f"{name}: forward call has more outputs than the op "
+                ), f"{name}: requires outputs number of fw_call == fw_op, but received {fw_call['outputs']} != {fw_op['outputs']}"
                 for output, output_ in zip(
                     fw_call["outputs"], fw_op["outputs"]
                 ):
 
@@ -20,26 +20,29 @@ KernelSignature ReshapeOpArgumentMapping(const ArgumentMappingContext& ctx) {
   if (ctx.HasOutput("XShape")) {
     if (ctx.InputSize("ShapeTensor") > 0) {
       return KernelSignature(
-          "reshape", {"X"}, {"ShapeTensor"}, {"Out", "XShape"});
+          "reshape_with_xshape", {"X"}, {"ShapeTensor"}, {"Out", "XShape"});
     } else if (ctx.HasInput("Shape")) {
-      return KernelSignature("reshape", {"X"}, {"Shape"}, {"Out", "XShape"});
+      return KernelSignature(
+          "reshape_with_xshape", {"X"}, {"Shape"}, {"Out", "XShape"});
     } else {
-      return KernelSignature("reshape", {"X"}, {"shape"}, {"Out", "XShape"});
+      return KernelSignature(
+          "reshape_with_xshape", {"X"}, {"shape"}, {"Out", "XShape"});
     }
   } else {
     if (ctx.InputSize("ShapeTensor") > 0) {
-      return KernelSignature("reshape_infer", {"X"}, {"ShapeTensor"}, {"Out"});
+      return KernelSignature("reshape", {"X"}, {"ShapeTensor"}, {"Out"});
     } else if (ctx.HasInput("Shape")) {
-      return KernelSignature("reshape_infer", {"X"}, {"Shape"}, {"Out"});
+      return KernelSignature("reshape", {"X"}, {"Shape"}, {"Out"});
     } else {
-      return KernelSignature("reshape_infer", {"X"}, {"shape"}, {"Out"});
+      return KernelSignature("reshape", {"X"}, {"shape"}, {"Out"});
     }
   }
 }
 
 KernelSignature ReshapeGradOpArgumentMapping(
     const ArgumentMappingContext& ctx UNUSED) {
-  return KernelSignature("reshape_grad", {"Out@GRAD"}, {}, {"X@GRAD"});
+  return KernelSignature(
+      "reshape_grad", {"XShape", "Out@GRAD"}, {}, {"X@GRAD"});
 }
 
 KernelSignature ReshapeDoubleGradOpArgumentMapping(
@@ -49,7 +52,7 @@ KernelSignature ReshapeDoubleGradOpArgumentMapping(
 
 }  // namespace phi
 
-PD_REGISTER_BASE_KERNEL_NAME(reshape2, reshape);
+PD_REGISTER_BASE_KERNEL_NAME(reshape2, reshape_with_xshape);
 PD_REGISTER_BASE_KERNEL_NAME(reshape2_grad, reshape_grad);
 PD_REGISTER_BASE_KERNEL_NAME(reshape2_grad_grad, reshape_double_grad);
 
 
@@ -421,27 +421,27 @@ class ReshapeKernel {
     }
     if (ctx.GetPlace().GetType() == phi::AllocationType::CPU) {
       auto &dev_ctx = ctx.device_context<phi::CPUContext>();
-      phi::ReshapeInferKernel(static_cast<const phi::CPUContext &>(dev_ctx),
-                              *in,
-                              pt_scalar_shape,
-                              out);
+      phi::ReshapeKernel(static_cast<const phi::CPUContext &>(dev_ctx),
+                         *in,
+                         pt_scalar_shape,
+                         out);
     }
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
     if (ctx.GetPlace().GetType() == phi::AllocationType::GPU) {
       auto &dev_ctx = ctx.device_context<phi::GPUContext>();
-      phi::ReshapeInferKernel(static_cast<const phi::GPUContext &>(dev_ctx),
-                              *in,
-                              pt_scalar_shape,
-                              out);
+      phi::ReshapeKernel(static_cast<const phi::GPUContext &>(dev_ctx),
+                         *in,
+                         pt_scalar_shape,
+                         out);
     }
 #endif
 #ifdef PADDLE_WITH_XPU
     if (ctx.GetPlace().GetType() == phi::AllocationType::XPU) {
       auto &dev_ctx = ctx.device_context<phi::XPUContext>();
-      phi::ReshapeInferKernel(static_cast<const phi::XPUContext &>(dev_ctx),
-                              *in,
-                              pt_scalar_shape,
-                              out);
+      phi::ReshapeKernel(static_cast<const phi::XPUContext &>(dev_ctx),
+                         *in,
+                         pt_scalar_shape,
+                         out);
     }
 #endif
   }
@@ -450,27 +450,28 @@ class ReshapeKernel {
 class ReshapeGradKernel {
  public:
   void operator()(const framework::ExecutionContext &ctx) const {
+    auto *x = ctx.Input<phi::DenseTensor>("X");
     auto *d_out = ctx.Input<phi::DenseTensor>(framework::GradVarName("Out"));
     auto *d_x = ctx.Output<phi::DenseTensor>(framework::GradVarName("X"));
     d_x->mutable_data(ctx.GetPlace(), d_out->type());
 
     if (ctx.GetPlace().GetType() == phi::AllocationType::CPU) {
       auto &dev_ctx = ctx.device_context<phi::CPUContext>();
       phi::ReshapeGradKernel(
-          static_cast<const phi::CPUContext &>(dev_ctx), *d_out, d_x);
+          static_cast<const phi::CPUContext &>(dev_ctx), *x, *d_out, d_x);
     }
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
     if (ctx.GetPlace().GetType() == phi::AllocationType::GPU) {
       auto &dev_ctx = ctx.device_context<phi::GPUContext>();
       phi::ReshapeGradKernel(
-          static_cast<const phi::GPUContext &>(dev_ctx), *d_out, d_x);
+          static_cast<const phi::GPUContext &>(dev_ctx), *x, *d_out, d_x);
     }
 #endif
 #ifdef PADDLE_WITH_XPU
     if (ctx.GetPlace().GetType() == phi::AllocationType::XPU) {
       auto &dev_ctx = ctx.device_context<phi::XPUContext>();
       phi::ReshapeGradKernel(
-          static_cast<const phi::XPUContext &>(dev_ctx), *d_out, d_x);
+          static_cast<const phi::XPUContext &>(dev_ctx), *x, *d_out, d_x);
     }
 #endif
   }
 
@@ -974,10 +974,6 @@ def gen_infermeta_func_str(args, op_info):
                 spmd_params = op_info.kernel_map['param']
         else:
             spmd_params = op_info.input_name_list
-        # TODO(GhostScreaming): specialized case for reshape_grad
-        # xshape is not kernel params, but inferspmd needs it.
-        if "reshape_grad" in op_info.kernel_map['func'][0]:
-            spmd_params = ["x"] + spmd_params
     op_info.spmd_params = spmd_params
 
     infermeta_inputs_str = get_infermeta_inputs_str(
 
@@ -402,12 +402,10 @@ void reduce_as_grad(const Tensor& x,
 }
 
 template <typename T>
-void reshape_grad(const Tensor& xshape,
-                  const Tensor& grad_out,
-                  Tensor* grad_x) {
+void reshape_grad(const Tensor& x, const Tensor& grad_out, Tensor* grad_x) {
   if (grad_x) {
-    auto xshape_dims = xshape.dims();
-    auto grad_x_tmp = reshape<T>(grad_out, common::vectorize(xshape_dims));
+    const auto& x_dims = x.dims();
+    auto grad_x_tmp = reshape<T>(grad_out, common::vectorize(x_dims));
     set_output<T>(grad_x_tmp, grad_x);
   }
 }
 
@@ -19,6 +19,7 @@
 
 PREFIX_TENSOR_NAME = 'input_'
 PREFIX_META_TENSOR_NAME = 'meta_'
+ORIGIN_PREFIX_TENSOR_NAME = 'origin_input_'
 
 
 def parse_plain_list(s: str, sep=",") -> list[str]:
@@ -55,6 +56,8 @@ def __init__(self, api_item_yaml):
 
         self.is_base_api = True
         self.is_only_composite_api = False
+        # Whether to generate code for inplace API
+        self.is_inplace_context = False
         if 'invoke' in api_item_yaml:
             self.is_base_api = False
             self.invoke = api_item_yaml['invoke']
@@ -93,6 +96,12 @@ def get_api_name(self, api_item_yaml):
     def get_api_func_name(self):
         return self.api
 
+    def is_inplace_input(self, input_name):
+        is_inplace_api = (
+            self.get_api_func_name()[-1] == "_" or self.is_inplace_context
+        )
+        return is_inplace_api and input_name in self.inplace_map.values()
+
     def get_input_tensor_args(self, inplace_flag=False):
         input_args = []
         inplace_type_map = {
@@ -599,13 +608,25 @@ def gene_infer_meta(self, kernel_output_names, code_indent) -> str:
         for param in infer_meta_params:
             if param in input_names:
                 if self.inputs['input_info'][param] == "const Tensor&":
-                    param_code = (
-                        param_code
-                        + "MakeMetaTensor(*"
-                        + PREFIX_TENSOR_NAME
-                        + param
-                        + "), "
-                    )
+                    if self.is_inplace_input(param):
+                        meta_tensor_code += f"""
+{code_indent}  auto {ORIGIN_PREFIX_TENSOR_NAME}{param} = *{PREFIX_TENSOR_NAME}{param};
+"""
+                        param_code = (
+                            param_code
+                            + "MakeMetaTensor("
+                            + ORIGIN_PREFIX_TENSOR_NAME
+                            + param
+                            + "), "
+                        )
+                    else:
+                        param_code = (
+                            param_code
+                            + "MakeMetaTensor(*"
+                            + PREFIX_TENSOR_NAME
+                            + param
+                            + "), "
+                        )
                 elif (
                     self.inputs['input_info'][param]
                     == "const std::vector<Tensor>&"
@@ -1166,7 +1187,12 @@ def get_kernel_args(self, kernel_tensor_type=None, code_indent=''):
                     kernel_args.append(PREFIX_TENSOR_NAME + param)
                 else:
                     if self.inputs['input_info'][param] == "const Tensor&":
-                        kernel_args.append("*" + PREFIX_TENSOR_NAME + param)
+                        if self.is_inplace_input(param):
+                            kernel_args.append(
+                                ORIGIN_PREFIX_TENSOR_NAME + param
+                            )
+                        else:
+                            kernel_args.append("*" + PREFIX_TENSOR_NAME + param)
                     elif (
                         self.inputs['input_info'][param]
                         == "const std::vector<Tensor>&"
@@ -1371,6 +1397,12 @@ def gene_dispatch_code(self, kernel_name, inplace_flag=False):
   }}
 """
 
+    def gene_base_api_code_for_inplace(self):
+        self.is_inplace_context = True
+        code = self.gene_base_api_code(inplace_flag=True)
+        self.is_inplace_context = False
+        return code
+
     def gene_base_api_code(self, inplace_flag=False):
         api_func_name = self.get_api_func_name()
         if inplace_flag and api_func_name[-1] != '_':
@@ -1416,7 +1448,7 @@ def gene_api_code(self):
             if len(self.inplace_map) > 0:
                 if self.api[-1] == '_':
                     api_code = ""
-                api_code = api_code + self.gene_base_api_code(inplace_flag=True)
+                api_code = api_code + self.gene_base_api_code_for_inplace()
             return api_code
         elif self.is_only_composite_api:
             # for composite and invoke api, dygraph use prim::xxx_grad method
 
@@ -866,11 +866,6 @@ def generate_specialized_infer_spmd_code(self) -> str:
         if kernel_params is None:
             kernel_params = input_names + attr_names
 
-        # TODO(GhostScreaming): specialized case for reshape_grad
-        # xshape is not kernel params, but inferspmd needs it.
-        if "reshape_grad" in self.kernel['func'][0]:
-            kernel_params = ["xshape"] + kernel_params
-
         input_decl_code = ""
         input_args_code = ""
         for param in kernel_params:
@@ -1275,6 +1270,7 @@ def generate_infer_global_shape_code(self) -> str:
                     and self.infer_meta['global_shape'] is not None
                     and self.outputs['names'][i]
                     == self.infer_meta['global_shape']
+                    and i > 0
                 ):
                     output_decl_code += (
                         SINGLE_GLOBAL_META_OUT_DECL_TEMPLATE.format(
@@ -1597,7 +1593,9 @@ def generate_infer_meta_code(self) -> str:
                 ] and self.need_to_generate_code_for_inplace_impl(i):
                     infer_meta_code += SET_DIMS_TEMPLATE.format(
                         dst=self.dist_output_args[i],
-                        src=self.dist_output_args[i] + '_tmp',
+                        src=self.dist_output_args[i] + '_tmp'
+                        if i > 0
+                        else self.dist_output_args[i],
                     )
 
         # TODO(GhostScreaming): kernel like reshape need calculate local_shape
 
@@ -962,12 +962,12 @@ void KernelWithXShapeInferMeta(const MetaTensor& xshape,
   dx->share_lod(xshape);
 }
 
-void ReshapeGradInferMeta(const MetaTensor& xshape,
-                          const MetaTensor& out,
-                          MetaTensor* dx) {
-  dx->set_dims(xshape.dims());
+void GradSameWithXInferMeta(const MetaTensor& x,
+                            const MetaTensor& out,
+                            MetaTensor* dx) {
+  dx->set_dims(x.dims());
   dx->set_dtype(out.dtype());
-  dx->share_lod(xshape);
+  dx->share_lod(x);
 }
 
 void LUGradInferMeta(const MetaTensor& x,
 
@@ -352,13 +352,13 @@ void InverseGradInferMeta(const MetaTensor& out,
                           const MetaTensor& dout,
                           MetaTensor* dx);
 
-void KernelWithXShapeInferMeta(const MetaTensor& xshape,
+void KernelWithXShapeInferMeta(const MetaTensor& x,
                                const MetaTensor& out,
                                MetaTensor* dx);
 
-void ReshapeGradInferMeta(const MetaTensor& xshape,
-                          const MetaTensor& out,
-                          MetaTensor* dx);
+void GradSameWithXInferMeta(const MetaTensor& xshape,
+                            const MetaTensor& out,
+                            MetaTensor* dx);
 
 void LUGradInferMeta(const MetaTensor& x,
                      const MetaTensor& out,
 
@@ -153,10 +153,6 @@ SpmdInfo ReshapeInferSpmd(const DistMetaTensor& x,
                           const std::vector<int64_t>& shape) {
   // Step0: Verify input args based on reshape logic
   auto x_shape = phi::vectorize(x.dims());
-  // For dynamic mode, deal with extra xshape dim.
-  if (x_shape[0] == 0) {
-    x_shape.erase(x_shape.begin());
-  }
 
   int x_ndim = static_cast<int>(x_shape.size());
   int out_ndim = static_cast<int>(shape.size());
@@ -317,6 +313,8 @@ SpmdInfo ReshapeInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr}, {out_dist_attr_dst}};
 }
 
+// FIXME(dev): XShape will be decprecated in the future, so we
+// need unify inferSpmd into ReshapeInferSpmd function.
 SpmdInfo ReshapeInferSpmdDynamic(const DistMetaTensor& x,
                                  const std::vector<int64_t>& shape) {
   auto spmd_info = ReshapeInferSpmd(x, shape);
@@ -328,26 +326,23 @@ SpmdInfo ReshapeInferSpmdDynamic(const DistMetaTensor& x,
   return spmd_info;
 }
 
-SpmdInfo ReshapeGradInferSpmd(const DistMetaTensor& x_shape,
+SpmdInfo ReshapeGradInferSpmd(const DistMetaTensor& x,
                               const DistMetaTensor& out_grad) {
   std::vector<int64_t> out_grad_shape = common::vectorize(out_grad.dims());
-  auto x_shape_dist_tmp = x_shape.dist_attr();
-  auto x_dims_mapping = x_shape_dist_tmp.dims_mapping();
-  x_dims_mapping.erase(x_dims_mapping.begin());
-  x_shape_dist_tmp.set_dims_mapping(x_dims_mapping);
-  auto tmp = ReshapeInferSpmd(DistMetaTensor(x_shape.dims(), x_shape_dist_tmp),
-                              out_grad_shape);
+  auto x_dist_tmp = x.dist_attr();
+  auto tmp =
+      ReshapeInferSpmd(DistMetaTensor(x.dims(), x_dist_tmp), out_grad_shape);
   // check no shard is needed
-  const auto& x_shape_dist_dst = PADDLE_GET_CONST(TensorDistAttr, tmp.first[0]);
+  const auto& x_dist_dst = PADDLE_GET_CONST(TensorDistAttr, tmp.first[0]);
   const auto& out_grad_dist_dst =
       PADDLE_GET_CONST(TensorDistAttr, tmp.second[0]);
-  PADDLE_ENFORCE_EQ(x_shape_dist_tmp.dims_mapping(),
-                    x_shape_dist_dst.dims_mapping(),
-                    common::errors::InvalidArgument(
-                        "x_shape should not be re shared: [%s] => [%s]",
-                        x_shape_dist_tmp.to_string(),
-                        x_shape_dist_dst.to_string()));
-  return {{out_grad_dist_dst}, {x_shape_dist_dst}};
+  PADDLE_ENFORCE_EQ(
+      x_dist_tmp.dims_mapping(),
+      x_dist_dst.dims_mapping(),
+      common::errors::InvalidArgument("x should not be re shared: [%s] => [%s]",
+                                      x_dist_tmp.to_string(),
+                                      x_dist_dst.to_string()));
+  return {{x_dist_tmp, out_grad_dist_dst}, {x_dist_dst}};
 }
 
 SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x,
Original file line number	Diff line number	Diff line change
`@@ -402,12 +402,10 @@ void reduce_as_grad(const Tensor& x,`
`402`	`402`	`}`
`403`	`403`
`404`	`404`	`template <typename T>`
`405`		`-void reshape_grad(const Tensor& xshape,`
`406`		`- const Tensor& grad_out,`
`407`		`- Tensor* grad_x) {`
	`405`	`+void reshape_grad(const Tensor& x, const Tensor& grad_out, Tensor* grad_x) {`
`408`	`406`	`if (grad_x) {`
`409`		`- auto xshape_dims = xshape.dims();`
`410`		`- auto grad_x_tmp = reshape<T>(grad_out, common::vectorize(xshape_dims));`
	`407`	`+ const auto& x_dims = x.dims();`
	`408`	`+ auto grad_x_tmp = reshape<T>(grad_out, common::vectorize(x_dims));`
`411`	`409`	`set_output<T>(grad_x_tmp, grad_x);`
`412`	`410`	`}`
`413`	`411`	`}`