[BugFix] Reject unsupported TMA atomic add dtypes

[morluto]

Fixes #2583.

Problem

The TMA atomic-add lowering accepted every dtype that could be represented by a TensorMap. That set is broader than the element types supported by cp.reduce.async.bulk.tensor.add, so unsupported types such as int16 reached tma_store_add and failed at kernel launch instead of producing a compile-time diagnostic.

Change

Validate the destination dtype in AtomicAdd::Lower before constructing the TMA descriptor. The accepted set now matches the tensor reduction contract: float16, bfloat16, float32, int32, uint32, and uint64.

The existing runtime test is also restricted to SM90 hardware, while compile-level coverage verifies that int16 and float64 are rejected and that float32 and int32 still lower to tma_store_add.

Regression coverage

Against the pre-fix native library, both new negative cases fail with DID NOT RAISE, demonstrating that the tests exercise the missing guard. The positive controls lower successfully and retain tma_store_add.

Validation

• python3 -m pre_commit run --files src/cuda/op/atomic_add.cc testing/python/language/test_tilelang_language_atomic.py — passed all hooks.
• Focused C++ -fsyntax-only compile of src/cuda/op/atomic_add.cc with the repository build flags — passed.
• python -m pytest -q testing/python/language/test_tilelang_language_atomic.py::test_tma_atomic_add_accepts_supported_dtype testing/python/language/test_tilelang_language_atomic.py::test_tma_atomic_add — 2 passed, 1 skipped on an SM86 GPU.
• Pre-fix proof: python -m pytest -q testing/python/language/test_tilelang_language_atomic.py::test_tma_atomic_add_rejects_unsupported_dtype — 2 failed with DID NOT RAISE, for int16 and float64.

Summary

• Added compile-time dtype validation for TMA atomic-add lowering.
• Accepted dtypes: float16, bfloat16, float32, int32, uint32, and uint64.
• Rejected unsupported dtypes before TMA descriptor generation.
• Added Hopper (sm_90) compile and lowering tests for supported and unsupported dtypes.
• Preserved plain TMA copy dtype handling.

C++ style / lint notes

• The PR changes C++ code but does not modify documented rules in docs/developer_guide/cpp_style.md.
• The C++ API Style Audit remains warning-only.
• No correctness or build issue is identified from the described changes.

[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

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: f1a2b4c9-9ff6-4e90-899f-e674221ab28b

📥 Commits

Reviewing files that changed from the base of the PR and between 53ff570 and 30c8dce.

📒 Files selected for processing (2)

• src/cuda/op/atomic_add.cc
• testing/python/language/test_tilelang_language_atomic.py

🚧 Files skipped from review as they are similar to previous changes (1)

• testing/python/language/test_tilelang_language_atomic.py

---

📝 Walkthrough

Walkthrough

TMA atomic-add lowering now validates destination dtypes before descriptor construction. Hopper-specific tests cover supported and unsupported dtypes during CUDA sm_90 lowering.

Changes

TMA atomic-add dtype validation

Layer / File(s)	Summary
TMA dtype validation src/cuda/op/atomic_add.cc	The TMA atomic-add path accepts bfloat16, float16/32, int32, and uint32. Unsupported destination dtypes trigger an ICHECK failure before descriptor construction.
Hopper-specific dtype tests testing/python/language/test_tilelang_language_atomic.py	The tests detect CUDA compute capability sm_90, compile parameterized TMA atomic-add programs, reject unsupported dtypes, and verify supported lowering with tma_store_add.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• tile-ai/tilelang#2743: Both PRs modify TMA atomic-add handling and CUDA architecture or dtype validation.
• tile-ai/tilelang#2780: Both PRs modify atomic-operation dtype handling and related language tests.
• tile-ai/tilelang#2833: Both PRs modify TMA atomic-add lowering in src/cuda/op/atomic_add.cc.

Suggested reviewers: penguin-wwy, leiwang1999

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Linked Issues check	⚠️ Warning	The PR adds compile-time validation, but it rejects uint64 and omits float64 support required by linked issue #2583.	Support all required valid reduction dtypes, including float64 and applicable uint64, while retaining compile-time rejection for unsupported types.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: rejecting unsupported TMA atomic-add destination dtypes.
Out of Scope Changes check	✅ Passed	The source and test changes directly address TMA atomic-add dtype validation and related regression coverage.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

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

341-352: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Cover every supported dtype branch.

The lowering accepts float16, bfloat16, float32, int32, uint32, and uint64. This test covers only float32 and int32. Add the other four supported dtypes to detect regressions in their validation branches.

Proposed test expansion

-@pytest.mark.parametrize("dtype", [T.float32, T.int32])
+@pytest.mark.parametrize(
+    "dtype",
+    [T.float16, T.bfloat16, T.float32, T.int32, T.uint32, T.uint64],
+)
 def test_tma_atomic_add_accepts_supported_dtype(dtype):

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@testing/python/language/test_tilelang_language_atomic.py` around lines 341 -
352, Expand the dtype parameter list in
test_tma_atomic_add_accepts_supported_dtype to include T.float16, T.bfloat16,
T.uint32, and T.uint64 alongside the existing T.float32 and T.int32 values,
while preserving the current lowering and kernel-source assertion.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@testing/python/language/test_tilelang_language_atomic.py`:
- Around line 12-16: Update _check_hopper() to query the currently selected CUDA
device rather than hardcoded device 0, using that device when calling
torch.cuda.get_device_properties. Preserve the existing CUDA availability guard
and Hopper capability check.

---

Nitpick comments:
In `@testing/python/language/test_tilelang_language_atomic.py`:
- Around line 341-352: Expand the dtype parameter list in
test_tma_atomic_add_accepts_supported_dtype to include T.float16, T.bfloat16,
T.uint32, and T.uint64 alongside the existing T.float32 and T.int32 values,
while preserving the current lowering and kernel-source assertion.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 0f1e0a3b-e29d-40d2-a863-25437a4a54bf

📥 Commits

Reviewing files that changed from the base of the PR and between bdb769a and 53ff570.

📒 Files selected for processing (2)

• src/cuda/op/atomic_add.cc
• testing/python/language/test_tilelang_language_atomic.py

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 53ff5705cf

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