[CINN][Backend Pass Update No.13] Update UpdateBufferAxis pass (#70533)

Albresky · web-flow · commit c5b95fed322d · 2025-01-12T12:36:31.000+08:00
* Update pass for update_buffer_axis

* Update codestyle

* Update update_buffer_axis_pass

* Traverse value expr of Store stmt; reformat code

* Fix Store traversal
diff --git a/paddle/cinn/optim/transform_gpu_forloop.cc b/paddle/cinn/optim/transform_gpu_forloop.cc
@@ -470,7 +470,12 @@ void OptimizeExprGPU(Expr *expr) {
   replace_index_to_bind_expr(expr);
 
   // resize buffer axis
-  UpdateBufferAxisPass(expr);
+  BlockPassManager pass_manager;
+  ir::stmt::BlockRef _block = ir::ConvertExprBlockToStmtBlock(*expr);
+  pass_manager.AddPass(optim::CreateUpdateBufferAxisPass());
+  pass_manager.Run(_block);
+  ir::Expr new_expr = ir::ConvertStmtBlockToExprBlock(_block);
+  *expr = new_expr;
 
   // Replace variables bound on block/thread to the actual blockIdx/threadIdx.
   ReplaceLoopVarToGpu replace_loop_var_to_gpu;
diff --git a/paddle/cinn/optim/update_buffer_axis_pass.cc b/paddle/cinn/optim/update_buffer_axis_pass.cc
@@ -20,13 +20,23 @@
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_mutator.h"
 #include "paddle/cinn/ir/ir_printer.h"
+#include "paddle/cinn/ir/stmt_visitors.h"
 #include "paddle/cinn/ir/utils/ir_copy.h"
 #include "paddle/cinn/ir/utils/ir_replace.h"
 #include "paddle/cinn/optim/replace_var_with_expr.h"
 #include "paddle/cinn/utils/string.h"
 
 namespace cinn {
 namespace optim {
+using ir::stmt::Alloc;
+using ir::stmt::BlockRef;
+using ir::stmt::Evaluate;
+using ir::stmt::For;
+using ir::stmt::Free;
+using ir::stmt::IfThenElse;
+using ir::stmt::Let;
+using ir::stmt::Schedule;
+using ir::stmt::Store;
 
 void FormalizeSingleIndex(const ir::Tensor& tensor,
                           std::vector<ir::Expr>* indices) {
@@ -43,34 +53,75 @@ void FormalizeSingleIndex(const ir::Tensor& tensor,
   }
 }
 
-class AnalyzeBufferAxis : public ir::IRMutator<> {
+class AnalyzeBufferAxis : public ir::IRMutator<>,
+                          public ir::stmt::StmtMutator<> {
  public:
   void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
 
-  void Visit(const ir::For* op, Expr* expr) override {
-    if (op->is_gpu_thread_binded()) {
-      var_bind_threads.insert(op->loop_var->name);
-      IRMutator::Visit(op, expr);
-      var_bind_threads.erase(op->loop_var->name);
+  void operator()(BlockRef block) {
+    ir::stmt::StmtMutator<>::VisitBlock(block);
+  }
+
+ private:
+  void VisitStmt(For stmt) override {
+    if (stmt->is_gpu_block_binded()) {
+      var_bind_threads.insert(stmt->loop_var()->name);
+      VisitBlock(stmt->body());
+      var_bind_threads.erase(stmt->loop_var()->name);
       return;
     }
-    IRMutator::Visit(op, expr);
+    VisitBlock(stmt->body());
   }
 
   // Analyze the buffer access inside store
-  void Visit(const ir::Store* op, Expr* expr) override {
-    ir::Store* store = expr->As<ir::Store>();
-    ir::Tensor tensor = store->tensor.as_tensor_ref();
+  void VisitStmt(Store stmt) override {
+    const ir::Tensor& tensor = stmt->tensor().as_tensor_ref();
+
     if (!tensor->buffer.defined() ||
         tensor->buffer->memory_type == ir::MemoryType::Heap) {
-      ir::IRMutator<>::Visit(op, expr);
+      ir::Expr value = stmt->value();
+      ir::IRMutator<>::Visit(&value, &value);
+      stmt->set_value(value);
       return;
     }
-    FormalizeSingleIndex(tensor, &(store->indices));
-    AnalyzeTensorAxis(store->indices, tensor);
-    ir::IRMutator<>::Visit(op, expr);
+
+    std::vector<ir::Expr> indices = stmt->indices();
+    FormalizeSingleIndex(tensor, &indices);
+    stmt->set_indices(indices);
+    AnalyzeTensorAxis(indices, tensor);
+    ir::Expr value = stmt->value();
+    ir::IRMutator<>::Visit(&value, &value);
+    stmt->set_value(value);
+  }
+
+  void VisitStmt(Schedule stmt) override {
+    const std::vector<ir::Var>& iter_vars = stmt->iter_vars();
+    const std::vector<ir::Expr>& iter_values = stmt->iter_values();
+    for (int i = 0; i < iter_vars.size(); ++i) {
+      iter_var_to_bind_expr_[iter_vars[i]->name] = iter_values[i];
+    }
+    VisitBlock(stmt->body());
   }
 
+  void VisitStmt(IfThenElse stmt) override {
+    VisitBlock(stmt->true_case());
+    if (stmt->false_case().defined()) {
+      VisitBlock(stmt->false_case());
+    }
+  }
+
+  void VisitStmt(Let stmt) override {
+    ir::Expr expr = stmt->body();
+    ir::IRMutator<>::Visit(&expr, &expr);
+    stmt->set_body(expr);
+  }
+
+  void VisitStmt(Alloc) override {}
+
+  void VisitStmt(Evaluate) override {}
+
+  void VisitStmt(Free) override {}
+
   // Analyze the buffer access inside load
   void Visit(const ir::Load* op, Expr* expr) override {
     ir::Load* load = expr->As<ir::Load>();
@@ -85,18 +136,6 @@ class AnalyzeBufferAxis : public ir::IRMutator<> {
     ir::IRMutator<>::Visit(op, expr);
   }
 
-  void Visit(const ir::ScheduleBlockRealize* x, Expr* expr) override {
-    const ir::ScheduleBlock* schedule_block =
-        x->schedule_block.As<ir::ScheduleBlock>();
-    const std::vector<ir::Var>& iter_vars = schedule_block->iter_vars;
-    const std::vector<ir::Expr>& iter_values = x->iter_values;
-    for (int i = 0; i < iter_vars.size(); ++i) {
-      iter_var_to_bind_expr_[iter_vars[i]->name] = iter_values[i];
-    }
-    ir::IRMutator<>::Visit(x, expr);
-  }
-
- private:
   void AnalyzeTensorAxis(const std::vector<Expr>& indices,
                          const ir::Tensor& tensor) {
     if (!tensor->buffer.defined() ||
@@ -164,7 +203,8 @@ class AnalyzeBufferAxis : public ir::IRMutator<> {
   std::unordered_set<std::string> var_bind_threads;
 };
 
-class ReplaceSameAxisToZero : public ir::IRMutator<> {
+class ReplaceSameAxisToZero : public ir::IRMutator<>,
+                              public ir::stmt::StmtMutator<> {
  public:
   ReplaceSameAxisToZero(
       const std::unordered_map<std::string, std::map<int, ir::Expr>>&
@@ -174,14 +214,42 @@ class ReplaceSameAxisToZero : public ir::IRMutator<> {
 
   void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
 
+  void operator()(BlockRef block) {
+    ir::stmt::StmtMutator<>::VisitBlock(block);
+  }
+
+ private:
   // Analyze the buffer access inside store
-  void Visit(const ir::Store* op, Expr* expr) override {
-    ir::Store* store = expr->As<ir::Store>();
-    ir::Tensor tensor = store->tensor.as_tensor_ref();
-    ReplaceIndices(tensor, &(store->indices));
-    ir::IRMutator<>::Visit(op, expr);
+  void VisitStmt(Store stmt) override {
+    ir::Tensor tensor = stmt->tensor().as_tensor_ref();
+    std::vector<Expr> expr = stmt->indices();
+    ReplaceIndices(tensor, &expr);
+    stmt->set_indices(expr);
+  }
+
+  void VisitStmt(IfThenElse stmt) override {
+    VisitBlock(stmt->true_case());
+    if (stmt->false_case().defined()) {
+      VisitBlock(stmt->false_case());
+    }
+  }
+
+  void VisitStmt(Let stmt) override {
+    ir::Expr expr = stmt->body();
+    ir::IRMutator<>::Visit(&expr, &expr);
+    stmt->set_body(expr);
   }
 
+  void VisitStmt(For stmt) override { VisitBlock(stmt->body()); }
+
+  void VisitStmt(Schedule stmt) override { VisitBlock(stmt->body()); }
+
+  void VisitStmt(Alloc) override {}
+
+  void VisitStmt(Evaluate) override {}
+
+  void VisitStmt(Free) override {}
+
   // Analyze the buffer access inside load
   void Visit(const ir::Load* op, Expr* expr) override {
     ir::Load* load = expr->As<ir::Load>();
@@ -190,15 +258,15 @@ class ReplaceSameAxisToZero : public ir::IRMutator<> {
     ir::IRMutator<>::Visit(op, expr);
   }
 
- private:
   void ReplaceIndices(const ir::Tensor& tensor, std::vector<Expr>* indices) {
     if (!tensor->buffer.defined() ||
         tensor->buffer->memory_type == ir::MemoryType::Heap) {
       return;
     }
     const std::string& buffer_name = tensor->buffer->name;
     if (buffer_name_access_same_index_expr_.count(buffer_name)) {
-      for (auto p : buffer_name_access_same_index_expr_.at(buffer_name)) {
+      for (const auto& p :
+           buffer_name_access_same_index_expr_.at(buffer_name)) {
         int r = p.first;
         // After optimization, some load indice may be removed, so we need this
         // condition
@@ -215,22 +283,32 @@ class ReplaceSameAxisToZero : public ir::IRMutator<> {
       buffer_name_access_same_index_expr_;
 };
 
-void UpdateBufferAxisPass(ir::Expr* expr) {
-  VLOG(6) << "Before UpdateBufferAxisPass, Expr = \n" << *expr;
+void UpdateBufferAxis(BlockRef block) {
+  VLOG(6) << "Before UpdateBufferAxisPass, Block = \n" << block;
 
   AnalyzeBufferAxis buffer_axis_analyzer;
-  buffer_axis_analyzer(expr);
-  for (auto p : buffer_axis_analyzer.buffer_name_access_same_index_expr) {
+  buffer_axis_analyzer(block);
+  for (const auto& p :
+       buffer_axis_analyzer.buffer_name_access_same_index_expr) {
     VLOG(6) << "Buffer name: " << p.first;
-    for (auto q : p.second) {
-      VLOG(6) << "  Index: " << q.first << " Expr: " << q.second;
+    for (const auto& q : p.second) {
+      VLOG(6) << "Index: " << q.first << " Expr: " << q.second;
     }
   }
 
   ReplaceSameAxisToZero replacer(
       buffer_axis_analyzer.buffer_name_access_same_index_expr);
-  replacer(expr);
-  VLOG(6) << "After UpdateBufferAxisPass, Expr = \n" << *expr;
+  replacer(block);
+  VLOG(6) << "After UpdateBufferAxisPass, Block = \n" << block;
+}
+
+LogicalResult UpdateBufferAxisPass::Run(BlockRef block) {
+  UpdateBufferAxis(block);
+  return LogicalResult::success();
+}
+
+std::unique_ptr<BlockPass> CreateUpdateBufferAxisPass() {
+  return std::make_unique<UpdateBufferAxisPass>();
 }
 
 }  // namespace optim
diff --git a/paddle/cinn/optim/update_buffer_axis_pass.h b/paddle/cinn/optim/update_buffer_axis_pass.h
@@ -16,6 +16,7 @@
 #include <string>
 
 #include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/pass/pass.h"
 
 namespace cinn {
 namespace optim {
@@ -58,7 +59,15 @@ namespace optim {
  *      B[i * n + 0] = ...
  *      ... = B[k * n + 0]
  */
-void UpdateBufferAxisPass(ir::Expr* expr);
+
+class UpdateBufferAxisPass : public BlockPass {
+ public:
+  UpdateBufferAxisPass() : BlockPass("update_buffer_axis_pass") {}
+
+  LogicalResult Run(ir::stmt::BlockRef) override;
+};
+
+std::unique_ptr<BlockPass> CreateUpdateBufferAxisPass();
 
 }  // namespace optim
 }  // namespace cinn
