Refactor test for InjectSetMaxNReg pass in TileLang

- Improved readability by restructuring conditional checks and assertions in the test cases.
- Enhanced clarity in the collection of `set_max_nreg` calls by simplifying the logic.
- Ensured consistent formatting and spacing throughout the test functions for better maintainability.

Author: LeiWang1999

src/transform/annotate_warp_group_reg_alloc.cc

@@ -9,16 +9,16 @@
 #include <unordered_set>
 #include <vector>
 
-#include "tir/transforms/ir_utils.h"
 #include "../op/builtin.h"
+#include "tir/transforms/ir_utils.h"
 
 namespace tvm {
 namespace tl {
 
 using namespace tir;
 
 class SetMaxNRegCollector : public StmtExprVisitor {
- public:
+public:
   static Array<IntImm> Collect(const PrimFunc &f) {
     SetMaxNRegCollector collector;
     collector(f->body);
@@ -28,7 +28,7 @@ class SetMaxNRegCollector : public StmtExprVisitor {
                : collector.nreg_;
   }
 
- private:
+private:
   void VisitStmt_(const EvaluateNode *op) final {
     if (const CallNode *call = op->value.as<CallNode>()) {
       if (call->op.same_as(set_max_nreg())) {
@@ -54,15 +54,15 @@ class SetMaxNRegCollector : public StmtExprVisitor {
 };
 
 class SetMaxNRegInjector : public StmtExprMutator {
- public:
+public:
   static PrimFunc Inject(PrimFunc f) {
     auto T = SetMaxNRegInjector();
     T.nreg_ = SetMaxNRegCollector::Collect(f);
     f.CopyOnWrite()->body = T(f->body);
     return f;
   }
 
- private:
+private:
   Stmt VisitStmt_(const EvaluateNode *op) final {
     if (const CallNode *call = op->value.as<CallNode>()) {
       if (call->op.same_as(set_max_nreg()) ||
@@ -97,7 +97,6 @@ class SetMaxNRegInjector : public StmtExprMutator {
       Optional<Stmt> consumer_body = if_then_else->else_case;
       ICHECK(consumer_body.defined()) << "Consumer body is undefined";
 
-
       int dec_reg = nreg_[0].as<IntImmNode>()->value;
       int inc_reg = nreg_[1].as<IntImmNode>()->value;
 
@@ -107,13 +106,13 @@ class SetMaxNRegInjector : public StmtExprMutator {
       // Only inject if we have valid register hints and no SIMT copy
       // For now, we assume no SIMT copy detection is available here
       // TODO: Add SIMT copy detection if needed
-      bool has_simt_copy = false;  // Placeholder
+      bool has_simt_copy = false; // Placeholder
 
       if (dec_reg >= 0 && inc_reg >= 0 && !has_simt_copy) {
         inc_reg_stmt = Evaluate(Call(DataType::Handle(), set_max_nreg(),
-                                   {inc_reg == 0 ? 240 : inc_reg, 1}));
+                                     {inc_reg == 0 ? 240 : inc_reg, 1}));
         dec_reg_stmt = Evaluate(Call(DataType::Handle(), set_max_nreg(),
-                                   {dec_reg == 0 ? 24 : dec_reg, 0}));
+                                     {dec_reg == 0 ? 24 : dec_reg, 0}));
       }
 
       // Inject register setting statements
@@ -127,7 +126,8 @@ class SetMaxNRegInjector : public StmtExprMutator {
       consumer_stmts.push_back(consumer_body.value());
       auto new_consumer_body = SeqStmt(consumer_stmts);
 
-      auto new_if_stmt = IfThenElse(if_then_else->condition, new_producer_body, new_consumer_body);
+      auto new_if_stmt = IfThenElse(if_then_else->condition, new_producer_body,
+                                    new_consumer_body);
       auto new_attr = AttrStmt(op->node, op->attr_key, op->value, new_if_stmt);
 
       return new_attr;
@@ -153,8 +153,9 @@ tvm::transform::Pass AnnotateWarpGroupRegAlloc() {
 
 TVM_FFI_STATIC_INIT_BLOCK({
   namespace refl = tvm::ffi::reflection;
-  refl::GlobalDef().def("tl.transform.AnnotateWarpGroupRegAlloc", AnnotateWarpGroupRegAlloc);
+  refl::GlobalDef().def("tl.transform.AnnotateWarpGroupRegAlloc",
+                        AnnotateWarpGroupRegAlloc);
 });
 
-}  // namespace tl
-}  // namespace tvm
+} // namespace tl
+} // namespace tvm

testing/python/transform/test_tilelang_transform_inject_set_max_nreg.py

@@ -1,12 +1,12 @@
 from tilelang import tvm as tvm
 import tilelang as tl
-from tilelang.utils.target import determine_target
 import tilelang.language as T
 import tilelang.testing
 from tvm import tir
 
 tilelang.disable_cache()
 
+
 def test_inject_set_max_nreg():
     """Test the InjectSetMaxNReg pass"""
 
@@ -37,21 +37,26 @@ def before(A: T.Tensor((512, 512), "float16"), B: T.Tensor((512, 512), "float16"
                     T.mbarrier_wait_parity(T.get_mbarrier(k % 3 + 3), T.bitwise_xor(k // 3 % 2, 1))
                     if v - 128 == 0:
                         T.tma_load(
-                            T.create_tma_descriptor(6, 2, A.data, 512, 512, 2, 1024, 32, 64, 1, 1, 0, 2, 2, 0),
-                            T.get_mbarrier(k % 3),
-                            T.tvm_access_ptr(T.type_annotation("float16"), A_shared.data, k % 3 * 2048, 2048, 2),
+                            T.create_tma_descriptor(6, 2, A.data, 512, 512, 2, 1024, 32, 64, 1, 1,
+                                                    0, 2, 2, 0), T.get_mbarrier(k % 3),
+                            T.tvm_access_ptr(
+                                T.type_annotation("float16"), A_shared.data, k % 3 * 2048, 2048, 2),
                             k * 32, by * 64)
-                    T.evaluate(tir.Call("handle", "tir.ptx_arrive_barrier", [T.get_mbarrier(k % 3)]))
+                    T.evaluate(
+                        tir.Call("handle", "tir.ptx_arrive_barrier", [T.get_mbarrier(k % 3)]))
             else:
                 # Consumer branch - should have set_max_nreg(240, 1)
                 for k in range(16):
                     T.mbarrier_wait_parity(T.get_mbarrier(k % 3), k // 3 % 2)
                     T.call_extern(
                         "handle", "tl::gemm_ss<64, 64, 32, 4, 1, 0, 0>",
-                        T.tvm_access_ptr(T.type_annotation("float16"), A_shared.data, k % 3 * 2048, 2048, 1),
-                        T.tvm_access_ptr(T.type_annotation("float16"), B_shared.data, k % 3 * 2048, 2048, 1),
+                        T.tvm_access_ptr(
+                            T.type_annotation("float16"), A_shared.data, k % 3 * 2048, 2048, 1),
+                        T.tvm_access_ptr(
+                            T.type_annotation("float16"), B_shared.data, k % 3 * 2048, 2048, 1),
                         T.tvm_access_ptr(T.type_annotation("float32"), C_local.data, 0, 32, 3))
-                    T.evaluate(tir.Call("handle", "tir.ptx_arrive_barrier", [T.get_mbarrier(k % 3 + 3)]))
+                    T.evaluate(
+                        tir.Call("handle", "tir.ptx_arrive_barrier", [T.get_mbarrier(k % 3 + 3)]))
 
     # Apply the InjectSetMaxNReg pass
     func = before
@@ -64,15 +69,15 @@ def before(A: T.Tensor((512, 512), "float16"), B: T.Tensor((512, 512), "float16"
     set_max_nreg_calls = []
 
     def collect_set_max_nreg(stmt):
-        if isinstance(stmt, tvm.tir.Evaluate):
-            if hasattr(stmt.value, 'op') and hasattr(stmt.value.op, 'name'):
-                if stmt.value.op.name == "tl.set_max_nreg":
-                    set_max_nreg_calls.append((stmt.value.args[0].value, stmt.value.args[1].value))
+        if (isinstance(stmt, tvm.tir.Evaluate) and hasattr(stmt.value, 'op') and
+                hasattr(stmt.value.op, 'name') and stmt.value.op.name == "tl.set_max_nreg"):
+            set_max_nreg_calls.append(stmt.value)
 
     tvm.tir.stmt_functor.post_order_visit(main_func.body, collect_set_max_nreg)
 
     # We should have at least 2 set_max_nreg calls (one for producer, one for consumer)
-    assert len(set_max_nreg_calls) >= 2, f"Expected at least 2 set_max_nreg calls, got {len(set_max_nreg_calls)}"
+    assert len(set_max_nreg_calls
+              ) >= 2, f"Expected at least 2 set_max_nreg calls, got {len(set_max_nreg_calls)}"
 
     # Check that we have the expected register values
     reg_values = [call[0] for call in set_max_nreg_calls]
@@ -118,15 +123,16 @@ def before_no_set_max_nreg(A: T.Tensor((512, 512), "float16")):
     set_max_nreg_calls = []
 
     def collect_set_max_nreg(stmt):
-        if isinstance(stmt, tvm.tir.Evaluate):
-            if hasattr(stmt.value, 'op') and hasattr(stmt.value.op, 'name'):
-                if stmt.value.op.name == "tl.set_max_nreg":
-                    set_max_nreg_calls.append(stmt.value)
+        if (isinstance(stmt, tvm.tir.Evaluate) and hasattr(stmt.value, 'op') and
+                hasattr(stmt.value.op, 'name') and stmt.value.op.name == "tl.set_max_nreg"):
+            set_max_nreg_calls.append(stmt.value)
 
     tvm.tir.stmt_functor.post_order_visit(main_func.body, collect_set_max_nreg)
 
     # Should have no set_max_nreg calls when no_set_max_nreg is present
-    assert len(set_max_nreg_calls) == 0, f"Expected 0 set_max_nreg calls when no_set_max_nreg is present, got {len(set_max_nreg_calls)}"
+    assert len(
+        set_max_nreg_calls
+    ) == 0, f"Expected 0 set_max_nreg calls when no_set_max_nreg is present, got {len(set_max_nreg_calls)}"
 
     print("InjectSetMaxNReg with no_set_max_nreg test passed!")