Skip to content

Conversation

@kurisu6912
Copy link
Collaborator

@kurisu6912 kurisu6912 commented Nov 4, 2025

This pr add support step and negative step in T.serial. But due to T.ceildiv allows non-negative variables only, negative step may cause undefined behavior.

I check whether it is a constant, and show warning when it may cause undefined behavior:

real_stop = tir.ceildiv(it.stop - it.start, it.step)
if isinstance(it.step, (int, IntImm)):
    value = it.step if isinstance(it.step, int) else it.step.value
    if value < 0:
        real_stop = tir.ceildiv(it.start - it.stop, -it.step)
else:
    logger.warning(
        f'Using a non-constant step `{it.step}` in stepped serial may lead to undefined behavior in tilelang'
    )

Summary by CodeRabbit

  • New Features

    • Support for stepped serial loops (positive and negative step values) in loop constructs.
  • Refactor

    • Loop-related constructors consolidated under a unified loop API, exposing a single entry for serial/parallel/pipelined/persistent patterns.
  • Tests

    • New tests validating stepped serial loop behavior and kernel results against expected references.

@github-actions
Copy link

github-actions bot commented Nov 4, 2025

👋 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
Copy link
Contributor

coderabbitai bot commented Nov 4, 2025

Note

Other AI code review bot(s) detected

CodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review.

Walkthrough

This PR consolidates loop constructors into a new tilelang/language/loop.py, exposes serial alongside Parallel/Persistent/Pipelined via tilelang.language, adds stepped-serial support (SerialForWithStep) in the v2 builder, and introduces tests for serial loops with non-unit steps.

Changes

Cohort / File(s) Summary
Loop consolidation (new)
tilelang/language/loop.py
New module implementing Parallel, Persistent, Pipelined, and serial helpers; normalizes args and delegates to underlying FFI or tir builders; serial routes to stepped or standard serial implementations based on step.
Top-level API
tilelang/language/__init__.py
Replaces separate imports with consolidated imports from .loop; adds serial to public exports and sources Parallel/Persistent/Pipelined from .loop.
Removed legacy modules
tilelang/language/parallel.py, tilelang/language/persistent.py, tilelang/language/pipeline.py
Deleted previous per-loop modules; their functionality moved into loop.py.
v2 builder: stepped serial
tilelang/language/v2/builder.py
Adds SerialForWithStep dataclass and extends ctx_for to handle stepped serial loops by computing iteration count via ceiling division (handles positive/negative/ symbolic steps); Range override now uses tilelang.language.serial.
Tests: serial step behavior
testing/python/language/test_tilelang_language_frontend_v2.py
Adds test_serial_for_with_step() defining kernels with positive and negative step serial loops (wrapped with tilelang.jit) and runtime assertions checking tensor results and IRBuilderFrame behavior for various T.serial call patterns.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant tilelang_language_serial as serial()
    participant v2_ctx_for as ctx_for
    participant tir_serial as tir.serial / tir.frame

    User->>tilelang_language_serial: call serial(start, stop, step, annotations)
    activate tilelang_language_serial

    alt step is 1 or None
        tilelang_language_serial->>tir_serial: return standard tir.serial frame
        tir_serial-->>tilelang_language_serial: ForFrame
    else step != 1
        tilelang_language_serial->>v2_ctx_for: SerialForWithStep(start,stop,step,annotations)
        activate v2_ctx_for
        v2_ctx_for->>v2_ctx_for: compute real_stop = ceildiv(|stop-start|, |step|)
        v2_ctx_for->>tir_serial: tir.serial(real_stop, annotations)
        tir_serial-->>v2_ctx_for: serial_frame
        v2_ctx_for-->>tilelang_language_serial: transformed loop yield (start + v*step)
        deactivate v2_ctx_for
    end

    tilelang_language_serial-->>User: loop frame (ForFrame or transformed serial)
    deactivate tilelang_language_serial
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

  • Pay attention to: ceiling-division correctness for positive/negative steps and zero-step handling.
  • Verify symbolic-step path and warning semantics in ctx_for.
  • Confirm re-exports in __init__.py don't introduce import cycles.
  • Review tests for completeness and correct CUDA invocation/assertions.

Possibly related PRs

