tile-ai / tile-ai/tilelang

[BUG][Fuzzer][diagnostic] `T.sts32` on a sub-4-byte-element shared buffer emits no compile-time diagnostic and crashes at launch with `CUDA_ERROR_MISALIGNED_ADDRESS`

Open
#2,980 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

### Required prerequisites

- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) and found no existing report.

### What version of TileLang are you using?

0.1.13.

### System information

NVIDIA L40S (sm_89), tilelang 0.1.13. Reproduced this session (CUDA_LAUNCH_BLOCKING=1).

### Problem description

`T.sts32` emits an unconditional 4-byte store (`*reinterpret_cast(ptr)=v`) at `&S[index]` with no compile-time diagnostic. On a `float8` (1-byte-element) shared buffer, `&S[pid]` is only 1-byte aligned for a non-multiple-of-4 `pid`, so the 4-byte store faults. The kernel **compiles clean** and only crashes at launch with `CUDA_ERROR_MISALIGNED_ADDRESS`. The identical pattern on a `float32` buffer (element naturally 4-byte aligned) runs fine.

This is filed as a **diagnostic** gap, not a request to prove per-`pid` alignment (which is a runtime value). The compiler has enough *static* information to warn here: the destination buffer's element width (fp8 = 1 byte) and the store width (`sts32` = 4 bytes) are both known at compile time, and the index is a non-constant. A 4-byte store into a 1-byte-element buffer at a runtime index *cannot* be guaranteed aligned — so a clear compile-time diagnostic ("`sts32` targets a sub-4-byte-element buffer at a non-aligned-constant index; the 4-byte store may fault") is possible and would replace the opaque launch-time crash. `T.sts32` is a low-level explicit-PTX primitive whose contract is "store 4 bytes" (its docstring says exactly that), so the caller owns alignment — but the missing diagnostic is the fileable point: the frontend has the dtype/width mismatch in hand and says nothing.

### Reproducible example code

```python
import torch, tilelang
import tilelang.language as T
N = 128

@tilelang.jit
def kern():
@T.prim_func
def main(X: T.Tensor((N,), "float8_e5m2"), Y: T.Tensor((N,), "float8_e5m2")):
with T.Kernel(N, threads=32) as pid:
S = T.alloc_shared((N,), "float8_e5m2")
T.sts32(S[pid], T.Cast("uint32", pid)) # 4-byte store at &S[pid]; pid not a multiple of 4 -> misaligned
T.sync_threads()
Y[pid] = S[pid]
return main

X = torch.zeros(N, dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
Y = torch.empty(N, dtype=torch.float8_e5m2, device="cuda")
kern()(X, Y); torch.cuda.synchronize()
# -> InternalError: CUDALaunch Error: CUDA_ERROR_MISALIGNED_ADDRESS
```

Control: the same pattern with a `float32` buffer runs OK — the fault is specific to the sub-4-byte element buffer's per-index pointer arithmetic.

### Traceback

```
InternalError: CUDALaunch Error: CUDA_ERROR_MISALIGNED_ADDRESS
```

The kernel compiles clean; the misaligned 4-byte store faults only at launch, with no compile-time diagnostic.

### Expected behavior

When `sts32` (a 4-byte store) targets a buffer whose element width is smaller than 4 bytes at a non-constant / non-4-byte-aligned index, the frontend should emit a compile-time diagnostic naming the alignment requirement (the dtype-width-vs-store-width mismatch is statically known). A clean compile followed by an opaque launch-time `CUDA_ERROR_MISALIGNED_ADDRESS` — with nothing pointing at the `sts32` call — is the worst outcome; even documenting the alignment precondition in the docstring would help. The `float32` control (element width == store width) shows the well-aligned case is fine, so the diagnostic need only fire on the width-mismatch case.

### Additional context

**Root cause.** [`sts32`, `builtin.py#L1845-L1851`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/language/builtin.py#L1845-L1851) retrieves the pointer and calls `tl.sts32(ptr, value)` with no dtype/width check; [codegen, `codegen_cuda.cc#L4261-L4269`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/cuda/codegen/codegen_cuda.cc#L4261-L4269) emits `tl::store_shared_32(ptr, value)`; the device helper does an unconditional `*reinterpret_cast(ptr)=v`. At the frontend `sts32` call the destination `Buffer`'s element dtype (fp8, 1 byte) is in hand and the store width (4 bytes) is fixed by the intrinsic — the width mismatch that guarantees possible misalignment at a non-constant index is statically visible there, but nothing checks it.

**Suggested fix.** In `sts32`/`sts64`/`sts128`, when the destination buffer's element width is smaller than the store width and the index is not a compile-time multiple of the store width, emit a compile-time diagnostic (warning or error) naming the alignment requirement — or at minimum document the precondition in the docstring. The check needs only statically-known values (element dtype width, store width, whether the index is an aligned constant), not a per-`pid` alignment proof.

**Provenance.** The unguarded pointer-and-store lowering of `sts32` is present in the tilelang `v0.1.13` sources (permalinks above); reproduced at runtime on 0.1.13 this session (fp8 buffer → launch-time `CUDA_ERROR_MISALIGNED_ADDRESS`; float32 control runs fine). Origin before 0.1.13 not bisected.

**Dedup.** No `sts32`/`sts64`/`sts128`/`store_shared` entry in the tracker. All catalogued misaligned-address bugs (TMA-bulk / cluster-copy / vectorizer / atomic lowering) are AUTOMATIC wide-store paths the user did not request; this is the explicit shared-store intrinsic's width-vs-element-alignment gap.

**Reach.** Triggers when `sts32`/`sts64`/`sts128` (explicit-PTX shared-store primitives) target a shared buffer whose element is smaller than the store width, at a non-aligned index. These are low-level primitives — the caller owns alignment, and the intrinsic's contract is "store N bytes" — so this is not a request to make the store itself safe; the fileable point is the missing compile-time diagnostic, since the frontend already holds the dtype-width mismatch. The shipped tests (`testing/python/language/test_tilelang_language_ldst.py`, grepped this session) exercise `sts32`/`sts64`/`sts128` only on `float32` buffers (element width == store width, naturally aligned), so the sub-store-width case that faults is never tested and CI stays green.

**Impact.** The trigger is narrow: an explicit `sts32`/`sts64`/`sts128` on a sub-store-width element buffer at a non-aligned index — an uncommon, low-level usage, and the aligned/matched-width cases are fine. When it does fire the consequence is a hard launch-time crash (`CUDA_ERROR_MISALIGNED_ADDRESS`) with no compile-time signal pointing at the `sts32` call, so the failure is deterministic but opaque — the cost is debuggability (an early diagnostic vs a raw CUDA fault), not silent wrong data.

Contributor guide

Open the contributing guide

Research direction

Start with tilelang/language/builtin.py around sts32, then trace its lowering in src/cuda/codegen/codegen_cuda.cc. Run testing/python/language/test_tilelang_language_ldst.py and inspect its existing sts32, sts64, and sts128 cases. Done means the sub-store-width, non-aligned-index case receives a clear compile-time diagnostic while matched-width cases remain valid, with regression coverage.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.