[DRR] Create PatternRewritePass in pass.h (#58655)

gongshaotian · web-flow · commit 03e59e4e0fc9 · 2023-11-15T15:44:49.000+08:00
* create class DrrBasePass

* add PatternPass in pass.h

* Change the RewritePattern type pass from inheriting Pass to inheriting PatternPass

* Rename PatternPass as PatternRewritePass and delete CanApplyOn() func

* fix bug

* Fix compilation bugs

* fix bug
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
@@ -143,17 +143,17 @@ class AddBrodcastToElementwisePattern : public pir::OpRewritePattern<OPTYPE> {
 };
 
 AddBroadcastToElementwisePass::AddBroadcastToElementwisePass()
-    : pir::Pass("add_broadcast_to_elementwise_pass", 1) {}
+    : pir::PatternRewritePass("add_broadcast_to_elementwise_pass", 1) {}
 
-bool AddBroadcastToElementwisePass::Initialize(pir::IrContext* context) {
+pir::RewritePatternSet AddBroadcastToElementwisePass::InitializePatterns(
+    pir::IrContext* context) {
   pir::RewritePatternSet ps(context);
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::AddOp>>(context);
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::SubtractOp>>(context);
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::MultiplyOp>>(context);
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::DivideOp>>(context);
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::ElementwisePowOp>>(
       context);
-
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::LessThanOp>>(context);
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::LessEqualOp>>(
       context);
@@ -164,15 +164,7 @@ bool AddBroadcastToElementwisePass::Initialize(pir::IrContext* context) {
   ps.Add<AddBrodcastToElementwisePattern<paddle::dialect::GreaterEqualOp>>(
       context);
 
-  patterns_ = ::pir::FrozenRewritePatternSet(std::move(ps));
-  return true;
-}
-
-void AddBroadcastToElementwisePass::Run(pir::Operation* op) {
-  pir::GreedyRewriteConfig cfg;
-  cfg.use_top_down_traversal = true;
-  cfg.max_iterations = 10;
-  pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
+  return ps;
 }
 
 bool AddBroadcastToElementwisePass::CanApplyOn(pir::Operation* op) const {
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h b/paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h
@@ -21,18 +21,13 @@ namespace cinn {
 namespace dialect {
 namespace ir {
 
-class AddBroadcastToElementwisePass : public pir::Pass {
+class AddBroadcastToElementwisePass : public pir::PatternRewritePass {
  public:
   AddBroadcastToElementwisePass();
 
-  bool Initialize(pir::IrContext *context) override;
-
-  void Run(pir::Operation *op) override;
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override;
 
   bool CanApplyOn(pir::Operation *op) const override;
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 }  // namespace ir
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
@@ -221,9 +221,11 @@ class UniformOpPattern : public pir::drr::DrrPatternBase<UniformOpPattern> {
   }
 };
 
-PdOpToCinnOpPass::PdOpToCinnOpPass() : pir::Pass("pd_to_cinn_pass", 1) {}
+PdOpToCinnOpPass::PdOpToCinnOpPass()
+    : pir::PatternRewritePass("pd_to_cinn_pass", 1) {}
 
-bool PdOpToCinnOpPass::Initialize(pir::IrContext *context) {
+pir::RewritePatternSet PdOpToCinnOpPass::InitializePatterns(
+    pir::IrContext *context) {
   pir::RewritePatternSet ps(context);
   ps.Add<ScaleOpPattern>(
       context);  // NOTE, scale op pattern should before AddBroadcastTo
@@ -232,15 +234,7 @@ bool PdOpToCinnOpPass::Initialize(pir::IrContext *context) {
   ps.Add<ReshapeOpPattern>(context);
   // ps.Add(UniformOpPattern().Build(context));
 
-  patterns_ = ::pir::FrozenRewritePatternSet(std::move(ps));
-  return true;
-}
-
-void PdOpToCinnOpPass::Run(pir::Operation *op) {
-  pir::GreedyRewriteConfig cfg;
-  cfg.use_top_down_traversal = true;
-  cfg.max_iterations = 10;
-  pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
+  return ps;
 }
 
 bool PdOpToCinnOpPass::CanApplyOn(pir::Operation *op) const {
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.h b/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.h
@@ -22,18 +22,13 @@ namespace cinn {
 namespace dialect {
 namespace ir {
 
-class PdOpToCinnOpPass : public pir::Pass {
+class PdOpToCinnOpPass : public pir::PatternRewritePass {
  public:
   PdOpToCinnOpPass();
 
-  bool Initialize(pir::IrContext *context) override;
-
-  void Run(pir::Operation *op) override;
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override;
 
   bool CanApplyOn(pir::Operation *op) const override;
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 void PdOp2CinnOpConverter(::pir::Program *program);
diff --git a/paddle/fluid/pir/transforms/constant_folding_pass.cc b/paddle/fluid/pir/transforms/constant_folding_pass.cc
@@ -172,7 +172,7 @@ class ConstantFoldingPattern : public pir::RewritePattern {
     // prepare op outputs
     std::vector<pir::Type> output_types;
     for (uint32_t i = 0; i < op->num_results(); i++) {
-      output_types.push_back(op->result_type(i));
+      output_types.push_back(op->result(i).type());
     }
 
     auto* temp_op =
diff --git a/paddle/fluid/pir/transforms/fusion/attention_fuse_pass.cc b/paddle/fluid/pir/transforms/fusion/attention_fuse_pass.cc
@@ -214,32 +214,17 @@ class MultiHeadMatmulFusePattern
   }
 };
 
-class AttentionFusePass : public pir::Pass {
+class AttentionFusePass : public pir::PatternRewritePass {
  public:
-  AttentionFusePass() : pir::Pass("attention_fuse_pass", 2) {}
+  AttentionFusePass() : pir::PatternRewritePass("attention_fuse_pass", 2) {}
 
-  bool Initialize(pir::IrContext *context) override {
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
     ps.Add(MultiHeadMatmulFusePattern().Build(context));
     // Add other attention variant fuse pattern.
 
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
+    return ps;
   }
-
-  void Run(pir::Operation *op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
-  }
-
-  bool CanApplyOn(pir::Operation *op) const override {
-    return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
-  }
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 }  // namespace
diff --git a/paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc
@@ -105,32 +105,18 @@ class FusedDropoutGradAddGradPattern
   }
 };
 
-class FusedDropoutAddPass : public pir::Pass {
+class FusedDropoutAddPass : public pir::PatternRewritePass {
  public:
-  FusedDropoutAddPass() : pir::Pass("fused_dropout_add_pass", 1) {}
+  FusedDropoutAddPass()
+      : pir::PatternRewritePass("fused_dropout_add_pass", 1) {}
 
-  bool Initialize(pir::IrContext *context) override {
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
-
     ps.Add(FusedDropoutAddPattern().Build(context));
     ps.Add(FusedDropoutGradAddGradPattern().Build(context));
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
-  }
 
-  void Run(pir::Operation *op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
+    return ps;
   }
-
-  bool CanApplyOn(pir::Operation *op) const override {
-    return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
-  }
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 }  // namespace
diff --git a/paddle/fluid/pir/transforms/fusion/fused_gemm_epilogue_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_gemm_epilogue_pass.cc
@@ -342,11 +342,12 @@ class FusedLinearReluGradPattern
   }
 };
 
-class FusedGemmEpiloguePass : public pir::Pass {
+class FusedGemmEpiloguePass : public pir::PatternRewritePass {
  public:
-  FusedGemmEpiloguePass() : pir::Pass("fused_gemm_epilogue_pass", 2) {}
+  FusedGemmEpiloguePass()
+      : pir::PatternRewritePass("fused_gemm_epilogue_pass", 2) {}
 
-  bool Initialize(pir::IrContext *context) override {
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
     ps.Add(FusedLinearGradPattern().Build(context));
     ps.Add(FusedLinearPattern().Build(context));
@@ -355,23 +356,8 @@ class FusedGemmEpiloguePass : public pir::Pass {
     ps.Add(FusedLinearGeluGradPattern().Build(context));
     ps.Add(FusedLinearReluGradPattern().Build(context));
 
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
+    return ps;
   }
-
-  void Run(pir::Operation *op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
-  }
-
-  bool CanApplyOn(pir::Operation *op) const override {
-    return op->name() == "builtin.module" && op->num_regions() > 0;
-  }
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 }  // namespace
diff --git a/paddle/fluid/pir/transforms/fusion/fused_linear_param_grad_add_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_linear_param_grad_add_pass.cc
@@ -316,36 +316,22 @@ class FusedMatmulAddGradAddbPattern
   }
 };
 
-class FusedLinearParamGradAddPass : public pir::Pass {
+class FusedLinearParamGradAddPass : public pir::PatternRewritePass {
  public:
   FusedLinearParamGradAddPass()
-      : pir::Pass("fused_linear_param_grad_add_pass", 1) {}
+      : pir::PatternRewritePass("fused_linear_param_grad_add_pass", 1) {}
 
-  bool Initialize(pir::IrContext *context) override {
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
     ps.Add(FusedMatmulAddGradAddPattern().Build(context));
     ps.Add(FusedMatmulGradAddPattern().Build(context));
     ps.Add(FusedMatmulAddaPattern().Build(context));
     ps.Add(FusedMatmulAddbPattern().Build(context));
     ps.Add(FusedMatmulAddGradAddaPattern().Build(context));
     ps.Add(FusedMatmulAddGradAddbPattern().Build(context));
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
-  }
-
-  void Run(pir::Operation *op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
-  }
 
-  bool CanApplyOn(pir::Operation *op) const override {
-    return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
+    return ps;
   }
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 }  // namespace
diff --git a/paddle/pir/pass/pass.h b/paddle/pir/pass/pass.h
@@ -18,8 +18,10 @@
 #include <string>
 #include <vector>
 
+#include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/core/enforce.h"
 #include "paddle/pir/pass/analysis_manager.h"
+#include "paddle/pir/pattern_rewrite/pattern_rewrite_driver.h"
 
 namespace pir {
 
@@ -105,4 +107,39 @@ class IR_API Pass {
   friend class detail::PassAdaptor;
 };
 
+class PatternRewritePass : public Pass {
+ public:
+  PatternRewritePass(const std::string& name,
+                     uint8_t opt_level,
+                     const std::vector<std::string>& dependents = {})
+      : Pass(name, opt_level, dependents) {}
+
+ protected:
+  virtual RewritePatternSet InitializePatterns(IrContext* context) = 0;
+
+  bool Initialize(IrContext* context) final {
+    RewritePatternSet ps = InitializePatterns(context);
+    IR_ENFORCE(ps.Empty() == false,
+               "Pass creation failed."
+               "When using PatternRewritePass to create a Pass, the number of "
+               "customized Patterns is required to be greater than zero."
+               "Suggested fix: Check whether Pattern is added to the "
+               "InitializePatterns() function of class [%s]",
+               name());
+    patterns_ = FrozenRewritePatternSet(std::move(ps));
+
+    return true;
+  }
+
+  void Run(Operation* op) override {
+    GreedyRewriteConfig cfg;
+    cfg.use_top_down_traversal = true;
+    cfg.max_iterations = 10;
+    ApplyPatternsGreedily(op->region(0), patterns_, cfg);
+  }
+
+ private:
+  FrozenRewritePatternSet patterns_;
+};
+
 }  // namespace pir
diff --git a/paddle/pir/pattern_rewrite/pattern_match.h b/paddle/pir/pattern_rewrite/pattern_match.h
@@ -346,6 +346,8 @@ class RewritePatternSet {
 
   void Clear() { native_patterns_.clear(); }
 
+  bool Empty() const { return native_patterns_.empty(); }
+
   // 'add' methods for adding patterns to the set.
   template <typename... Ts,
             typename ConstructorArg,
diff --git a/test/cpp/pir/pattern_rewrite/drr_same_type_binding_test.cc b/test/cpp/pir/pattern_rewrite/drr_same_type_binding_test.cc
@@ -284,31 +284,17 @@ void BuildProgram(pir::Builder &builder) {  // NOLINT
   builder.Build<paddle::dialect::FetchOp>(matmul_op6.out(), "out3", 2);
 }
 
-class DrrPatternRewritePass : public pir::Pass {
+class DrrPatternRewritePass : public pir::PatternRewritePass {
  public:
-  DrrPatternRewritePass() : pir::Pass("DrrPatternRewritePass", 1) {}
+  DrrPatternRewritePass()
+      : pir::PatternRewritePass("drr_pattern_rewrite_pass", 1) {}
 
-  bool Initialize(pir::IrContext *context) override {
+  pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
     ps.Add(SameTypeBindingTestPattern().Build(context));
 
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
+    return ps;
   }
-
-  void Run(pir::Operation *op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
-  }
-
-  bool CanApplyOn(pir::Operation *op) const override {
-    return op->name() == "builtin.module" && op->num_regions() > 0;
-  }
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 TEST(DrrTest, drr_demo) {
diff --git a/test/cpp/pir/pattern_rewrite/drr_test.cc b/test/cpp/pir/pattern_rewrite/drr_test.cc
diff --git a/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc b/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc

Original file line number	Diff line number	Diff line change
`@@ -172,7 +172,7 @@ class ConstantFoldingPattern : public pir::RewritePattern {`
`172`	`172`	`// prepare op outputs`
`173`	`173`	`std::vector<pir::Type> output_types;`
`174`	`174`	`for (uint32_t i = 0; i < op->num_results(); i++) {`
`175`		`- output_types.push_back(op->result_type(i));`
	`175`	`+ output_types.push_back(op->result(i).type());`
`176`	`176`	`}`
`177`	`177`
`178`	`178`	`auto* temp_op =`