[Bug] TMA descriptors ignore Buffer.elem_offset in global base addresses
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 745
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 104
Description
### Observed behavior
CUDA TMA descriptor lowering uses `Buffer::data` directly as the descriptor's global base address. When a global buffer view has a nonzero `Buffer.elem_offset`, the descriptor therefore addresses the backing allocation from element zero instead of the start of the view.
The same omission is present in the descriptor paths for:
- bulk TMA load/store;
- `tile::gather4` / `tile::scatter4`;
- im2col TMA loads;
- `cp.reduce.async.bulk.tensor` atomic adds.
For example, a `float16` buffer with `elem_offset=16` should encode a descriptor base 32 bytes past `buffer.data`, but the current lowering passes `buffer.data` unchanged.
### Expected behavior
The descriptor base should include the buffer view's element offset converted to the TMA payload byte representation. Descriptor coordinates should remain relative to that adjusted view base.
```text
Before: Buffer.data ----------------------> TMA descriptor base
elem_offset (dropped)
After: Buffer.data -> + elem_offset bytes -> TMA descriptor base
Range::min remains a descriptor coordinate
```
### Analysis
This is separate from #2681. That PR checks whether the innermost `Range::min` coordinate is suitably aligned before selecting descriptor TMA. It does not incorporate `Buffer.elem_offset` into the descriptor base address. The two offsets enter the address flow at different stages and both need to be preserved.
#2803 is also adjacent but distinct: it rejects descriptor lowering when a base pointer is bound inside the device body and cannot be host-encoded. A regular host-encodable buffer view with nonzero `elem_offset` should remain eligible; its descriptor base just needs the view displacement.
### Reproduction evidence
A lowering regression can construct buffer views with constant `elem_offset` values and inspect the stable `tl.create_tma_descriptor` / `tl.create_tma_im2col_descriptor` calls before host extraction. On the affected revision, argument 2 is the raw buffer data variable. The expected result is `tirx.handle_add_byte_offset(buffer.data, byte_offset)` for all four descriptor families.
The source-level reproductions cover SM90 bulk copy, im2col, and atomic paths plus SM100a gather/scatter. Hopper/Blackwell runtime execution is still useful follow-up evidence, but is not required to observe the incorrect descriptor address in lowered IR.
Contributor guide
Research direction
Locate the lowering paths for bulk TMA load/store, tile gather/scatter, im2col TMA loads, and cp.reduce.async.bulk.tensor atomic adds. Run the lowering reproductions and inspect the stable tl.create_tma_descriptor and tl.create_tma_im2col_descriptor calls; done means each descriptor base uses handle_add_byte_offset with the converted elem_offset while coordinates remain relative to the adjusted view base.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, hpc, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100