PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt
Copy file name to clipboard
+2-11Lines changed: 2 additions & 11 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt
Copy file name to clipboard
+2-11Lines changed: 2 additions & 11 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/CMakeLists.txt
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/CMakeLists.txt
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc
Copy file name to clipboard
+98-98Lines changed: 98 additions & 98 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.cc renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.cc
Copy file name to clipboard
+98-98Lines changed: 98 additions & 98 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.h
Copy file name to clipboard
+3-4Lines changed: 3 additions & 4 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.h
Copy file name to clipboard
+3-4Lines changed: 3 additions & 4 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_pass.cc renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass.cc
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_pass.cc renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass.cc
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_pass_utils.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass_utils.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_pass_utils.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass_utils.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_util.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_util.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_util.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_util.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/op_group.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_group.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/op_group.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_group.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/op_node.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_node.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/op_node.h renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_node.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_pass.cc renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_pass.cc renamed to ‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
@@ -1,15 +1,6 @@
+add_subdirectory(group_merge)
+
 if(NOT CINN_ONLY)
-  cinn_cc_library(
-    op_with_group_merge_pass
-    SRCS
-    group_with_group_merge_pass.cc
-    op_with_group_merge_pass.cc
-    cinn_group_lowering_pass.cc
-    tensor_node.cc
-    DEPS
-    op_dialect_vjp
-    pir_compiler
-    cinn_runtime_dialect)
 
   cinn_cc_library(
     pd_to_cinn_pass
 
@@ -0,0 +1,13 @@
+if(NOT CINN_ONLY)
+  cinn_cc_library(
+    op_with_group_merge_pass
+    SRCS
+    group_with_group_merge_pass.cc
+    op_with_group_merge_pass.cc
+    cinn_group_lowering_pass.cc
+    tensor_node.cc
+    DEPS
+    op_dialect_vjp
+    pir_compiler
+    cinn_runtime_dialect)
+endif()
@@ -14,27 +14,28 @@
 
 #pragma once
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/cinn_group_lowering_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/cinn_group_lowering_pass.h"
 
 #include <unordered_map>
 
 #include "paddle/cinn/adt/generate_map_expr.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/op_attribute.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.h"
 #include "paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h"
 #include "paddle/cinn/hlir/dialect/runtime/ir/runtime_dialect.h"
 #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/pir/core/program.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
+#include "paddle/pir/pass/pass_registry.h"
+#include "paddle/pir/pattern_rewrite/frozen_rewrite_pattern_set.h"
 
 PD_DECLARE_bool(cinn_enable_map_expr);
 
-namespace cinn {
-namespace dialect {
-namespace ir {
+namespace {
 
 std::vector<pir::Value> GetBlockOutsideInput(
     const std::vector<pir::Operation*> op_list) {
@@ -127,115 +128,114 @@ std::vector<pir::Operation*> GetOutputOpList(
   return vec_res;
 }
 
-std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program) {
-  ::pir::IrContext* ctx = ::pir::IrContext::Instance();
-
-  ctx->GetOrRegisterDialect<cinn::dialect::RuntimeDialect>();
-  ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
-  ctx->GetOrRegisterDialect<paddle::dialect::KernelDialect>();
-
-  std::string jit_op_name = cinn::dialect::JitKernelOp::name();
-  ::pir::OpInfo op_info = ctx->GetRegisteredOpInfo(jit_op_name);
-
-  auto ir_program = std::make_unique<::pir::Program>(ctx);
-  std::unordered_map<pir::Value, pir::Value> value_map;
-
-  auto target = cinn::common::DefaultNVGPUTarget();
-  auto scope = cinn::hlir::framework::BuildScope(target, *program);
-
-  for (auto it = program->block()->begin(); it != program->block()->end();
-       ++it) {
-    if (it->isa<cinn::dialect::GroupOp>()) {
-      // GetOpList and Call cinn CodeGen
-      auto group_op = it->dyn_cast<cinn::dialect::GroupOp>();
-
-      // op fusion
-      auto op_fusion = cinn::dialect::ir::OpFusionPassInternal(
-          GetOpListNotIncludeYield(group_op.ops()),
-          GetOutputOpList(group_op.ops()));
-
-      // fusion merge
-      auto group_list =
-          cinn::dialect::ir::GeneralFusionMergePassInternal(op_fusion);
-
-      // using yield op to sort
-      std::unordered_map<::pir::Value, size_t> value2id;
-      auto yeild_op = group_op.ops().back();
-      for (size_t i = 0; i < yeild_op->num_operands(); ++i) {
-        value2id[yeild_op->operand_source(i)] = i;
+class GroupOpPattern : public pir::OpRewritePattern<cinn::dialect::GroupOp> {
+ public:
+  using pir::OpRewritePattern<cinn::dialect::GroupOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(cinn::dialect::GroupOp group_op,
+                       pir::PatternRewriter& rewriter) const override {
+    ::pir::IrContext* ctx = ::pir::IrContext::Instance();
+    auto target = cinn::common::DefaultNVGPUTarget();
+    auto* program = group_op->GetParentProgram();
+    VLOG(4) << "Before GroupOpPattern: " << *program;
+    // TODO(Aurelius84): Remove scope after cleaning PirCompiler usless Build
+    // Interface
+    auto scope = std::make_shared<cinn::hlir::framework::Scope>();
+
+    VLOG(4) << "start Lowering Group Op: " << group_op;
+    // using yield op to sort
+    std::unordered_map<::pir::Value, size_t> value2id;
+    auto yeild_op = group_op.ops().back();
+    for (size_t i = 0; i < yeild_op->num_operands(); ++i) {
+      value2id[yeild_op->operand_source(i)] = i;
+    }
+    std::unordered_map<pir::Value, pir::Value> value_map;
+
+    // op fusion
+    auto op_fusion = cinn::dialect::ir::OpFusionPassInternal(
+        GetOpListNotIncludeYield(group_op.ops()),
+        GetOutputOpList(group_op.ops()));
+
+    // fusion merge
+    auto group_list =
+        cinn::dialect::ir::GeneralFusionMergePassInternal(op_fusion);
+
+    for (auto group : group_list) {
+      auto ir_compiler = std::make_shared<cinn::hlir::framework::PirCompiler>(
+          *program, target, scope);
+      if (FLAGS_cinn_enable_map_expr) {
+        cinn::adt::TryGenerateMapExprFromGroup(group);
       }
 
-      for (auto group : group_list) {
-        auto ir_compiler = std::make_shared<cinn::hlir::framework::PirCompiler>(
-            *program, target, scope);
-        hlir::framework::PirCompilerManager::Instance().insert(ir_compiler);
-        if (FLAGS_cinn_enable_map_expr) {
-          adt::TryGenerateMapExprFromGroup(group);
-        }
-        auto fn_ptr_res = ir_compiler->BuildCUDAJITInfo({group});
-        std::unordered_map<std::string, ::pir::Attribute> op_attrs{
-            {cinn::dialect::JitKernelOp::kAttrName,
-             cinn::dialect::CUDAJITInfoAttribute::get(ctx, fn_ptr_res[0])},
-        };
-
-        // Generate jit kernel op input and output
-        auto vec_ins = GetBlockOutsideInput(group->ops);
-
-        std::vector<pir::Value> vec_new_ins;
-        for (size_t i = 0; i < vec_ins.size(); ++i) {
-          vec_new_ins.push_back(value_map.at(vec_ins[i]));
+      auto fn_ptr_res = ir_compiler->BuildCUDAJITInfo({group});
+      std::unordered_map<std::string, ::pir::Attribute> op_attrs{
+          {cinn::dialect::JitKernelOp::kAttrName,
+           cinn::dialect::CUDAJITInfoAttribute::get(ctx, fn_ptr_res[0])},
+      };
+
+      // Generate jit kernel op input and output
+      auto vec_ins = GetBlockOutsideInput(group->ops);
+      for (size_t i = 0; i < vec_ins.size(); ++i) {
+        if (value_map.find(vec_ins[i]) != value_map.end()) {
+          vec_ins[i] = value_map.at(vec_ins[i]);
         }
+      }
 
-        std::unordered_map<size_t, size_t> codegen2orig;
+      std::vector<pir::Type> vec_types;
+      for (size_t i = 0; i < group->output_values.size(); ++i) {
+        vec_types.push_back(group->output_values[i].type());
+      }
 
-        std::vector<pir::Type> vec_types;
-        for (size_t i = 0; i < group->output_values.size(); ++i) {
-          vec_types.push_back(group->output_values[i].type());
+      auto jit_kernel_op = rewriter.Build<cinn::dialect::JitKernelOp>(
+          vec_ins, op_attrs, vec_types);
+      for (size_t i = 0; i < jit_kernel_op.num_results(); ++i) {
+        auto find_it = value2id.find(group->output_values[i]);
+        if (find_it != value2id.end()) {
+          rewriter.ReplaceAllUsesWith(group_op.result(find_it->second),
+                                      jit_kernel_op.result(i));
         }
+        value_map[group->output_values[i]] = jit_kernel_op.result(i);
+      }
+    }
+    value_map.clear();
+    VLOG(4) << "Before GroupOpPattern.EraseOp: " << *program;
+    rewriter.EraseOp(group_op);
+    return true;
+  }
+};
 
-        ::pir::Operation* cinn_op =
-            ::pir::Operation::Create(vec_new_ins, op_attrs, vec_types, op_info);
+class CinnGroupLoweringPass : public pir::PatternRewritePass {
+ public:
+  CinnGroupLoweringPass() : pir::PatternRewritePass("cinn_group_lowering", 1) {}
 
-        for (size_t i = 0; i < cinn_op->num_results(); ++i) {
-          auto find_it = value2id.find(group->output_values[i]);
-          value_map[group->output_values[i]] = cinn_op->result(i);
-          if (find_it != value2id.end()) {
-            value_map[group_op.result(find_it->second)] = cinn_op->result(i);
-          }
-        }
+  pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
+    context->GetOrRegisterDialect<cinn::dialect::RuntimeDialect>();
+    context->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
+    context->GetOrRegisterDialect<paddle::dialect::KernelDialect>();
 
-        ir_program->block()->push_back(cinn_op);
-      }
+    pir::RewritePatternSet ps(context);
+    ps.Add<GroupOpPattern>(context);
 
-    } else {
-      std::vector<pir::Value> vec_ins;
+    return ps;
+  }
 
-      for (size_t i = 0; i < it->num_operands(); ++i) {
-        if (it->operand_source(i)) {
-          vec_ins.push_back(value_map.at(it->operand_source(i)));
-        } else {
-          vec_ins.push_back(it->operand_source(i));
-        }
-      }
+  bool CanApplyOn(pir::Operation* op) const override {
+    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+  }
+};
 
-      std::vector<pir::Type> vec_types;
-      for (size_t i = 0; i < it->num_results(); ++i) {
-        vec_types.push_back(it->result(i).type());
-      }
+}  // namespace
 
-      ::pir::OpInfo info1 = ctx->GetRegisteredOpInfo(it->name());
-      ::pir::Operation* op =
-          ::pir::Operation::Create(vec_ins, it->attributes(), vec_types, info1);
+namespace cinn {
+namespace dialect {
+namespace ir {
 
-      ir_program->block()->push_back(op);
-      for (size_t i = 0; i < it->num_results(); ++i) {
-        value_map[it->result(i)] = op->result(i);
-      }
-    }
-  }
-  return ir_program;
+std::unique_ptr<::pir::Pass> CreateCinnGroupLoweringPass() {
+  return std::make_unique<CinnGroupLoweringPass>();
 }
 
 }  // namespace ir
 }  // namespace dialect
 }  // namespace cinn
+
+REGISTER_IR_PASS(cinn_group_lowering, CinnGroupLoweringPass);
@@ -14,14 +14,13 @@
 
 #pragma once
 
-#include "paddle/pir/core/program.h"
+#include <memory>
+#include "paddle/pir/pass/pass.h"
 
 namespace cinn {
 namespace dialect {
 namespace ir {
-
-std::unique_ptr<pir::Program> CINNGroupLoweringPass(::pir::Program* program);
-
+std::unique_ptr<::pir::Pass> CreateCinnGroupLoweringPass();
 }  // namespace ir
 }  // namespace dialect
 }  // namespace cinn
@@ -16,15 +16,15 @@
 #include <set>
 #include <unordered_map>
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_group.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_group.h"
 #include "paddle/pir/core/ir_printer.h"
 #include "paddle/pir/core/value.h"
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_pass_utils.h"
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass_utils.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.h"
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/group_with_group_merge_util.h"
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 #include "paddle/phi/core/flags.h"
 
 #include "paddle/cinn/common/is_reachable_predicator.h"
 
@@ -14,8 +14,8 @@
 
 #pragma once
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_group.h"
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_group.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 
 namespace cinn {
 namespace dialect {
 
@@ -27,7 +27,7 @@
 #include "paddle/pir/core/value.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 
 namespace cinn {
 namespace dialect {
 
@@ -16,8 +16,8 @@
 
 #include <memory>
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_node.h"
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_node.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 
 namespace cinn {
 namespace dialect {
 
@@ -15,8 +15,8 @@
 #pragma once
 
 #include <memory>
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_util.h"
-#include "paddle/cinn/hlir/dialect/operator/transforms/tensor_node.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/tensor_node.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/core/operation.h"
 
 
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/cinn/hlir/dialect/operator/transforms/op_with_group_merge_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.h"
 
 #include <limits.h>
 #include <memory>