PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/CMakeLists.txt
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/fluid/pir/dialect/CMakeLists.txt
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc
Copy file name to clipboard
+63-54Lines changed: 63 additions & 54 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc
Copy file name to clipboard
+63-54Lines changed: 63 additions & 54 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h
Copy file name to clipboard
+10-1Lines changed: 10 additions & 1 deletion b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h
Copy file name to clipboard
+10-1Lines changed: 10 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc
Copy file name to clipboard
+21Lines changed: 21 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc
Copy file name to clipboard
+21Lines changed: 21 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc
Copy file name to clipboard
+52-45Lines changed: 52 additions & 45 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc
Copy file name to clipboard
+52-45Lines changed: 52 additions & 45 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
@@ -246,7 +246,8 @@ set(op_dialect_srcs
     ${pir_op_source_file}
     ${pir_bwd_op_source_file}
     ${pir_update_op_source_file}
-    ${api_source_file})
+    ${api_source_file}
+    ${PADDLE_SOURCE_DIR}/paddle/fluid/pir/transforms/shape_optimization_pass.cc)
 
 if(WITH_MKLDNN)
   set(op_dialect_srcs ${op_dialect_srcs} ${onednn_op_source_file}
 
@@ -36,74 +36,32 @@ bool BroadcastOpInferSymbolicShape(
   return true;
 }
 
-bool SliceOpInferSymbolicShape(pir::Operation *op,
-                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  // TODO(zhangbopd): Not implemented yet, different from the one in paddle
-  // dialect. And Currently only support start/end/axis with single value.
-  pir::AttributeMap attributes = op->attributes();
-
-  auto GetAttrInt64Value = [&](const std::string &name) -> int64_t {
-    std::vector<pir::Attribute> attr =
-        attributes[name].dyn_cast<pir::ArrayAttribute>().AsVector();
-    PADDLE_ENFORCE_GT(
-        attr.size(),
-        0,
-        phi::errors::PreconditionNotMet(
-            "Only Support [%s] op len(%s) == 1 , but received %d.",
-            op->name(),
-            name,
-            attr.size()));
-    return attr[0].dyn_cast<pir::Int64Attribute>().data();
-  };
-
-  const int64_t start = GetAttrInt64Value("starts");
-  const int64_t end = GetAttrInt64Value("ends");
-  const int64_t axis = GetAttrInt64Value("axes");
-
-  const pir::Value operand_source = op->operand_source(0);
-  const auto &operand_shape_or_data =
-      shape_analysis->GetShapeOrDataForValue(operand_source);
-
-  const auto GetOutDimExprs = [&]() -> symbol::TensorShapeOrDataDimExprs {
-    std::vector<symbol::DimExpr> out_sym_shape = operand_shape_or_data.shape();
-    if (end == std::numeric_limits<int>::max()) {
-      out_sym_shape[axis] = out_sym_shape[axis] - start;
-    } else {
-      out_sym_shape[axis] = end - start;
-    }
-    symbol::TensorShapeOrDataDimExprs shape_dim_expr(out_sym_shape);
-    if (operand_shape_or_data.data().has_value()) {
-      std::vector<symbol::DimExpr> out_data;
-      for (int64_t i = start; i < end; i++) {
-        out_data.push_back(operand_shape_or_data.data().value()[i]);
-      }
-      shape_dim_expr.SetData(out_data);
-    }
-    return shape_dim_expr;
-  };
-  symbol::ShapeOrDataDimExprs shape_data{GetOutDimExprs()};
-
-  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
-  return true;
-}
-
 bool ConcatOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const auto input_values = op->operands_source();
   const auto input_size = input_values.size();
 
-  const int axis =
-      op->attributes().at("axis").dyn_cast<pir::Int32Attribute>().data();
+  int axis = op->attributes().at("axis").dyn_cast<pir::Int32Attribute>().data();
 
-  // TODO(zhangbopd): Need support GetShapeOrDataForValue().data() case.
   const auto &GetOutDimExprs = [&]() -> std::vector<symbol::DimExpr> {
     std::vector<symbol::DimExpr> out_dims =
         shape_analysis->GetShapeOrDataForValue(input_values[0]).shape();
+
+    size_t rank = out_dims.size();
+    axis = axis >= 0 ? axis : std::max(int64_t(0), int64_t(axis + rank));
+
     for (size_t i = 1; i < input_size; ++i) {
       const auto &operand_shape_or_data =
           shape_analysis->GetShapeOrDataForValue(input_values[i]);
       out_dims[axis] = out_dims[axis] + operand_shape_or_data.shape()[axis];
     }
+
+    for (size_t i = 1; i < rank; ++i) {
+      if (i == static_cast<size_t>(axis)) continue;
+      paddle::dialect::details::BuildCstrEqForTensorListAlongAxis(
+          shape_analysis, input_values, i);
+    }
+
     return out_dims;
   };
 
@@ -164,4 +122,55 @@ bool ReshapeOpInferSymbolicShape(
   return true;
 }
 
+bool SliceOpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  // TODO(zhangbopd): Not implemented yet, different from the one in paddle
+  // dialect. And Currently only support start/end/axis with single value.
+  pir::AttributeMap attributes = op->attributes();
+
+  auto GetAttrInt64Value = [&](const std::string &name) -> int64_t {
+    std::vector<pir::Attribute> attr =
+        attributes[name].dyn_cast<pir::ArrayAttribute>().AsVector();
+    PADDLE_ENFORCE_GT(
+        attr.size(),
+        0,
+        phi::errors::PreconditionNotMet(
+            "Only Support [%s] op len(%s) == 1 , but received %d.",
+            op->name(),
+            name,
+            attr.size()));
+    return attr[0].dyn_cast<pir::Int64Attribute>().data();
+  };
+
+  const int64_t start = GetAttrInt64Value("starts");
+  const int64_t end = GetAttrInt64Value("ends");
+  const int64_t axis = GetAttrInt64Value("axes");
+
+  const pir::Value operand_source = op->operand_source(0);
+  const auto &operand_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(operand_source);
+
+  const auto GetOutDimExprs = [&]() -> symbol::TensorShapeOrDataDimExprs {
+    std::vector<symbol::DimExpr> out_sym_shape = operand_shape_or_data.shape();
+    if (end == std::numeric_limits<int>::max()) {
+      out_sym_shape[axis] = out_sym_shape[axis] - start;
+    } else {
+      out_sym_shape[axis] = end - start;
+    }
+    symbol::TensorShapeOrDataDimExprs shape_dim_expr(out_sym_shape);
+    if (operand_shape_or_data.data().has_value()) {
+      std::vector<symbol::DimExpr> out_data;
+      for (int64_t i = start; i < end; i++) {
+        out_data.push_back(operand_shape_or_data.data().value()[i]);
+      }
+      shape_dim_expr.SetData(out_data);
+    }
+    return shape_dim_expr;
+  };
+  symbol::ShapeOrDataDimExprs shape_data{GetOutDimExprs()};
+
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+  return true;
+}
+
 }  // namespace cinn::dialect
