mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Chaofan LinandLei Wang
9dda774aff
* [BugFix] Use BufferRegion in tl.cumsum to infer buffer shape * remove debug lines * remove rubbish * Fix decorator syntax for atomic_different_memory_orders_program --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
34 lines
891 B
Python
34 lines
891 B
Python
import torch
|
|
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T
|
|
|
|
|
|
@tilelang.jit(
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
},)
|
|
def _cumsum_view_infer_layout(hidden):
|
|
num_tokens = T.dynamic('num_tokens')
|
|
|
|
@T.prim_func
|
|
def buggy_kernel(x: T.Tensor[(num_tokens, hidden), 'float']):
|
|
with T.Kernel(num_tokens, threads=128) as pid:
|
|
smem = T.alloc_shared((hidden,), dtype='float')
|
|
T.copy(x[pid, :], smem)
|
|
T.cumsum(T.view(smem, (1, hidden)), dim=1)
|
|
|
|
return buggy_kernel
|
|
|
|
|
|
def test_cumsum_view_infer_layout():
|
|
hidden = 128
|
|
x = torch.randn(1, hidden, device='cuda', dtype=torch.float)
|
|
kernel = _cumsum_view_infer_layout(hidden)
|
|
kernel(x)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tilelang.testing.main()
|