Suggested reviewers

  • LeiWang1999
  • XuehaiPan

Poem

🐰
I hopped through loops both short and long,
Collected steps to craft this song.
From start to stop, in twos or one,
I spin the code until it's done —
A rabbit's cheer for loops made strong.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% 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 '[Feat] Add support for T.serial with step and negative step' clearly and concisely describes the main feature addition in the changeset.
✨ 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

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

@kurisu6912 kurisu6912 linked an issue Nov 4, 2025 that may be closed by this pull request
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (2)
testing/python/language/test_tilelang_language_frontend_v2.py (1)

284-284: Use integer literals for int32 tensor assignments.

Lines 284 and 287 assign float literals (1.0, 2.0) to an int32 tensor. While TVM likely handles the implicit conversion, using integer literals (1, 2) improves clarity and avoids potential type confusion.

Apply this diff:

             for i in range(0, 10, 2):
                 T.device_assert(0 <= i < 10 and i % 2 == 0, "i out of range")
-                A[i] = 1.0
+                A[i] = 1
             for i in range(1, 10, 2):
                 T.device_assert(1 <= i < 10 and i % 2 == 1, "i out of range")
-                A[i] = 2.0
+                A[i] = 2

Also applies to: 287-287

tilelang/language/__init__.py (1)

26-26: Remove unused noqa directive.

The noqa: F401 directive is unnecessary as the F401 warning is not being raised for this import.

Apply this diff:

-from .loop import serial, Parallel, Persistent, Pipelined  # noqa: F401
+from .loop import serial, Parallel, Persistent, Pipelined

Based on static analysis.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 778b97d and 5a6204e.

📒 Files selected for processing (7)
  • testing/python/language/test_tilelang_language_frontend_v2.py (1 hunks)
  • tilelang/language/__init__.py (1 hunks)
  • tilelang/language/loop.py (1 hunks)
  • tilelang/language/parallel.py (0 hunks)
  • tilelang/language/persistent.py (0 hunks)
  • tilelang/language/pipeline.py (0 hunks)
  • tilelang/language/v2/builder.py (3 hunks)
💤 Files with no reviewable changes (3)
  • tilelang/language/parallel.py
  • tilelang/language/pipeline.py
  • tilelang/language/persistent.py
🧰 Additional context used
🧬 Code graph analysis (4)
testing/python/language/test_tilelang_language_frontend_v2.py (3)
tilelang/jit/__init__.py (3)
  • jit (275-276)
  • jit (280-291)
  • jit (294-361)
tilelang/language/v2/builder.py (2)
  • prim_func (144-148)
  • prim_func (594-687)
tilelang/language/print.py (1)
  • device_assert (144-155)
tilelang/language/__init__.py (3)
tilelang/language/loop.py (1)
  • serial (97-108)
tilelang/language/tir/ir.py (1)
  • serial (10-32)
tilelang/language/ast/ir.py (1)
  • serial (672-700)
tilelang/language/v2/builder.py (1)
tilelang/language/loop.py (1)
  • serial (97-108)
tilelang/language/loop.py (1)
tilelang/language/v2/builder.py (1)
  • SerialForWithStep (99-103)
🪛 Ruff (0.14.3)
tilelang/language/__init__.py

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

Remove unused noqa directive

(RUF100)

tilelang/language/v2/builder.py

263-265: 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 (6)
tilelang/language/v2/builder.py (3)

98-103: LGTM!

The SerialForWithStep dataclass is well-defined with appropriate types for representing a serial loop with a custom step.


247-267: LGTM!

The ctx_for logic correctly handles SerialForWithStep for both positive and negative stepping:

  • For positive steps: real_stop = ceildiv(stop - start, step) correctly computes iteration count
  • For negative constant steps: adjusts real_stop = ceildiv(start - stop, -step)
  • Yields the transformed value start + loop_var * step to produce the expected iteration sequence

The warning for non-constant steps is also appropriate.


475-478: LGTM!

The updated override method correctly routes Python's range to T.serial, which now supports stepped iteration. The local import avoids potential circular dependency issues.

tilelang/language/loop.py (3)

12-32: LGTM!

The Parallel function is well-implemented with clear documentation and appropriate handling of the optional coalesced_width annotation.


35-54: LGTM!

