[CINN] Support dynamic shape broadcast in cinn (#61156)

zyfncg · web-flow · commit 5c08810ab4db · 2024-01-27T17:15:07.000+08:00
* support dynamic shape broadcast in cinn

* fix some infer symbolic shape

* fix complie

* fix complie

* fix ShapeBroadcastOp InferSymbolicShape

* clone TransBufferWithDynamicShape
diff --git a/paddle/cinn/hlir/dialect/operator/ir/manual_op.cc b/paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
@@ -434,8 +434,8 @@ bool GenerateShapeOp::InferSymbolicShape(
   }();
 
   // TODO(HongyuJia): use op->result(0) to infer the shape
-  std::vector<symbol::DimExpr> shape(
-      std::int64_t(substituted_dim_exprs.size()));
+  std::vector<symbol::DimExpr> shape{
+      std::int64_t(substituted_dim_exprs.size())};
   symbol::ShapeOrDataDimExprs shape_or_data_dim_exprs{
       symbol::TensorShapeOrDataDimExprs(shape, substituted_dim_exprs)};
 
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/lower_cinn_fusion_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/lower_cinn_fusion_op_pass.cc
@@ -65,11 +65,6 @@ bool SameInputOutputShape(
   return x.shape() == out.shape();
 }
 
-void ReplaceAllUsesWithInput(paddle::dialect::ExpandOp expand) {
-  pir::Value x = expand.x();
-  expand.out().ReplaceAllUsesWith(x);
-}
-
 // Returns true if success
 bool EraseOneExpand(
     pir::Block* block,
@@ -83,9 +78,11 @@ bool EraseOneExpand(
     auto generate_shape_op =
         expand.shape().defining_op<cinn::dialect::GenerateShapeOp>();
     CHECK_NOTNULL(generate_shape_op);
-    ReplaceAllUsesWithInput(expand);
+    rewriter.ReplaceAllUsesWith(expand.out(), expand.x());
     rewriter.EraseOp(expand);
-    rewriter.EraseOp(generate_shape_op);
+    if (generate_shape_op->use_empty()) {
+      rewriter.EraseOp(generate_shape_op);
+    }
     return true;
   }
   return false;
@@ -448,7 +445,6 @@ void SimplyConditionBlock(
 
 void CompileGroupToJitKernelOp(
     const std::vector<pir::Value>& group_inputs,
-    const std::vector<pir::Type>& output_types,
     const std::shared_ptr<cinn::hlir::framework::PirCompiler>& pir_compiler,
     pir::PatternRewriter& rewriter,  // NOLINT
     std::unordered_map<pir::Block*, cinn::dialect::ir::GroupPtr>* group_map) {
@@ -461,12 +457,16 @@ void CompileGroupToJitKernelOp(
   auto op_attr_map = ComplieGroupAsOpAttribute(pir_compiler, group_list);
   VLOG(4) << "The size of group_map is : " << group_map->size();
   for (auto& [block, group] : *group_map) {
+    std::vector<pir::Type> output_types;
+    const auto& group_output_values = group->output_values;
+    for (size_t i = 0; i < group_output_values.size(); ++i) {
+      output_types.push_back(group_output_values[i].type());
+    }
     auto& yeild_op = block->back();
     CHECK(yeild_op.isa<pir::YieldOp>()) << "Last op of block should be yield";
     rewriter.set_insertion_point(&yeild_op);
     auto jit_kernel_op = rewriter.Build<cinn::dialect::JitKernelOp>(
         group_inputs, op_attr_map.at(group), output_types);
-    auto group_output_values = group->GetGroupOutputValues();
     CHECK(jit_kernel_op.num_results() == group_output_values.size());
     for (size_t i = 0; i < jit_kernel_op.num_results(); ++i) {
       rewriter.ReplaceAllUsesWith(group_output_values[i],
@@ -517,8 +517,7 @@ pir::Operation* ComplieBroadcastTreeToConditionBlock(
   VLOG(6) << "After simply condition block: " << *program;
 
   // 3. complie condition block to jit_kernel_op
-  CompileGroupToJitKernelOp(
-      group_inputs, output_types, pir_compiler, rewriter, &group_map);
+  CompileGroupToJitKernelOp(group_inputs, pir_compiler, rewriter, &group_map);
   VLOG(6) << "complie condition block to jit_kernel_op: " << *program;
 
   return cond_op;
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
@@ -356,7 +356,7 @@ class OpFusionPassHelper {
 
  private:
   void DoOpFusion() {
-    auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(
+    const auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(
         ops_.front()->GetParentProgram());
     for (auto consumer : ops_) {
       auto consumer_kind =
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h
@@ -27,6 +27,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/pir/core/operation.h"
 #include "paddle/pir/core/value.h"
+#include "paddle/pir/dialect/shape/utils/shape_analysis.h"
 
 namespace cinn {
 namespace dialect {
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc
@@ -52,33 +52,42 @@ using TensorDim = std::variant<TensorDimInShape, TensorDimInData>;
 using TensorDim4SymbolNameT =
     std::function<std::optional<TensorDim>(const std::string& symbol_name)>;
 
+using SymbolName2CachedValue = std::unordered_map<symbol::DimExpr, pir::Value>;
+
 struct CachedDimExprToValueConverter {
   CachedDimExprToValueConverter(
+      const std::shared_ptr<SymbolName2CachedValue>& symbol_names2cached_value,
       const TensorDim4SymbolNameT& TensorDim4SymbolNameVal,
       pir::PatternRewriter* rewriter_val)
-      : TensorDim4SymbolName(TensorDim4SymbolNameVal), rewriter(rewriter_val) {}
+      : symbol_names2cached_value_(symbol_names2cached_value),
+        TensorDim4SymbolName(TensorDim4SymbolNameVal),
+        rewriter(rewriter_val) {}
 
   TensorDim4SymbolNameT TensorDim4SymbolName;
   pir::PatternRewriter* rewriter;
 
-  // TODO(): Refactor to cached version if std::hash<symbol::DimExpr>() is
-  // ready. std::unordered_map<symbol::DimExpr, pir::Value>
-  // symbol_names2cached_value_;
-
   pir::Value ConvertToValue(const symbol::DimExpr& dim_expr) {
-    // TODO():  cache the returned value if std::hash<symbol::DimExpr>() is
-    // ready
-    return std::visit(
-        [&](const auto& impl) { return ConvertToValueImpl(impl); },
-        dim_expr.variant());
+    auto iter = symbol_names2cached_value_->find(dim_expr);
+    if (iter == symbol_names2cached_value_->end()) {
+      pir::Value value =
+          std::visit([&](const auto& impl) { return ConvertToValueImpl(impl); },
+                     dim_expr.variant());
+      iter = symbol_names2cached_value_->emplace(dim_expr, value).first;
+    }
+    return iter->second;
   }
 
+  std::shared_ptr<SymbolName2CachedValue> symbol_names2cached_value_;
+
   pir::Value GetInputShapeByInputTensor(pir::Value input_tensor) {
     auto iter = tensor2shape_.find(input_tensor);
     if (iter == tensor2shape_.end()) {
-      pir::Value input_shape =
+      pir::Value shape =
           rewriter->Build<paddle::dialect::ShapeOp>(input_tensor).out();
-      iter = tensor2shape_.emplace(input_tensor, input_shape).first;
+      pir::Value cast_shape =
+          rewriter->Build<paddle::dialect::CastOp>(shape, phi::DataType::INT64)
+              .out();
+      iter = tensor2shape_.emplace(input_tensor, cast_shape).first;
     }
     return iter->second;
   }
@@ -219,7 +228,11 @@ struct CachedDimExprToValueConverter {
 class SplitGenerateShapeIntoShapeOps
     : public pir::OpRewritePattern<cinn::dialect::GenerateShapeOp> {
  public:
-  using pir::OpRewritePattern<cinn::dialect::GenerateShapeOp>::OpRewritePattern;
+  SplitGenerateShapeIntoShapeOps(
+      pir::IrContext* context,
+      const std::shared_ptr<SymbolName2CachedValue>& symbol_names2cached_value)
+      : OpRewritePattern(context),
+        symbol_names2cached_value_(symbol_names2cached_value) {}
 
   bool MatchAndRewrite(cinn::dialect::GenerateShapeOp op,
                        pir::PatternRewriter& rewriter) const override {
@@ -236,10 +249,13 @@ class SplitGenerateShapeIntoShapeOps
     TensorDim4SymbolNameT TensorDim4SymbolName =
         MakeGetterTensorDim4SymbolName(op);
     if (!TensorDim4SymbolName) return std::nullopt;
-    CachedDimExprToValueConverter converter{TensorDim4SymbolName, rewriter};
+    CachedDimExprToValueConverter converter{
+        symbol_names2cached_value_, TensorDim4SymbolName, rewriter};
     return GetValueOfRewritedOps(dim_exprs, &converter);
   }
 
+  std::shared_ptr<SymbolName2CachedValue> symbol_names2cached_value_;
+
   TensorDim4SymbolNameT MakeGetterTensorDim4SymbolName(
       cinn::dialect::GenerateShapeOp op) const {
     std::unordered_map<std::string, TensorDim> symbol_name2tenso_dim{};
@@ -330,8 +346,10 @@ class SplitGenerateShapeIntoShapeOps
       CachedDimExprToValueConverter* converter) const {
     const std::vector<pir::Value>& values_from_dim_exprs =
         GetValuesOfRewritedOps(dim_exprs, converter);
-    return converter->rewriter->Build<pir::CombineOp>(values_from_dim_exprs)
-        .out();
+    if (values_from_dim_exprs.size() == 1) return values_from_dim_exprs.at(0);
+    pir::Value vec =
+        converter->rewriter->Build<pir::CombineOp>(values_from_dim_exprs).out();
+    return converter->rewriter->Build<paddle::dialect::ConcatOp>(vec).out();
   }
 
   std::vector<pir::Value> GetValuesOfRewritedOps(
@@ -353,8 +371,8 @@ SplitGenerateShapeIntoShapeOpsPass::SplitGenerateShapeIntoShapeOpsPass()
 pir::RewritePatternSet SplitGenerateShapeIntoShapeOpsPass::InitializePatterns(
     pir::IrContext* context) {
   pir::RewritePatternSet ps(context);
-  // elementwise ops
-  ps.Add<SplitGenerateShapeIntoShapeOps>(context);
+  ps.Add<SplitGenerateShapeIntoShapeOps>(
+      context, std::make_shared<SymbolName2CachedValue>());
   return ps;
 }
 
diff --git a/paddle/cinn/optim/optimize.cc b/paddle/cinn/optim/optimize.cc
@@ -65,7 +65,7 @@ Expr Optimize(Expr e,
     RemoveGpuForloopsAxis(&copied);
   }
   CudaSyncThreadsDropIfThenElse(&copied);
-  TransBufferWithDynamicShape(&copied);
+  // TransBufferWithDynamicShape(&copied);
 #endif
 
   SimplifyBlocks(&copied);
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc
@@ -257,15 +257,14 @@ bool ShapeSrOpInferSymbolicShape(
 bool StackOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source = op->operand_source(0);
-  symbol::ShapeOrDataDimExprs operand_shape_or_data =
+  const auto &operand_shape_or_data =
       shape_analysis->GetShapeOrDataForValue(operand_source);
 
   std::vector<symbol::DimExpr> out_dims;
   std::vector<symbol::DimExpr> out_dims_data;
   if (operand_shape_or_data.data().has_value()) {
     out_dims_data = operand_shape_or_data.data().value();
-    out_dims.emplace_back(
-        static_cast<std::int64_t>(operand_shape_or_data.shape().size()));
+    out_dims = operand_shape_or_data.shape();
   }
   // else : pir::VectorType x =
   // operand_source.type().dyn_cast<pir::VectorType>();
@@ -740,19 +739,17 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
   std::vector<symbol::DimExpr> out_dims;
   if (operand_shape_or_data.data().has_value()) {
     out_dims.push_back(operand_shape_or_data.data().value()[start]);
+    symbol::ShapeOrDataDimExprs shape_data{symbol::TensorShapeOrDataDimExprs(
+        std::vector<symbol::DimExpr>{out_dims.size()}, out_dims)};
+    shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+    op->set_attribute("symbolic_shape",
+                      pir::shape::SymbolAttribute::get(
+                          pir::IrContext::Instance(), shape_data));
+    return true;
   }
 
-  symbol::ShapeOrDataDimExprs shape_data{
-      symbol::TensorShapeOrDataDimExprs(out_dims)};
-  if (operand_shape_or_data.data().has_value()) {
-    shape_data.SetData(operand_shape_or_data.shape());
-  }
-  op->set_attribute(
-      "symbolic_shape",
-      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
-
-  pir::Value res = op->result(0);
-  shape_analysis->SetShapeOrDataForValue(res, shape_data);
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
   return true;
 }
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -4134,7 +4134,7 @@ bool ShapeBroadcastOp::InferSymbolicShape(
 
   pir::Value res = result(0);
   // TODO(HongyuJia): use op->result(0) to infer the shape
-  std::vector<symbol::DimExpr> shape(std::int64_t(output_data.size()));
+  std::vector<symbol::DimExpr> shape{std::int64_t(output_data.size())};
 
   symbol::ShapeOrDataDimExprs output_data_shape{
       symbol::TensorShapeOrDataDimExprs(shape, output_data)};
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc
@@ -73,22 +73,23 @@ struct CombineOpInferSymbolicShapeInterfaceModel
       }
     }
 
-    // TODO(zhangbopd): use op->result(0) to infer the shape
-    symbol::ShapeOrDataDimExprs shape_data{
-        symbol::TensorShapeOrDataDimExprs(out_dims)};
-
     if (operand_source_1st_data.has_value()) {
-      std::vector<symbol::DimExpr> tmp_shape(std::int64_t(out_dims.size()));
-      symbol::ShapeOrDataDimExprs temp_shape_data{
-          symbol::TensorShapeOrDataDimExprs(tmp_shape, out_dims)};
-      shape_data = temp_shape_data;
+      std::vector<symbol::DimExpr> tmp_shape{int64_t(out_dims.size())};
+      shape_analysis->SetShapeOrDataForValue(
+          op->result(0),
+          symbol::TensorShapeOrDataDimExprs(tmp_shape, out_dims));
+    } else {
+      // TODO(zhangbopd): use op->result(0) to infer the shape
+      shape_analysis->SetShapeOrDataForValue(
+          op->result(0), symbol::TensorShapeOrDataDimExprs(out_dims));
     }
 
-    op->set_attribute("symbolic_shape",
-                      pir::shape::SymbolAttribute::get(
-                          pir::IrContext::Instance(), shape_data));
-    auto res = op->result(0);
-    shape_analysis->SetShapeOrDataForValue(res, shape_data);
+    op->set_attribute(
+        "symbolic_shape",
+        pir::shape::SymbolAttribute::get(
+            pir::IrContext::Instance(),
+            shape_analysis->GetShapeOrDataForValue(op->result(0))));
+
     return true;
   }
 
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -89,7 +89,9 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/divide_group_op_to_fusion_op_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/lower_cinn_fusion_op_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/move_generate_shape_ops_to_prologue_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.h"
 #include "paddle/cinn/hlir/framework/pir_compiler.h"
 #include "paddle/fluid/pir/transforms/build_cinn_pass.h"
 #endif
@@ -1555,23 +1557,25 @@ void AddCinnPass(std::shared_ptr<PassManager> &pass_manager,  // NOLINT
   bool has_dynamic_shape = HasDynamicShape(program);
   pass_manager->EnableIRPrinting();
 
+  pass_manager->AddPass(cinn::dialect::ir::CreatePdOpToCinnOpPass());
+  pass_manager->AddPass(
+      std::make_unique<cinn::dialect::ir::AddBroadcastToElementwisePass>());
+
   if (has_dynamic_shape) {
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+    pass_manager->AddPass(cinn::dialect::ir::CreateInsertBroadcastPass());
     pass_manager->AddPass(pir::CreateShapeOptimizationPass());
     pass_manager->AddPass(
         std::make_unique<
             cinn::dialect::ir::FuseShapeOpsIntoGenerateShapeOpPass>());
     pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
-  }
-  pass_manager->EnableIRPrinting();
-  pass_manager->AddPass(cinn::dialect::ir::CreatePdOpToCinnOpPass());
-  if (has_dynamic_shape) {
     pass_manager->AddPass(pir::CreateShapeOptimizationPass());
   }
-  pass_manager->AddPass(
-      std::make_unique<cinn::dialect::ir::AddBroadcastToElementwisePass>());
-  pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
+
   pass_manager->AddPass(pir::CreateBuildCinnPass());
 
+  pass_manager->AddPass(
+      cinn::dialect::ir::CreateMoveGenerateShapeOpsToProloguePass());
   pass_manager->AddPass(cinn::dialect::ir::CreateDivideGroupOpToFusionOpPass());
   pass_manager->AddPass(cinn::dialect::ir::CreateDynamicReshapeOpPass());
   pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
@@ -1590,6 +1594,9 @@ void AddCinnPass(std::shared_ptr<PassManager> &pass_manager,  // NOLINT
         cinn::dialect::ir::CreateLowerCinnDyShapeFusionOpPass());
   }
   pass_manager->AddPass(cinn::dialect::ir::CreateLowerCinnFusionOpPass());
+  pass_manager->AddPass(
+      std::make_unique<
+          cinn::dialect::ir::SplitGenerateShapeIntoShapeOpsPass>());
 #else
   PADDLE_THROW(platform::errors::Unimplemented(
       "Currently we only support CINN Pass for Pir under @to_static, please "
diff --git a/paddle/pir/dialect/shape/utils/dim_expr.h b/paddle/pir/dialect/shape/utils/dim_expr.h
@@ -224,7 +224,26 @@ class ShapeOrData {
   explicit ShapeOrData(const std::vector<T>& shape)
       : shape_(shape), data_(std::nullopt) {}
   explicit ShapeOrData(const std::vector<T>& shape, const std::vector<T>& data)
-      : shape_(shape), data_(data) {}
+      : shape_(shape), data_(data) {
+    // Vaild check
+    if (shape.size() == 0) {
+      IR_ENFORCE(data.size() == 1,
+                 "When shape is 0-D, size of data shoubld be 1, but got %d.",
+                 data.size());
+    } else if (shape.size() == 1) {
+      IR_ENFORCE(shape[0].template Has<int64_t>(),
+                 "When shape is 1-D, value of shape shoubld be int");
+      IR_ENFORCE(
+          shape[0].template Get<int64_t>() == static_cast<int64_t>(data.size()),
+          "When shape is 1-D, size of data shoubld be the same as "
+          "value[%d] of shape, but got [%d].",
+          shape[0].template Get<std::int64_t>(),
+          data.size());
+    } else {
+      IR_THROW("Size of shape shoubld be 0 or 1, but got %d", shape.size());
+    }
+  }
+
   ShapeOrData() = default;
   ShapeOrData(const ShapeOrData&) = default;
   ShapeOrData(ShapeOrData&&) = default;
diff --git a/test/ir/pir/cinn/symbolic/test_cinn_broadcast_symbolic.py b/test/ir/pir/cinn/symbolic/test_cinn_broadcast_symbolic.py
diff --git a/test/ir/pir/cinn/symbolic/test_cinn_sub_graph_symbolic.py b/test/ir/pir/cinn/symbolic/test_cinn_sub_graph_symbolic.py

Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ Expr Optimize(Expr e,`
`65`	`65`	`RemoveGpuForloopsAxis(&copied);`
`66`	`66`	`}`
`67`	`67`	`CudaSyncThreadsDropIfThenElse(&copied);`
`68`		`- TransBufferWithDynamicShape(&copied);`
	`68`	`+ // TransBufferWithDynamicShape(&copied);`
`69`	`69`	`#endif`
`70`	`70`
`71`	`71`	`SimplifyBlocks(&copied);`