[Enhancement] Implement dynamic unroll factor in CUDA code generation

[LeiWang1999]

This commit introduces support for specifying a dynamic unroll factor in the CUDA code generation. The unroll_factor map is added to store unroll factors for loop variables, allowing for more flexible and optimized loop unrolling. Additionally, the unroll function is integrated into the loop language, enabling users to define unroll factors directly in their code. This enhancement improves performance by allowing tailored unrolling strategies based on specific loop characteristics.

Summary by CodeRabbit

• New Features

  • Support for per-loop unroll factors in CUDA output, allowing explicit pragma unroll N.
  • Public unroll API added to the language for explicit unrolled loops and aliases for serial/unroll variants.

• Tests

  • New tests verifying unroll behavior and that custom unroll factors produce the expected CUDA pragmas.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Walkthrough

This PR adds per-loop unroll factor support end-to-end: new TileLang unroll API, builder support for UnrollForWithStep, storage of unroll factors in CUDA codegen, and emission of #pragma unroll or #pragma unroll <factor> in generated CUDA kernels; tests validate both behaviors.

Changes

Cohort / File(s)	Summary
CUDA codegen src/target/codegen_cuda.h, src/target/codegen_cuda.cc	Added a private unroll_factor map to CodeGenTileLangCUDA; codegen now emits #pragma unroll <factor> when a factor is registered for a loop variable, otherwise emits default #pragma unroll.
Language API — loop construction tilelang/language/loop.py	Added unroll() function (params: explicit, unroll_factor, annotations), returned frame.ForFrame types, and exported Serial and Unroll aliases; serial() gains an explicit return annotation.
Language exports tilelang/language/__init__.py	Export surface expanded to include unroll, Serial, and Unroll from tilelang.language.loop.
Loop builder framework tilelang/language/v2/builder.py	Added @dataclass UnrollForWithStep (inherits SerialForWithStep); ctx_for now branches on UnrollForWithStep vs SerialForWithStep and constructs tir.unroll or tir.serial respectively.
Tests testing/python/language/test_tilelang_language_unroll.py	New test module with test_unroll_with_step (checks #pragma unroll) and test_unroll_with_unroll_factor (checks #pragma unroll 4).
Examples / kernels examples/gemm_sp/example_custom_compress.py, tilelang/language/experimental/gemm_sp.py	Minor behavior-preserving changes: reset semantics in compress example; intrinsic name changed to tl.tileop.gemm_sp_py in two call sites.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant TileLangAPI as TileLang API (unroll)
    participant Builder as v2.builder (UnrollForWithStep)
    participant CodeGen as CodeGenTileLangCUDA
    participant CUDA as CUDA Kernel Source

    User->>TileLangAPI: write loop using unroll(...) (optional unroll_factor)
    TileLangAPI->>Builder: create ForFrame / UnrollForWithStep
    Builder->>CodeGen: emit IR frames (tir.unroll or tir.serial) with metadata (unroll_factor)
    CodeGen->>CodeGen: store unroll_factor map per loop var
    CodeGen->>CUDA: generate kernel source, emit:
    alt factor present
        CodeGen->>CUDA: "#pragma unroll <factor>"
    else
        CodeGen->>CUDA: "#pragma unroll"
    end
    CUDA-->>User: compiled kernel source (checked by tests)

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

• Review tilelang/language/loop.py for annotation handling and unroll_factor validation vs explicit pragma flag.
• Inspect tilelang/language/v2/builder.py branching to ensure correct IR (tir.unroll vs tir.serial) and preserved step validation.
• Verify src/target/codegen_cuda.cc logic that maps loop variables to factors and emits correct pragmas.
• Check new tests for adequate coverage and any brittle string-matching of pragmas.

Suggested reviewers

• tzj-fxz

Poem

🐰 I nudge each loop to dance and play,
A factor tucked in every way.
Pragmas sing in CUDA bright,
Threads hop, unrolled just right—
Hooray for kernels in bunny light!

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 37.50% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and clearly summarizes the main change: implementing support for dynamic unroll factors in CUDA code generation, which is the primary objective across all modified files.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[LeiWang1999]

@codex review

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (3)

src/target/codegen_cuda.cc (1)

2669-2673: Verify cleanup of unroll_factor map entries after use.

The unroll_factor map stores entries keyed by VarNode*, but these entries are never removed after the corresponding ForNode is processed. While this doesn't cause incorrect behavior (the map is scoped to a single function's codegen), it could lead to stale entries if the same VarNode* address is reused in a different context.

