tile-ai / tile-ai/tilelang

[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

Open
#3,017 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 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.