PaddlePaddle
diff --git a/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+13-9Lines changed: 13 additions & 9 deletions b/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+13-9Lines changed: 13 additions & 9 deletions
diff --git a/‎paddle/fluid/operators/ops_signature/flatten_sig.cc
Copy file name to clipboard
+6-4Lines changed: 6 additions & 4 deletions b/‎paddle/fluid/operators/ops_signature/flatten_sig.cc
Copy file name to clipboard
+6-4Lines changed: 6 additions & 4 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
Copy file name to clipboard
-6Lines changed: 0 additions & 6 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
Copy file name to clipboard
-6Lines changed: 0 additions & 6 deletions
diff --git a/‎paddle/fluid/pir/transforms/gpu/transpose_flatten_concat_fuse_pass.cc
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions b/‎paddle/fluid/pir/transforms/gpu/transpose_flatten_concat_fuse_pass.cc
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions
diff --git a/‎paddle/fluid/primitive/composite/composite.h
Copy file name to clipboard
+6-19Lines changed: 6 additions & 19 deletions b/‎paddle/fluid/primitive/composite/composite.h
Copy file name to clipboard
+6-19Lines changed: 6 additions & 19 deletions
diff --git a/‎paddle/phi/backends/xpu/xpu2_op_list.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/backends/xpu/xpu2_op_list.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/backends/xpu/xpu3_op_list.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/backends/xpu/xpu3_op_list.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/flatten.cc
Copy file name to clipboard
+6-10Lines changed: 6 additions & 10 deletions b/‎paddle/phi/infermeta/spmd_rules/flatten.cc
Copy file name to clipboard
+6-10Lines changed: 6 additions & 10 deletions
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+16-15Lines changed: 16 additions & 15 deletions b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+16-15Lines changed: 16 additions & 15 deletions
diff --git a/‎paddle/phi/kernels/flatten_grad_kernel.cc
Copy file name to clipboard
+5-3Lines changed: 5 additions & 3 deletions b/‎paddle/phi/kernels/flatten_grad_kernel.cc
Copy file name to clipboard
+5-3Lines changed: 5 additions & 3 deletions
@@ -3591,12 +3591,13 @@ struct QuantizeLinearOpTranscriber : public OpTranscriber {
   }
 };
 
-struct Reshape2GradOpTranscriber : public OpTranscriber {
+template <typename OpT>
+struct WithXShapeGradOpTranscriber : public OpTranscriber {
   pir::Operation* operator()(pir::IrContext* ctx,
                              TranslationContext* param_map,
                              const OpDesc& op_desc,
                              pir::Block* block) override {
-    VLOG(4) << "Translate Reshape2Grad...";
+    VLOG(4) << "Translate " << op_desc.Type() << ".....";
     pir::Builder builder(ctx, block);
     auto& input_xshape_name = op_desc.Input("XShape")[0];
     auto& input_outgrad_name = op_desc.Input("Out@GRAD")[0];
@@ -3645,13 +3646,12 @@ struct Reshape2GradOpTranscriber : public OpTranscriber {
         ::common::errors::InvalidArgument(
             "input type must be DenseTensorType, but received: %s.",
             input_outgrad_value.type()));
+    // NOTE(Aurelius84): Even though we use xshape to construct grad op,
+    // but in GradKernel we still use dx->dims by default.
+    OpT grad_op = builder.Build<OpT>(xshape_value, input_outgrad_value);
+    param_map->PushValue(out_name, grad_op.result(0));
 
-    dialect::ReshapeGradOp reshape_grad_op =
-        builder.Build<dialect::ReshapeGradOp>(xshape_value,
-                                              input_outgrad_value);
-    param_map->PushValue(out_name, reshape_grad_op.result(0));
-
-    return reshape_grad_op.operation();
+    return grad_op.operation();
   }
 };
 
@@ -3747,7 +3747,11 @@ OpTranslator::OpTranslator() {
   special_handlers["c_embedding"] = CEmbeddingOpTranscriber();
   special_handlers["quantize_linear"] = QuantizeLinearOpTranscriber();
   special_handlers["dequantize_linear"] = QuantizeLinearOpTranscriber();
-  special_handlers["reshape2_grad"] = Reshape2GradOpTranscriber();
+  // To process Op with XShape output in old IR
+  special_handlers["reshape2_grad"] =
+      WithXShapeGradOpTranscriber<dialect::ReshapeGradOp>();
+  special_handlers["flatten_contiguous_range_grad"] =
+      WithXShapeGradOpTranscriber<dialect::FlattenGradOp>();
 }
 
 }  // namespace translator
 