Consider clearing the entry after use in VisitStmt_(const ForNode*):

 void CodeGenTileLangCUDA::VisitStmt_(const tir::ForNode *op) {
   if (op->kind == tir::ForKind::kUnrolled) {
     PrintIndent();
     if (unroll_factor.count(op->loop_var.get())) {
       stream << "#pragma unroll "
              << PrintExpr(unroll_factor[op->loop_var.get()]) << "\n";
+      unroll_factor.erase(op->loop_var.get());
     } else {
       stream << "#pragma unroll\n";
     }
   }

tilelang/language/v2/builder.py (1)

293-300: Dead code: the else branch is unreachable.

The else branch (lines 297-300) can never be executed because:

1. Line 278 already validates that it is an instance of SerialForWithStep or UnrollForWithStep
2. UnrollForWithStep extends SerialForWithStep, so lines 293-296 cover all cases

Consider removing the unreachable branch:

             if isinstance(it, UnrollForWithStep):
                 real_frame = tir.unroll(real_stop, annotations=it.annotations)
             elif isinstance(it, SerialForWithStep):
                 real_frame = tir.serial(real_stop, annotations=it.annotations)
-            else:
-                raise TypeError(
-                    f"Invalid for loop, got {it}({type(it)}), expect one of the following: "
-                    "range, T.serial, T.unroll, T.grid, T.parallel, T.vectorized, T.thread_binding")

Alternatively, if you want to keep defensive programming for future extensibility, use an assertion:

else:
    assert False, f"Unexpected for loop type: {type(it)}"

testing/python/language/test_tilelang_language_unroll.py (1)

20-32: Consider adding a test for combined step and unroll_factor.

Based on the unroll function in tilelang/language/loop.py, using both step and unroll_factor together should work (when step is not None). Consider adding a test case to verify this combination produces the expected #pragma unroll <factor>:

def test_unroll_with_step_and_factor():

    @T.prim_func
    def main(A_ptr: T.handle):
        A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)

        for _blockIdx in T.thread_binding(1, thread="blockIdx.x"):
            for _threadIdx in T.thread_binding(128, thread="threadIdx.x"):
                for i in T.unroll(0, 16, step=2, unroll_factor=4):
                    A[0, i] = 1.0

    kernel = tilelang.compile(main, target="cuda")
    assert "#pragma unroll 4" in kernel.get_kernel_source()

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b10ef75 and b42ec64.

📒 Files selected for processing (6)

• src/target/codegen_cuda.cc (2 hunks)
• src/target/codegen_cuda.h (1 hunks)
• testing/python/language/test_tilelang_language_unroll.py (1 hunks)
• tilelang/language/__init__.py (1 hunks)
• tilelang/language/loop.py (3 hunks)
• tilelang/language/v2/builder.py (3 hunks)

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-09-12T09:47:46.474Z

Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.

Applied to files:

• src/target/codegen_cuda.cc

🧬 Code graph analysis (3)

tilelang/language/v2/builder.py (2)

tilelang/language/loop.py (2)

• unroll (114-174)
• serial (97-111)

tilelang/language/tir/ir.py (2)

• unroll (84-113)
• serial (9-31)

tilelang/language/__init__.py (3)

tilelang/language/loop.py (5)

• Parallel (12-32)
• Persistent (35-54)
• Pipelined (57-94)
• serial (97-111)
• unroll (114-174)

tilelang/language/tir/ir.py (2)

• serial (9-31)
• unroll (84-113)

tilelang/language/ast/ir.py (2)

• serial (672-700)
• unroll (765-793)

tilelang/language/loop.py (2)

tilelang/language/v2/builder.py (2)

