cinn(backends): fuse last dim reduce loop and set bound of reduce thread (#61576)

6clc · web-flow · commit e53a04e87a72 · 2024-02-07T10:28:45.000+08:00
diff --git a/paddle/cinn/ir/group_schedule/tactic/align_iter_space_tactic.cc b/paddle/cinn/ir/group_schedule/tactic/align_iter_space_tactic.cc
@@ -30,9 +30,6 @@ void AlignIterSpaceTactic::Init(ScheduleContext* context) {
 void AlignIterSpaceTactic::Apply(ir::IRSchedule* sch,
                                  const std::string& block_id) {
   ir::Expr block = sch->GetBlock(block_id);
-  if (analyzer::IsReductionSBlock(block)) {
-    return;
-  }
 
   std::vector<ir::Expr> loops = sch->GetLoops(block_id);
   ir::Expr src_total_extent{1};
@@ -66,16 +63,20 @@ void AlignIterSpaceTactic::Apply(ir::IRSchedule* sch,
     }
     if (context_->iter_space_info.sp_space.size() < loops.size() - 1) {
       loops = sch->GetLoops(block_id);
+
+      // Align the loop in the current block that needs to be aligned with the
+      // reduce loop in iter_space_info
       std::vector<ir::Expr> rb_loops(
-          loops.begin() + context_->iter_space_info.sp_space.size(),
-          loops.end());
+          loops.end() - context_->iter_space_info.rb_space.size(), loops.end());
       sch->Fuse(rb_loops);
     }
     if (context_->iter_space_info.sp_space.size() > 1) {
+      // Align the loop in the current block that needs to be aligned with the
+      // spatial loop in iter_space_info
       loops = sch->GetLoops(block_id);
       std::vector<ir::Expr> sp_loops(
           loops.begin(),
-          loops.begin() + context_->iter_space_info.sp_space.size());
+          loops.end() - context_->iter_space_info.rb_space.size());
       sch->Fuse(sp_loops);
     }
   } else {
diff --git a/paddle/cinn/ir/group_schedule/tactic/tile_tactic.cc b/paddle/cinn/ir/group_schedule/tactic/tile_tactic.cc
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/cinn/ir/group_schedule/tactic/tile_tactic.h"
+#include "paddle/cinn/common/target.h"
 #include "paddle/cinn/ir/ir.h"
 
 namespace cinn {
@@ -32,11 +33,17 @@ void TileTactic::Init(ScheduleContext* context) {
     }
   };
   auto GetTreeReduceSize = [&](const ir::Expr& total_rb_extent) -> int64_t {
+    const int64_t max_num_threads =
+        common::DefaultNVGPUTarget().max_num_threads();
+    int64_t nums_thread_per_block = max_num_threads;
     if (total_rb_extent.is_constant()) {
       int64_t extent = static_cast<int64_t>(total_rb_extent.get_constant());
-      return GetFirstFactor(extent);
+      nums_thread_per_block = GetFirstFactor(extent);
+    } else {
+      nums_thread_per_block = context_->bucket_info.rb_lower_bound;
     }
-    return context_->bucket_info.rb_lower_bound;
+    return nums_thread_per_block > max_num_threads ? max_num_threads
+                                                   : nums_thread_per_block;
   };
   auto GetNumThreadPerBlock = [&](int64_t lower_bound) -> int64_t {
     // When designing the tile config, we can further subdivided.