@@ -16,23 +16,32 @@
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
 namespace cinn::dialect {
-// using paddle::dialect::ScaleOpInferSymbolicShape;
 
 bool BroadcastOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool SliceOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool ConcatOpInferSymbolicShape(pir::Operation *op,
                                 pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool ReduceMaxOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool ReduceMinOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool ReduceProdOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool ReduceSumOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 bool ReshapeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
 
+bool SliceOpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 }  // namespace cinn::dialect
@@ -78,4 +78,25 @@ bool ReduceInferDim(pir::Operation *op,
   return true;
 }
 
+void BuildCstrEqForTensorListAlongAxis(
+    pir::ShapeConstraintIRAnalysis *shape_analysis,
+    const symbol::TensorListShapeOrDataDimExprs &shape_data_list,
+    int axis) {
+  for (size_t i = 1; i < shape_data_list.size(); ++i) {
+    shape_analysis->CreateDimExprBuilder().CstrEq(
+        shape_data_list[0].shape()[axis], shape_data_list[i].shape()[axis]);
+  }
+}
+
+void BuildCstrEqForTensorListAlongAxis(
+    pir::ShapeConstraintIRAnalysis *shape_analysis,
+    const std::vector<pir::Value> &values,
+    int axis) {
+  for (size_t i = 1; i < values.size(); ++i) {
+    shape_analysis->CreateDimExprBuilder().CstrEq(
+        shape_analysis->GetShapeOrDataForValue(values[0]).shape()[axis],
+        shape_analysis->GetShapeOrDataForValue(values[i]).shape()[axis]);
+  }
+}
+
 }  // namespace paddle::dialect::details
