PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc
Copy file name to clipboard
+26-9Lines changed: 26 additions & 9 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc
Copy file name to clipboard
+26-9Lines changed: 26 additions & 9 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass.cc
Copy file name to clipboard
+14-3Lines changed: 14 additions & 3 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass.cc
Copy file name to clipboard
+14-3Lines changed: 14 additions & 3 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
Copy file name to clipboard
+14-3Lines changed: 14 additions & 3 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
Copy file name to clipboard
+14-3Lines changed: 14 additions & 3 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.h
Copy file name to clipboard
+8-2Lines changed: 8 additions & 2 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.h
Copy file name to clipboard
+8-2Lines changed: 8 additions & 2 deletions
diff --git a/‎paddle/cinn/hlir/framework/pir/group.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/cinn/hlir/framework/pir/group.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/transforms/CMakeLists.txt
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/fluid/pir/transforms/CMakeLists.txt
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/transforms/infer_symbolic_shape_pass.cc
Copy file name to clipboard
+131Lines changed: 131 additions & 0 deletions b/‎paddle/fluid/pir/transforms/infer_symbolic_shape_pass.cc
Copy file name to clipboard
+131Lines changed: 131 additions & 0 deletions
@@ -28,6 +28,7 @@
 #include "paddle/cinn/hlir/framework/pir_compiler.h"
 #include "paddle/cinn/runtime/flags.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/pir/core/program.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
 #include "paddle/pir/pass/pass_registry.h"
