mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Lei Wang
a03df604e4
* [Feature] Enhance fill operation to support various buffer types - Added support for `BufferLoad` in the `fill` function to handle different buffer types. - Updated `Fill` class to process region descriptors and buffer regions, improving flexibility in buffer handling. - Introduced checks for static bounds in region definitions to ensure safety during operations. - Refactored loop induction variable handling in `FillNode` to accommodate sliced regions. * lint fix
54 lines
1.4 KiB
Python
54 lines
1.4 KiB
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 _fill_with_static_region_kernel():
|
|
num_tokens = T.symbolic('num_tokens')
|
|
|
|
@T.prim_func
|
|
def buggy_kernel(x: T.Tensor[(num_tokens,), 'int64']): # noqa: F821
|
|
with T.Kernel(num_tokens, threads=128) as _:
|
|
T.fill(x[0:128], 0)
|
|
|
|
return buggy_kernel
|
|
|
|
|
|
@tilelang.jit(
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
},)
|
|
def _fill_with_dynamic_region_kernel():
|
|
num_tokens = T.symbolic('num_tokens')
|
|
|
|
@T.prim_func
|
|
def buggy_kernel(x: T.Tensor[(num_tokens,), 'int64']): # noqa: F821
|
|
with T.Kernel(num_tokens, threads=128) as _:
|
|
a, b = T.alloc_var('int'), T.alloc_var('int')
|
|
T.fill(x[a:b], 0)
|
|
|
|
return buggy_kernel
|
|
|
|
|
|
def test_fill_with_static_region_kernel():
|
|
kernel = _fill_with_static_region_kernel()
|
|
x = torch.zeros((256,), dtype=torch.int64, device='cuda')
|
|
kernel(x)
|
|
|
|
|
|
def test_fill_with_dynamic_region_kernel():
|
|
kernel = _fill_with_dynamic_region_kernel()
|
|
x = torch.zeros((256,), dtype=torch.int64, device='cuda')
|
|
kernel(x)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tilelang.testing.main()
|