[Feat] Add support for `T.serial` with step and negative step

[kurisu6912]

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'
    )

[github-actions[bot]]

👋 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[bot]]

[!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

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

• tile-ai/tilelang#1120 — touches tilelang/language/v2/builder.py and is likely related to SerialForWithStep integration and ctx_for changes.

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[kurisu6912]

@codex review

[chatgpt-codex-connector[bot]]

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

[kurisu6912]

@codex review

[chatgpt-codex-connector[bot]]

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

[LeiWang1999]

@codex review

[LeiWang1999]

@codex review