@@ -22,11 +22,13 @@ KernelSignature FlattenOpArgumentMapping(const ArgumentMappingContext& ctx) {
         "flatten", {"X"}, {"start_axis", "stop_axis"}, {"Out", "XShape"});
   }
   if (ctx.HasOutput("XShape")) {
-    return KernelSignature(
-        "flatten", {"X"}, {"start_axis", "stop_axis"}, {"Out", "XShape"});
+    return KernelSignature("flatten_with_xshape",
+                           {"X"},
+                           {"start_axis", "stop_axis"},
+                           {"Out", "XShape"});
   } else {
     return KernelSignature(
-        "flatten_infer", {"X"}, {"start_axis", "stop_axis"}, {"Out"});
+        "flatten", {"X"}, {"start_axis", "stop_axis"}, {"Out"});
   }
 }
 
@@ -39,7 +41,7 @@ KernelSignature FlattenGradOpArgumentMapping(
 }  // namespace phi
 
 PD_REGISTER_BASE_KERNEL_NAME(flatten_contiguous_range, flatten);
-PD_REGISTER_BASE_KERNEL_NAME(flatten_contiguous_range_grad, flatten_grad);
+PD_REGISTER_BASE_KERNEL_NAME(flatten_contiguous_range, flatten_with_xshape);
 
 PD_REGISTER_ARG_MAPPING_FN(flatten_contiguous_range,
                            phi::FlattenOpArgumentMapping);
 
@@ -968,12 +968,6 @@ bool FlattenOpInferSymbolicShape(
   symbol::ShapeOrDataDimExprs out_shape_data{
       symbol::TensorShapeOrDataDimExprs(out_shape)};
   infer_context->SetShapeOrDataForValue(op->result(0), out_shape_data);
-
-  std::vector<symbol::DimExpr> xshape_shape = x_shape;
-  xshape_shape.insert(xshape_shape.begin(), symbol::DimExpr{0});
-  symbol::ShapeOrDataDimExprs xshape_shape_data{
-      symbol::TensorShapeOrDataDimExprs(xshape_shape)};
-  infer_context->SetShapeOrDataForValue(op->result(1), xshape_shape_data);
   return true;
 }
 
 
@@ -52,8 +52,7 @@ class NTransposeFlattenConcatFusePattern : public paddle::drr::DrrPatternBase {
       transpose_op({&pat.Tensor("transpose_in_" + std::to_string(i))},
                    {&pat.Tensor("transpose_out_" + std::to_string(i))});
       flatten_op({&pat.Tensor("transpose_out_" + std::to_string(i))},
-                 {&pat.Tensor("flatten_out_" + std::to_string(i)),
-                  &pat.Tensor("flatten_xshape_" + std::to_string(i))});
+                 {&pat.Tensor("flatten_out_" + std::to_string(i))});
       combine_in.push_back(&pat.Tensor("flatten_out_" + std::to_string(i)));
     }
     const auto &combine_op = pat.Op(pir::CombineOp::name());
 
@@ -1094,9 +1094,7 @@ std::tuple<Tensor, Tensor, Tensor> instance_norm_decomp(
 }
 
 template <typename T>
