tile-ai / tile-ai/tilelang

[BUG] Invalid layout

Open
#2,158 0 comments 0 reactions 0 assignees View on GitHub
bug
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.