mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Lei Wang
4fa2d24d27
* [Refactor] Simplify guard construction in loop partitioning - Removed the need for a separate guard check based on bijectivity, allowing the analyzer to determine the necessity of the guard. - Streamlined the guard construction process by consolidating index bounds checks into a single loop, improving code clarity and maintainability. - Enhanced handling of non-bijective layouts and cases where loop extents differ from layout input shapes. * add a reproducable test
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
import torch
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit(
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
},
|
|
)
|
|
def _buggy_kernel(S: T.Tensor((8), T.bfloat16), D: T.Tensor((4, 64), T.bfloat16)):
|
|
with T.Kernel(1, threads=128):
|
|
S_shared = T.alloc_shared((8), T.bfloat16)
|
|
S_fragment = T.alloc_fragment((8), T.float32)
|
|
D_shared = T.alloc_shared((4, 64), T.bfloat16)
|
|
|
|
T.copy(S, S_shared)
|
|
T.copy(S_shared, S_fragment)
|
|
for k in T.serial(64):
|
|
for i in T.Parallel(4):
|
|
D_shared[i, k] = S_fragment[i]
|
|
T.copy(D_shared, D)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test():
|
|
test_S = torch.randn((8), dtype=torch.bfloat16, device="cuda")
|
|
test_D = torch.zeros((4, 64), dtype=torch.bfloat16, device="cuda")
|
|
_buggy_kernel(test_S, test_D)
|
|
ref_D = test_S[:4].view(4, 1).repeat(1, 64)
|
|
torch.testing.assert_close(test_D, ref_D)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test()
|