The Persistent function is straightforward and well-documented.


57-94: LGTM!

The Pipelined function correctly normalizes start/stop and handles optional list parameters with appropriate defaults.

@kurisu6912
Copy link
Collaborator Author

@codex review

@chatgpt-codex-connector
Copy link

To use Codex here, create a Codex account and connect to github.

@kurisu6912
Copy link
Collaborator Author

@codex review

@chatgpt-codex-connector
Copy link

To use Codex here, create a Codex account and connect to github.

LeiWang1999
LeiWang1999 previously approved these changes Nov 5, 2025
@LeiWang1999
Copy link
Member

@codex review

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ 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".

… step values and prevent undefined behavior. Added error handling for zero step values and improved logging for non-constant steps.
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

♻️ Duplicate comments (1)
tilelang/language/v2/builder.py (1)

254-267: [Duplicate] Zero step validation still missing.

This issue was flagged in a previous review. The code computes tir.ceildiv(it.stop - it.start, it.step) at line 255 without verifying that it.step is non-zero. When step=0, this causes division by zero before any loop is created. Python's range(0, 10, 0) raises ValueError: range() arg 3 must not be zero, so the DSL should validate and reject zero step similarly.

Apply this diff to add zero-step validation:

 def ctx_for(self, it):
     self.check_continue_break()
     it = unwrap_expr(it)
     if isinstance(it, SerialForWithStep):
+        # Validate non-zero step
+        if isinstance(it.step, (int, IntImm)):
+            value = it.step if isinstance(it.step, int) else it.step.value
+            if value == 0:
+                raise ValueError("Serial loop step must not be zero")
         real_stop = tir.ceildiv(it.stop - it.start, it.step)
         if isinstance(it.step, (int, IntImm)):
             value = it.step if isinstance(it.step, int) else it.step.value

Additionally, verify the ceil division logic for edge cases:

#!/bin/bash
# Check for any existing zero-step validation in the codebase
rg -nP "step.*==\s*0|step.*zero|ValueError.*step" --type=py

# Search for range() calls that might be affected
rg -nP "range\s*\([^)]*,\s*[^)]*,\s*[^)]*\)" --type=py -A2 -B2
🧹 Nitpick comments (2)
tilelang/language/v2/builder.py (1)

269-272: Consider extracting the error message to a module-level constant.

The error message is quite long and static analysis (Ruff TRY003) suggests extracting it to improve maintainability.

+_INVALID_FOR_LOOP_ERROR = (
+    "Invalid for loop, got {it}({type}), expect one of the following: "
+    "range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding"
+)
+
 def ctx_for(self, it):
     ...
     else:
         if not isinstance(it, tir.frame.ForFrame):
-            raise TypeError(
-                f"Invalid for loop, got {it}({type(it)}), expect one of the following: "
-                "range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding")
+            raise TypeError(_INVALID_FOR_LOOP_ERROR.format(it=it, type=type(it)))
testing/python/language/test_tilelang_language_frontend_v2.py (1)

278-313: Good test coverage, but consider adding zero-step validation test.

The tests effectively cover positive steps, negative steps, and frame type assertions. The use of device_assert within loops validates that iteration variables stay in expected ranges.

Consider adding a test case for step=0, which should raise a ValueError similar to Python's range(0, 10, 0):

def test_serial_zero_step_error():
    """Verify that zero step is rejected."""
    import pytest
    
    @T.prim_func
    def invalid_zero_step(A: T.Tensor((10,), T.int32)):
        with T.Kernel(1) as _:
            for i in range(0, 10, 0):  # Should raise ValueError
                A[i] = i
    
    with pytest.raises(ValueError, match="step must not be zero"):
        invalid_zero_step
📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5a6204e and c7c9a49.

