PaddlePaddle
diff --git a/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
Copy file name to clipboard
+5-3Lines changed: 5 additions & 3 deletions b/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
Copy file name to clipboard
+5-3Lines changed: 5 additions & 3 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/group_tile_util.cc
Copy file name to clipboard
+1-6Lines changed: 1 addition & 6 deletions b/‎paddle/cinn/ir/group_schedule/config/group_tile_util.cc
Copy file name to clipboard
+1-6Lines changed: 1 addition & 6 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/config/group_tile_util.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/cinn/ir/group_schedule/config/group_tile_util.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc
Copy file name to clipboard
+49-5Lines changed: 49 additions & 5 deletions b/‎paddle/cinn/ir/group_schedule/tactic/tile_first_general_tactic.cc
Copy file name to clipboard
+49-5Lines changed: 49 additions & 5 deletions
@@ -36,7 +36,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
 
-PD_DECLARE_bool(group_schedule_tiling_first);
+PD_DECLARE_bool(cinn_enable_grid_reduce);
 
 namespace cinn {
 namespace hlir {
@@ -808,8 +808,10 @@ std::shared_ptr<FusionGroupInfo> GetFusionGroupInfo(
                    });
   }
 
-  group_info->can_apply_grid_reduce =
-      GetCanApplyGridReduce(op_compute_bodies, group_info->reduce_axis);
+  if (FLAGS_cinn_enable_grid_reduce) {
+    group_info->can_apply_grid_reduce =
+        GetCanApplyGridReduce(op_compute_bodies, group_info->reduce_axis);
+  }
 
   VLOG(4) << group_info->DebugPrint();
   return group_info;
 
@@ -34,7 +34,6 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
   // A tensor is downstream of reduce either if it is produced by a reduce, or
   // if it has data dependency on another tensor that is downstream of reduce.
   std::unordered_set<std::string> reduce_downstream_tensor_names;
-  int reduce_count = 0;
 
   const auto IsReduceDownstream = [&](const ir::Expr& expr_block) {
     for (auto& expr_load : ChildTensorLoads(expr_block)) {
@@ -90,9 +89,6 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
     bool is_reduce_downstream = IsReduceDownstream(expr_block);
     bool output_has_reduce_axis = CheckOutputHasReduceAxis(body, expr_block);
 
-    if (is_reduce) {
-      ++reduce_count;
-    }
     if (is_reduce_downstream || is_reduce) {
       AddReduceDownstream(expr_block);
     }
@@ -105,8 +101,7 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
       return false;
     }
   }
-
-  return reduce_count == 1;
+  return true;
 }
 
 }  // namespace ir
 
