[PIR Pass] Support hook of value replace in pir pass (#69806)

zyfncg · web-flow · commit 51dfb95eb753 · 2024-12-02T19:14:12.000+08:00
* support hook of value replace for pir pass

* fix bug

* refine code

* fix bug
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -85,6 +85,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_util.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/pir_to_py_code_converter.h"
 #include "paddle/cinn/hlir/framework/pir_compiler.h"
+#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 #endif
 
 using paddle::dialect::ApiBuilder;
@@ -2392,22 +2393,23 @@ void BindUtils(pybind11::module *m) {
 
 namespace {
 
-#ifdef PADDLE_WITH_CINN
-std::shared_ptr<pir::PassManager> CreatePassManager() {
-  pir::IrContext *ctx = pir::IrContext::Instance();
-  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
-  ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
-  ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
-  auto pass_manager = std::make_shared<pir::PassManager>(ctx);
-  if (FLAGS_print_ir) {
-    pass_manager->EnableIRPrinting();
-  }
-  return pass_manager;
-}
-#endif
-
 void ApplyCinnPass(Program &program) {  // NOLINT
 #ifdef PADDLE_WITH_CINN
+  auto CreatePassManager = [&]() -> std::shared_ptr<pir::PassManager> {
+    pir::IrContext *ctx = pir::IrContext::Instance();
+    ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+    ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
+    ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
+    auto pass_manager = std::make_shared<pir::PassManager>(ctx);
+    if (FLAGS_print_ir) {
+      pass_manager->EnableIRPrinting();
+    }
+    auto &shape_analysis = pir::ShapeAnalysisManager::Instance().Get(&program);
+    pass_manager->SetValueReplacedHook([&](pir::Value from, pir::Value to) {
+      shape_analysis.ShareShapeOrData(from, to);
+    });
+    return pass_manager;
+  };
   cinn::dialect::ir::ApplyCinnPass(&program, CreatePassManager);
 #else
   PADDLE_THROW(common::errors::Unimplemented(
@@ -2418,6 +2420,17 @@ void ApplyCinnPass(Program &program) {  // NOLINT
 
 void CheckInferSymbolicIfNeed(Program &program) {  // NOLINT
 #ifdef PADDLE_WITH_CINN
+  auto CreatePassManager = [&]() -> std::shared_ptr<pir::PassManager> {
+    pir::IrContext *ctx = pir::IrContext::Instance();
+    ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+    ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
+    ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
+    auto pass_manager = std::make_shared<pir::PassManager>(ctx);
+    if (FLAGS_print_ir) {
+      pass_manager->EnableIRPrinting();
+    }
+    return pass_manager;
+  };
   cinn::dialect::ir::CheckInferSymbolicIfNeed(&program, CreatePassManager);
 #else
   // Do nothing.
diff --git a/paddle/pir/include/dialect/shape/utils/shape_analysis.h b/paddle/pir/include/dialect/shape/utils/shape_analysis.h
@@ -196,6 +196,9 @@ class IR_API ShapeConstraintIRAnalysis final
   void SetShapeOrDataForValue(Value val,
                               const symbol::ShapeOrDataDimExprs& shape_or_data);
 
+  // Set ShapeOrData of `to` value by ShapeOrData of `from` value.
+  void ShareShapeOrData(Value from, Value to);
+
   bool IsEqual(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs) const;
 
   bool IsGreatThanOne(const symbol::DimExpr& dim_expr) const;
diff --git a/paddle/pir/include/pass/pass.h b/paddle/pir/include/pass/pass.h
@@ -76,6 +76,7 @@ class IR_API Pass {
  public:
   inline static const char kParamScopeAttr[] = "__param_scope__";
   inline static const char kPlaceAttr[] = "__place__";
+  inline static const char kValueReplaceHookAttr[] = "__value_replaced_hook__";
 
   explicit Pass(const std::string& name,
                 uint8_t opt_level,
diff --git a/paddle/pir/include/pass/pass_manager.h b/paddle/pir/include/pass/pass_manager.h
@@ -115,6 +115,10 @@ class IR_API PassManager {
 
   void AddInstrumentation(std::unique_ptr<PassInstrumentation> pi);
 
+  void SetValueReplacedHook(const VALUE_REPLACED_HOOK_FUNC &hook) {
+    value_replaced_hook_ = hook;
+  }
+
  private:
   bool Initialize(IrContext *context);
 
@@ -135,6 +139,8 @@ class IR_API PassManager {
 
   std::unique_ptr<PassInstrumentor> instrumentor_;
 
+  VALUE_REPLACED_HOOK_FUNC value_replaced_hook_ = nullptr;
+
   // For access member of pass_adaptor_.
   friend class detail::PassAdaptor;
 };
diff --git a/paddle/pir/include/pattern_rewrite/pattern_match.h b/paddle/pir/include/pattern_rewrite/pattern_match.h
@@ -302,6 +302,8 @@ class RewriterBase : public Builder {
 
   virtual void NotifyOperationInserted(Operation* op) {}
 
+  virtual void NotifyValueReplaced(Value from, Value to) {}
+
   virtual void StartRootUpdate(Operation* op) {}
 
   virtual void FinalizeRootUpdate(Operation* op) {}
diff --git a/paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h b/paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h
@@ -19,6 +19,8 @@
 
 namespace pir {
 
+using VALUE_REPLACED_HOOK_FUNC = std::function<void(pir::Value, pir::Value)>;
+
 class FrozenRewritePatternSet;
 
 /// This enum will control which ops will be added to the worklist during the
@@ -57,6 +59,9 @@ class IR_API GreedyRewriteConfig {
   /// - ExistingOps: only pre-existing ops are added to the worklist.
   GreedyRewriteStrictness strict_mode = GreedyRewriteStrictness::AnyOp;
 
+  // Hook function for replacing the value.
+  VALUE_REPLACED_HOOK_FUNC value_replaced_hook = nullptr;
+
   static constexpr int64_t kNoLimit = -1;
 };
 
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -613,6 +613,12 @@ void ShapeConstraintIRAnalysis::SetShapeOrDataForValue(
   context_.SetShapeOrDataForValue(val, shape_or_data);
 }
 
+void ShapeConstraintIRAnalysis::ShareShapeOrData(Value from, Value to) {
+  if (context_.HasShapeOrDataForValue(from)) {
+    context_.SetShapeOrDataForValue(to, context_.GetShapeOrDataForValue(from));
+  }
+}
+
 bool ShapeConstraintIRAnalysis::IsEqual(const symbol::DimExpr& lhs,
                                         const symbol::DimExpr& rhs) const {
   return context_.IsEqual(lhs, rhs);
diff --git a/paddle/pir/src/pass/pass.cc b/paddle/pir/src/pass/pass.cc
@@ -89,8 +89,12 @@ GreedyRewriteConfig PatternRewritePass::InitializeConfig() {
 
 void PatternRewritePass::Run(Operation* op) {
   VLOG(4) << "Run PatternRewritePass: " << name();
-  auto [_, num_rewrites] =
-      ApplyPatternsGreedily(op, patterns_, InitializeConfig());
+  GreedyRewriteConfig config = InitializeConfig();
+  if (Has(kValueReplaceHookAttr)) {
+    config.value_replaced_hook =
+        Get<VALUE_REPLACED_HOOK_FUNC>(kValueReplaceHookAttr);
+  }
+  auto [_, num_rewrites] = ApplyPatternsGreedily(op, patterns_, config);
   AddStatistics(num_rewrites);
 }
 
@@ -202,6 +206,10 @@ bool PassManager::Run(Operation* op) {
 bool PassManager::Initialize(IrContext* context) {
   for (auto& pass : passes()) {
     if (!pass->Initialize(context)) return false;
+    if (value_replaced_hook_) {
+      pass->SetNotOwned<VALUE_REPLACED_HOOK_FUNC>(Pass::kValueReplaceHookAttr,
+                                                  &value_replaced_hook_);
+    }
   }
 
   return true;
diff --git a/paddle/pir/src/pattern_rewrite/pattern_match.cc b/paddle/pir/src/pattern_rewrite/pattern_match.cc
@@ -126,6 +126,9 @@ void RewriterBase::ReplaceOp(Operation* op,
       new_values.size(),
       common::errors::InvalidArgument("incorrect # of replacement values"));
   op->ReplaceAllUsesWith(new_values);
+  for (uint32_t i = 0; i < op->num_results(); ++i) {
+    NotifyValueReplaced(op->result(i), new_values[i]);
+  }
 
   NotifyOperationRemoved(op);
   op->Erase();
@@ -144,8 +147,10 @@ void RewriterBase::EraseOp(Operation* op) {
 
 // Find uses of `from` and replace it with `to`.
 void RewriterBase::ReplaceAllUsesWith(Value from, Value to) {
-  for (auto it = from.use_begin(); it != from.use_end();)
+  for (auto it = from.use_begin(); it != from.use_end();) {
     UpdateRootInplace(it.owner(), [&]() { (it++)->set_source(to); });
+  }
+  NotifyValueReplaced(from, to);
 }
 
 // Find uses of `from` and replace them with `to` if the `functor` returns true.
@@ -155,9 +160,15 @@ void RewriterBase::ReplaceUseIf(Value from,
   // Use post-increment operator for iterator since set_source() will change
   // `it`.
   // TODO(zhangbopd): Add unit test for this.
+  bool replaced = false;
   for (auto it = from.use_begin(); it != from.use_end();) {
-    if (functor(*it))
+    if (functor(*it)) {
       UpdateRootInplace(it.owner(), [&]() { (it++)->set_source(to); });
+      replaced = true;
+    }
+  }
+  if (replaced) {
+    NotifyValueReplaced(from, to);
   }
 }
 
diff --git a/paddle/pir/src/pattern_rewrite/pattern_rewrite_driver.cc b/paddle/pir/src/pattern_rewrite/pattern_rewrite_driver.cc
@@ -53,6 +53,9 @@ class GreedyPatternRewriteDriver : public pir::PatternRewriter {
         }
       }
     }
+    if (config.value_replaced_hook) {
+      value_replaced_hook_fn_ = config.value_replaced_hook;
+    }
   }
 
   std::pair<bool, int64_t> Simplify() {
@@ -156,6 +159,12 @@ class GreedyPatternRewriteDriver : public pir::PatternRewriter {
     AddToWorklist(op);
   }
 
+  void NotifyValueReplaced(pir::Value from, pir::Value to) override {
+    if (value_replaced_hook_fn_) {
+      value_replaced_hook_fn_(from, to);
+    }
+  }
+
   /// Add the given operation to the worklist.
   void AddToWorklist(pir::Operation* op) {
     if (config_.strict_mode == pir::GreedyRewriteStrictness::AnyOp ||
@@ -207,6 +216,7 @@ class GreedyPatternRewriteDriver : public pir::PatternRewriter {
   std::unordered_set<pir::Operation*> strict_mode_filtered_ops_;
   pir::Region& region_;
   pir::PatternApplicator matcher_;
+  pir::VALUE_REPLACED_HOOK_FUNC value_replaced_hook_fn_ = nullptr;
 };
 
 }  // namespace

Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,9 @@ class GreedyPatternRewriteDriver : public pir::PatternRewriter {`
`53`	`53`	`}`
`54`	`54`	`}`
`55`	`55`	`}`
	`56`	`+ if (config.value_replaced_hook) {`
	`57`	`+ value_replaced_hook_fn_ = config.value_replaced_hook;`
	`58`	`+ }`
`56`	`59`	`}`
`57`	`60`
`58`	`61`	`std::pair<bool, int64_t> Simplify() {`
`@@ -156,6 +159,12 @@ class GreedyPatternRewriteDriver : public pir::PatternRewriter {`
`156`	`159`	`AddToWorklist(op);`
`157`	`160`	`}`
`158`	`161`
	`162`	`+ void NotifyValueReplaced(pir::Value from, pir::Value to) override {`
	`163`	`+ if (value_replaced_hook_fn_) {`
	`164`	`+ value_replaced_hook_fn_(from, to);`
	`165`	`+ }`
	`166`	`+ }`
	`167`	`+`
`159`	`168`	`/// Add the given operation to the worklist.`
`160`	`169`	`void AddToWorklist(pir::Operation* op) {`
`161`	`170`	`if (config_.strict_mode == pir::GreedyRewriteStrictness::AnyOp \|\|`
`@@ -207,6 +216,7 @@ class GreedyPatternRewriteDriver : public pir::PatternRewriter {`
`207`	`216`	`std::unordered_set<pir::Operation*> strict_mode_filtered_ops_;`
`208`	`217`	`pir::Region& region_;`
`209`	`218`	`pir::PatternApplicator matcher_;`
	`219`	`+ pir::VALUE_REPLACED_HOOK_FUNC value_replaced_hook_fn_ = nullptr;`
`210`	`220`	`};`
`211`	`221`
`212`	`222`	`} // namespace`