-std::tuple<Tensor, Tensor> flatten_decomp(const Tensor& x,
-                                          int start_axis,
-                                          int end_axis) {
+Tensor flatten_decomp(const Tensor& x, int start_axis, int end_axis) {
   auto x_dim = x.shape();
   if (x_dim.size() == 0) {
     start_axis = 0;
@@ -1109,17 +1107,8 @@ std::tuple<Tensor, Tensor> flatten_decomp(const Tensor& x,
 
   if (has_dynamic_shape(x.shape())) {
     auto x_shape = shape<T>(x);
-    Tensor x_shape_tensor = full<T>({1}, 0, x_shape.dtype());
-    std::vector<Tensor> tmp_shape;
-    tmp_shape.push_back(x_shape_tensor);
-    for (size_t i = 0; i < x_dim.size(); i++) {
-      tmp_shape.push_back(get_slice<T>(x_shape, i));
-    }
-    x_shape_tensor = concat<T>(tmp_shape);
-    x_shape_tensor =
-        backend::full_with_tensor<T>(x_shape_tensor, 0.0, DataType::FLOAT32);
     if (end_axis == start_axis) {
-      return std::make_tuple(backend::reshape<T>(x, x_shape), x_shape_tensor);
+      return backend::reshape<T>(x, x_shape);
     }
     std::vector<Tensor> out_shape;
 
@@ -1139,18 +1128,16 @@ std::tuple<Tensor, Tensor> flatten_decomp(const Tensor& x,
     }
 
     Tensor out_shape_tensor = concat<T>(out_shape);
-    return std::make_tuple(backend::reshape<T>(x, out_shape_tensor),
-                           x_shape_tensor);
+    return backend::reshape<T>(x, out_shape_tensor);
   } else {
     std::vector<int64_t> tmp_shape(x_dim);
     tmp_shape.insert(tmp_shape.begin(), 0);
-    auto xshape = full<T>(tmp_shape, 0.0, DataType::FLOAT32);
     if (x_dim.size() == 0) {
       std::vector<int64_t> res_shape(1, 1);
-      return std::make_tuple(reshape<T>(x, res_shape), xshape);
+      return reshape<T>(x, res_shape);
     }
     if (end_axis == start_axis) {
-      return std::make_tuple(reshape<T>(x, x_dim), xshape);
+      return reshape<T>(x, x_dim);
     }
 
     int slice_numel = 1;
@@ -1166,7 +1153,7 @@ std::tuple<Tensor, Tensor> flatten_decomp(const Tensor& x,
       out_shape.push_back(x_dim[i]);
     }
 
-    return std::make_tuple(reshape<T>(x, out_shape), xshape);
+    return reshape<T>(x, out_shape);
   }
 }
 
 
@@ -445,6 +445,8 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::INT8,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32})},
       {"flatten",
        XPUKernelSet({phi::DataType::INT64,
 
@@ -468,6 +468,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::INT8,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32})},
       {"flatten",
        XPUKernelSet({phi::DataType::INT64,
 
@@ -105,7 +105,7 @@ SpmdInfo FlattenInferSpmd(const DistMetaTensor& x,
                                       x_ndim,
                                       x_dims_mapping.size()));
 
-  // obtain target shape and use ReshapeInferSpmdDynamic to infer
+  // obtain target shape and use ReshapeInferSpmd to infer
   start_axis = PreprocessAxis(start_axis, x_ndim);
   stop_axis = PreprocessAxis(stop_axis, x_ndim);
   std::vector<int64_t> dst_shape;
@@ -125,8 +125,8 @@ SpmdInfo FlattenInferSpmd(const DistMetaTensor& x,
   VLOG(4) << "Start_axis: " << start_axis;
   VLOG(4) << "Stop_axis: " << stop_axis;
   VLOG(4) << "FlattenInferSpmd: output shape: [" << str_join(dst_shape) << "]";
-  VLOG(4) << "use ReshapeInferSpmdDynamic to infer distributed attribute";
-  return ReshapeInferSpmdDynamic(x, dst_shape);
+  VLOG(4) << "use ReshapeInferSpmd to infer distributed attribute";
+  return ReshapeInferSpmd(x, dst_shape);
 }
 
 // TODO(jeff41404): consider xshape and use ReshapeInferSpmdReverse in future
@@ -188,15 +188,11 @@ SpmdInfo FlattenInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr}, {out_dist_attr_dst}};
 }
 
