tile-ai / tile-ai/tilelang

[BUG] Empty stepped unrolled loops execute once

Open
#2,886 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 searched the issue tracker for an existing report.

### TileLang version

Reproduced on `0.1.13` at commit `1d155f4b80865edfe0009ad952135b7afbd4f05a`.

### System information

- Linux x86-64
- Python 3.12.13
- TVM 0.25.dev0
- TileLang built from source

### Problem description

Explicitly unrolling an empty loop with a non-unit step executes its body once.

A TIRX loop with `extent=0` has no legal loop-variable values, regardless of its step. After `tl.transform.UnrollLoop`, however, the empty loop body is emitted once. This can introduce writes or other side effects that were unreachable in the original program.

The trip-count calculation in `src/transform/unroll_loop.cc` is:

```cpp
1 + (extent - 1) / step
```

For `extent=0, step=2`, C++ integer division truncates `-1 / 2` toward zero, so the expression evaluates to `1`.

### Reproducible example

```python
import numpy as np

import tilelang as tl
from tilelang import tvm

output = tvm.tirx.decl_buffer((1,), "int32", name="output")
i = tvm.tirx.Var("i", "int32")

loop = tvm.tirx.For(
i,
0,
0,
tvm.tirx.ForKind.UNROLLED,
tvm.tirx.BufferStore(output, 1, [0]),
step=tvm.tirx.IntImm("int32", 2),
annotations={"pragma_unroll_explicit": True},
)

func = tvm.tirx.PrimFunc(
[output.data],
loop,
buffer_map={output.data: output},
).with_attr("global_symbol", "main")

mod = tl.transform.UnrollLoop()(tvm.IRModule.from_expr(func))
executable = tvm.compile(mod["main"], target="c").jit(options=["-std=c++17"])

result = tvm.runtime.tensor(np.zeros(1, dtype="int32"))
executable["main"](result)
print(result.numpy())
```

Observed output:

```text
[1]
```

### Expected behavior

The loop body should not execute, and the output should remain:

```text
[0]
```

An empty loop must retain zero-trip semantics after explicit unrolling.

Contributor guide

Open the contributing guide

Research direction

Start in src/transform/unroll_loop.cc, where tl.transform.UnrollLoop calculates the trip count, and compare the extent=0, step=2 case with normal loop semantics. Run the provided Python reproducer and confirm that the transformed loop leaves output at [0] rather than [1].

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.