[BUG] Invalid layout
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 742
- 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) that this hasn't already been reported. (comment there if it has.)
### What version of TileLang are you using?
0.1.9
### System information
0.1.9
### Problem description
TVM internal error on normal layout
### Reproducible example code
The Python snippets:
```python
import tilelang
import tilelang.language as T
import torch
@tilelang.jit
def get_kernel():
num_threads = 256
block_m, block_k = 32, 256
in_block_m, in_block_k = 1, 32
out_block_m, out_block_k = 1, 128
num_in_sf_per_block_m, num_in_sf_per_block_k = block_m // in_block_m, block_k // in_block_k # 32, 8
num_in_sf_per_out_sf_m, num_in_sf_per_out_sf_k = out_block_m // in_block_m, out_block_k // in_block_k # 1, 4
num_out_sf_per_block_m, num_out_sf_per_block_k = block_m // out_block_m, block_k // out_block_k # 32, 2
num_in_sf_per_out_sf = num_in_sf_per_out_sf_m * num_in_sf_per_out_sf_k # 4
@T.prim_func
def kernel():
with T.Kernel(1, threads=num_threads) as pid:
x_fragment = T.alloc_fragment((num_in_sf_per_block_m, num_in_sf_per_block_k), T.uint32) # 32, 8
x_fragment_reshaped = T.alloc_fragment((num_out_sf_per_block_m, num_out_sf_per_block_k, num_in_sf_per_out_sf), T.uint32) # (32, 2, 4)
T.fill(x_fragment, 0)
T.fill(x_fragment_reshaped, 0)
# Reshape input scaling factors to match output scaling factor layout
for i, j in T.Parallel(num_in_sf_per_block_m, num_in_sf_per_block_k):
out_sf_m_idx = i // num_in_sf_per_out_sf_m
out_sf_k_idx = j // num_in_sf_per_out_sf_k
in_sf_idx = (i % num_in_sf_per_out_sf_m) * num_in_sf_per_out_sf_k + (j % num_in_sf_per_out_sf_k)
if j < num_in_sf_per_block_k // 2:
x_fragment_reshaped[out_sf_m_idx, out_sf_k_idx, in_sf_idx] = x_fragment[i, j]
return kernel
kernel = get_kernel()
```
### Traceback
```pytb
tvm.error.InternalError: Loop layout is not injective: Fragment([32, 2, 4] -> [2], replicate: 1, thread: 252, forward_thread: _i * 8 + _k, forward_index: [_j], thread_range: I.Range(0, 256))
errors: ["Could not normalize iterators"]
loop AST: for i in T.parallel(32):
for j in T.parallel(2):
for k in T.parallel(4):
x_fragment_reshaped = T.Buffer((32, 2, 4), "uint32", scope="local.fragment")
x_fragment_reshaped[i, j, k] = T.Cast("uint32", 0)
```
### Expected behavior
_No response_
### Additional context
_No response_
Contributor guide
Research direction
Start by running the provided Python reproducer with TileLang 0.1.9 and inspect the traceback around loop layout normalization and the non-injective fragment layout. The fix is complete when this normal-layout example no longer raises the TVM InternalError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100