support more than 2 warp groups

examples/gemm_sm100/gemm_auto_tcgen5mma.py

@@ -40,7 +40,7 @@ def main(
             for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
                 T.copy(A[by * block_M, k * block_K], A_shared)  # not trans_A
                 T.copy(B[bx * block_N, k * block_K], B_shared)  # trans_B
-                T.gemm(A_shared, B_shared, C_tmem, trans_A, trans_B, wg_wait=-1, clear_accum=k == 0)
+                T.gemm(A_shared, B_shared, C_tmem, trans_A, trans_B, clear_accum=k == 0)
 
             T.copy(C_tmem, C_local)
             T.copy(C_local, C_shared)

src/transform/auto_schedule.cc

@@ -39,12 +39,11 @@
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
-#include <unordered_set>
-
 #include <algorithm>
 #include <cmath>
 #include <iostream>
 #include <memory>
+#include <numeric>
 #include <optional>
 #include <queue>
 #include <sstream>
@@ -567,8 +566,6 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) {
     extractor(func->body);
     Stmt body_to_schedule;
     bool has_tilelang_root = false;
-    PrimExpr updated_thread_extent; // Will be set if warpgroup partition
-                                    // doubles thread extent
     IterVar thread_var; // Thread index variable for warpgroup partition
 
     if (extractor.body.defined()) {
@@ -614,22 +611,15 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) {
 
     // Build ScheduleUnits from IRStructure
     ScheduleUnitBuilder unit_builder;
-    // Get thread index variable for warpgroup partition
-    // First try to get from body_to_schedule, if not found, try from the entire
-    // function body
-    thread_var = ThreadTagChecker::GetThreadVar(body_to_schedule);
-    if (!thread_var.defined()) {
-      thread_var = ThreadTagChecker::GetThreadVar(func->body);
-    }
     if (thread_var.defined()) {
       unit_builder.SetThreadVar(thread_var);
     } else {
       LOG(FATAL) << "Could not find thread index variable, warpgroup "
                     "partition will use default";
     }
-    unit_builder.SetEnableWarpPartition(config.enable_warp_partition);
-    unit_builder.SetSharedMemoryLimit(config.shared_memory_limit);
-    bool double_thread = unit_builder.Build(ir_structure);
+    unit_builder.SetWarpSpeicializeConfig(config);
+    unit_builder.SetSharedMemoryLimit(GetSharedMemoryLimit(target));
+    std::vector<PrimExpr> thread_count = unit_builder.Build(ir_structure);
 
     if (!config.enable_warpgroup_partition) {
       Stmt new_body = ConvertIRStructureToStmt(ir_structure.get(), enable_epi);
@@ -655,28 +645,13 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) {
     int next_barrier_id = 1;
     std::vector<Buffer> barrier_buffers;
     Map<ObjectRef, ObjectRef> barrier_map;
-    // Determine thread count for barrier arrive_count calculations
-    PrimExpr thread_count[2];
-    if (!config.enable_thread_extend) {
-      ICHECK(config.enable_warp_partition);
-      // sm_100: use fixed warp size (32) for both partitions
-      thread_count[0] = IntImm(DataType::Int(32), 32);
-      thread_count[1] = IntImm(DataType::Int(32), 32);
-    } else {
-      // sm_90: original behavior
-      thread_count[0] = thread_var->dom->extent;
-      thread_count[1] = double_thread ? thread_var->dom->extent
-                                      : IntImm(DataType::Int(32),
-                                               config.producer_thread_count);
-    }
     LoopNestingInfo loop_info;
     std::vector<MultiVersionBufferInfo> buffer_infos;
-    PrimExpr barrier_count = config.enable_thread_extend
-                                 ? thread_count[0] + thread_count[1]
-                                 : thread_var->dom->extent;
+    PrimExpr updated_thread_extent = std::accumulate(
+        thread_count.begin() + 1, thread_count.end(), thread_count[0]);
     Buffer neutral_sync_shared_barrier =
-        makeBarrierBuffer(barrier_count, "neutral_sync_shared_barrier", 1,
-                          barrier_buffers, barrier_map);
+        makeBarrierBuffer(updated_thread_extent, "neutral_sync_shared_barrier",
+                          1, barrier_buffers, barrier_map);
     AnalyzeAndInsertBarriers(
         ir_structure.get(), next_barrier_id, barrier_buffers, barrier_map,
         thread_count, loop_info, buffer_infos, neutral_sync_shared_barrier);
@@ -687,19 +662,7 @@ tvm::transform::Pass AutoSchedule(const bool enable_epi) {
     // Apply warpgroup partition to entire IRStructure
     Stmt new_body = ApplyWarpgroupPartitionToIRStructure(
         ir_structure.get(), thread_var, barrier_buffers, barrier_map,
-        enable_epi, thread_count, double_thread, config,
-        neutral_sync_shared_barrier);
-
-    if (config.enable_thread_extend) {
-      // sm_90: may need to update thread extent
-      if (double_thread) {
-        updated_thread_extent = thread_var->dom->extent * 2;
-      } else {
-        updated_thread_extent =
-            thread_var->dom->extent +
-            IntImm(DataType::Int(32), config.producer_thread_count);
-      }
-    }
+        enable_epi, thread_count, config, neutral_sync_shared_barrier);
 
     // If we extracted from tilelang_root block, replace the body
     Stmt final_body;

src/transform/auto_schedule.h

@@ -85,58 +85,25 @@ struct ComponentInfo {
   bool uses_tensor_core_{false};
 };
 
-// Warp specialization architecture enum
-enum class WarpSpecializeArch : uint8_t {
-  kHopper = 0,
-  kBlackwell = 1,
-  kUnsupported = 2,
-};
-
-// Configuration for warp specialization
-struct WarpSpecializeConfig {
-  WarpSpecializeArch arch = WarpSpecializeArch::kUnsupported;
-  int consumer_max_nreg = 0;
-  int producer_max_nreg = 0;
-  int producer_thread_count = 0;
-  bool enable_set_max_nreg = false;
-  bool enable_warpgroup_partition = false;
-  bool enable_thread_extend = false;
-  bool enable_warp_partition = false;
-  int shared_memory_limit = 0;
-};
-
 // Factory function to get warp specialization configuration for a target
 inline WarpSpecializeConfig GetWarpSpecializeConfig(Target target) {
   if (TargetIsHopper(target)) {
-    return {WarpSpecializeArch::kHopper,
-            240,
-            24,
-            128,
-            true,
-            true,
-            true,
-            false,
-            228 * 1024};
+    return {WarpSpecializeArch::kHopper, 240, 24, 128, true, true, true, false};
+  } else if (TargetIsSm100(target)) {
+    return {WarpSpecializeArch::kBlackwell, 0, 0, 32, false, true, false, true};
+  } else {
+    return {
+        WarpSpecializeArch::kUnsupported, 0, 0, 0, false, false, false, false};
+  }
+}
+
+inline int64_t GetSharedMemoryLimit(Target target) {
+  if (TargetIsHopper(target)) {
+    return 228 * 1024;
   } else if (TargetIsSm100(target)) {
-    return {WarpSpecializeArch::kBlackwell,
-            0,
-            0,
-            32,
-            false,
-            true,
-            false,
-            true,
-            228 * 1024};
+    return 228 * 1024;
   } else {
-    return {WarpSpecializeArch::kUnsupported,
-            0,
-            0,
-            0,
-            false,
-            false,
-            false,
-            false,
-            0};
+    return 48 * 1024;
   }
 }
 

src/transform/auto_schedule/barrier.h

@@ -167,22 +167,24 @@ static void
 AnalyzeAndInsertBarriers(IRStructure *node, int &next_barrier_id,
                          std::vector<Buffer> &barrier_buffers,
                          Map<ObjectRef, ObjectRef> &barrier_map,
-                         PrimExpr thread_count[2], LoopNestingInfo &loop_info,
+                         const std::vector<PrimExpr> &thread_count,
+                         LoopNestingInfo &loop_info,
                          std::vector<MultiVersionBufferInfo> &buffer_infos,
                          Buffer neutral_sync_shared_barrier);
 static void
 AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id,
                             std::vector<Buffer> &barrier_buffers,
                             Map<ObjectRef, ObjectRef> &barrier_map,
-                            PrimExpr thread_count[2],
+                            const std::vector<PrimExpr> &thread_count,
                             LoopNestingInfo &loop_info,
                             std::vector<MultiVersionBufferInfo> &buffer_infos,
                             Buffer neutral_sync_shared_barrier);
 static void
 AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id,
                            std::vector<Buffer> &barrier_buffers,
                            Map<ObjectRef, ObjectRef> &barrier_map,
-                           PrimExpr thread_count[2], LoopNestingInfo &loop_info,
+                           const std::vector<PrimExpr> &thread_count,
+                           LoopNestingInfo &loop_info,
                            std::vector<MultiVersionBufferInfo> &buffer_infos,
                            Buffer neutral_sync_shared_barrier);
 
@@ -543,7 +545,8 @@ static void
 AnalyzeAndInsertBarriers(IRStructure *node, int &next_barrier_id,
                          std::vector<Buffer> &barrier_buffers,
                          Map<ObjectRef, ObjectRef> &barrier_map,
-                         PrimExpr thread_count[2], LoopNestingInfo &loop_info,
+                         const std::vector<PrimExpr> &thread_count,
+                         LoopNestingInfo &loop_info,
                          std::vector<MultiVersionBufferInfo> &buffer_infos,
                          Buffer neutral_sync_shared_barrier) {
   if (!node)
@@ -575,7 +578,7 @@ static void
 AnalyzeSequenceNodeBarriers(SequenceNode *seq, int &next_barrier_id,
                             std::vector<Buffer> &barrier_buffers,
                             Map<ObjectRef, ObjectRef> &barrier_map,
-                            PrimExpr thread_count[2],
+                            const std::vector<PrimExpr> &thread_count,
                             LoopNestingInfo &loop_info,
                             std::vector<MultiVersionBufferInfo> &buffer_infos,
                             Buffer neutral_sync_shared_barrier) {
@@ -786,7 +789,8 @@ static void
 AnalyzeControlNodeBarriers(ControlNode *ctrl, int &next_barrier_id,
                            std::vector<Buffer> &barrier_buffers,
                            Map<ObjectRef, ObjectRef> &barrier_map,
-                           PrimExpr thread_count[2], LoopNestingInfo &loop_info,
+                           const std::vector<PrimExpr> &thread_count,
+                           LoopNestingInfo &loop_info,
                            std::vector<MultiVersionBufferInfo> &buffer_infos,
                            Buffer neutral_sync_shared_barrier) {
   if (!ctrl || !ctrl->child)

src/transform/auto_schedule/schedule_builder.cc

@@ -237,7 +237,9 @@ void CollectSuffixTasks(IRStructure *root,
   }
 }
 
-bool AssignWarpgroupIdsGlobal(IRStructure *root, bool enable_warp_partition) {
+std::vector<PrimExpr>
+AssignWarpgroupIdsGlobal(IRStructure *root, const WarpSpecializeConfig &config,
+                         PrimExpr thread_count) {
   if (!root) {
     LOG(FATAL) << "Empty root";
   }
@@ -269,7 +271,7 @@ bool AssignWarpgroupIdsGlobal(IRStructure *root, bool enable_warp_partition) {
   CollectPrefixTasks(root, prefix_tasks, prefix_valid);
 
   std::unordered_set<TaskNode *> suffix_tasks;
-  if (enable_warp_partition) {
+  if (config.enable_warp_partition) {
     CollectSuffixTasks(root, all_tasks, uf, suffix_tasks);
   }
 
@@ -360,7 +362,12 @@ bool AssignWarpgroupIdsGlobal(IRStructure *root, bool enable_warp_partition) {
         }
       }
     }
-    return true;
+    if (config.enable_thread_extend) {
+      return {thread_count, thread_count};
+    } else {
+      return {IntImm(DataType::Int(32), 32), IntImm(DataType::Int(32), 32),
+              thread_count - IntImm(DataType::Int(32), 64)};
+    }
   } else {
     int64_t warpgroup0_latency = 0;
     int64_t warpgroup1_latency = 0;
@@ -387,7 +394,13 @@ bool AssignWarpgroupIdsGlobal(IRStructure *root, bool enable_warp_partition) {
         }
       }
     }
-    return false;
+    if (config.enable_thread_extend) {
+      return {thread_count,
+              IntImm(DataType::Int(32), config.producer_thread_count)};
+    } else {
+      return {IntImm(DataType::Int(32), 32), IntImm(DataType::Int(32), 32),
+              thread_count - IntImm(DataType::Int(32), 64)};
+    }
   }
 }
 

src/transform/auto_schedule/schedule_builder.h

@@ -29,7 +29,28 @@ using namespace tir;
 class TaskUnionFind;
 struct ComponentInfo;
 
-bool AssignWarpgroupIdsGlobal(IRStructure *root, bool enable_warp_partition);
+// Warp specialization architecture enum
+enum class WarpSpecializeArch : uint8_t {
+  kHopper = 0,
+  kBlackwell = 1,
+  kUnsupported = 2,
+};
+
+// Configuration for warp specialization
+struct WarpSpecializeConfig {
+  WarpSpecializeArch arch = WarpSpecializeArch::kUnsupported;
+  int consumer_max_nreg = 0;
+  int producer_max_nreg = 0;
+  int producer_thread_count = 0;
+  bool enable_set_max_nreg = false;
+  bool enable_warpgroup_partition = false;
+  bool enable_thread_extend = false;
+  bool enable_warp_partition = false;
+};
+
+std::vector<PrimExpr>
+AssignWarpgroupIdsGlobal(IRStructure *root, const WarpSpecializeConfig &config,
+                         PrimExpr thread_count);
 
 // Extract all sequential task nodes from the IR structure tree
 void GatherTaskNodes(const std::vector<std::shared_ptr<IRStructure>> &nodes,
@@ -54,11 +75,12 @@ void CollectSuffixTasks(IRStructure *root,
 // Builder that collects ScheduleUnits from IRStructure
 class ScheduleUnitBuilder {
 public:
-  bool Build(std::shared_ptr<IRStructure> &root) {
+  std::vector<PrimExpr> Build(std::shared_ptr<IRStructure> &root) {
     ScheduleRecursive(root, {});
 
     // Global warpgroup id assignment from the top level
-    return AssignWarpgroupIdsGlobal(root.get(), enable_warp_partition_);
+    return AssignWarpgroupIdsGlobal(root.get(), config_,
+                                    thread_var_->dom->extent);
   }
 
   // New recursive scheduling function that replaces Collect method
@@ -589,14 +611,16 @@ class ScheduleUnitBuilder {
   void SetThreadVar(IterVar thread_var) { thread_var_ = thread_var; }
 
   // Set enable_warp_partition flag
-  void SetEnableWarpPartition(bool enable) { enable_warp_partition_ = enable; }
+  void SetWarpSpeicializeConfig(const WarpSpecializeConfig &config) {
+    config_ = config;
+  }
 
   // Set shared memory limit for pipeline (in bytes)
   void SetSharedMemoryLimit(int64_t bytes) { shared_memory_limit_ = bytes; }
 
 private:
-  IterVar thread_var_; // Thread index variable for warpgroup partition
-  bool enable_warp_partition_ = false;
+  IterVar thread_var_;          // Thread index variable for warpgroup partition
+  WarpSpecializeConfig config_; // Configuration for warp specialization
   int64_t shared_memory_limit_ = 48 * 1024;
 
   // Check if two regions refer to the same buffer