[DRR] Create PatternRewritePass in pass.h

paddle/fluid/pir/drr/api/drr_base_pass.h

@@ -0,0 +1,59 @@
+// 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
+//
+//     http://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/pir/pass/pass.h"
+#include "paddle/pir/pattern_rewrite/pattern_rewrite_driver.h"
+
+#pragma once
+
+namespace pir {
+namespace drr {
+
+class DrrBasePass : public pir::Pass {
+ public:
+  DrrBasePass(const std::string &name = "DrrBasePass",
+              uint8_t opt_level = 1,
+              const std::vector<std::string> &dependents = {})
+      : pir::Pass(name, opt_level, dependents) {}
+
+  virtual void DrrPassInitialize(pir::RewritePatternSet *ps,
+                                 pir::IrContext *context) {
+    return;
+  }
+
+  bool Initialize(pir::IrContext *context) final {
+    pir::RewritePatternSet ps(context);
+    DrrPassInitialize(&ps, context);
+
+    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
+    return true;
+  }
+
+  void Run(pir::Operation *op) final {
+    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 final {
+    return op->name() == "builtin.module" && op->num_regions() > 0;
+  }
+
+ private:
+  pir::FrozenRewritePatternSet patterns_;
+};
+
+}  // namespace drr
+}  // namespace pir

test/cpp/pir/pattern_rewrite/drr_test.cc

@@ -18,11 +18,10 @@
 
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/drr/api/drr_base_pass.h"
 #include "paddle/fluid/pir/drr/api/drr_pattern_base.h"
 #include "paddle/pir/core/builtin_dialect.h"
-#include "paddle/pir/pass/pass.h"
 #include "paddle/pir/pass/pass_manager.h"
-#include "paddle/pir/pattern_rewrite/pattern_rewrite_driver.h"
 #include "paddle/pir/transforms/dead_code_elimination_pass.h"
 
 class RemoveRedundentReshapePattern
@@ -187,35 +186,18 @@ void BuildProgram(pir::Builder &builder) {  // NOLINT
   builder.Build<paddle::dialect::FetchOp>(relu_op_second.out(), "out", 0);
 }
 
-class DrrPatternRewritePass : public pir::Pass {
+class DrrPatternRewritePass : public pir::drr::DrrBasePass {
  public:
-  DrrPatternRewritePass() : pir::Pass("DrrPatternRewritePass", 1) {}
-
-  bool Initialize(pir::IrContext *context) override {
-    pir::RewritePatternSet ps(context);
-    ps.Add(RemoveRedundentReshapePattern().Build(context));
-    ps.Add(RemoveRedundentTransposePattern().Build(context));
-    ps.Add(RemoveRedundentCastPattern().Build(context));
-    ps.Add(RemoveUselessCastPattern().Build(context));
-    ps.Add(FoldExpandToConstantPattern().Build(context));
-
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
+  DrrPatternRewritePass() : pir::drr::DrrBasePass("DrrPatternRewritePass", 1) {}
+
+  void DrrPassInitialize(pir::RewritePatternSet *ps,
+                         pir::IrContext *context) override {
+    ps->Add(RemoveRedundentReshapePattern().Build(context));
+    ps->Add(RemoveRedundentTransposePattern().Build(context));
+    ps->Add(RemoveRedundentCastPattern().Build(context));
+    ps->Add(RemoveUselessCastPattern().Build(context));
+    ps->Add(FoldExpandToConstantPattern().Build(context));
   }
-
-  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) {