📒 Files selected for processing (3)
  • testing/python/language/test_tilelang_language_frontend_v2.py (2 hunks)
  • tilelang/language/loop.py (1 hunks)
  • tilelang/language/v2/builder.py (3 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
testing/python/language/test_tilelang_language_frontend_v2.py (2)
tilelang/language/print.py (1)
  • device_assert (144-155)
tilelang/language/loop.py (1)
  • serial (97-111)
tilelang/language/loop.py (1)
tilelang/language/v2/builder.py (1)
  • SerialForWithStep (104-108)
tilelang/language/v2/builder.py (1)
tilelang/language/loop.py (1)
  • serial (97-111)
🪛 Ruff (0.14.3)
tilelang/language/v2/builder.py

270-272: 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 (3)
tilelang/language/v2/builder.py (2)

103-108: LGTM! Clean dataclass definition.

The SerialForWithStep dataclass is well-structured and provides the necessary fields to represent a serial loop with a custom step parameter.


485-487: LGTM! Correct routing to the new serial API.

The override correctly routes range() to T.serial, which now supports the step parameter. The import is appropriately scoped within the method.

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

6-7: LGTM! Necessary imports for serial step testing.

The imports support runtime assertions about frame types and symbolic variables in the new test.

Comment on lines +97 to +111
def serial(start: tir.PrimExpr,
stop: tir.PrimExpr | None = None,
step: tir.PrimExpr | None = None,
*,
annotations: dict[str, Any] | None = None):
step_is_one = False
step_is_one |= isinstance(step, int) and step == 1
step_is_one |= isinstance(step, IntImm) and step.value == 1
if step is None or step_is_one:
return tb_tir.serial(start, stop, annotations=annotations)
else:
if stop is None:
stop = start
start = IntImm(start.dtype, 0) if hasattr(start, "dtype") else 0
return SerialForWithStep(start, stop, step, annotations=annotations)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Missing zero-step validation.

The function correctly routes step values to either tb_tir.serial (for step=None or step=1) or SerialForWithStep, and properly normalizes start/stop. However, zero step is not validated before creating SerialForWithStep, which will cause division by zero in builder.py::ctx_for line 255.

Apply this diff to validate step:

 def serial(start: tir.PrimExpr,
            stop: tir.PrimExpr | None = None,
            step: tir.PrimExpr | None = None,
            *,
            annotations: dict[str, Any] | None = None):
+    # Validate non-zero step for constant values
+    if isinstance(step, int) and step == 0:
+        raise ValueError("Serial loop step must not be zero")
+    if isinstance(step, IntImm) and step.value == 0:
+        raise ValueError("Serial loop step must not be zero")
+    
     step_is_one = False
     step_is_one |= isinstance(step, int) and step == 1
     step_is_one |= isinstance(step, IntImm) and step.value == 1

Minor: Consider using logical or instead of bitwise |=.

While lines 103-104 work correctly, the bitwise OR pattern is unconventional for boolean accumulation. Consider:

-    step_is_one = False
-    step_is_one |= isinstance(step, int) and step == 1
-    step_is_one |= isinstance(step, IntImm) and step.value == 1
+    step_is_one = (isinstance(step, int) and step == 1) or \
+                  (isinstance(step, IntImm) and step.value == 1)
📝 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
def serial(start: tir.PrimExpr,
stop: tir.PrimExpr | None = None,
step: tir.PrimExpr | None = None,
*,
annotations: dict[str, Any] | None = None):
step_is_one = False
step_is_one |= isinstance(step, int) and step == 1
step_is_one |= isinstance(step, IntImm) and step.value == 1
if step is None or step_is_one:
return tb_tir.serial(start, stop, annotations=annotations)
else:
if stop is None:
stop = start
start = IntImm(start.dtype, 0) if hasattr(start, "dtype") else 0
return SerialForWithStep(start, stop, step, annotations=annotations)
def serial(start: tir.PrimExpr,
stop: tir.PrimExpr | None = None,
step: tir.PrimExpr | None = None,
*,
annotations: dict[str, Any] | None = None):
# Validate non-zero step for constant values
if isinstance(step, int) and step == 0:
raise ValueError("Serial loop step must not be zero")
if isinstance(step, IntImm) and step.value == 0:
raise ValueError("Serial loop step must not be zero")
step_is_one = (isinstance(step, int) and step == 1) or \
(isinstance(step, IntImm) and step.value == 1)
if step is None or step_is_one:
return tb_tir.serial(start, stop, annotations=annotations)
else:
if stop is None:
stop = start
start = IntImm(start.dtype, 0) if hasattr(start, "dtype") else 0
return SerialForWithStep(start, stop, step, annotations=annotations)
🤖 Prompt for AI Agents
In tilelang/language/loop.py around lines 97 to 111, the function lacks
validation for a zero step which will cause a division-by-zero downstream; add a
guard before creating SerialForWithStep that checks if step represents zero
(e.g., isinstance(step, int) and step == 0 or isinstance(step, IntImm) and
step.value == 0) and raise a ValueError with a clear message if so; also replace
the bitwise accumulation using '|=' on step_is_one with conventional boolean
logic using 'or' (e.g., step_is_one = (isinstance(step, int) and step == 1) or
(isinstance(step, IntImm) and step.value == 1)) so the intent is clearer.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tilelang/language/v2/builder.py (1)

265-268: Optional: Clarify why non-constant steps cause undefined behavior.

The warning correctly alerts users but could be more explicit about the root cause.

Consider enhancing the warning message:

                logger.warning(
-                   f'Using a non-constant step `{it.step}` in stepped serial may lead to undefined behavior in tilelang'
+                   f'Using a non-constant step `{it.step}` in stepped serial may lead to undefined behavior in tilelang. '
+                   'T.ceildiv requires non-negative operands; negative runtime step values will produce incorrect iteration counts.'
                )
📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c7c9a49 and 3d10d99.

📒 Files selected for processing (1)
  • tilelang/language/v2/builder.py (3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tilelang/language/v2/builder.py (1)
tilelang/language/loop.py (1)
  • serial (97-111)
🪛 Ruff (0.14.3)
tilelang/language/v2/builder.py

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

(TRY003)


275-277: Avoid specifying long messages outside the exception class

(TRY003)

🔇 Additional comments (4)
tilelang/language/v2/builder.py (4)

103-109: LGTM: Clean dataclass definition for stepped serial loops.

The SerialForWithStep dataclass appropriately encapsulates the loop parameters needed for serial iteration with custom step values.


258-259: LGTM: Zero-step validation correctly prevents division by zero.

The check properly rejects zero steps for constant values, addressing the concern from previous reviews.


274-279: LGTM: Existing for-loop handling preserved.

The else branch correctly maintains backward compatibility for non-stepped serial loops and other loop constructs.


490-493: Import path verified and correct.

The from tilelang.language import serial at line 490 correctly imports the serial function from tilelang/language/loop.py, which is properly exported through tilelang/language/__init__.py. The change integrates the stepped serial loop support as intended.

Comment on lines +256 to +263
if isinstance(it.step, (int, IntImm)):
step_value = it.step if isinstance(it.step, int) else it.step.value
if step_value == 0:
raise ValueError('Invalid stepped serial: step must be non-zero')
if step_value > 0:
real_stop = tir.ceildiv(it.stop - it.start, step_value)
else:
real_stop = tir.ceildiv(it.start - it.stop, -step_value)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Critical: Handle empty iteration range for negative steps.

When step < 0 and start <= stop (e.g., T.serial(0, 10, -3)), line 263 computes ceildiv(0 - 10, 3) = -3, producing a negative real_stop. This causes invalid IR generation when passed to tir.serial at line 269.

Python's range(0, 10, -3) produces an empty sequence; the DSL should yield 0 iterations in this case.

Apply this diff to clamp real_stop to non-negative values:

                if step_value > 0:
-                   real_stop = tir.ceildiv(it.stop - it.start, step_value)
+                   real_stop = tir.max(tir.ceildiv(it.stop - it.start, step_value), 0)
                else:
-                   real_stop = tir.ceildiv(it.start - it.stop, -step_value)
+                   real_stop = tir.max(tir.ceildiv(it.start - it.stop, -step_value), 0)
📝 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
if isinstance(it.step, (int, IntImm)):
step_value = it.step if isinstance(it.step, int) else it.step.value
if step_value == 0:
raise ValueError('Invalid stepped serial: step must be non-zero')
if step_value > 0:
real_stop = tir.ceildiv(it.stop - it.start, step_value)
else:
real_stop = tir.ceildiv(it.start - it.stop, -step_value)
if isinstance(it.step, (int, IntImm)):
step_value = it.step if isinstance(it.step, int) else it.step.value
if step_value == 0:
raise ValueError('Invalid stepped serial: step must be non-zero')
if step_value > 0:
real_stop = tir.max(tir.ceildiv(it.stop - it.start, step_value), 0)
else:
real_stop = tir.max(tir.ceildiv(it.start - it.stop, -step_value), 0)
🧰 Tools
🪛 Ruff (0.14.3)

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

(TRY003)

🤖 Prompt for AI Agents
In tilelang/language/v2/builder.py around lines 256 to 263, the computation of
real_stop for negative step values can produce a negative count (e.g., start <=
stop with step < 0), causing invalid IR; after computing real_stop for the
negative-step branch, clamp it to a non-negative integer (e.g., real_stop =
max(0, real_stop)) so empty iteration ranges yield zero iterations before it is
passed to tir.serial.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

♻️ Duplicate comments (1)
tilelang/language/v2/builder.py (1)

256-263: Empty iteration range still produces negative real_stop.

When step < 0 and start <= stop (e.g., T.serial(0, 10, -3)), line 263 computes ceildiv(0 - 10, 3) = -3, producing a negative real_stop that causes invalid IR.

This issue was already flagged in previous reviews and remains unresolved. The suggested fix is to clamp real_stop to non-negative values after the computation.

🧹 Nitpick comments (1)
tilelang/language/v2/builder.py (1)

264-268: Clarify the undefined behavior for non-constant steps.

The warning mentions "undefined behavior" but doesn't explain that line 268 assumes a positive step. If the step expression evaluates to a negative value at runtime, ceildiv(it.stop - it.start, it.step) will produce incorrect iteration counts.

Consider clarifying the warning message:

             else:
                 logger.warning(
-                    f'Using a non-constant step `{it.step}` in stepped serial may lead to undefined behavior in tilelang'
+                    f'Using a non-constant step `{it.step}` in stepped serial may lead to undefined behavior in tilelang. '
+                    f'The implementation assumes a positive step; negative runtime values will produce incorrect iteration counts.'
                 )
📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3d10d99 and 54f2da3.

📒 Files selected for processing (1)
  • tilelang/language/v2/builder.py (3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tilelang/language/v2/builder.py (2)
tilelang/carver/roller/hint.py (1)
  • step (240-243)
tilelang/language/loop.py (1)
  • serial (97-111)
🪛 Ruff (0.14.3)
tilelang/language/v2/builder.py

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

(TRY003)


275-277: 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). (2)
  • GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
  • GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
🔇 Additional comments (3)
tilelang/language/v2/builder.py (3)

103-108: LGTM!

The dataclass structure is clean and appropriate for representing stepped serial loops.


269-272: LGTM!

The frame creation and value transformation (it.start + v * it.step) correctly implement the stepped loop logic.


490-493: LGTM!

The override correctly returns tilelang.language.serial for range, enabling stepped serial loops in the DSL.

@LeiWang1999
Copy link
Member

@codex review

Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ 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".

Comment on lines +256 to +269
if isinstance(it.step, (int, IntImm)):
step_value = it.step if isinstance(it.step, int) else it.step.value
if step_value == 0:
raise ValueError('Invalid stepped serial: step must be non-zero')
if step_value > 0:
real_stop = tir.ceildiv(it.stop - it.start, step_value)
else:
real_stop = tir.ceildiv(it.start - it.stop, -step_value)
else:
logger.warning(
f'Using a non-constant step `{it.step}` in stepped serial may lead to undefined behavior in tilelang'
)
real_stop = tir.ceildiv(it.stop - it.start, it.step)
real_frame = tir.serial(real_stop, annotations=it.annotations)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Guard trip count when step cannot reach stop

The new stepped loop path computes real_stop with tir.ceildiv and passes it directly to tir.serial. When the step sign is incompatible with the start/stop ordering (e.g. range(10, 0, 2) or range(0, 10, -2)), (stop - start) and step have opposite signs, so real_stop becomes negative and we attempt to build a serial loop with a negative extent. TVM’s serial builder rejects negative extents, raising an error for loops that should simply execute zero iterations per Python semantics. Consider clamping the trip count to zero when no iterations are expected before creating the tir.serial frame.

Useful? React with 👍 / 👎.

@LeiWang1999 LeiWang1999 merged commit 777881e into tile-ai:main Nov 6, 2025
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature request] Support strides for T.serial

3 participants