[CINN] Fix the bug in the preloaded tensor which has a write dependency. (#72365)

ZhangX-21 · web-flow · commit 9a05216dfa0c · 2025-05-07T22:34:00.000+08:00
diff --git a/paddle/cinn/ir/group_schedule/config/group_tile_config.cc b/paddle/cinn/ir/group_schedule/config/group_tile_config.cc
@@ -308,7 +308,7 @@ inline bool CheckThreadDimensionCanVectorize(int threads,
                                              int nums,
                                              int factor) {
   const int deal_elements_in_warp = threads * factor;
-  if (nums % deal_elements_in_warp == 0) {
+  if (nums == deal_elements_in_warp) {
     return true;
   }
   return false;
diff --git a/paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc b/paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc
@@ -372,7 +372,8 @@ bool ContainsVectorizableAxis(const ir::IRSchedule* sch,
 
   VLOG(4) << "Checking ContainsVectorizableAxis on block: [" << block_id
           << "], loop:\n"
-          << sch->GetModule().GetExprs().front() << "\n vectorize expr:\n"
+          << sch->GetModule().GetExprs().front() << "\n"
+          << "vectorize expr:\n"
           << vectorize_expr;
 
   // Get all the lter values in the axis bind that contain a loop var and the
@@ -416,23 +417,18 @@ void ReduceRegionWithReduceBlockVectorizeTilingSchedule(
     const std::string& block_id,
     const int rd_thread,
     const int vectorize_factor) {
-  int loop_axis = 1;
-  int threads_axis = 2;
-  int vectorize_axis = 3;
+  int threads_axis = 1;
+  int vectorize_axis = 2;
   auto loops = sch->GetLoops(block_id);
   if (ContainsVectorizableAxis(sch, loops.size() - 1, block_id)) {
-    sch->Split(loops[1], {-1, rd_thread, vectorize_factor});
+    sch->Split(loops[1], {rd_thread, vectorize_factor});
     loops = sch->GetLoops(block_id);
     sch->Vectorize(loops[vectorize_axis], vectorize_factor);
   } else {
     sch->Split(loops[1], {-1, rd_thread});
     loops = sch->GetLoops(block_id);
   }
 
-  loops = sch->GetLoops(block_id);
-  sch->Reorder({loops[threads_axis], loops[loop_axis]});
-  threads_axis = 1;
-  loops = sch->GetLoops(block_id);
   if (IsReductionSBlock(sch->GetBlock(block_id))) {
     ir::Expr rf_tensor =
         sch->FactorizeReduction(loops[threads_axis],
@@ -459,7 +455,7 @@ void ReduceRegionWithSpatialBlockVectorizeTilingSchedule(
     const int vectorize_factor) {
   auto loops = sch->GetLoops(block_id);
   if (ContainsVectorizableAxis(sch, loops.size() - 1, block_id)) {
-    sch->Split(loops[1], std::vector<int>{-1, rd_thread, vectorize_factor});
+    sch->Split(loops[1], std::vector<int>{rd_thread, vectorize_factor});
 
     // set vectorize schedule primitives
     loops = sch->GetLoops(block_id);
diff --git a/paddle/cinn/optim/vectorize_for_trans.cc b/paddle/cinn/optim/vectorize_for_trans.cc
@@ -477,15 +477,24 @@ class VectorizeForTransMutator : public ir::IRMutator<ir::Expr *> {
   void Visit(const ir::Load *op, ir::Expr *expr) override {
     auto *node = expr->As<ir::Load>();
     auto *tensor = node->tensor.As<ir::_Tensor_>();
+    PADDLE_ENFORCE_NOT_NULL(
+        tensor,
+        ::common::errors::InvalidArgument(
+            "Expected _Tensor_ node in Load, but received nullptr."));
+
     if (in_vectorize_ && node->is_addr_tensor() &&
-        scalar_tensor_without_vectorize_axis_.count(tensor->name)) {
-      PreLoadScalarTensorWithoutVectorizeAxis(node, &node->indices, expr);
+        tensor_can_vectorized_.count(tensor->name)) {
+      TensorVectorized(node, &node->indices, false);
+      return;
+    }
+
+    if (in_vectorize_ && schedule_block_write_dependency_.count(tensor->name)) {
       return;
     }
 
     if (in_vectorize_ && node->is_addr_tensor() &&
-        tensor_can_vectorized_.count(tensor->name)) {
-      TensorVectorized(node, &node->indices, false);
+        scalar_tensor_without_vectorize_axis_.count(tensor->name)) {
+      PreLoadScalarTensorWithoutVectorizeAxis(node, &node->indices, expr);
       return;
     }
 
@@ -502,6 +511,7 @@ class VectorizeForTransMutator : public ir::IRMutator<ir::Expr *> {
         tensor,
         ::common::errors::InvalidArgument(
             "Expected _Tensor_ node in Store, but received nullptr."));
+    schedule_block_write_dependency_.insert(tensor->name);
 
     if (in_vectorize_ && node->is_addr_tensor() &&
         tensor_can_vectorized_.count(tensor->name)) {
@@ -585,6 +595,7 @@ class VectorizeForTransMutator : public ir::IRMutator<ir::Expr *> {
     tensor_to_vectorized_vars_.clear();
     tensor_can_vectorized_.clear();
     scalar_tensor_without_vectorize_axis_.clear();
+    schedule_block_write_dependency_.clear();
     scalar_tensor_to_local_var_.clear();
     scalar_tensor_to_local_buffer_.clear();
     preload_scalar_tensor_stmts_.clear();
@@ -888,6 +899,8 @@ class VectorizeForTransMutator : public ir::IRMutator<ir::Expr *> {
 
   std::unordered_set<std::string> tensor_can_vectorized_;
   std::unordered_set<std::string> scalar_tensor_without_vectorize_axis_;
+  // avoid to preload tensor which is written in schedule block.
+  std::unordered_set<std::string> schedule_block_write_dependency_;
 
   paddle::flat_hash_map<std::string, ir::Var> tensor_to_vectorized_vars_;
   paddle::flat_hash_map<std::string, ir::Var> scalar_tensor_to_local_var_;

Original file line number	Diff line number	Diff line change
`@@ -308,7 +308,7 @@ inline bool CheckThreadDimensionCanVectorize(int threads,`
`308`	`308`	`int nums,`
`309`	`309`	`int factor) {`
`310`	`310`	`const int deal_elements_in_warp = threads * factor;`
`311`		`- if (nums % deal_elements_in_warp == 0) {`
	`311`	`+ if (nums == deal_elements_in_warp) {`
`312`	`312`	`return true;`
`313`	`313`	`}`
`314`	`314`	`return false;`