mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Lei Wang
b01cb93586
* [Feature] Implement ProveFragmentContains Function for Fragment Thread Validation - Added the ProveFragmentContains function to check if the threads accessing elements of a smaller fragment are a subset of those accessing a larger fragment. - This function ensures valid access when transitioning from a smaller to a larger fragment layout. - Updated layout.cc and utils.cc to incorporate this new functionality, enhancing the layout validation process. - Removed the previous implementation of ProveFragmentContains from parallel.cc to streamline the codebase. * fix * Refactor ParallelOpNode Layout Handling - Removed the initial DeReplicate attempt from InferLayout to streamline layout inference. - Added DeReplicate logic to ComputeLoopLayoutFromBuffer to reduce replication when validating layout candidates. - Updated test cases to disable caching and ensure proper functionality of loop layout kernels. * fix * Refactor Test Cases for Loop Layout - Removed caching disablement and print statements from the loop layout identity test for cleaner output. - Updated the main execution block to directly call the testing framework, enhancing test execution flow.
110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.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
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn)
|
|
kernel = loop_layout_kernel.compile(M=M, N=N, 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
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn)
|
|
kernel = loop_layout_kernel.compile(M=M, N=N, 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.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
|
|
loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn, replicate=2)
|
|
kernel = copy_with_layout_kernel.compile(M=M, N=N, 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
|
|
)
|
|
|
|
|
|
@tilelang.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
|
|
@tilelang.testing.requires_cuda_compute_version(9, 0)
|
|
def test_annotate_replicate_loop_layout_vec4():
|
|
M, N = 128, 32
|
|
|
|
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(M=M, N=N, 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()
|