mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
Lei Wang
9936636438
* Add unroll loop functionality and corresponding test case - Introduced a new `UnrollLoop` function in the transform module to unroll loops based on various configuration options. - Added a test case in `test_tilelang_language_unroll.py` to validate the behavior of `T.unroll` with only the extent parameter, ensuring correct kernel generation with unroll pragmas. * Refactor unroll kernel implementation and update test case - Changed the kernel function in `test_tilelang_language_unroll.py` to use a new `unroll_kernel` function that compiles and returns the output tensor, improving clarity and structure. - Updated the `OptimizeForTarget` function in `phase.py` to ensure the `UnrollLoop` transformation is applied correctly, maintaining consistency in optimization phases. * lint fix * lint fix
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import tvm as tvm
|
|
from tilelang import language as T
|
|
|
|
|
|
def test_unroll_with_step():
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (16, 16), dtype=T.float32, align=16)
|
|
|
|
for _blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for _threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
for i in T.unroll(0, 16, step=4):
|
|
A[0, i] = 1.0
|
|
|
|
kernel = tilelang.compile(main)
|
|
assert "#pragma unroll" in kernel.get_kernel_source()
|
|
|
|
|
|
# TODO: unroll factor is not supported on hip, skip.
|
|
@tilelang.testing.requires_cuda
|
|
def test_unroll_with_unroll_factor():
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (16, 16), dtype=T.float32, align=16)
|
|
|
|
for _blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for _threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
for i in T.unroll(0, 16, unroll_factor=4):
|
|
A[0, i] = 1.0
|
|
|
|
kernel = tilelang.compile(main)
|
|
assert "#pragma unroll 4" in kernel.get_kernel_source()
|
|
|
|
|
|
def test_unroll_with_extent_only():
|
|
"""Test T.unroll with only extent parameter."""
|
|
|
|
@tilelang.jit
|
|
def unroll_kernel():
|
|
out = T.empty((512,), dtype=T.float32)
|
|
with T.Kernel(1, threads=512):
|
|
tid = T.get_thread_binding()
|
|
for i in T.unroll(tid % 32):
|
|
out[i] = i
|
|
return out
|
|
|
|
kernel = unroll_kernel.compile()
|
|
source = kernel.get_kernel_source()
|
|
assert "#pragma unroll" in source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|