-SpmdInfo FlattenGradInferSpmd(const DistMetaTensor& xshape,
+SpmdInfo FlattenGradInferSpmd(const DistMetaTensor& x,
                               const DistMetaTensor& out_grad) {
-  // TODO(jeff41404): when ReshapeInferSpmd and ReshapeGradInferSpmd can deliver
-  // distributed attribute of xshape, we will use ReshapeGradInferSpmd directly
-  // in future return ReshapeGradInferSpmd(xshape, out_grad);
-  auto shape = phi::vectorize(xshape.dims());
-  shape = std::vector<int64_t>(shape.begin() + 1, shape.end());
+  auto shape = phi::vectorize(x.dims());
   const auto& spmd = ReshapeInferSpmd(out_grad, shape);
-  return {{xshape.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
+  return {{x.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
 }
 
 }  // namespace phi::distributed
@@ -1562,18 +1562,27 @@ void FFTR2CInferMeta(const MetaTensor& x,
   }
 }
 
-void FlattenInferMeta(const MetaTensor& x,
-                      int start_axis,
-                      int stop_axis,
-                      MetaTensor* out) {
-  FlattenWithXShapeInferMeta(x, start_axis, stop_axis, out, nullptr);
-}
-
 void FlattenWithXShapeInferMeta(const MetaTensor& x,
                                 int start_axis,
                                 int stop_axis,
                                 MetaTensor* out,
                                 MetaTensor* xshape) {
+  FlattenInferMeta(x, start_axis, stop_axis, out);
+  if (xshape == nullptr) return;
+  const auto& x_dims = x.dims();
+  std::vector<int64_t> xshape_dims(x_dims.size() + 1);
+  xshape_dims[0] = 0;
+  for (int i = 0; i < x_dims.size(); ++i) {
+    xshape_dims[i + 1] = x_dims[i];
+  }
+  xshape->set_dims(common::make_ddim(xshape_dims));
+  xshape->share_lod(x);
+}
+
+void FlattenInferMeta(const MetaTensor& x,
+                      int start_axis,
+                      int stop_axis,
+                      MetaTensor* out) {
   auto x_dims = x.dims();
   int in_dims_size = x_dims.size();
 
@@ -1635,14 +1644,6 @@ void FlattenWithXShapeInferMeta(const MetaTensor& x,
     // are the same.
     out->share_lod(x);
   }
-  if (xshape == nullptr) return;
-  std::vector<int64_t> xshape_dims(x_dims.size() + 1);
-  xshape_dims[0] = 0;
-  for (int i = 0; i < x_dims.size(); ++i) {
-    xshape_dims[i + 1] = x_dims[i];
-  }
-  xshape->set_dims(common::make_ddim(xshape_dims));
-  xshape->share_lod(x);
 }
 
 void Flatten2InferMeta(const MetaTensor& x,
 
@@ -22,12 +22,14 @@ namespace phi {
 
 template <typename T, typename Context>
 void FlattenGradKernel(const Context& dev_ctx,
-                       const DenseTensor& xshape,
+                       const DenseTensor& x,
                        const DenseTensor& out_grad,
                        DenseTensor* x_grad) {
-  auto xshape_dims = xshape.dims();
+  // NOTE: [Why not to use x.dims() ?]
+  // Because inplace strategy is different between old IR and PIR,
+  // we need fix it into x.dims() after cleaning old IR system.
+  auto x_dims = x_grad->dims();
   dev_ctx.Alloc(x_grad, out_grad.dtype());
-  auto x_dims = common::slice_ddim(xshape_dims, 1, xshape_dims.size());
   phi::Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
   x_grad->Resize(x_dims);
 }