[CINN][Backend Pass Update No.10] Update ReplaceCrossThreadReduction pass (#70592)

LittleHeroZZZX · web-flow · commit 5718f746d945 · 2025-01-07T19:47:14.000+08:00
* Update replaceCrossThreadReduction

* Add visit logics for IfThenElse stmt

* Refine test and exception message

* Leverage help function Mutate to refactor CrossThreadReductionReplacer
diff --git a/paddle/cinn/ir/schedule/ir_schedule_util.cc b/paddle/cinn/ir/schedule/ir_schedule_util.cc
@@ -110,6 +110,18 @@ int GetLoopExtent(const Expr& loop) {
   return static_cast<int>(loop.As<ir::For>()->extent.get_constant());
 }
 
+int GetLoopExtent(const ir::stmt::For loop) {
+  PADDLE_ENFORCE_EQ(
+      cinn::common::is_zero(loop->min()),
+      true,
+      ::common::errors::InvalidArgument("For node's min should be zero."));
+  PADDLE_ENFORCE_EQ(loop->extent().is_constant(),
+                    true,
+                    ::common::errors::InvalidArgument(
+                        "For node's extent should be constant."));
+  return static_cast<int>(loop->extent().get_constant());
+}
+
 void SetCudaAxisInfo(ir::LoweredFunc lowered_func) {
   auto CannotProveLT = [](const ir::Expr& lhs, const ir::Expr& rhs) -> bool {
     std::vector<ir::Expr> exprs{rhs, lhs};
diff --git a/paddle/cinn/ir/schedule/ir_schedule_util.h b/paddle/cinn/ir/schedule/ir_schedule_util.h
@@ -65,6 +65,13 @@ Tensor GetReadTensor(const Expr& block, int index);
  */
 int GetLoopExtent(const Expr& loop);
 
+/**
+ * \brief Given a For node, return its extent as int.
+ * @param loop The given For node
+ * @return The extent of For node
+ */
+int GetLoopExtent(const ir::stmt::For loop);
+
 /**
  * \brief Given a vector of Exprs, return whether they contain a var with
  * specific name.
diff --git a/paddle/cinn/optim/replace_cross_thread_reduction.cc b/paddle/cinn/optim/replace_cross_thread_reduction.cc
@@ -26,7 +26,9 @@
 #include "paddle/cinn/ir/ir_mutator.h"
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/ir/schedule/ir_schedule_util.h"
+#include "paddle/cinn/ir/utils/stmt_converter.h"
 #include "paddle/cinn/lang/compute.h"
+#include "paddle/cinn/pass/pass_manager.h"
 
 namespace cinn {
 namespace optim {
@@ -40,27 +42,17 @@ struct BufferCmp {
 };
 
 thread_local std::set<ir::Buffer, BufferCmp> shm_buffer_;
-struct CrossThreadReductionReplacer : public ir::IRMutator<> {
+struct CrossThreadReductionReplacer {
   void operator()(ir::LoweredFunc fn) { Visit(fn.As<ir::_LoweredFunc_>()); }
 
  private:
-  bool CanReplace(const ir::ScheduleBlockRealize* block_realize) {
-    const ir::ScheduleBlock* schedule_block =
-        block_realize->schedule_block.As<ir::ScheduleBlock>();
-
-    PADDLE_ENFORCE_NOT_NULL(
-        schedule_block,
-        ::common::errors::PreconditionNotMet(
-            "The schedule block pointer in CanReplace must not be null."));
-
-    if (block_realize->schedule_block.As<ir::ScheduleBlock>()->name.substr(
-            0, 4) == "root") {
+  bool CanReplace(const ir::stmt::Schedule block) {
+    if (block->name().substr(0, 4) == "root") {
       return false;
     }
 
-    const std::vector<ir::Expr>& iter_values = block_realize->iter_values;
-    const std::vector<ir::Var>& iter_vars = schedule_block->iter_vars;
-    ir::Expr body = schedule_block->body;
+    const std::vector<ir::Expr>& iter_values = block->iter_values();
+    const std::vector<ir::Var>& iter_vars = block->iter_vars();
 
     std::unordered_set<std::string> reduce_var_names;
     for (int i = 0; i < iter_values.size(); ++i) {
@@ -72,23 +64,22 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<> {
             if (x->as_var()) {
               reduce_var_names.insert(x->as_var()->name);
             }
-
             return false;
           });
     }
 
-    auto IsThreadBindOnReduceAxis = [&](const ir::For* for_node) {
-      return reduce_var_names.count(for_node->loop_var->name) > 0 &&
+    auto IsThreadBindOnReduceAxis = [&](const ir::stmt::For& for_node) {
+      return reduce_var_names.count(for_node->loop_var()->name) > 0 &&
              for_node->is_gpu_thread_binded();
     };
 
     std::vector<int> thread_binded_reduce_loop_indices;
     bool is_thread_binded_inner_loop = false;
     for (int i = 0; i < cur_loops_.size(); ++i) {
-      bool is_thread_bind_on_reduce =
-          IsThreadBindOnReduceAxis(cur_loops_[i].As<ir::For>());
-      if (is_thread_bind_on_reduce && ir::GetLoopExtent(cur_loops_[i]) == 1)
+      bool is_thread_bind_on_reduce = IsThreadBindOnReduceAxis(cur_loops_[i]);
+      if (is_thread_bind_on_reduce && ir::GetLoopExtent(cur_loops_[i]) == 1) {
         return false;
+      }
       if (is_thread_binded_inner_loop || is_thread_bind_on_reduce) {
         if (ir::GetLoopExtent(cur_loops_[i]) > 1024) {
           return false;
@@ -115,21 +106,22 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<> {
   int GetBlockSize() const {
     int block_size = 1;
     for (auto& loop : cur_loops_) {
-      if (loop->as<ir::For>()->is_gpu_thread_binded()) {
+      if (loop->is_gpu_thread_binded()) {
         block_size *= ir::GetLoopExtent(loop);
       }
     }
     return block_size;
   }
 
   template <typename OpT>
-  void ReplaceByContinuousReduceExternCall(ir::Expr* store, bool return_warp) {
-    auto* node = store->As<ir::Store>()->value.As<OpT>();
+  void ReplaceByContinuousReduceExternCall(ir::stmt::Store store,
+                                           bool return_warp) {
+    auto* node = store->value().As<OpT>();
     PADDLE_ENFORCE_NOT_NULL(
         node, ::common::errors::InvalidArgument("The node must not be null."));
     auto& operand = node->b();
     std::string reduce_func_name = hlir::pe::CrossThreadReduceExternalFuncName(
-        store->As<ir::Store>()->value, operand.template As<ir::Load>()->tensor);
+        store->value(), operand.template As<ir::Load>()->tensor);
     auto tmp_dtype =
         operand.template As<ir::Load>()->tensor.as_tensor()->type();
     auto tmp_buffer = ir::_Buffer_::Make(
@@ -138,18 +130,18 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<> {
     tmp_buffer->dtype = tmp_dtype;
     tmp_buffer->memory_type = ir::MemoryType::GPUShared;
     shm_buffer_.insert(tmp_buffer);
-    store->As<ir::Store>()->value = lang::CallExtern(
-        reduce_func_name, {node->b(), tmp_buffer, ir::Expr(return_warp)});
+    store->set_value(lang::CallExtern(
+        reduce_func_name, {node->b(), tmp_buffer, ir::Expr(return_warp)}));
   }
 
   template <typename OpT>
-  void ReplaceByDiscreteReduceExternCall(ir::Expr* store) {
-    auto* node = store->As<ir::Store>()->value.As<OpT>();
+  void ReplaceByDiscreteReduceExternCall(ir::stmt::Store store) {
+    auto* node = store->value().As<OpT>();
     PADDLE_ENFORCE_NOT_NULL(
         node, ::common::errors::InvalidArgument("The node must not be null."));
     auto& operand = node->b();
     std::string reduce_func_name = hlir::pe::DiscreteReduceExternalFuncName(
-        store->As<ir::Store>()->value, operand.template As<ir::Load>()->tensor);
+        store->value(), operand.template As<ir::Load>()->tensor);
     auto tmp_dtype =
         operand.template As<ir::Load>()->tensor.as_tensor()->type();
     auto tmp_buffer = ir::_Buffer_::Make(
@@ -158,12 +150,12 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<> {
     tmp_buffer->dtype = tmp_dtype;
     tmp_buffer->memory_type = ir::MemoryType::GPUShared;
     shm_buffer_.insert(tmp_buffer);
-    store->As<ir::Store>()->value =
-        lang::CallExtern(reduce_func_name, {node->b(), tmp_buffer});
+    store->set_value(
+        lang::CallExtern(reduce_func_name, {node->b(), tmp_buffer}));
   }
 
   template <typename OpT>
-  void ReplaceByReduceExternCall(ir::Expr* store,
+  void ReplaceByReduceExternCall(ir::stmt::Store store,
                                  const ir::ReduceMethod& method) {
     std::visit(cinn::adt::match{
                    [&](const ir::NoneReduceMethod&) {
@@ -181,10 +173,11 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<> {
                method);
   }
 
-  void Visit(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
-
-  void Visit(ir::_LoweredFunc_* fn) override {
-    ir::IRMutator<>::Visit(fn);
+  void Visit(ir::_LoweredFunc_* fn) {
+    ir::stmt::Mutate(
+        fn->body_block,
+        [&](ir::stmt::StmtRef stmt) { PreCall(stmt); },
+        [&](ir::stmt::StmtRef stmt) { PostCall(stmt); });
     if (std::find_if(fn->temp_bufs.begin(),
                      fn->temp_bufs.end(),
                      [&](const ir::Buffer& buf) -> bool {
@@ -198,74 +191,98 @@ struct CrossThreadReductionReplacer : public ir::IRMutator<> {
     shm_buffer_.clear();
   }
 
-  void Visit(const ir::ScheduleBlockRealize* expr, ir::Expr* op) override {
-    if (!CanReplace(expr)) {
-      VLOG(6) << "Can't replace cross thread reduction: " << *op;
-      IRMutator::Visit(expr, op);
-      return;
+  void PreCall(ir::stmt::StmtRef stmt) {
+    switch (stmt->stmt_type()) {
+      case ir::StmtNodeTy::Schedule:
+        VisitStmt(stmt.as<ir::stmt::Schedule>());
+        break;
+      case ir::StmtNodeTy::For:
+        cur_loops_.push_back(stmt.as<ir::stmt::For>());
+        break;
+      default:
+        break;
     }
-    VLOG(6) << "Can replace cross thread reduction: " << *op;
+  }
 
-    const ir::ScheduleBlock* schedule_block =
-        expr->schedule_block.As<ir::ScheduleBlock>();
-    PADDLE_ENFORCE_NOT_NULL(
-        schedule_block,
-        ::common::errors::PreconditionNotMet(
-            "The schedule block pointer in Visit must not be null."));
-    ir::Expr original_update_body = schedule_block->body;
-    ir::Expr original_update_stmt;
-    PADDLE_ENFORCE_EQ(original_update_body.As<ir::Block>() ||
-                          original_update_body.As<ir::Store>(),
-                      true,
-                      ::common::errors::InvalidArgument(
-                          "The type of original_update_body is incorrect."
-                          "Expected type is Block or Store."));
-    if (original_update_body.As<ir::Block>()) {
-      PADDLE_ENFORCE_EQ(
-          original_update_body.As<ir::Block>()->stmts.size(),
-          1,
-          ::common::errors::InvalidArgument(
-              "The size of stmts is incorrect."
-              "Expected size is 1, but receive %d.",
-              original_update_body.As<ir::Block>()->stmts.size()));
-      original_update_stmt = original_update_body.As<ir::Block>()->stmts[0];
-    } else if (original_update_body.As<ir::Store>()) {
-      original_update_stmt = original_update_body;
+  void PostCall(ir::stmt::StmtRef stmt) {
+    switch (stmt->stmt_type()) {
+      case ir::StmtNodeTy::For:
+        cur_loops_.pop_back();
+        break;
+      default:
+        break;
     }
-
-#define REPLACE_TO_EXTERNAL_CALL(Op)                              \
-  if (original_update_stmt.As<ir::Store>()->value.As<Op>()) {     \
-    ReplaceByReduceExternCall<Op>(&original_update_stmt,          \
-                                  schedule_block->reduce_method); \
   }
 
-    REPLACE_TO_EXTERNAL_CALL(ir::Add)
-    REPLACE_TO_EXTERNAL_CALL(ir::Mul)
-    REPLACE_TO_EXTERNAL_CALL(ir::Max)
-    REPLACE_TO_EXTERNAL_CALL(ir::Min)
-    REPLACE_TO_EXTERNAL_CALL(ir::And)
-    REPLACE_TO_EXTERNAL_CALL(ir::Or)
-#undef REPLACE_TO_EXTERNAL_CALL
-
-    VLOG(6) << "Replace cross thread reduction: " << *op;
-
-    IRMutator::Visit(expr, op);
-  }
+  void VisitStmt(ir::stmt::Schedule stmt) {
+    if (!CanReplace(stmt)) {
+      return;
+    }
+    ir::stmt::BlockRef original_update_body = stmt->body();
 
-  void Visit(const ir::For* expr, ir::Expr* op) override {
-    cur_loops_.push_back(*op);
-    IRMutator::Visit(expr, op);
-    cur_loops_.pop_back();
+    ir::stmt::Store original_update_stmt;
+    PADDLE_ENFORCE_EQ(original_update_body->stmts().size(),
+                      1,
+                      ::common::errors::InvalidArgument(
+                          "The size of statements is incorrect."
+                          "Expected size is 1, but receive %d.",
+                          original_update_body->stmts().size()));
+    PADDLE_ENFORCE_EQ(original_update_body->stmts()[0].isa<ir::stmt::Store>(),
+                      true,
+                      ::common::errors::InvalidArgument(
+                          "The stmt in schedule's body should be store "
+                          "statement, but get %s.",
+                          original_update_body->stmts()[0]->stmt_type()));
+    original_update_stmt =
+        original_update_body->stmts()[0].as<ir::stmt::Store>();
+
+    switch (original_update_stmt->value()->node_type()) {
+      case cinn::ir::IrNodeTy::Add:
+        ReplaceByReduceExternCall<ir::Add>(original_update_stmt,
+                                           stmt->reduce_method());
+        break;
+      case cinn::ir::IrNodeTy::Mul:
+        ReplaceByReduceExternCall<ir::Mul>(original_update_stmt,
+                                           stmt->reduce_method());
+        break;
+      case cinn::ir::IrNodeTy::Max:
+        ReplaceByReduceExternCall<ir::Max>(original_update_stmt,
+                                           stmt->reduce_method());
+        break;
+      case cinn::ir::IrNodeTy::Min:
+        ReplaceByReduceExternCall<ir::Min>(original_update_stmt,
+                                           stmt->reduce_method());
+        break;
+      case cinn::ir::IrNodeTy::And:
+        ReplaceByReduceExternCall<ir::And>(original_update_stmt,
+                                           stmt->reduce_method());
+        break;
+      case cinn::ir::IrNodeTy::Or:
+        ReplaceByReduceExternCall<ir::Or>(original_update_stmt,
+                                          stmt->reduce_method());
+        break;
+      default:
+        PADDLE_THROW(::common::errors::InvalidArgument(
+            "The node type is not supported in cross thread reduction."));
+    }
   }
 
  private:
-  std::vector<ir::Expr> cur_loops_;
+  std::vector<ir::stmt::For> cur_loops_;
 };
 
 }  // namespace
 
 void ReplaceCrossThreadReduction(ir::LoweredFunc fn) {
-  CrossThreadReductionReplacer()(fn);
+  FuncPassManager manager;
+  manager.AddPass(std::make_unique<ReplaceCrossThreadReductionPass>());
+  manager.Run(fn);
+}
+
+LogicalResult ReplaceCrossThreadReductionPass::Run(ir::LoweredFunc func) {
+  CrossThreadReductionReplacer replacer;
+  replacer(func);
+  return LogicalResult::success();
 }
 
 }  // namespace optim
diff --git a/paddle/cinn/optim/replace_cross_thread_reduction.h b/paddle/cinn/optim/replace_cross_thread_reduction.h
@@ -16,14 +16,21 @@
  * This file implements the strategy to remove the unnecessary nested block.
  */
 #pragma once
-#include <vector>
 
-#include "paddle/cinn/common/common.h"
 #include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/pass/pass.h"
 
 namespace cinn {
 namespace optim {
 
+class ReplaceCrossThreadReductionPass : public FuncPass {
+ public:
+  ReplaceCrossThreadReductionPass()
+      : FuncPass("replace_cross_thread_reduction") {}
+
+  LogicalResult Run(ir::LoweredFunc func) override;
+};
+
 /**
  * Replace cross thread reduction to external call.
  */
diff --git a/paddle/cinn/optim/replace_cross_thread_reduction_test.cc b/paddle/cinn/optim/replace_cross_thread_reduction_test.cc
