mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
import pytest
|
|
from tilelang import language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def _tilelang_issue_layout_free_inference_choose_smallest_replication():
|
|
@T.prim_func
|
|
def main(A: T.Tensor((128, 4), T.float), B: T.Tensor((128, 4), T.float)):
|
|
with T.Kernel(1, threads=128) as _:
|
|
A_frag = T.alloc_fragment((128, 4), T.float)
|
|
B_frag = T.alloc_fragment((128, 4), T.float)
|
|
S_frag = T.alloc_fragment((4,), T.float)
|
|
T.annotate_layout(
|
|
{
|
|
A_frag: T.Fragment(A_frag.shape, lambda i, j: (i, j)),
|
|
}
|
|
)
|
|
for i, j in T.Parallel(128, 4):
|
|
A_frag[i, j] = S_frag[j]
|
|
for i, j in T.Parallel(128, 4):
|
|
B_frag[i, j] = S_frag[j]
|
|
|
|
return main
|
|
|
|
|
|
def test_tilelang_issue_layout_free_inference_choose_smallest_replication():
|
|
kernel = _tilelang_issue_layout_free_inference_choose_smallest_replication()
|
|
source = kernel.get_kernel_source()
|
|
assert "float S_frag[4];" in source, "S_frag is not in the source"
|
|
assert "float B_frag[4];" in source, "B_frag is not in the source"
|
|
assert "float A_frag[4];" in source, "A_frag is not in the source"
|
|
|
|
|
|
@tilelang.jit
|
|
def _invalid_fragment_write_owner_layout():
|
|
def token_loop_layout_fn(i, j, rep):
|
|
return i * 4 + rep % 4 + (rep // 4) * 32, j
|
|
|
|
def scalar_buffer_layout_fn(i, j, rep):
|
|
return i * 4 + j + rep * 32, 0
|
|
|
|
loop_layout = T.Fragment((8, 4), forward_fn=token_loop_layout_fn, replicate=8)
|
|
buffer_layout = T.Fragment((8, 4), forward_fn=scalar_buffer_layout_fn, replicate=2)
|
|
|
|
@T.prim_func
|
|
def main():
|
|
with T.Kernel(1, threads=64):
|
|
frag = T.alloc_fragment((8, 4), T.float32)
|
|
T.annotate_layout({frag: buffer_layout})
|
|
for i, j in T.Parallel(8, 4, loop_layout=loop_layout):
|
|
frag[i, j] = T.float32(1.0)
|
|
|
|
return main
|
|
|
|
|
|
def test_reject_fragment_write_from_non_owner_threads():
|
|
with pytest.raises(Exception, match="Layout infer conflict"):
|
|
_invalid_fragment_write_owner_layout()
|
|
|
|
|
|
@tilelang.jit
|
|
def _scalar_reduce_accumulator_owner_layout():
|
|
@T.prim_func
|
|
def main(out: T.Tensor((1,), T.float32)):
|
|
with T.Kernel(1, threads=128):
|
|
acc = T.alloc_fragment((128,), T.float32)
|
|
total = T.alloc_fragment((1,), T.float32)
|
|
total[0] = T.float32(0.0)
|
|
for k in T.Parallel(128):
|
|
acc[k] = T.float32(1.0)
|
|
T.reduce_sum(acc, total, clear=False)
|
|
if T.get_thread_binding() == 0:
|
|
out[0] = total[0]
|
|
|
|
return main
|
|
|
|
|
|
def test_scalar_reduce_accumulator_owner_layout():
|
|
kernel = _scalar_reduce_accumulator_owner_layout()
|
|
source = kernel.get_kernel_source()
|
|
assert "AllReduce" in source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|