[PIR][CUSTOM PASS] Part1: Reconstruct drr

paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.cc

@@ -20,7 +20,6 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
-#include "paddle/fluid/pir/drr/api/match_context.h"
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/pass/pass.h"
 #include "paddle/pir/pattern_rewrite/pattern_applicator.h"

paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc

@@ -25,7 +25,6 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
-#include "paddle/fluid/pir/drr/api/match_context.h"
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/dialect/shape/utils/shape_utils.h"
 #include "paddle/pir/pass/pass.h"

paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.cc

@@ -21,7 +21,6 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
-#include "paddle/fluid/pir/drr/api/match_context.h"
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/pass/pass.h"
 #include "paddle/pir/pattern_rewrite/pattern_applicator.h"

paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc

@@ -19,8 +19,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
-#include "paddle/fluid/pir/drr/api/drr_pattern_base.h"
-#include "paddle/fluid/pir/drr/api/match_context.h"
+#include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/pass/pass.h"
@@ -31,7 +30,7 @@ namespace cinn {
 namespace dialect {
 namespace ir {
 
-class SumOpPattern : public paddle::drr::DrrPatternBase<SumOpPattern> {
+class SumOpPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // Source Pattern
@@ -55,9 +54,11 @@ class SumOpPattern : public paddle::drr::DrrPatternBase<SumOpPattern> {
                 {"keep_dim", pattern.Attr("keep_dim")}});
     res.Tensor("ret") = cinn_reduce_sum(res.Tensor("arg0"));
   }
+
+  std::string name() const override { return "SumOpPattern"; }
 };
 
-class MaxOpPattern : public paddle::drr::DrrPatternBase<MaxOpPattern> {
+class MaxOpPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // Source Pattern
@@ -80,9 +81,11 @@ class MaxOpPattern : public paddle::drr::DrrPatternBase<MaxOpPattern> {
                 {"keep_dim", pattern.Attr("keep_dim")}});
     res.Tensor("ret") = cinn_reduce_max(res.Tensor("arg0"));
   }
+
+  std::string name() const override { return "MaxOpPattern"; }
 };
 
-class MinOpPattern : public paddle::drr::DrrPatternBase<MinOpPattern> {
+class MinOpPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // Source Pattern
@@ -105,9 +108,11 @@ class MinOpPattern : public paddle::drr::DrrPatternBase<MinOpPattern> {
                 {"keep_dim", pattern.Attr("keep_dim")}});
     res.Tensor("ret") = cinn_reduce_max(res.Tensor("arg0"));
   }
+
+  std::string name() const override { return "MinOpPattern"; }
 };
 
-class ProdOpPattern : public paddle::drr::DrrPatternBase<ProdOpPattern> {
+class ProdOpPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // Source Pattern
@@ -130,6 +135,8 @@ class ProdOpPattern : public paddle::drr::DrrPatternBase<ProdOpPattern> {
                 {"keep_dim", pattern.Attr("keep_dim")}});
     res.Tensor("ret") = cinn_reduce_max(res.Tensor("arg0"));
   }
+
+  std::string name() const override { return "ProdOpPattern"; }
 };
 
 class ScaleOpPattern : public pir::OpRewritePattern<paddle::dialect::ScaleOp> {
@@ -586,7 +593,7 @@ class ExpandOpPattern
   }
 };
 
-class UniformOpPattern : public paddle::drr::DrrPatternBase<UniformOpPattern> {
+class UniformOpPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // Source Pattern
@@ -632,6 +639,8 @@ class UniformOpPattern : public paddle::drr::DrrPatternBase<UniformOpPattern> {
                 {"diag_val", pattern.Attr("min_value")}});
     res.Tensor("ret") = cinn_uniform();
   }
+
+  std::string name() const override { return "ProdOpPattern"; }
 };
 
 PdOpToCinnOpPass::PdOpToCinnOpPass()

paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc

@@ -24,7 +24,6 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
-#include "paddle/fluid/pir/drr/api/match_context.h"
 #include "paddle/pir/core/builtin_dialect.h"
 #include "paddle/pir/dialect/shape/utils/dim_expr.h"
 #include "paddle/pir/pass/pass.h"

paddle/fluid/pir/drr/CMakeLists.txt

