[BUG][Fuzzer][wrong-code] `buf[(i+1)%N]` fragment read in `T.Parallel` silently returns the local element instead of the indexed value
- 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 similar report.
### What version of TileLang are you using?
0.1.13 (latest release)
### System information
TileLang 0.1.13; CUDA 12.8; NVIDIA L40S (sm_89), single GPU.
### Problem description
A `local.fragment` READ whose index is a non-identity affine function of the `T.Parallel`
loop variable — for example a rotation `buf[(i+1)%N]` or a shifted read `buf[(i+OFF)%N]` —
is silently miscompiled: the index expression is discarded and each thread reads *its own*
local element instead. The kernel compiles with no error and runs to completion; the result
is deterministically wrong.
Concretely, for `B[i] = buf[(i+1) % N]` (a rotate-by-1) over an `int32` fragment of `N=128`
with `threads=128`, the emitted CUDA is:
```
buf[0] = A[((int)threadIdx.x)];
B[((int)threadIdx.x)] = buf[0]; // the (i+1)%N index is gone — no cross-thread read
```
so `B == A` (identity) instead of the rotation. All 128 elements are wrong. The same kernel
over a `shared` buffer emits the real indexed read and is correct (see control below).
Trigger boundary (offset and elements-per-thread)
The whole index expression is dropped, so the *kind* of wrongness depends only on the
loop→thread partition, not on the offset value:
| config | epr | got | expected | wrong |
|---|---|---|---|---|
| `threads=128`, `OFF=1` | 1 | `[0,1,2,3,4,5,6,7]` | `[1,2,3,4,5,6,7,8]` | 128/128 |
| `threads=128`, `OFF=3` | 1 | `[0,1,2,3,4,5,6,7]` | `[3,4,5,6,7,8,9,10]` | 128/128 |
| `threads=128`, `OFF=5` | 1 | `[0,1,2,3,4,5,6,7]` | `[5,6,7,8,9,10,11,12]` | 128/128 |
| `threads=64`, `OFF=3` | 2 | `[1,0,3,2,5,4,7,6]` | `[3,4,5,6,7,8,9,10]` | 128/128 |
| `threads=32`, `OFF=1` | 4 | `[1,2,3,0,5,6,7,4]` | `[1,2,3,4,5,6,7,8]` | 32/128 |
At one element per thread (`threads==N`) the read degenerates to plain identity `buf[i]`;
with more than one element per thread the intra-thread element index is permuted within each
per-thread group. In every case the offset/modulo never reaches the emitted address.
Not a regression — see Provenance.
### Reproducible example code
```python
import tilelang, tilelang.language as T, torch, numpy as np
N, OFF = 128, 1
def build(shared):
@T.prim_func
def main(A: T.Tensor((N,), "int32"), B: T.Tensor((N,), "int32")):
with T.Kernel(1, threads=128) as bx:
buf = T.alloc_shared((N,), "int32") if shared else T.alloc_fragment((N,), "int32")
T.copy(A, buf)
for i in T.Parallel(N):
B[i] = buf[(i + OFF) % N] # rotate-by-1 read
return main
a = torch.arange(N, dtype=torch.int32, device="cuda")
exp = a.cpu().numpy()[(np.arange(N) + OFF) % N]
frag = tilelang.compile(build(shared=False), out_idx=[1])(a).cpu().numpy()
shrd = tilelang.compile(build(shared=True), out_idx=[1])(a).cpu().numpy()
print("fragment:", "PASS" if np.array_equal(frag, exp) else f"FAIL ({(frag!=exp).sum()}/{N})")
# -> fragment: FAIL (128/128) got [0,1,2,...], expected [1,2,3,...]
print("shared :", "PASS" if np.array_equal(shrd, exp) else f"FAIL ({(shrd!=exp).sum()}/{N})")
# -> shared : PASS (same kernel, shared buffer, computes the rotation)
```
### Traceback
No traceback — the kernel compiles and runs to completion; the result is silently wrong and
deterministic.
### Expected behavior
The read `buf[(i+1) % N]` should return element `(i+1) % N` of the fragment, the same value the
identical `shared`-buffer kernel returns. When an element lives on another thread, honoring the
index requires a cross-thread exchange (a shuffle or a shared staging buffer); the current code
emits neither and silently substitutes the thread's own local slot. The same result is already
computed correctly on the `shared` path over identical inputs, so the operation is computable —
computing it (rather than silently reading the local element) seems the natural expectation.
Alternatively, if a non-identity fragment access index inside `T.Parallel` is meant to be
unsupported, rejecting it at compile time would at least surface the problem instead of
returning wrong data — which direction is appropriate depends on whether this access pattern is
meant to be supported.
### Additional context
**Root cause.** Layout inference for a `T.Parallel` loop over a `local.fragment` conflates three
things that are only equal in the identity case, and never reconciles them when they diverge: (1)
the **loop partition** — which `T.Parallel` iteration runs on which thread; (2) the fragment's
**physical distribution** — which logical element lives in which thread's register slot; and (3)
the **access index expression** actually written at each fragment reference. `ParallelOpNode`
picks a source fragment access and, from the fragment shape and thread bounds, derives a single
partition map that folds (1) and (2) together — e.g. for `N` elements over `N` threads, logical
element `j → (thread=j, slot=0)`. It then *reuses that same map, via its `Forward`, to rewrite
every fragment access index* — implicitly treating (3) as if it were the map's own input, i.e.
assuming the index written is exactly the loop-iteration variable `i`.
That assumption is unchecked. When the access index is a non-identity affine function of `i`
(`(i+OFF)%N`, a rotation, a shifted stencil), the offset/modulo is not evaluated against the
partition at all — the `Forward` rewrite substitutes the thread's *own* slot, so `buf[(i+OFF)%N]`
lowers to the local element (`buf[0]` at one element per thread) and the `+OFF`/`%N` never reaches
the emitted address. The information needed to honor the index — element `(i+OFF)%N` lives on a
*different* thread, so the read requires a cross-thread exchange (a shuffle, or staging through
shared memory) — is silently discarded. There is no step that (a) evaluates the index expression
to the (thread, slot) it actually names, then (b) emits the exchange when that thread ≠ the
current one, nor (c) rejects the access when no such exchange is planned. Note the inference
*does* handle the constant-index case explicitly — a non-zero **constant** fragment index is
rejected with `LOG(FATAL)` ("Only fragment[0] access is allowed within T.Parallel loop") — so the
gap is specifically the non-constant, `i`-dependent affine index, which reaches the unconditional
`Forward` rewrite instead of either being computed or being rejected the way the constant case is.
Where
The loop→(thread, slot) partition is built in
[`ParallelOpNode::InferLayout`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/op/parallel.cc#L325)
(the constant-index branch `LOG(FATAL)`s a non-zero constant index; the non-constant branch keys
the partition on the loop variable and does not reconcile a non-identity index against it). The
access index is then rewritten by applying that layout's `Forward` **unconditionally** — with no
inspection of the index form — to the load at
[`parallel.cc:50`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/op/parallel.cc#L50)
and the store at
[`parallel.cc:61`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/op/parallel.cc#L61)
(`new_indices = layout_map_[buffer]->Forward(indices)`). Because that map is the loop-iteration
partition, not an evaluation of `(i+OFF)%N`, the offset/modulo is lost and the load collapses to
the local slot. The `shared` path does not go through this partition-based remap — it takes
ordinary indexed lowering, keeps the real address, and is correct (the control above).
**Suggested fix.** Before the unconditional `Forward` rewrite at `parallel.cc:50/61`, reconcile the access index against the partition: evaluate the index expression `(i+OFF)%N` to the (thread, slot) it names, and — when that thread differs from the current one — emit the cross-thread exchange the read requires (a shuffle, or staging through shared memory) instead of substituting the local slot. If honoring non-identity fragment indices inside `T.Parallel` is out of scope, reject them at compile time — the same treatment the constant-index case already gets via the `LOG(FATAL)` in `InferLayout` — so the non-constant affine index surfaces an error rather than silently reading the local element. This is a nontrivial change to layout inference, not a one-liner (not verified end-to-end).
**Provenance.** Reproduced on 0.1.13; the fragment layout-inference and buffer-remap code in
`src/op/parallel.cc` is unchanged at the v0.1.9 tag (`441c3b0`) and since. I did not bisect an introducing
PR; origin before 0.1.9 is unverified.
**Dedup.** I searched the open and closed tracker and found no existing report of this defect. Distinct from #2395/#2396 (fixed by #2462), which hardened a *serial-loop* fallback
to reject a dynamic non-zero fragment write index; that guard converts the bad case to a compile
error. This repro is a `T.Parallel` data-parallel *read*, it compiles without triggering that
guard, and it returns wrong data rather than crashing or being rejected.
**Reach.** The triggering ingredient is a fragment access index that is a non-identity affine
function of the `T.Parallel` loop variable (offset, rotation, or modulo). The access *pattern*
itself — ring-buffer / rotate reads (`buf[(i+1)%N]`), shifted stencils — is a common kernel idiom,
but in TileLang a cross-element read like this is normally expressed through a `shared` buffer
(which supports arbitrary cross-thread addresses, and is the correct control here); doing it on a
`local.fragment` is uncommon, and is not documented as either supported or rejected for
`T.Parallel` fragments. Grepping `examples/` and `testing/` this session (0.1.13) I found `%` used
on *shared* buffer indices and inside `T.Fragment(forward_fn=...)` layout lambdas, but no shipped
example uses a modulo/offset *access* index on a `local.fragment` inside `T.Parallel` — so the
pattern is one step from existing kernels and no current test exercises it, which is why CI is
green. The bug fires on any non-identity index at any `threads`/tile; the most common config
(`threads==extent`, one element per thread) degenerates to a clean identity read.
**Impact.** The trigger is narrow: it needs a non-identity affine access index (offset,
rotation, modulo) on a `local.fragment` read inside `T.Parallel`, not a plain `buf[i]`. When
it fires the failure is silent-wrong-code with no crash or diagnostic — the whole output
region is corrupted (128/128 elements above), and the wrong value is deterministic (the index
expression is dropped regardless of runtime data), so a rotate/stencil kernel would always
return the wrong permutation rather than failing only on some inputs. Fixing it either closes
this access class by emitting the required cross-thread exchange or converts the silent wrong
result into a compile-time rejection.
Contributor guide
Research direction
Run the reported Python reproducer first, comparing the fragment and shared-buffer results. Read src/op/parallel.cc, especially ParallelOpNode::InferLayout and the Forward rewrites at the load and store sites. Done means non-identity fragment indices no longer silently produce the local element, either by correct cross-thread handling or explicit compile-time rejection, with a regression test covering the repro.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100