[BugFix] Include buffer offsets in TMA descriptor bases - #2833
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughTMA descriptor lowering now includes ChangesTMA descriptor offset handling
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7a7a4ef054
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6af6bfa6bb
ℹ️ 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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/cuda/test_tilelang_tma_descriptor_elem_offset.py`:
- Around line 161-168: Update both pytest.raises match patterns in
test_fp4_unpacked_copy_rejects_16_byte_offset and
test_fp4_unpacked_copy_rejects_subbyte_offset to escape the dot in the literal
T.tma_copy text, using the existing regex-safe approach so near-miss messages
are rejected.
🪄 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: c78e7e22-24ea-41d6-89b0-a95ddf745749
📒 Files selected for processing (9)
src/cuda/op/atomic_add.ccsrc/cuda/op/copy.ccsrc/cuda/op/copy.hsrc/cuda/op/copy_analysis.ccsrc/cuda/transform/lower_hopper_intrin.cctesting/python/cuda/test_tilelang_tma_descriptor_elem_offset.pytilelang/jit/adapter/nvrtc/wrapper.pytilelang/jit/adapter/utils.pytilelang/jit/adapter/wrapper.py
🚧 Files skipped from review as they are similar to previous changes (2)
- src/cuda/op/atomic_add.cc
- src/cuda/op/copy.h
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f9ddd04332
ℹ️ 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".
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/cuda/test_tilelang_tma_descriptor_elem_offset.py`:
- Around line 224-227: Update CanProveTMADescriptorBaseAligned to prove
alignment for symbolic elem_offset expressions whose multiplication by the
element size is guaranteed to be a 16-byte multiple, including elem_offset * 8
for the float16 case; then update
test_auto_copy_falls_back_for_symbolically_aligned_device_offset to expect
descriptor lowering rather than fallback.
🪄 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: 92492f59-289d-403e-856f-11bec7669a6e
📒 Files selected for processing (3)
src/cuda/op/copy_analysis.cctesting/python/cuda/test_tilelang_tma_descriptor_elem_offset.pytilelang/jit/adapter/cutedsl/wrapper.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/cuda/op/copy_analysis.cc
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9c6fa8917a
ℹ️ 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".
1 participant
Fixes #2832.
Problem
Descriptor-based TMA lowering used
Buffer::dataas the global base address and droppedBuffer.elem_offset. Coordinates were therefore relative to the backing allocation rather than the buffer view for bulk load/store, gather/scatter, im2col, and TMA atomic paths.This is orthogonal to #2681: that PR validates the alignment of the
Range::mincoordinate used by bulk copies, while this change preserves the buffer view displacement before descriptor coordinates are applied.Change
Add one CUDA lowering helper that converts
elem_offsetthrough the existing TMA payload-byte calculation and emitshandle_add_byte_offsetwhen the offset is nonzero. Use that adjusted base for tiled bulk copy, gather/scatter, im2col, and TMA atomic descriptors. Zero-offset buffers retain the existing raw-data expression.Regression coverage
The new lowering test inspects the stable descriptor-creation IR before host extraction. It covers:
The bulk fixture uses a non-contiguous 2D view so it cannot take the descriptorless 1D fast path. On the affected source, descriptor argument 2 is the raw buffer data variable; the regression requires the corresponding
handle_add_byte_offsetexpression.Validation
ninja -C build tilelang -j$(nproc)with the pip CUDA 13.0 toolchain — CUDA-enabledlibtilelang.solinked successfully..venv/bin/python -m pytest testing/python/cuda/test_tilelang_tma_descriptor_elem_offset.py -q— 4 passed.bash format.sh --files src/cuda/op/copy.h src/cuda/op/copy.cc src/cuda/op/atomic_add.cc testing/python/cuda/test_tilelang_tma_descriptor_elem_offset.py— all configured checks passed.Proof gap
The available GPU is SM86, so Hopper/Blackwell runtime execution was not available. The regression validates the emitted descriptor address flow for all affected lowering paths, but it does not execute TMA instructions on SM90/SM100 hardware.
Summary
Buffer.elem_offsetin payload bytes.handle_add_byte_offsetfor nonzero offsets.Validation
C++ style / lint notes
docs/developer_guide/cpp_style.md.