tile-ai / tile-ai/tilelang

[Bug] TMA descriptors ignore Buffer.elem_offset in global base addresses

Open
#2,832 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.