@@ -65,4 +65,15 @@ bool ReduceInferDim(pir::Operation *op,
                     const std::vector<int64_t> &axis,
                     bool keep_dim,
                     bool reduce_all);
+
+void BuildCstrEqForTensorListAlongAxis(
+    pir::ShapeConstraintIRAnalysis *shape_analysis,
+    const symbol::TensorListShapeOrDataDimExprs &shape_data_list,
+    int axis);
+
+void BuildCstrEqForTensorListAlongAxis(
+    pir::ShapeConstraintIRAnalysis *shape_analysis,
+    const std::vector<pir::Value> &values,
+    int axis);
+
 }  // namespace paddle::dialect::details
@@ -80,16 +80,6 @@ bool ShapeSrOpInferSymbolicShape(
   return ShapeOpInferSymbolicShape(op, shape_analysis);
 }
 
-void BuildCstrEqForTensorListAlongAxis(
-    pir::ShapeConstraintIRAnalysis *shape_analysis,
-    const symbol::TensorListShapeOrDataDimExprs &shape_data_list,
-    int axis) {
-  for (size_t i = 1; i < shape_data_list.size(); ++i) {
-    shape_analysis->CreateDimExprBuilder().CstrEq(
-        shape_data_list[0].shape()[axis], shape_data_list[i].shape()[axis]);
-  }
-}
-
 bool StackOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source = op->operand_source(0);