@@ -19,8 +19,8 @@ namespace cinn {
 namespace ir {
 
 // Check whether we can apply grid reduce in this group.
-// We can apply grid reduce if there is exactly one reduce, and whose result is
-// not broadcasted before output.
+// We can apply grid reduce if there is no reduce-then-broadcast dependency
+// in this group.
 bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,
                            const std::vector<int64_t>& reduce_axis);
 
 
@@ -79,6 +79,7 @@ class TileFirstGeneralTactic final : public ScheduleTactic {
   std::vector<int32_t> vec_flatten_axis_;
   std::vector<int32_t> vec_reduce_axis_;
   std::unordered_map<std::string, std::string> map_rf_block_;
+  std::unordered_map<std::string, std::string> map_global_rf_block_;
 };
 
 void TileFirstGeneralTactic::Init(ScheduleContext* context) {
@@ -381,19 +382,46 @@ void TileFirstGeneralTactic::SplitSptialInner(ir::IRSchedule* sch,
 
 void TileFirstGeneralTactic::SplitReduceInner(ir::IRSchedule* sch,
                                               const std::string& block_id) {
+  const int64_t rd_block = context_->config.tile_config.grid_reduce_num;
+  const int64_t rd_thread = 16;
+  const int cur_reduce_axis = 2;
+
+  // [ R ] => [ rd_block*rd_thread, rd_inner ]
   auto loops = sch->GetLoops(block_id);
-  // [S(-1), S(32), R] => [S(-1), S(32), R(16), R(-1)]
-  sch->Split(loops[2], std::vector<int>{16, -1});
+  sch->Split(loops[cur_reduce_axis],
+             std::vector<int>{-1, rd_block * rd_thread});
+  loops = sch->GetLoops(block_id);
+  sch->Reorder({loops[cur_reduce_axis + 1], loops[cur_reduce_axis]});
 
   loops = sch->GetLoops(block_id);
   if (IsReductionSBlock(sch->GetBlock(block_id)) &&
       ir::GetLoopExtent(loops[2]) != 1) {
     ir::Expr rf_tensor =
-        sch->FactorizeReduction(loops[2],
-                                0,
+        sch->FactorizeReduction(loops[cur_reduce_axis],
+                                /* rf_axis = */ 0,
                                 /* with_write_back_block_init = */ false);
     map_rf_block_[block_id] = rf_tensor.as_tensor_ref()->name;
   }
+
+  // [ rd_block*rd_thread ] => [ rd_block, rd_thread ]
+  if (rd_block > 1) {
+    loops = sch->GetLoops(block_id);
+    sch->Split(loops[cur_reduce_axis], {rd_block, rd_thread});
+
+    if (IsReductionSBlock(sch->GetBlock(block_id))) {
+      loops = sch->GetLoops(map_rf_block_[block_id]);
+      sch->Split(loops[cur_reduce_axis], {rd_block, rd_thread});
+
+      loops = sch->GetLoops(block_id);
+      ir::Expr rf_tensor =
+          sch->FactorizeReduction(loops[cur_reduce_axis],
+                                  /* rf_axis = */ 0,
+                                  /* with_write_back_block_init = */ false);
+      std::string tensor_name = rf_tensor.as_tensor_ref()->name;
+      map_global_rf_block_[block_id] = tensor_name;
+      rf_tensor.as_tensor_ref()->WithBuffer("global", "_" + tensor_name);
+    }
+  }
 }
 
 void TileFirstGeneralTactic::VariableTypeAssignment(
@@ -435,6 +463,12 @@ void TileFirstGeneralTactic::SetDiscreteReduceType(
                      ->schedule_block.As<ir::ScheduleBlock>();
     block->reduce_method = cinn::ir::DiscreteReduceMethod();
   }
+  if (map_global_rf_block_.count(block_id) > 0) {
+    auto block = sch->GetBlock(map_global_rf_block_[block_id])
+                     .As<ir::ScheduleBlockRealize>()
+                     ->schedule_block.As<ir::ScheduleBlock>();
+    block->reduce_method = cinn::ir::DiscreteReduceMethod();
+  }
 }
 
 void TileFirstGeneralTactic::BindCudaInfo(ir::IRSchedule* sch,
@@ -446,14 +480,24 @@ void TileFirstGeneralTactic::BindCudaInfo(ir::IRSchedule* sch,
   const auto DoBind = [&](const std::vector<ir::Expr>& loops) {
     sch->Bind(loops[0], "blockIdx.x");
     sch->Bind(loops[1], "threadIdx.x");
-    sch->Bind(loops[2], "threadIdx.y");
+    if (context_->config.tile_config.grid_reduce_num > 1) {
+      sch->Bind(loops[2], "blockIdx.y");
+      if (loops.size() > 3) {
+        sch->Bind(loops[3], "threadIdx.y");
+      }
+    } else {
+      sch->Bind(loops[2], "threadIdx.y");
+    }
   };
 
   DoBind(sch->GetLoops(block_id));
 
   if (map_rf_block_.count(block_id) > 0) {
     DoBind(sch->GetLoops(map_rf_block_[block_id]));
   }
+  if (map_global_rf_block_.count(block_id) > 0) {
+    DoBind(sch->GetLoops(map_global_rf_block_[block_id]));
+  }
 }
 
 std::unique_ptr<ScheduleTactic> CreateTileFirstGeneralTactic() {
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,6 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,`
`34`	`34`	`// A tensor is downstream of reduce either if it is produced by a reduce, or`
`35`	`35`	`// if it has data dependency on another tensor that is downstream of reduce.`
`36`	`36`	`std::unordered_set<std::string> reduce_downstream_tensor_names;`
`37`		`- int reduce_count = 0;`
`38`	`37`
`39`	`38`	`const auto IsReduceDownstream = [&](const ir::Expr& expr_block) {`
`40`	`39`	`for (auto& expr_load : ChildTensorLoads(expr_block)) {`
`@@ -90,9 +89,6 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,`
`90`	`89`	`bool is_reduce_downstream = IsReduceDownstream(expr_block);`
`91`	`90`	`bool output_has_reduce_axis = CheckOutputHasReduceAxis(body, expr_block);`
`92`	`91`
`93`		`- if (is_reduce) {`
`94`		`- ++reduce_count;`
`95`		`- }`
`96`	`92`	`if (is_reduce_downstream \|\| is_reduce) {`
`97`	`93`	`AddReduceDownstream(expr_block);`
`98`	`94`	`}`
`@@ -105,8 +101,7 @@ bool GetCanApplyGridReduce(const std::vector<ir::Expr>& op_compute_bodies,`
`105`	`101`	`return false;`
`106`	`102`	`}`
`107`	`103`	`}`
`108`		`-`
`109`		`- return reduce_count == 1;`
	`104`	`+ return true;`
`110`	`105`	`}`
`111`	`106`
`112`	`107`	`} // namespace ir`