[BUG][Fuzzer][ice-on-valid-code] Dynamic-shape output placed before the input supplying its symbolic dim crashes at call time (`IndexError`) instead of allocating the output
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 745
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 104
Description
### Required prerequisites
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) and [Discussions](https://github.com/tile-ai/tilelang/discussions) that this hasn't already been reported. (comment there if it has.)
### What version of TileLang are you using?
0.1.13 (latest release)
### System information
TileLang 0.1.13 / PyTorch 2.8 / CUDA 12.8, single NVIDIA L40S (sm_89). Backend `tvm_ffi` (the default) and `cython` both affected.
### Problem description
A `@tilelang.jit` kernel whose **output** tensor has a symbolic (dynamic) dimension crashes at call time with `IndexError: list index out of range` whenever that output parameter appears in the signature *before* the input parameter that supplies the symbolic size. The kernel compiles fine; the crash is in the host wrapper that allocates the output.
The size resolution assumes every symbolic output dim is backed by an *already-allocated* tensor. The wrapper builds its tensor list in parameter order in a single pass, so when the output comes first it looks up a slot that has not been filled yet.
The precise root is in `_process_dynamic_symbolic`: it records each symbolic dim against the **first parameter (in signature order) whose buffer shape mentions it**, and does **not exclude output buffers**. When the output is the first param to mention the dim, the recorded owner *is the output itself*, and the allocation loop then reads the output's own not-yet-appended slot (`tensor_list[ref_tensor_idx]` where `ref_tensor_idx` is the output's own index). It is a self-reference, not (as one might assume) a reference to a later input. Verified by inspecting the map: for `main(B(N,), A(N,))` with `out_idx=[0]` the map is `{'N': (0, 0, 0, 1)}` — owner index 0 = B, the output.
The same kernel with the output parameter placed *after* the input (the shape it copies from) works correctly, which is why every shipped example — and the dynamic-symbolic test, which uses `out_idx=[-1]` — dodges it.
Console output (repro below): crash vs working control
```
output-before-input: CRASH IndexError: list index out of range
output-after-input : PASS
```
Traceback for the crashing case:
```
File ".../tilelang/jit/adapter/tvm_ffi.py", line 255, in func
shape.append(tensor_list[ref_tensor_idx].shape[ref_shape_idx])
~~~~~~~~~~~^^^^^^^^^^^^^^^^
IndexError: list index out of range
```
### Reproducible example code
```python
import tilelang, tilelang.language as T, torch
N = T.symbolic("N")
# out_idx=[0]: output B is param 0, input A is param 1.
# B's shape (N,) is resolved from a LATER input param -> crash.
@tilelang.jit(out_idx=[0])
def out_first():
@T.prim_func
def main(B: T.Tensor((N,), "float32"), A: T.Tensor((N,), "float32")):
with T.Kernel(1, threads=128):
for i in T.Parallel(N):
B[i] = A[i] + T.float32(9)
return main
# Control: same kernel, output B placed AFTER input A.
@tilelang.jit(out_idx=[1])
def out_last():
@T.prim_func
def main(A: T.Tensor((N,), "float32"), B: T.Tensor((N,), "float32")):
with T.Kernel(1, threads=128):
for i in T.Parallel(N):
B[i] = A[i] + T.float32(9)
return main
a = torch.randn(64, device="cuda")
ref = a.double() + 9
try:
b = out_first()(a)
print("output-before-input:", "PASS" if torch.allclose(b.double(), ref, atol=1e-4) else "FAIL")
except Exception as e:
print(f"output-before-input: CRASH {type(e).__name__}") # -> IndexError
b = out_last()(a)
print("output-after-input :", "PASS" if torch.allclose(b.double(), ref, atol=1e-4) else "FAIL") # -> PASS
```
### Traceback
```
Traceback (most recent call last):
File ".../tilelang/jit/kernel.py", line 204, in __call__
return self.torch_function(*args, **kwds)
File ".../tilelang/jit/adapter/tvm_ffi.py", line 255, in func
shape.append(tensor_list[ref_tensor_idx].shape[ref_shape_idx])
~~~~~~~~~~~^^^^^^^^^^^^^^^^
IndexError: list index out of range
```
### Expected behavior
The kernel should allocate the output and run regardless of where the output parameter sits in the signature — as it already does when the output is placed last. `out_idx` is documented as "Index(es) of the output tensors to return" with no ordering restriction, so an output at index 0 whose dim is shared with a later input is a legal signature that the compiler already compiles; only the runtime output-allocation step fails on it.
### Additional context
**Root cause.** The symbolic-dim owner recording does not exclude output buffers, so when the output is the first parameter to name a shared dim it records itself as that dim's owner, and the output-allocation loop then reads the output's own not-yet-appended slot.
Mechanism
Two-level:
- **Source root (where the wrong index is recorded).** [`_process_dynamic_symbolic`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/jit/adapter/tvm_ffi.py#L145-L175) records each symbolic shape dim against the **first parameter, in signature order, whose buffer shape mentions it** — [`dynamic_symbolic_map[shape] = (0, i, j, 1)`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/jit/adapter/tvm_ffi.py#L166), gated only by `shape not in dynamic_symbolic_map` (first-wins) — with **no exclusion of `result_idx` (output) parameters**. So if the output is the first param to name the dim, the owner index it stores is the output's own index.
- **Consumption root (where it crashes).** In [`_convert_torch_func.func`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/jit/adapter/tvm_ffi.py#L224-L282) the allocation loop walks parameters in order and for an output param resolves each symbolic dim via [`tensor_list[ref_tensor_idx].shape[ref_shape_idx]`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/jit/adapter/tvm_ffi.py#L255). `tensor_list` only holds the params processed *before* the current one, so a `ref_tensor_idx` at-or-after the current output — including the **self-reference** where the owner is the output itself — indexes past the end and raises `IndexError`.
The self-reference case is the common one: whenever the output precedes every input that shares the dim, `_process_dynamic_symbolic` points the dim at the output's own slot.
Why the control passes, and why it hits both backends
When the output is last (`out_idx=[-1]`, the form used everywhere in `examples/`/`testing/`), an input already names the shared dim, so `_process_dynamic_symbolic` records that input as the owner; by the time the output is allocated the input is already in `tensor_list`, so the lookup succeeds. The `cython` backend carries the identical single-pass resolution at [`cython_wrapper.pyx#L202`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/jit/adapter/cython/cython_wrapper.pyx#L202) and crashes the same way (tested: `IndexError` on the repro kernel with `execution_backend="cython"`).
**Suggested fix.** Two independent points would each close it: (1) in `_process_dynamic_symbolic`, prefer a non-output (`result_idx`) parameter as the recorded owner of a shape dim, so a shared dim is always bound to an input that is populated before any output; and/or (2) in the allocation loop of both adapters, split into two passes (bind all inputs first, then allocate outputs) so `tensor_list[ref_tensor_idx]` is always populated, or resolve output shapes from a position-keyed inputs map rather than the `tensor_list` under construction. Both adapters carry the same allocation-loop shape (not verified end-to-end).
**Provenance.** The single-pass `tensor_list[ref_tensor_idx]` resolution plus the output-inclusive owner recording are present in both adapters at the 0.1.13 release; I did not pin the introducing PR. Not established as a regression of previously-correct behavior.
**Dedup.** Searched the open and closed tracker and found no existing report of this defect.
**Tested boundary.** The trigger is exactly whether the **recorded owner index** of a symbolic dim (the first buffer param, in order, that mentions it — recorded in [`_process_dynamic_symbolic`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/jit/adapter/tvm_ffi.py#L145-L175)) is >= the index of the output being allocated. Because output buffers are not excluded from that recording, the common case is a **self-reference**: the output is the first param to name the dim. Cells run on 0.1.13/L40S (full sweep in the Generalization table below): (a) 2-param output-first (`out_idx=[0]`) — `IndexError`, map `N:(0,0,0)` (owner = the output itself); (b) same kernel `execution_backend="cython"` — identical `IndexError` (`cython_wrapper.pyx:202`); (c) 3-param middle output (`out_idx=[1]`) whose dim is owned by input idx 0 (processed first) — runs correctly, map `N:(0,0,0)`. So (a)/(b) and (c) share one root; the boundary is owner-index-vs-output-index, not the output's absolute position.
**Generalization — two-level root.** (a) *Source-level:* `_process_dynamic_symbolic` (`tvm_ffi.py:161-166`, twin in `cython_wrapper.pyx`) records a symbolic dim against the first buffer param that mentions it, first-wins, **not excluding output (`result_idx`) buffers**. (b) *Operator-level:* every path that returns a symbolic-shaped output through the same allocation loop — single output, multi-output, any output dtype/rank — since the symbol resolution runs before the output tensor is created and is dtype/rank-agnostic. 4-axis sweep, all cells run on 0.1.13 / L40S:
| axis | cell tested (input) | observed result | same-root? |
|---|---|---|---|
| repro (baseline) | `out_idx=[0]`, `main(B(N,),A(N,))`, call `(a[64])` | `IndexError` @ `tvm_ffi.py:255`; map `N:(0,0,0)` (owner=output) | yes |
| related-source | same kernel, `execution_backend="cython"` | `IndexError` @ `cython_wrapper.pyx:202` | yes (identical logic, twin file) |
| related-source | `out_idx=[1]`, `main(X(8,)static, B(M,)out, A(M,)in)` — var first named by output | `IndexError`; map `M:(0,1,0)` (owner=output idx1, **self-ref**, not the later input idx2) | yes — sharpens root to self-reference |
| related-operator | multi-output `out_idx=[0,1]`, `main(B(N,),C(N,),A(N,))` | `IndexError` | yes |
| related-operator | scalar-owner: `out_idx=[0]`, `main(B(N,), n:int32, A(N,))` | `IndexError` @ 255 — N still recorded as buffer-shape (ref_id 0) against output B, NOT the scalar (ref_id 2); scalar `ref_id==2` branch (`inputs[...]`) never reached | yes (same shape-loop path) |
| related-type | 2D int8 output `out_idx=[0]`, `main(B(M,4)int8, A(M,4)int8)` | `IndexError` | yes (resolution precedes tensor creation; dtype/rank orthogonal) |
| control (dodge) | `out_idx=[1]`, `main(A(N,)in, B(N,)out)` — input names dim first | PASS, correct value | n/a (owner idx0 < output idx1) |
| control (dodge) | 3-param `out_idx=[1]`, `main(A(N,)in idx0, B out, C in)` | PASS | n/a (owner idx0 processed first) |
**Reframe.** Kept the specific title/framing but the Problem now names the class: *output is the first param to name a shared symbolic dim* (subsumes both "output before its input" and the self-reference case). The `ref_id==1` stride branch (`tvm_ffi.py:257`, `cython_wrapper.pyx` twin) is code-identical to the crashing shape branch and shares the exact `tensor_list[ref_tensor_idx]` lookup, so it is **same-root by inspection** (symbolic output strides are uncommon and awkward to express via `T.Tensor`, not run here). No distinct-adjacent bug and no new bug surfaced during the sweep. The scalar-owner cell revealed the draft's original "later input" narrative was imprecise — the recorded owner is a *buffer* (here the output itself), and the `ref_id==2` scalar branch that reads the fully-available `inputs` list is not the crashing path.
**Impact.** The trigger is narrow: a dynamic-shape kernel whose symbolic output parameter is the first param to name a shared dim (e.g. declared before the input it borrows the dim from) — an ordering the shipped `out_idx=[-1]` convention never uses. When it fires it is a deterministic call-time `IndexError` in the host output-allocation wrapper: the kernel never runs, nothing is computed, and the failure is loud and immediate rather than a silent wrong value, so it cannot reach a workload's results undetected. Fixing it closes an output-parameter-ordering class that `out_idx` documents as legal but the runtime allocator rejects, aligning the two adapters' allocation path with what the compiler already accepts.
**Reach.** (Run, not asserted.) The trigger is a legal signature (`out_idx` is documented with no ordering constraint) with two ingredients: a **symbolic output dimension** shared with an input (common in dynamic-shape kernels) and the output being the **first parameter to name that dim** (e.g. placed before the input). Placing the output last is the overwhelmingly common convention. I ran the shipped citations on 0.1.13/L40S: the dynamic-symbolic test `testing/python/profiler/test_tilelang_profiler_dynamic_symbolic.py` (`out_idx=[-1]`) — **11 passed** (output C last, dim owned by input A idx 0, dodges); the canonical `examples/dynamic_shape/example_dynamic.py` (eager `C = T.empty((M,N),...)`, `out_idx=[-1]`) — **runs, matches torch** (output constructed after inputs, dodges). A grep of `examples/` at v0.1.13 found no `out_idx=[0]` and no symbolic kernel whose output precedes its owning input; the only symbolic example with a non-trailing out_idx is `deepseek_v32/inference/kernel.py` (`out_idx=[4]`, output after its inputs, dodges). So no shipped example or test exercises the output-first ordering that trips it, which is why CI is green.
Contributor guide
Research direction
Start with _process_dynamic_symbolic and _convert_torch_func.func in tilelang/jit/adapter/tvm_ffi.py, then compare the matching allocation logic in cython_wrapper.pyx. Run the provided output-first and output-last reproductions on both backends; done means the output-first kernel allocates successfully and matches the expected values without IndexError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100