support ExpandAsOpInferSymbolicShape (#63576)

yulangz · web-flow · commit 0ae676f130e2 · 2024-04-19T16:52:24.000+08:00
* support ExpandAsOpInferSymbolicShape;
in add_broadcast_to_elementwise_pass, insert shape data after create op;
allow Concht when input size is 1.

* fix compile error
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.cc
@@ -107,6 +107,14 @@ bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
 
   if (x_dims != y_dims) {
     auto output_shape = GetOutputShape(x_dims, y_dims);
+    pir::ShapeConstraintIRAnalysis& shape_analysis =
+        pir::ShapeAnalysisManager::Instance().Get(op->GetParentProgram());
+    std::vector<symbol::DimExpr> out_dim;
+    out_dim.reserve(output_shape.size());
+    for (auto d : output_shape) {
+      out_dim.emplace_back(d);
+    }
+
     if (!IsSameDim(x_dims, output_shape)) {
       // add broadcast to input 0
       if (auto full_op = op->operand_source(0)
@@ -122,13 +130,18 @@ bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
                 .dyn_cast<paddle::dialect::PlaceAttribute>()
                 .data());
         op->operand(0).set_source(new_full->result(0));
+        shape_analysis.SetShapeOrDataForValue(
+            new_full.result(0), symbol::TensorShapeOrDataDimExprs(out_dim));
       } else {
         auto new_transpose_op = rewriter->Build<cinn::dialect::BroadcastOp>(
             op->operand_source(0),
             cinn::hlir::framework::pir::GetBroadcastAxis(x_dims, output_shape),
             output_shape);
 
         op->operand(0).set_source(new_transpose_op->result(0));
+        shape_analysis.SetShapeOrDataForValue(
+            new_transpose_op.result(0),
+            symbol::TensorShapeOrDataDimExprs(out_dim));
       }
     }
 
@@ -147,13 +160,18 @@ bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
                 .data());
 
         op->operand(1).set_source(new_full->result(0));
+        shape_analysis.SetShapeOrDataForValue(
+            new_full.result(0), symbol::TensorShapeOrDataDimExprs(out_dim));
       } else {
         auto new_transpose_op = rewriter->Build<cinn::dialect::BroadcastOp>(
             op->operand_source(1),
             cinn::hlir::framework::pir::GetBroadcastAxis(y_dims, output_shape),
             output_shape);
 
         op->operand(1).set_source(new_transpose_op->result(0));
+        shape_analysis.SetShapeOrDataForValue(
+            new_transpose_op.result(0),
+            symbol::TensorShapeOrDataDimExprs(out_dim));
       }
     }
 
diff --git a/paddle/cinn/hlir/pe/transform.cc b/paddle/cinn/hlir/pe/transform.cc
@@ -426,8 +426,9 @@ ir::Tensor Concat(const ir::Tensor& A,
 ir::Tensor Concat(const std::vector<ir::Tensor>& input_tensors,
                   int axis,
                   const std::string& name) {
+  // input size 1 is valid for Concat
   int input_size = input_tensors.size();
-  CHECK_GE(input_size, 2U) << "Concat should have at least 2 input tensors";
+  CHECK_GE(input_size, 1U) << "Concat should have at least 1 input tensors";
   std::vector<Expr> output_shape = input_tensors[0]->shape;
   int input_dim = output_shape.size();
   CHECK(axis >= -input_dim && axis < input_dim)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -204,10 +204,20 @@ bool SparseWeightEmbeddingOpInferSymbolicShape(
 
 bool ExpandAsOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() +
-      " 's InferSymbolicShape interface is NOT implemented "
-      "now because of the lack of necessary information."));
+  std::vector<int> target_shape =
+      paddle::dialect::details::GetVectorAttr<int>(op, "target_shape");
+  const std::vector<symbol::DimExpr> &output_dims = [&] {
+    std::vector<symbol::DimExpr> output_dims;
+    output_dims.reserve(target_shape.size());
+    for (int shape : target_shape) {
+      output_dims.push_back(shape);
+    }
+    return output_dims;
+  }();
+
+  shape_analysis->SetShapeOrDataForValue(
+      op->result(0), symbol::TensorShapeOrDataDimExprs(output_dims));
+
   return true;
 }
 