• SerialForWithStep (108-112)
• UnrollForWithStep (116-117)

tilelang/language/tir/ir.py (2)

• unroll (84-113)
• serial (9-31)

🪛 Ruff (0.14.6)

tilelang/language/v2/builder.py

298-300: Avoid specifying long messages outside the exception class

(TRY003)

tilelang/language/__init__.py

28-28: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

29-29: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

30-30: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

31-31: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

32-32: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

33-33: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

34-34: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

tilelang/language/loop.py

168-168: Avoid specifying long messages outside the exception class

(TRY003)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Quick Lint

🔇 Additional comments (7)

src/target/codegen_cuda.cc (1)

315-320: LGTM! Per-loop unroll factor emission is correctly implemented.

The logic correctly checks if the loop variable has an entry in unroll_factor and emits #pragma unroll <factor> when present, otherwise falls back to the plain #pragma unroll. The use of PrintExpr handles the IntImm value correctly.

tilelang/language/v2/builder.py (1)

115-118: LGTM! Clean extension for unroll support.

The UnrollForWithStep dataclass correctly extends SerialForWithStep without adding new fields, which allows sharing the stepped-loop handling logic while enabling type-based dispatch for frame creation.

tilelang/language/__init__.py (1)

27-35: LGTM! Public API surface correctly expanded.

The new exports (unroll, Serial, Unroll) are properly imported from .loop and follow the established pattern in this file. The # noqa: F401 comments are consistent with the rest of the file's style.

Note: The static analysis hints about unused noqa directives are because F401 checks are not enabled in the linter config. This is a false positive and can be safely ignored since the pattern is consistent throughout the file.

testing/python/language/test_tilelang_language_unroll.py (1)

1-35: LGTM! Good test coverage for the new unroll functionality.

The tests correctly verify:

1. T.unroll with step emits a plain #pragma unroll
2. T.unroll with unroll_factor emits #pragma unroll 4

Both tests follow the project's testing patterns and provide adequate coverage for the new feature.

tilelang/language/loop.py (3)

7-9: New imports and frame return typing look consistent

The new imports for SerialForWithStep, UnrollForWithStep, and frame align with their usage below (step-based control flow and frame.ForFrame return types). No issues here.

---

97-112: serial step-dispatch logic is reasonable

The serial helper cleanly preserves the default path (step is None or 1 → tb_tir.serial) while routing other cases through SerialForWithStep. The step_is_one detection for both int and IntImm is a nice touch, and reusing the existing start/stop semantics in the else branch keeps behavior predictable.

---

177-178: Aliases Serial/Unroll improve ergonomics

Re-exporting serial/unroll as Serial and Unroll is a nice ergonomic touch and keeps the public API consistent with the other loop helpers (e.g., Parallel, Persistent, Pipelined).

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Duplicate friend declaration detected.

The friend void PrintConst(...) declaration appears twice in this class (lines 81-82 and lines 144-145). The duplicate on lines 144-145 should be removed.

   std::unordered_map<const VarNode *, std::string> fragment_shapes;
   std::unordered_map<const VarNode *, std::string> fragment_layouts;
   std::unordered_map<const VarNode *, IntImm> unroll_factor;
-  friend void PrintConst(const FloatImmNode *op, std::ostream &os,
-                         CodeGenTileLangCUDA *p);
   void PrintWmmaScope(const std::string &scope, DataType t,

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	std::unordered_map<const VarNode *, IntImm> unroll_factor;
	friend void PrintConst(const FloatImmNode *op, std::ostream &os,
	CodeGenTileLangCUDA *p);
	std::unordered_map<const VarNode *, IntImm> unroll_factor;

🤖 Prompt for AI Agents

In src/target/codegen_cuda.h around lines 143-145, there is a duplicate friend
declaration "friend void PrintConst(const FloatImmNode *op, std::ostream &os,
CodeGenTileLangCUDA *p);" that already exists earlier (lines 81-82); delete the
duplicate declaration at lines 144-145 so the class only contains the single
friend declaration at its original location.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Fix step_is_one handling and align pragma_unroll_explicit validation/message