@@ -130,7 +131,11 @@ std::vector<pir::Operation*> GetOutputOpList(
 
 class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {
  public:
-  using pir::OpRewritePattern<cinn::dialect::GroupOp>::OpRewritePattern;
+  GroupOpPattern(
+      ::pir::IrContext* context,
+      const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis)
+      : pir::OpRewritePattern<cinn::dialect::GroupOp>(context),
+        shape_analysis_(shape_analysis) {}
 
   bool MatchAndRewrite(cinn::dialect::GroupOp group_op,
                        pir::PatternRewriter& rewriter) const override {
@@ -154,15 +159,17 @@ class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {
     // op fusion
     auto op_fusion = cinn::dialect::ir::OpFusionPassInternal(
         GetOpListNotIncludeYield(group_op.ops()),
-        GetOutputOpList(group_op.ops()));
+        GetOutputOpList(group_op.ops()),
+        shape_analysis_);
 
     // fusion merge
-    auto group_list =
-        cinn::dialect::ir::GeneralFusionMergePassInternal(op_fusion);
+    auto group_list = cinn::dialect::ir::GeneralFusionMergePassInternal(
+        op_fusion, shape_analysis_);
 
     for (auto group : group_list) {
       auto ir_compiler = cinn::hlir::framework::PirCompilerManager::Create(
           *program, target, scope);
+      group->shape_analysis = shape_analysis_;
       if (FLAGS_cinn_enable_map_expr) {
         cinn::adt::TryGenerateMapExprFromGroup(group);
       }
@@ -202,26 +209,35 @@ class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {
     rewriter.EraseOp(group_op);
     return true;
   }
+
+ private:
+  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_;
 };
 
 class CinnGroupLoweringPass : public pir::PatternRewritePass {
  public:
-  CinnGroupLoweringPass() : pir::PatternRewritePass("cinn_group_lowering", 1) {}
+  CinnGroupLoweringPass(
+      const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis)
+      : pir::PatternRewritePass("cinn_group_lowering", 1),
+        shape_analysis_(shape_analysis) {}
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
     context->GetOrRegisterDialect<cinn::dialect::RuntimeDialect>();
     context->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
     context->GetOrRegisterDialect<paddle::dialect::KernelDialect>();
 
     pir::RewritePatternSet ps(context);
-    ps.Add<GroupOpPattern>(context);
+    ps.Add<GroupOpPattern>(context, shape_analysis_);
 
     return ps;
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
     return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
   }
+
+ private:
+  const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis_;
 };
 
 }  // namespace
@@ -230,12 +246,13 @@ namespace cinn {
 namespace dialect {
 namespace ir {
 
-std::unique_ptr<::pir::Pass> CreateCinnGroupLoweringPass() {
-  return std::make_unique<CinnGroupLoweringPass>();
+std::unique_ptr<::pir::Pass> CreateCinnGroupLoweringPass(
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis) {
+  return std::make_unique<CinnGroupLoweringPass>(shape_analysis);
 }
 
 }  // namespace ir
 }  // namespace dialect
 }  // namespace cinn
 
-REGISTER_IR_PASS(cinn_group_lowering, CinnGroupLoweringPass);
+// REGISTER_IR_PASS(cinn_group_lowering, CinnGroupLoweringPass);
@@ -15,12 +15,15 @@
 #pragma once
 
 #include <memory>
+#include "paddle/pir/dialect/shape/utils/shape_utils.h"
 #include "paddle/pir/pass/pass.h"
 
 namespace cinn {
 namespace dialect {
 namespace ir {
-std::unique_ptr<::pir::Pass> CreateCinnGroupLoweringPass();
+std::unique_ptr<::pir::Pass> CreateCinnGroupLoweringPass(
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis =
+        nullptr);
 }  // namespace ir
 }  // namespace dialect
 }  // namespace cinn
@@ -1033,7 +1033,10 @@ class FusionPassRegistrar final : public Registrar {
 // code generation.
 class GeneralFusionMergePassHelper {
  public:
-  explicit GeneralFusionMergePassHelper(const GroupList& group_list) {
+  explicit GeneralFusionMergePassHelper(
+      const GroupList& group_list,
+      const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis)
+      : shape_analysis_(shape_analysis) {
     fusion_groups_ = group_list;
     // init input to consumers.
     InitInputToConsumers();
@@ -2123,21 +2126,29 @@ class GeneralFusionMergePassHelper {
     }
   }
 
+  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis() const {
+    return CHECK_NOTNULL(shape_analysis_.lock());
+  }
+
   GroupList fusion_groups_;
   std::unordered_map<GroupPtr, int> fusion_groups_index_;
   std::unordered_set<const ::pir::Operation*> output_ops_set_;
   std::unordered_map<::pir::Value,
                      std::unordered_set<GroupPtr, Hasher, Comparator>>
       input_to_consumers_;
+  std::weak_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_;
 };
 
-GroupList GeneralFusionMergePassInternal(const GroupList& group_list) {
+GroupList GeneralFusionMergePassInternal(
+    const GroupList& group_list,
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis) {
   if (group_list.size() <= 1) {
     VLOG(3) << "Don't do Fusoin Merge Pass...!";
     return group_list;
   }
 
-  GeneralFusionMergePassHelper fusion_merge_pass_helper(group_list);
+  GeneralFusionMergePassHelper fusion_merge_pass_helper(group_list,
+                                                        shape_analysis);
   auto res = fusion_merge_pass_helper();
 
   if (VLOG_IS_ON(6)) {
 
@@ -245,7 +245,10 @@ class OpFusionPassHelper {
  public:
   explicit OpFusionPassHelper(
       const std::vector<pir::Operation*>& op_list,
-      const std::vector<pir::Operation*>& output_op_list = {}) {
+      const std::vector<pir::Operation*>& output_op_list = {},
+      const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis =
+          nullptr)
+      : shape_analysis_(shape_analysis) {
     // init fusion relation
     InitFusionRelation();
     // filter op data, create group for each op
@@ -608,6 +611,11 @@ class OpFusionPassHelper {
 
     return false;
   }
+
+  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis() const {
+    return CHECK_NOTNULL(shape_analysis_.lock());
+  }
+
   std::vector<::pir::Operation*> ops_;
   std::unordered_map<const ::pir::Operation*, GroupPtr> fusion_groups_;
   std::unordered_set<const ::pir::Operation*> output_ops_set_;
@@ -625,14 +633,17 @@ class OpFusionPassHelper {
     std::unordered_map<OpPatternKind, ConditionFunction> fusion_op_kind = {};
   };
   std::unordered_map<OpPatternKind, FusionRelation> fusion_relation_map_;
+  std::weak_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_;
 };
 
 GroupList OpFusionPassInternal(
     const std::vector<pir::Operation*>& op_list,
-    const std::vector<pir::Operation*>& output_op_list) {
+    const std::vector<pir::Operation*>& output_op_list,
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis) {
   VLOG(3) << "OpFusionPass...!";
 
-  auto op_fusion_helper = OpFusionPassHelper(op_list, output_op_list);
+  auto op_fusion_helper =
+      OpFusionPassHelper(op_list, output_op_list, shape_analysis);
   auto res = op_fusion_helper();
 
   if (VLOG_IS_ON(6)) {
 
@@ -16,6 +16,7 @@
 
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 #include "paddle/pir/core/program.h"
+#include "paddle/pir/dialect/shape/utils/shape_utils.h"
 
 namespace cinn {
 namespace dialect {
@@ -26,9 +27,14 @@ using GroupList = std::vector<GroupPtr>;
 
 GroupList OpFusionPassInternal(
     const std::vector<pir::Operation*>& op_list,
-    const std::vector<pir::Operation*>& output_op_list = {});
+    const std::vector<pir::Operation*>& output_op_list = {},
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis =
+        nullptr);
 
-GroupList GeneralFusionMergePassInternal(const GroupList& group_list);
+GroupList GeneralFusionMergePassInternal(
+    const GroupList& group_list,
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis =
+        nullptr);
 
 }  // namespace ir
 }  // namespace dialect
 
@@ -20,6 +20,7 @@
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/pir/core/operation.h"
 #include "paddle/pir/core/value.h"
+#include "paddle/pir/dialect/shape/utils/shape_utils.h"
 
 namespace cinn {
 
@@ -73,6 +74,8 @@ struct Group {
   // if as sub-group, used for belong groups.
   std::unordered_set<std::shared_ptr<Group>> belong_groups;
 
+  std::shared_ptr<::pir::ShapeConstraintIRAnalysis> shape_analysis = nullptr;
+
   // for op lowering.
   std::vector<std::string> input_names;
   std::vector<std::string> output_names;
 
@@ -1,9 +1,12 @@
 file(GLOB_RECURSE transforms_srcs "*.cc")
 if(NOT WITH_CINN)
   list(
-    REMOVE_ITEM transforms_srcs ${CMAKE_CURRENT_SOURCE_DIR}/build_cinn_pass.cc
+    REMOVE_ITEM
+    transforms_srcs
+    ${CMAKE_CURRENT_SOURCE_DIR}/build_cinn_pass.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/sub_graph_extract_pass.cc
-    ${CMAKE_CURRENT_SOURCE_DIR}/sub_graph_detector.cc)
+    ${CMAKE_CURRENT_SOURCE_DIR}/sub_graph_detector.cc
+    ${CMAKE_CURRENT_SOURCE_DIR}/infer_symbolic_shape_pass.cc)
 endif()
 
 set(transforms_deps drr op_dialect op_dialect_vjp standalone_executor pir)
 
@@ -0,0 +1,131 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pir/transforms/infer_symbolic_shape_pass.h"
+
+#include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/pir/pass/pass.h"
+#include "paddle/pir/pass/pass_registry.h"
+
+namespace {
+
+void InferUnaryElementwiseSymbolicShape(
+    const pir::Operation& op,
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis) {
+  auto input = op.operand_source(0);
+  auto output = op.result(0);
+  const auto& in_sym_dims =
+      shape_analysis->GetOrCreateSymbolicDimsForRankedValue(input);
+  const auto& out_sym_dims =
+      shape_analysis->GetOrCreateSymbolicDimsForRankedValue(output);
+  pir::SymbolicDimMgr& sym_dim_mgr = shape_analysis->symbolicDimMgr();
+  for (auto i = 0; i < out_sym_dims.size(); ++i) {
+    if (in_sym_dims[i].IsDynamic() || out_sym_dims[i].IsDynamic()) {
+      sym_dim_mgr.MapSymbolicDimEqual(in_sym_dims[i], out_sym_dims[i]);
+    } else {
+      // do nothing
+    }
+  }
+}
+
+// TODO(zyfncg): support broadcast for elementwise ops.
+void InferBinaryElementwiseSymbolicShape(
+    const pir::Operation& op,
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis) {
+  auto input0 = op.operand_source(0);
+  auto input1 = op.operand_source(1);
+  auto output = op.result(0);
+  const auto& in_sym_dims0 =
+      shape_analysis->GetOrCreateSymbolicDimsForRankedValue(input0);
+  const auto& in_sym_dims1 =
+      shape_analysis->GetOrCreateSymbolicDimsForRankedValue(input1);
+  const auto& out_sym_dims =
+      shape_analysis->GetOrCreateSymbolicDimsForRankedValue(output);
+  pir::SymbolicDimMgr& sym_dim_mgr = shape_analysis->symbolicDimMgr();
+  for (auto i = 0; i < out_sym_dims.size(); ++i) {
+    if (in_sym_dims0[i].IsDynamic() || in_sym_dims1[i].IsDynamic() ||
+        out_sym_dims[i].IsDynamic()) {
+      sym_dim_mgr.MapSymbolicDimEqual(in_sym_dims0[i], out_sym_dims[i]);
+      sym_dim_mgr.MapSymbolicDimEqual(in_sym_dims1[i], out_sym_dims[i]);
+    } else {
+      // do nothing
+    }
+  }
+}
+
+class InferSymbolicShapePass : public pir::Pass {
+ public:
+  InferSymbolicShapePass(
+      const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis)
+      : pir::Pass("infer_symbolic_shape_pass", /*opt_level=*/1),
+        shape_analysis_(shape_analysis) {}
+
+  void Run(pir::Operation* op) override {
+    auto module_op = op->dyn_cast<pir::ModuleOp>();
+    IR_ENFORCE(module_op, "infer_symbolic_shape_pass should run on module op.");
+
+    for (auto& op : *module_op.block()) {
+      if (op.isa<cinn::dialect::GroupOp>()) {
+        for (auto* local_op : op.dyn_cast<cinn::dialect::GroupOp>().ops()) {
+          InferSymbolicShape(*local_op);
+        }
+      } else {
+        InferSymbolicShape(op);
+      }
+    }
+  }
+
+  bool CanApplyOn(pir::Operation* op) const override {
+    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+  }
+
+ private:
+  typedef void (*InferSymShapeFunc)(
+      const pir::Operation&,
+      const std::shared_ptr<pir::ShapeConstraintIRAnalysis>&);
+  void InferSymbolicShape(const pir::Operation& op) {
+    thread_local static std::unordered_map<std::string, InferSymShapeFunc>
+        infer_sym_shape_map(GetInferSymShapeMap());
+    auto it = infer_sym_shape_map.find(op.name());
+
+    if (it != infer_sym_shape_map.end()) {
+      it->second(op, shape_analysis_);
+    } else {
+      LOG(WARNING) << "[" << op.name()
+                   << "] is not supported for infer_symbolic_shape pass.";
+    }
+  }
+
+  static std::unordered_map<std::string, InferSymShapeFunc>
+  GetInferSymShapeMap() {
+    return std::unordered_map<std::string, InferSymShapeFunc>{
+        {paddle::dialect::ExpOp::name(), &InferUnaryElementwiseSymbolicShape},
+        {paddle::dialect::SubtractOp::name(),
+         &InferBinaryElementwiseSymbolicShape}};
+  }
+
+  std::shared_ptr<pir::ShapeConstraintIRAnalysis> shape_analysis_;
+};
+
+}  // namespace
+
+namespace pir {
+
+std::unique_ptr<Pass> CreateInferSymbolicShapePass(
+    const std::shared_ptr<pir::ShapeConstraintIRAnalysis>& shape_analysis) {
+  return std::make_unique<InferSymbolicShapePass>(shape_analysis);
+}
+
+}  // namespace pir