[dajiaohuang]

Summary

• Convert Python integer row indices to int32 TIR expressions in the gather4/scatter4 frontend helpers.
• Add a Blackwell-gated compile regression covering literal indices in both operations.

Fixes #2601.

Validation

• Ruff check and format, Python syntax compile, and git diff check pass.
• Focused pytest collection is blocked by the local Windows Python 3.13 environment: Home Assistant's pytest plugin imports cntl, which is unavailable on Windows. The TileLang test itself did not run.

Summary

• Convert Python integer row indices to int32 TIR constants in tma_gather4 and tma_scatter4. Other row expressions remain unchanged.
• Add a Blackwell-gated regression test that compiles both operations with literal row indices and checks the generated CUDA source for gather and scatter operations.
• Address the constant-row compiler crash reported in issue #2601. The new test checks compilation, not runtime numeric correctness.

Validation

• Ruff checks and formatting, Python syntax compilation, and git diff checks reportedly pass.
• Focused pytest collection did not complete in the reported Windows Python 3.13 environment. A Home Assistant pytest plugin could not import fcntl, so the TileLang test did not run.

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

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: tile-ai/tilelang
• Branch: main
• File: .coderabbit.yaml
• Commit: 424792e
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: c40ec0fa-cc82-46b7-a9ae-69192e9d6416

Commits

Reviewing files that changed from the base of the PR and between 2e4b4fd and 6aa331b.

Files selected for processing: 1

• tilelang/language/copy_op.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

Walkthrough

tma_gather4 and tma_scatter4 now convert Python integer row indices to int32 TIR constants while preserving other row values. A CUDA compile test exercises both operations with constant rows and checks the generated source.

Changes: Constant TMA Row Indices

Layer / File(s)	Summary
Row-index normalization and compile test tilelang/language/copy_op.py, testing/python/language/test_tilelang_language_tma_gather_scatter.py	Both APIs convert integer row indices to int32 TIR constants and leave other values unchanged. A CUDA compute capability 10 test compiles gather and scatter with constant rows and checks that the generated source contains gather4 and scatter4 operations.

Priority: Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Possibly related PRs

• tile-ai/tilelang#3685: Adds CuTeDSL TMA gather4 and scatter4 operations and tests row ordering and write behavior.

Suggested reviewers: yongqi-zhuo, leiwang1999

Merge Risk: Minimal · up to 6aa33

No actionable merge-blocking issue is identified in this change.

Security architecture risk: Minimal · up to 2e4b4

This is a narrow input-compatibility fix. Existing buffer checks and caller synchronization responsibilities remain intact, and no material security risk was identified in the changed behavior.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The identified change affects the representation of caller-supplied row indices for operations on caller-supplied buffers. The inspected test adds fixed literal indices rather than a new externally supplied input channel or privileged execution endpoint.

Trust Boundaries and Controls

• observed — Both operations retain checks for buffer types and memory scopes, rank-two tensors, matching dtypes, four-row shared tiles, unit inner global stride, and exactly four row indices. Integer normalization does not replace these checks.

Pre-merge checks: 5

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title "Fix constant row indices for TMA gather and scatter" matches the main change. The PR converts Python integer row indices to int32 TIR constants in the gather4 and scatter4 helpers, and adds…
Linked Issues check		Directly linked issue #2601 requires that constant row indices in T.tma_gather4 and T.tma_scatter4 either compile or raise an actionable error. The reported crash occurs in LowerBulkGather4. The issue…
Out of Scope Changes check		The whole-PR diff from the merge base changes two files. In tilelang/language/copy_op.py, the row normalization in tma_gather4 and tma_scatter4 changes. In testing/python/language/test_tilelang_langua…
Docstring Coverage		Docstring coverage is inconclusive: files from earlier reviews still need to be rechecked within the check's analysis limits. Not blocking.

Full details: Title check

Explanation

The title "Fix constant row indices for TMA gather and scatter" matches the main change. The PR converts Python integer row indices to int32 TIR constants in the gather4 and scatter4 helpers, and adds a Blackwell-gated compile test for both operations. The title is specific and describes the primary fix. It does not mention the test addition, which is expected for a title.

Full details: Linked Issues check

Explanation

Directly linked issue #2601 requires that constant row indices in T.tma_gather4 and T.tma_scatter4 either compile or raise an actionable error. The reported crash occurs in LowerBulkGather4. The issue states that the only difference between a crash and a clean lowering is whether each rows element is a constant or a tensor-loaded value. The diff at head changes both frontend helpers. Each Python int row becomes tirx.const(row, dtype="int32") before the annotation is built. Non-int rows pass through unchanged. The count check for exactly four rows still follows. The change since the previous review only inlines the former normalization helper. Its behavior is the same. The added test, test_gather_scatter_constant_rows_compile, compiles a kernel with literal rows [5, 17, 42, 9] for both operations. It asserts that the generated source contains tma_load_gather4 and tma_store_scatter4. The test is gated to CUDA compute capability 10 and below 11. The PR reports that the test did not run. The crash mechanism is inferred from the issue analysis and was not reproduced here. The code evidence shows that constant rows now reach the lowering path as int32 TIR objects. This addresses the reported trigger. The issue does not require runtime numeric validation, so that is not assessed.

Full details: Out of Scope Changes check

Explanation

The whole-PR diff from the merge base changes two files. In tilelang/language/copy_op.py, the row normalization in tma_gather4 and tma_scatter4 changes. In testing/python/language/test_tilelang_language_tma_gather_scatter.py, one program helper and one compile test are added. Both changes directly support issue #2601. No unrelated behavior, refactoring, or files appear in the diff. The inlining since the previous review does not change behavior and does not add scope.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.

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.}

[LeiWang1999]

A small simplification to the row normalization.