There are a couple of small but important inconsistencies in unroll:

1. step_is_one is never updated

   • step_is_one is initialized to False and then never modified, so step == 1 is never treated specially.
   • Given the symmetry with serial and the if step is None or step_is_one: branch, it looks like you intended to treat step == 1 as equivalent to the default and go through tb_tir.unroll.
   • Either compute step_is_one like in serial or drop it entirely. If you want the symmetry, something like this would fix it:

   def unroll(start: tir.PrimExpr,
              stop: tir.PrimExpr | None = None,
              step: tir.PrimExpr | None = None,
              *,
              explicit: bool = False,
              unroll_factor: int | None = None,
              annotations: dict[str, Any] | None = None) -> frame.ForFrame:
   @@

• step_is_one = False

• step_is_one = False
• step_is_one |= isinstance(step, int) and step == 1
• step_is_one |= isinstance(step, IntImm) and getattr(step, "value", None) == 1

2. **Comment vs error message for `pragma_unroll_explicit` contradict**  
- Line 166 comment says: `# check pragma_unroll_explicit must be False`, but the error message says `"pragma_unroll_explicit must be True when unroll_factor is not None"`.  
- The condition `if annotations.get("pragma_unroll_explicit", True): raise ...` currently *forbids* `unroll_factor` when `pragma_unroll_explicit` is `True`, which matches the comment (must be False to be allowed) but contradicts the error text.  
- To avoid confusing users, make the condition and message agree and describe the actual rule. For example, if the intent is “you cannot combine an explicit unroll with `unroll_factor`”, consider:

```diff
-    if unroll_factor is not None:
-        # check pragma_unroll_explicit must be False
-        if annotations.get("pragma_unroll_explicit", True):
-            raise ValueError("pragma_unroll_explicit must be True when unroll_factor is not None")
-        annotations.update({"pragma_unroll_factor": unroll_factor})
+    if unroll_factor is not None:
+        # Disallow combining explicit unrolling with an unroll factor.
+        if annotations.get("pragma_unroll_explicit", False):
+            raise ValueError(
+                "unroll_factor cannot be used when pragma_unroll_explicit is True"
+            )
+        annotations.update({"pragma_unroll_factor": unroll_factor})

This also keeps Ruff’s TRY003 warning in check by using a concise, targeted message.

3. Minor doc / comment mismatch (optional)
   • The comment above the annotations block still mentions “{"pragma_unroll_explicit": True} by default” while the code actually uses the explicit flag (defaulting to False). Updating the comment to match the current behavior would avoid confusion for future readers.

Overall structure of unroll looks good; tightening up these details will make the behavior and error reporting much clearer.

🧰 Tools

🪛 Ruff (0.14.6)

168-168: Avoid specifying long messages outside the exception class

(TRY003)

🤖 Prompt for AI Agents

In tilelang/language/loop.py around lines 114-175, compute step_is_one (e.g.,
set step_is_one = (step == 1 or (hasattr(step, "dtype") and int(getattr(step,
"value", 1)) == 1)) or mirror the logic used in serial) so that step == 1 is
treated like step is None and routed to tb_tir.unroll; change the
pragma_unroll_explicit validation so the condition and message agree — if the
intent is to forbid combining explicit unrolling with an unroll_factor keep the
existing check but change the ValueError text to "pragma_unroll_explicit must be
False when unroll_factor is not None" (or invert both condition and message if
you prefer the opposite rule), and update the preceding comment to correctly
state the default annotation behavior (use explicit param) so comments,
condition, and error message are consistent.

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Another round soon, please!

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (4)

src/target/codegen_cuda.h (1)

141-145: Duplicate friend declaration detected.

The friend void PrintConst(...) declaration appears twice in this class—once at lines 81-82 and again at lines 144-145. This is redundant and likely a merge artifact.

Consider removing the duplicate declaration at lines 144-145:

  std::unordered_map<const VarNode *, std::string> fragment_shapes;
  std::unordered_map<const VarNode *, std::string> fragment_layouts;
  std::unordered_map<const VarNode *, IntImm> unroll_factor;
-  friend void PrintConst(const FloatImmNode *op, std::ostream &os,
-                         CodeGenTileLangCUDA *p);
  void PrintWmmaScope(const std::string &scope, DataType t,

tilelang/language/v2/builder.py (1)

293-300: Unreachable else branch - consider simplifying.

The else branch at lines 297-300 is logically unreachable. The outer isinstance check at line 278 guarantees it is either SerialForWithStep or UnrollForWithStep. Since UnrollForWithStep inherits from SerialForWithStep, the inner checks at lines 293-296 will always match one of these types.

While this serves as defensive coding, you could simplify by removing the else branch or restructuring:

             if isinstance(it, UnrollForWithStep):
                 real_frame = tir.unroll(real_stop, annotations=it.annotations)
-            elif isinstance(it, SerialForWithStep):
+            else:
+                # Must be SerialForWithStep (the base type)
                 real_frame = tir.serial(real_stop, annotations=it.annotations)
-            else:
-                raise TypeError(
-                    f"Invalid for loop, got {it}({type(it)}), expect one of the following: "
-                    "range, T.serial, T.unroll, T.grid, T.parallel, T.vectorized, T.thread_binding")

testing/python/language/test_tilelang_language_unroll.py (1)

1-37: Consider adding edge case tests.

The current tests cover the happy path. Consider adding tests for:

1. unroll_factor=1 to verify single-iteration unroll
2. Negative test: verify that explicit=True with unroll_factor raises ValueError (as per loop.py validation)

def test_unroll_factor_with_explicit_should_raise():
    """Verify that explicit=True with unroll_factor raises ValueError."""
    import pytest
    
    with pytest.raises(ValueError, match="pragma_unroll_explicit"):
        @T.prim_func
        def main(A_ptr: T.handle):
            A = T.match_buffer(A_ptr, (16, 16), dtype="float32")
            for _blockIdx in T.thread_binding(1, thread="blockIdx.x"):
                for _threadIdx in T.thread_binding(128, thread="threadIdx.x"):
                    for i in T.unroll(0, 16, explicit=True, unroll_factor=4):
                        A[0, i] = 1.0

Would you like me to generate additional edge case tests or open an issue to track this?

tilelang/language/loop.py (1)

137-146: Consider validating unroll_factor and documenting the constraint

Right now any int is accepted for unroll_factor, including 0 or negative values, which will likely be meaningless or rejected downstream by CUDA pragmas.

Consider adding a lightweight check and reflecting it in the docstring:

-    unroll_factor : int
-        The unroll factor of the loop.
+    unroll_factor : int
+        The unroll factor of the loop. Must be a positive integer.
@@
-    if unroll_factor is not None:
+    if unroll_factor is not None:
+        if unroll_factor <= 0:
+            raise ValueError("unroll_factor must be a positive integer")
         # require non‑explicit unroll when using a factor
         ...

This keeps misuse localized at the API boundary instead of surfacing later in codegen.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b10ef75 and d1330f7.

📒 Files selected for processing (6)

• src/target/codegen_cuda.cc (2 hunks)
• src/target/codegen_cuda.h (1 hunks)
• testing/python/language/test_tilelang_language_unroll.py (1 hunks)
• tilelang/language/__init__.py (1 hunks)
• tilelang/language/loop.py (3 hunks)
• tilelang/language/v2/builder.py (3 hunks)

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-09-12T09:47:46.474Z

Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.

Applied to files:

• src/target/codegen_cuda.cc

🧬 Code graph analysis (4)

tilelang/language/v2/builder.py (2)

tilelang/language/loop.py (2)

• unroll (114-174)
• serial (97-111)

tilelang/language/tir/ir.py (2)

• unroll (84-113)
• serial (9-31)

testing/python/language/test_tilelang_language_unroll.py (1)

tilelang/language/loop.py (1)

• unroll (114-174)

tilelang/language/__init__.py (2)

tilelang/language/loop.py (5)

• Parallel (12-32)
• Persistent (35-54)
• Pipelined (57-94)
• serial (97-111)
• unroll (114-174)

tilelang/language/tir/ir.py (2)

• serial (9-31)
• unroll (84-113)

tilelang/language/loop.py (2)

tilelang/language/v2/builder.py (1)

• SerialForWithStep (108-112)

tilelang/language/tir/ir.py (2)

• unroll (84-113)
• serial (9-31)

🪛 Ruff (0.14.6)

tilelang/language/v2/builder.py

298-300: Avoid specifying long messages outside the exception class

(TRY003)

tilelang/language/__init__.py

28-28: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

29-29: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

30-30: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

31-31: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

32-32: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

33-33: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

34-34: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

tilelang/language/loop.py

168-168: Avoid specifying long messages outside the exception class

(TRY003)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (8)

src/target/codegen_cuda.cc (2)

2669-2675: LGTM! The pragma_unroll_factor attribute handling is correctly implemented.

The flow correctly captures the unroll factor from the attribute and stores it in the map before the corresponding ForNode is visited. The ICHECK validates the factor is an IntImmNode before the downcast.

---

312-321: LGTM! Per-loop unroll pragma emission looks correct.

The conditional emission of #pragma unroll with or without a factor correctly handles both cases. The logic properly uses PrintExpr to serialize the IntImm factor value.

tilelang/language/__init__.py (1)

27-35: LGTM! The new loop exports are correctly added.

The consolidated import block properly exports the new serial, unroll, Serial, and Unroll symbols alongside the existing Parallel, Persistent, and Pipelined exports.

Regarding the static analysis hints about unused # noqa: F401 directives—these are false positives. The directives are valid for __init__.py re-exports, but Ruff may not have the F401 rule enabled in the project configuration. You can safely ignore these warnings or remove the noqa comments if F401 is not enforced.

tilelang/language/v2/builder.py (1)

115-118: LGTM! Clean dataclass inheritance.

The UnrollForWithStep dataclass correctly extends SerialForWithStep, inheriting all fields (start, stop, step, annotations). This is a clean way to differentiate unroll loops from serial loops while reusing the same structure.

testing/python/language/test_tilelang_language_unroll.py (2)

6-18: LGTM! Test correctly validates step-based unrolling.

The test verifies that T.unroll(0, 16, step=4) generates a loop with #pragma unroll in the CUDA output. The assertion correctly checks for the presence of the unroll pragma without a specific factor, which aligns with the implementation in loop.py where step creates an UnrollForWithStep without setting pragma_unroll_factor.

---

21-33: LGTM! Test correctly validates factor-based unrolling.

The test verifies that T.unroll(0, 16, unroll_factor=4) generates a loop with #pragma unroll 4 in the CUDA output. This exercises the full flow through pragma_unroll_factor annotation → codegen attribute handling → pragma emission.

tilelang/language/loop.py (2)

97-112: serial changes and type annotation look consistent

The serial(...) helper’s signature and return type annotation (-> frame.ForFrame) align with the new unroll(...) function and with the SerialForWithStep path for stepped loops. The step_is_one handling matches the intended semantics and doesn’t need further changes.

---

7-10: Imports and public aliases are coherent

Importing UnrollForWithStep alongside SerialForWithStep, introducing frame.ForFrame for type hints, and exposing Serial/Unroll aliases all look consistent with the existing language API surface.

Also applies to: 177-178

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Fix step_is_one handling so step=1 uses the simple unroll path

Right now step_is_one is always False, so step=1 will go through UnrollForWithStep instead of the plain tb_tir.unroll path, unlike serial(...). This is almost certainly unintended and diverges from the serial helper’s behavior.

You can mirror the serial implementation by actually computing step_is_one:

-    step_is_one = False
-    if stop is None:
-        stop = start
-        if hasattr(start, "dtype"):
-            start = IntImm(start.dtype, 0)
-        else:
-            start = 0
-
-    # Ensure annotations has {"pragma_unroll_explicit": True} by default
+    if stop is None:
+        stop = start
+        if hasattr(start, "dtype"):
+            start = IntImm(start.dtype, 0)
+        else:
+            start = 0
+
+    step_is_one = False
+    step_is_one |= isinstance(step, int) and step == 1
+    step_is_one |= isinstance(step, IntImm) and getattr(step, "value", None) == 1
+
+    # Ensure annotations has {"pragma_unroll_explicit": explicit} by default
@@
-    if step is None or step_is_one:
-        return tb_tir.unroll(start, stop, annotations=annotations)
-    else:
-        return UnrollForWithStep(start, stop, step, annotations=annotations)
+    if step is None or step_is_one:
+        return tb_tir.unroll(start, stop, annotations=annotations)
+    else:
+        return UnrollForWithStep(start, stop, step, annotations=annotations)

This keeps step=None/step=1 on the simple path and only uses UnrollForWithStep when the step is genuinely non‑unit.

🧰 Tools

🪛 Ruff (0.14.6)

168-168: Avoid specifying long messages outside the exception class

(TRY003)

🤖 Prompt for AI Agents

In tilelang/language/loop.py around lines 149 to 175, the variable step_is_one
is left False so step=1 incorrectly bypasses the simple unroll path; compute
step_is_one based on the passed step (mirror serial: treat None or a unit
constant/IntImm of value 1 as "one") before the branch so that step=None or
step==1 use tb_tir.unroll, and only non‑unit steps go to UnrollForWithStep; also
ensure the dtype/IntImm case is handled consistently with how start was
normalized.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Clarify pragma_unroll_explicit constraint and align error message/comment

The comment says pragma_unroll_explicit “must be False” when unroll_factor is set, but:

• The code enforces exactly that (it raises when the flag is True), and
• The error message says “must be True”, which contradicts both the comment and the actual condition and triggers TRY003 due to length.

Consider tightening this block:

-    # Ensure annotations has {"pragma_unroll_explicit": True} by default
+    # Ensure annotations has {"pragma_unroll_explicit": explicit} by default
@@
-    if unroll_factor is not None:
-        # check pragma_unroll_explicit must be False
-        if annotations.get("pragma_unroll_explicit", True):
-            raise ValueError("pragma_unroll_explicit must be True when unroll_factor is not None")
-        annotations.update({"pragma_unroll_factor": unroll_factor})
+    if unroll_factor is not None:
+        # require non‑explicit unroll when using a factor
+        if annotations.get("pragma_unroll_explicit", False):
+            raise ValueError("unroll_factor requires pragma_unroll_explicit=False")
+        annotations["pragma_unroll_factor"] = unroll_factor

This makes the requirement obvious, fixes the contradictory message, and shortens the exception text to satisfy the TRY003 hint.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	# Ensure annotations has {"pragma_unroll_explicit": True} by default
	if annotations is None:
	annotations = {"pragma_unroll_explicit": explicit}
	else:
	# Add "pragma_unroll_explicit": True if not already present
	annotations = dict(annotations)
	annotations.setdefault("pragma_unroll_explicit", explicit)
	if unroll_factor is not None:
	# check pragma_unroll_explicit must be False
	if annotations.get("pragma_unroll_explicit", True):
	raise ValueError("pragma_unroll_explicit must be True when unroll_factor is not None")
	annotations.update({"pragma_unroll_factor": unroll_factor})
	# Ensure annotations has {"pragma_unroll_explicit": explicit} by default
	if annotations is None:
	annotations = {"pragma_unroll_explicit": explicit}
	else:
	# Add "pragma_unroll_explicit": True if not already present
	annotations = dict(annotations)
	annotations.setdefault("pragma_unroll_explicit", explicit)
	if unroll_factor is not None:
	# require non‑explicit unroll when using a factor
	if annotations.get("pragma_unroll_explicit", False):
	raise ValueError("unroll_factor requires pragma_unroll_explicit=False")
	annotations["pragma_unroll_factor"] = unroll_factor

🧰 Tools

🪛 Ruff (0.14.6)

168-168: Avoid specifying long messages outside the exception class

(TRY003)

🤖 Prompt for AI Agents

In tilelang/language/loop.py around lines 157 to 170, the inline comment and the
raised ValueError contradict each other: the code currently raises when
annotations.get("pragma_unroll_explicit", True) is True but the error text says
"must be True" and the comment wording is unclear; update the comment to state
clearly that pragma_unroll_explicit must be False when unroll_factor is
provided, and change the ValueError text to a short, correct message like
"pragma_unroll_explicit must be False" (keep it brief to satisfy TRY003).

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

tilelang/language/__init__.py (1)

27-35: Ruff RUF100: unused # noqa: F401 directives on new loop imports

Static analysis reports RUF100 on these lines because F401 is disabled in the Ruff config, so # noqa: F401 is considered unnecessary here.

To keep linters clean, either remove the # noqa: F401 from these new imports or adjust the Ruff configuration if you do want to suppress F401 globally for re-export imports. For this PR, the minimal change is to drop the noqa comments on lines 28–34:

-from .loop import (
-    Parallel,  # noqa: F401
-    Persistent,  # noqa: F401
-    Pipelined,  # noqa: F401
-    serial,  # noqa: F401
-    unroll,  # noqa: F401
-    Serial,  # noqa: F401
-    Unroll,  # noqa: F401
-)
+from .loop import (
+    Parallel,
+    Persistent,
+    Pipelined,
+    serial,
+    unroll,
+    Serial,
+    Unroll,
+)

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d1330f7 and 9a1a674.

📒 Files selected for processing (3)

• examples/gemm_sp/example_custom_compress.py (1 hunks)
• tilelang/language/__init__.py (1 hunks)
• tilelang/language/experimental/gemm_sp.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

tilelang/language/__init__.py (3)

tilelang/language/loop.py (5)

• Parallel (12-32)
• Persistent (35-54)
• Pipelined (57-94)
• serial (97-111)
• unroll (114-174)

tilelang/language/tir/ir.py (2)

• serial (9-31)
• unroll (84-113)

tilelang/language/ast/ir.py (2)

• serial (672-700)
• unroll (765-793)

🪛 Ruff (0.14.6)

tilelang/language/__init__.py

28-28: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

29-29: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

30-30: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

31-31: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

32-32: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

33-33: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

---

34-34: Unused noqa directive (non-enabled: F401)

Remove unused noqa directive

(RUF100)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: Test for Python 3.12 with ROCm-6.3 (on self-hosted-amd)
• GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (3)

tilelang/language/__init__.py (1)

27-35: Public re-exports of serial/unroll/Serial/Unroll look consistent, but they now shadow tir.ir versions

Importing serial and unroll from .loop here will override the serial/unroll symbols brought in earlier via from .tir.ir import *, and adds Serial/Unroll to the public surface. That matches the new higher-level loop wrappers, but it is a behavior change for tilelang.language.serial/unroll (new signature and semantics like step, explicit, unroll_factor).

Please double‑check that all internal and external call sites are expected to use the new loop-level behavior and that there are no users depending on the previous tir.ir implementation details.

tilelang/language/experimental/gemm_sp.py (1)

201-201: Verify the renamed intrinsic exists and all references are updated.

The intrinsic operator name has been changed to "tl.tileop.gemm_sp_py" in the gemm_sp_v2 function. Please confirm:

1. The new intrinsic "tl.tileop.gemm_sp_py" is registered and implemented
2. All references to the previous intrinsic name have been updated consistently across the codebase
3. The naming difference between gemm_sp (uses "tl.tileop.gemm_sp") and gemm_sp_v2 (uses "tl.tileop.gemm_sp_py") is intentional

Additionally, verify that the enriched summary accurately reflects the changes—only gemm_sp_v2 was modified, not both functions.

examples/gemm_sp/example_custom_compress.py (1)

272-300: Per‑group reset of non_zero_cnt / non_zero_elt_log_idx looks correct

Re‑initializing non_zero_cnt[0] and the full non_zero_elt_log_idx vector once per (tm, g_i) group (before scanning the 4 values) is the right scope: it prevents stale indices from leaking across groups and guarantees defined values for the later fixup/metadata logic even when there are 0 or 1 non‑zeros in the group. Given elem = 2, the explicit loop is also trivial cost-wise.

No issues from my side here.