Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into custom_stride

Author: Tongkaio

paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc

@@ -21,6 +21,7 @@
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/common/ddim.h"
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
@@ -107,8 +108,12 @@ bool MakeGenerateShapeOpAttribute(
     std::vector<pir::Attribute>* output_dim_expr_attrs,
     GenerateShapeOp::SymbolBindings* symbol_bindings) {
   const auto& shape_or_data_dim_exprs = ShapeOrDataDimExprs4Value(output_shape);
-  CHECK(shape_or_data_dim_exprs.data().has_value());
-  const auto& out_dim_exprs = shape_or_data_dim_exprs.data().value();
+  ExprVec data_vec =
+      paddle::dialect::details::GetExprVecFromData(shape_or_data_dim_exprs);
+  // CHECK(shape_or_data_dim_exprs.data().has_value());
+  CHECK(data_vec.size());
+  // const auto& out_dim_exprs = shape_or_data_dim_exprs.data().value();
+  const auto& out_dim_exprs = data_vec;
   return MakeGenerateShapeOpAttribute(ir_context,
                                       ShapeOrDataDimExprs4Value,
                                       out_dim_exprs,

paddle/cinn/hlir/framework/pir/group.cc

@@ -52,7 +52,6 @@ std::shared_ptr<Group> Group::Clone(::pir::Block* target_block,
 
   new_group->input_names = this->input_names;
   new_group->output_names = this->output_names;
-  new_group->output_values = this->output_values;
   new_group->fn_name = this->fn_name;
   new_group->int_args_map = this->int_args_map;
   new_group->alignment_schedule_info = this->alignment_schedule_info;

paddle/fluid/framework/operator.cc

@@ -1704,6 +1704,11 @@ void OperatorWithKernel::RunImpl(const Scope& scope,
     all_kernels_must_compute_runtime_shape_ = true;
   const Scope* cur_scope = &scope;
   CheckWhetherPreparePhiData(Inputs(), Outputs(), scope);
+#if defined(PADDLE_WITH_XPU)
+  if (std::getenv("XPU_NEED_PREPARE_PHI_DATA") != nullptr) {
+    need_prepare_phi_data_ = atoi(std::getenv("XPU_NEED_PREPARE_PHI_DATA"));
+  }
+#endif
   if (!enable_cache_runtime_context_) {
     RuntimeContext ctx(Inputs(), Outputs(), scope);
     RunImpl(scope, place, &ctx);

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.cc

@@ -23,7 +23,9 @@ bool InferSymbolicShapeElementWiseBinary(
   // For ElementWiseBinary ops, if the input tensor is from full op, the value
   // of fullop is useless, only the shape need doing broadcast
   bool x_from_fullop =
-      op->operand_source(0).defining_op()->isa<paddle::dialect::FullOp>();
+      op->operand_source(0).defining_op()
+          ? op->operand_source(0).defining_op()->isa<paddle::dialect::FullOp>()
+          : false;
   if (!x_from_fullop && x_shapeordata.data().has_value()) {
     shape_0 = x_shapeordata.data().value();
   } else {
@@ -34,7 +36,9 @@ bool InferSymbolicShapeElementWiseBinary(
       shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
   std::vector<symbol::DimExpr> shape_1;
   bool y_from_fullop =
-      op->operand_source(1).defining_op()->isa<paddle::dialect::FullOp>();
+      op->operand_source(1).defining_op()
+          ? op->operand_source(1).defining_op()->isa<paddle::dialect::FullOp>()
+          : false;
   if (!y_from_fullop && y_shapeordata.data().has_value()) {
     shape_1 = y_shapeordata.data().value();
   } else {
@@ -224,4 +228,12 @@ bool Remainder_OpInferSymbolicShape(
   return InferSymbolicShapeElementWiseBinary(op, shape_analysis);
 }
 
+bool SubtractOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return InferSymbolicShapeElementWiseBinary(op, shape_analysis);
+}
+bool Subtract_OpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return InferSymbolicShapeElementWiseBinary(op, shape_analysis);
+}
 }  // namespace paddle::dialect

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h

@@ -53,5 +53,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(NotEqual)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(NotEqual_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Remainder)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Remainder_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract_)
 
 }  // namespace paddle::dialect

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h

@@ -75,6 +75,22 @@ std::vector<T> GetVectorAttr(const ::pir::Operation *op,
   return vec_res;
 }
 
+inline ExprVec GetExprVecFromData(const ShapeOrData &shapeordata) {
+  if (shapeordata.isa<TensorListExprs>()) {
+    ExprVec result;
+    TensorListExprs list =
+        shapeordata.dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
+    for (size_t i = 0; i < list.size(); i++) {
+      for (auto expr : list[i].data().value()) {
+        result.emplace_back(expr);
+      }
+    }
+    return result;
+  } else {
+    return shapeordata.data().value();
+  }
+}
+
 std::optional<std::vector<int64_t>> VecExpr2Int64(const ExprVec &expr_vec);
 
 ExprVec VecInt642Expr(const std::vector<int64_t> &int_vec);

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc

@@ -289,6 +289,21 @@ bool ConcatOpInferSymbolicShape(
   axis = axis >= 0 ? axis : std::max(int64_t(0), int64_t(axis + rank));
 
   if (shape_data_list[0].data().has_value()) {
+    if (rank == 1) {
+      ExprVec data = details::GetExprVecFromData(
+          shape_analysis->GetShapeOrDataForValue(operand_source));
+      const std::vector<symbol::DimExpr> shape{std::int64_t(data.size())};
+      symbol::ShapeOrDataDimExprs shape_data{
+          symbol::TensorShapeOrDataDimExprs(shape, data)};
+      pir::Value res = op->result(0);
+      shape_analysis->SetShapeOrDataForValue(res, shape_data);
+
+      return true;
+    } else {
+      PADDLE_THROW(phi::errors::Unimplemented(
+          op->name() +
+          " 's InferSymbolicShape can NOT deal with rank > 1 now."));
+    }
     std::vector<symbol::DimExpr> data;
     data.reserve(shape_data_list.size());
     for (auto &data_elem : shape_data_list) {
@@ -436,9 +451,9 @@ bool SqueezeOpInferSymbolicShape(
         if (in_dims_sym[current] == 1) {
           should_squeeze[current] = true;
         } else if (!in_dims_sym[current].Has<std::int64_t>()) {
-          PADDLE_THROW(
-              phi::errors::Unimplemented("SqueezeOpInferSymbolicShape CAN NOT "
-                                         "deal with symbol in axis now"));
+          should_squeeze[current] = true;
+        } else {
+          should_squeeze[current] = true;
         }
       }
     }

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.cc

@@ -379,14 +379,7 @@ bool Sinh_OpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
 }
-bool SubtractOpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  return SameOperandsAndResultShape(op, shape_analysis);
-}
-bool Subtract_OpInferSymbolicShape(
-    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  return SameOperandsAndResultShape(op, shape_analysis);
-}
+
 bool TanOpInferSymbolicShape(pir::Operation *op,
                              pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h

@@ -105,8 +105,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sin_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sinh)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sinh_)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Subtract_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tan)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tan_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tanh)

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc

@@ -165,6 +165,7 @@ bool Cumsum_OpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return CumsumOpInferSymbolicShape(op, shape_analysis);
 }
+
 bool DiagEmbedOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source = op->operand_source(0);
@@ -280,6 +281,7 @@ bool KthvalueOpInferSymbolicShape(
   shape_analysis->SetShapeOrDataForValue(op->result(1), shape_data);
   return true;
 }
+
 bool ReshapeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source = op->operand_source(0);
@@ -329,10 +331,11 @@ bool ReshapeOpInferSymbolicShape(
     const auto &numel =
         GetProduct(original_shape, [](const auto &) { return true; });
 
+    ExprVec target_shape = details::GetExprVecFromData(operand_shape_or_data);
     const auto &product_exclude_minus_one =
-        GetProduct(operand_shape_or_data.data().value(), IsNotMinusOne);
+        GetProduct(target_shape, IsNotMinusOne);
 
-    const auto &input_dims = operand_shape_or_data.data().value();
+    const auto &input_dims = target_shape;
 
     std::vector<symbol::DimExpr> out_dims;
     out_dims.reserve(input_dims.size());