mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Lei Wang
5685c56a45
* [Enhancement] Improve layout annotation handling in parallel loops * Updated the handling of layout annotations in parallel loops to ensure they are only attached to the outermost loop. This change clarifies the rationale behind layout management, allowing for better reasoning and transformation of nested parallel regions. * Introduced new parameters in the `Parallel` function to accept layout annotations directly, enhancing usability and validation during layout inference. * Refactored the layout inference logic to prioritize annotated layouts over inferred ones, improving the efficiency of layout management in parallel operations. This update aims to strengthen the integrity of layout annotations and streamline the handling of nested parallel loops in the system. * lint fix * lint fix * [Enhancement] Update test and builder for improved debugging and functionality * Added print statements in `test_tilelang_language_let.py` to output the `main` function and the source of the compiled module, aiding in debugging. * Modified the `bind_immutable` method in `builder.py` to include logging of the bound value type, enhancing traceability during the binding process. * Commented out the call to `tilelang.testing.main()` in favor of directly invoking the test function, streamlining the testing workflow. * replicate example * [Enhancement] Add loop layout support to T.copy and improve parallel loop handling * Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify a layout hint for the outermost parallel loop in SIMT copies. * Enhanced the handling of layout annotations in the CopyNode class to enforce compatibility with SIMT copy operations, ensuring that layout hints are only applied when appropriate. * Added tests to validate the functionality of loop layout annotations in copy operations, improving the robustness of the layout management system. * [Enhancement] Refine loop annotation handling in AtomicAdd and Copy operations * Updated the handling of loop annotations in both AtomicAddNode and CopyNode classes to ensure that coalesced width and parallel loop layout annotations are only applied to the outermost loop. * This change improves the clarity and correctness of loop management in SIMT operations, aligning with the recent enhancements made to layout annotation handling in parallel loops.
118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
# TODO(lei): replicate loop layout and more complicated layout cases
|
|
@tilelang.lazy_jit
|
|
def loop_layout_kernel(A, B, loop_layout):
|
|
M, N = T.const("M, N")
|
|
A: T.Tensor[(M, N), T.float32]
|
|
B: T.Tensor[(M, N), T.float32]
|
|
|
|
with T.Kernel(1, threads=128):
|
|
for i, j in T.Parallel(M, N, loop_layout=loop_layout):
|
|
B[i, j] = A[i, j]
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_loop_layout_fragment_vec4():
|
|
def loop_layout_fn(i, j):
|
|
elems = i * 32 + j
|
|
forward_thread = (elems // 4) % 128
|
|
forward_local = elems % 4 + (elems // 512) * 4
|
|
return forward_thread, forward_local
|
|
|
|
M, N = 128, 32
|
|
A = T.Tensor((M, N), T.float32)
|
|
B = T.Tensor((M, N), T.float32)
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn)
|
|
kernel = loop_layout_kernel.compile(A=A, B=B, loop_layout=loop_layout)
|
|
code = kernel.get_kernel_source()
|
|
|
|
# Expect vectorized copy along innermost dimension (float4)
|
|
assert "*(float4*)(B + ((i * 512) + (((int)threadIdx.x) * 4))) = *(float4*)(A + ((i * 512) + (((int)threadIdx.x) * 4)));" in code
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_loop_layout_identity():
|
|
def loop_layout_fn(i, j):
|
|
forward_thread = i
|
|
forward_local = j
|
|
return forward_thread, forward_local
|
|
|
|
M, N = 128, 32
|
|
A = T.Tensor((M, N), T.float32)
|
|
B = T.Tensor((M, N), T.float32)
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn)
|
|
kernel = loop_layout_kernel.compile(A=A, B=B, loop_layout=loop_layout)
|
|
code = kernel.get_kernel_source()
|
|
|
|
assert "*(float4*)(B + ((((int)threadIdx.x) * 32) + (i * 4))) = *(float4*)(A + ((((int)threadIdx.x) * 32) + (i * 4)));" in code
|
|
|
|
|
|
@tilelang.lazy_jit
|
|
def copy_with_layout_kernel(A, B, loop_layout):
|
|
M, N = T.const("M, N")
|
|
A: T.Tensor[(M, N), T.float32]
|
|
B: T.Tensor[(M, N), T.float32]
|
|
|
|
with T.Kernel(1, threads=128):
|
|
T.copy(A, B, loop_layout=loop_layout)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_copy_loop_layout_annotated_replicate_vec4():
|
|
def loop_layout_fn(i, j, rep):
|
|
elems = i * 32 + j
|
|
fth = (elems // 4) % 64 + rep * 64
|
|
floc = elems % 4 + (elems // (64 * 4)) * 4
|
|
return fth, floc
|
|
|
|
M, N = 128, 32
|
|
A = T.Tensor((M, N), T.float32)
|
|
B = T.Tensor((M, N), T.float32)
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn, replicate=2)
|
|
kernel = copy_with_layout_kernel.compile(A=A, B=B, loop_layout=loop_layout)
|
|
code = kernel.get_kernel_source()
|
|
|
|
assert "*(float4*)(B + ((i * 512) + (((int)threadIdx.x) * 4))) = *(float4*)(A + ((i * 512) + (((int)threadIdx.x) * 4)));" in code
|
|
|
|
|
|
@tilelang.lazy_jit
|
|
def replicate_loop_layout_kernel(A, B, loop_layout):
|
|
M, N = T.const("M, N")
|
|
A: T.Tensor[(M, N), T.float32]
|
|
B: T.Tensor[(M, N), T.float32]
|
|
|
|
with T.Kernel(1, threads=128):
|
|
for i, j in T.Parallel(M, N, loop_layout=loop_layout):
|
|
B[i, j] = A[i, j]
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_annotate_replicate_loop_layout_vec4():
|
|
M, N = 128, 32
|
|
A = T.Tensor((M, N), T.float32)
|
|
B = T.Tensor((M, N), T.float32)
|
|
|
|
def loop_layout_fn(i, j, rep):
|
|
elems = i * 32 + j
|
|
forward_thread = (elems // 4) % 64 + rep * 64
|
|
forward_local = elems % 4 + (elems // (64 * 4)) * 4
|
|
return forward_thread, forward_local
|
|
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn, replicate=2)
|
|
|
|
kernel = replicate_loop_layout_kernel.compile(A, B, loop_layout=loop_layout)
|
|
code = kernel.get_kernel_source()
|
|
|
|
assert (
|
|
"*(float4*)(B + ((i * 256) + ((((int)threadIdx.x) & 63) * 4))) = *(float4*)(A + ((i * 256) + ((((int)threadIdx.x) & 63) * 4)));"
|
|
in code
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|