@@ -1,4 +1,4 @@
-file(GLOB DRR_SRCS "*.cc" "api/*.cc")
+file(GLOB DRR_SRCS "*.cc" "include/*.cc")
 
 set(op_creator_gen_file
     ${PADDLE_SOURCE_DIR}/paddle/fluid/pir/dialect/op_generator/op_creator_drr_gen.py

paddle/fluid/pir/drr/README.md

@@ -8,9 +8,8 @@ DRR can reduce the development cost of PASS, allowing developers to focus on pro
 
 Taking PASS to eliminate redundant CastOp as an example, the code example developed using DRR is as follows:
 ~~~ c++
-// 1. Inherit specialized template class from DrPatternBase
-class RemoveRedundentCastPattern
-    : public paddle::drr::DrrPatternBase<RemoveRedundentCastPattern> {
+// 1. Inherit class from DrPatternBase
+class RemoveRedundentCastPattern : public paddle::drr::DrrPatternBase {
   // 2. Overload operator()
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // 3. Define a SourcePattern containing two consecutive CastOps using Op, Tensor, and Attribute
@@ -32,6 +31,8 @@ class RemoveRedundentCastPattern
         res.Op(paddle::dialect::CastOp::name(),
                {{"dtype", pat.Attr("dtype2")}})(res.Tensor("arg0"));
   }
+
+  std::string name() const override { return "RemoveRedundentCastPattern"; }
 };
 ~~~
 
@@ -165,7 +166,7 @@ Attribute Attr(const AttrComputeFunc& attr_compute_func) const</pre></td>
 ## 3 Example
 Example 1: Matmul + Add -> FusedGemmEpilogue
 ~~~ c++
-class FusedLinearPattern : public paddle::drr::DrrPatternBase<FusedLinearPattern> {
+class FusedLinearPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
 	// Define SourcePattern
@@ -193,13 +194,14 @@ class FusedLinearPattern : public paddle::drr::DrrPatternBase<FusedLinearPattern
         {&res.Tensor("x"), &res.Tensor("w"), &res.Tensor("bias")},
         {&res.Tensor("out")});
   }
+
+  std::string name() const override { return "FusedLinearPattern"; }
 };
 ~~~
 
 Example 2: Full + Expand -> Full
 ~~~ c++
-class FoldExpandToConstantPattern
-    : public paddle::drr::DrrPatternBase<FoldExpandToConstantPattern> {
+class FoldExpandToConstantPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // Define SourcePattern
@@ -226,5 +228,7 @@ class FoldExpandToConstantPattern
                                 {"place", pat.Attr("place_1")}});
     res.Tensor("ret") = full2();
   }
+
+  std::string name() const override { return "FoldExpandToConstantPattern"; }
 };
 ~~~

paddle/fluid/pir/drr/README_cn.md

@@ -8,9 +8,8 @@ DRR ( Declarative Rewrite Rule ) 是来处理这种 DAG-to-DAG 类型的一套 P
 
 以消除冗余 CastOp 的 PASS 为例，使用 DRR 的代码开发示例如下：
 ~~~ c++
-// 1. 继承 DrrPatternBase 的特化模板类
-class RemoveRedundentCastPattern
-    : public paddle::drr::DrrPatternBase<RemoveRedundentCastPattern> {
+// 1. 继承 DrrPatternBase 类
+class RemoveRedundentCastPattern : public paddle::drr::DrrPatternBase {
   // 2. 重载 operator()
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // 3. 使用 Op、Tensor 和 Attribute 定义一个包含两个连续 CastOp 的 SourcePattern
@@ -32,6 +31,8 @@ class RemoveRedundentCastPattern
         res.Op(paddle::dialect::CastOp::name(),
                {{"dtype", pat.Attr("dtype2")}})(res.Tensor("arg0"));
   }
+
+  std::string name() const override { return "RemoveRedundentCastPattern"; }
 };
 ~~~
 
@@ -168,7 +169,7 @@ Attribute Attr(const AttrComputeFunc& attr_compute_func) const</pre></td>
 ## 3 使用示例
 Example 1: Matmul + Add -> FusedGemmEpilogue
 ~~~ c++
-class FusedLinearPattern : public paddle::drr::DrrPatternBase<FusedLinearPattern> {
+class FusedLinearPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // 定义 Source Pattern
@@ -196,13 +197,14 @@ class FusedLinearPattern : public paddle::drr::DrrPatternBase<FusedLinearPattern
         {&res.Tensor("x"), &res.Tensor("w"), &res.Tensor("bias")},
         {&res.Tensor("out")});
   }
+
+  std::string name() const override { return "FusedLinearPattern"; }
 };
 ~~~
 
 Example 2: Full + Expand -> Full
 ~~~ c++
-class FoldExpandToConstantPattern
-    : public paddle::drr::DrrPatternBase<FoldExpandToConstantPattern> {
+class FoldExpandToConstantPattern : public paddle::drr::DrrPatternBase {
  public:
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     // 定义 Source Pattern
@@ -229,5 +231,7 @@ class FoldExpandToConstantPattern
                                 {"place", pat.Attr("place_1")}});
     res.Tensor("ret") = full2();
   }
+
+  std::string name() const override { return "FoldExpandToConstantPattern"; }
 };
 ~~~

paddle/fluid/pir/drr/api/match_context.h → paddle/fluid/pir/drr/include/drr_match_context.h

@@ -17,8 +17,9 @@
 #include <memory>
 #include <string>
 
-#include "paddle/fluid/pir/drr/api/tensor_interface.h"
-#include "paddle/fluid/pir/drr/ir_operation.h"
+namespace pir {
+class Value;
+}
 
 namespace paddle {
 namespace drr {
@@ -30,7 +31,7 @@ class MatchContext final {
  public:
   MatchContext(std::shared_ptr<const MatchContextImpl> impl);
 
-  const TensorInterface& Tensor(const std::string& tensor_name) const;
+  const pir::Value& Tensor(const std::string& tensor_name) const;
 
   template <typename T>
   T Attr(const std::string& attr_name) const;