@@ -119,7 +109,8 @@ bool StackOpInferSymbolicShape(pir::Operation *op,
     } else {
       for (int i = 0; i < rank; ++i) {
         if (i == axis) continue;
-        BuildCstrEqForTensorListAlongAxis(shape_analysis, shape_data_list, i);
+        details::BuildCstrEqForTensorListAlongAxis(
+            shape_analysis, shape_data_list, i);
       }
       shape_dim_exprs.insert(shape_dim_exprs.begin() + axis,
                              static_cast<std::int64_t>(shape_data_list.size()));
@@ -194,31 +185,42 @@ bool ReshapeOpInferSymbolicShape(
   const symbol::ShapeOrDataDimExprs &operand_shape_or_data =
       shape_analysis->GetShapeOrDataForValue(operand_source_shape);
 
-  const std::vector<symbol::DimExpr> out_dims = [&] {
-    std::vector<symbol::DimExpr> out_dims;
-    out_dims = operand_shape_or_data.data().value();
+  const auto &GetProduct = [&](const auto &dim_exprs, const auto &Filter) {
+    symbol::DimExpr product{1};
+    for (const auto &dim_expr : dim_exprs) {
+      if (Filter(dim_expr)) {
+        product = product * dim_expr;
+      }
+    }
+    return product;
+  };
 
-    symbol::DimExpr product = symbol::DimExpr(1);
-    symbol::DimExpr numel = symbol::DimExpr(1);
+  const auto &IsNotMinusOne = [&](const symbol::DimExpr &dim_expr) {
+    if (dim_expr.isa<int64_t>()) {
+      return dim_expr.dyn_cast<int64_t>() != static_cast<int64_t>(-1);
+    }
+    return true;
+  };
 
+  const std::vector<symbol::DimExpr> out_dims = [&] {
     const auto &original_shape =
         shape_analysis->GetShapeOrDataForValue(op->operand_source(0)).shape();
-    for (auto &dim_expr : original_shape) {
-      numel = numel * dim_expr;
-    }
 
-    for (size_t i = 0; i < out_dims.size(); i++) {
-      if (out_dims[i].isa<int64_t>()) {
-        if (out_dims[i].dyn_cast<int64_t>() != static_cast<int64_t>(-1)) {
-          product = product * out_dims[i];
-        } else if (i == out_dims.size() - 1) {
-          out_dims[i] = numel / product;
-        } else {
-          // doing nothing
-        }
-      } else {
-        product = product * out_dims[i];
-      }
+    const auto &numel =
+        GetProduct(original_shape, [](const auto &) { return true; });
+
+    const auto &product_exclude_minus_one =
+        GetProduct(operand_shape_or_data.data().value(), IsNotMinusOne);
+
+    const auto &input_dims = operand_shape_or_data.data().value();
+
+    std::vector<symbol::DimExpr> out_dims;
+    out_dims.reserve(input_dims.size());
+    for (const auto &dim_expr : input_dims) {
+      const auto &out_dim_expr = IsNotMinusOne(dim_expr)
+                                     ? dim_expr
+                                     : (numel / product_exclude_minus_one);
+      out_dims.emplace_back(out_dim_expr);
     }
 
     return out_dims;
@@ -352,15 +354,20 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
       int64_t axis = axes[i];
       auto end =
           IsMaxInt(dim_expr_ends[i]) ? out_shape[axis] : dim_expr_ends[i];
-      if ((starts[i] >= 0 && ends[i] >= 0) ||
-          (starts[i] <= 0 && ends[i] <= 0)) {  // both negtive or positive.
+
+      bool both_negative_or_positive =
+          (starts[i] >= 0 && ends[i] >= 0) || (starts[i] <= 0 && ends[i] <= 0);
+      bool start_negative_end_positive = starts[i] <= 0 && ends[i] >= 0;
+      bool start_positive_end_negative = starts[i] >= 0 && ends[i] <= 0;
+
+      if (both_negative_or_positive) {
         out_shape[axis] = end - dim_expr_starts[i];
-      } else if (starts[i] <= 0 &&
-                 ends[i] >= 0) {  // negtive start, positive end
+      } else if (start_negative_end_positive) {
         out_shape[axis] = end - dim_expr_starts[i] - out_shape[axis];
-      } else if (starts[i] >= 0 &&
-                 ends[i] <= 0) {  // positive start, negtive end
+      } else if (start_positive_end_negative) {
         out_shape[axis] = out_shape[axis] - dim_expr_starts[i] + end;
+      } else {
+        LOG(FATAL) << "Dead code";
       }
     }
 
@@ -429,14 +436,14 @@ bool ConcatOpInferSymbolicShape(
 
   const std::vector<symbol::DimExpr> &out_dims = [&] {
     std::vector<symbol::DimExpr> out_dims = shape_data_list[0].shape();
-    for (size_t i = 1; i < shape_data_list.size(); ++i) {
-      for (size_t j = 0; j < rank; ++j) {
-        if (j != static_cast<size_t>(axis)) {
-          // This func have bug
-          BuildCstrEqForTensorListAlongAxis(shape_analysis, shape_data_list, i);
-          continue;
-        }
-        out_dims[axis] = out_dims[axis] + shape_data_list[i].shape()[axis];
+    for (size_t i = 0; i < rank; ++i) {
+      if (i != static_cast<size_t>(axis)) {
+        details::BuildCstrEqForTensorListAlongAxis(
+            shape_analysis, shape_data_list, i);
+        continue;
+      }
+      for (size_t j = 1; j < shape_data_list.size(); ++j) {
+        out_dims[axis] = out_dims[axis] + shape_data_list[j].shape()[axis];
       }
     }
     return out_dims;
 
@@ -95,3 +95,7 @@ bool Tril_OpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis);
 
 }  // namespace paddle::dialect
+
+namespace cinn::dialect {
+using paddle::dialect::ScaleOpInferSymbolicShape;
+}