mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
Lei Wang
6e439530b5
* Refactor: Replace @tilelang.lazy_jit with @tilelang.jit in examples and tests Updated all instances of the @tilelang.lazy_jit decorator to @tilelang.jit in the lazyjit example notebooks and related test files. This change aligns the code with the new JIT compilation approach, enhancing consistency across the codebase. Additionally, removed the lazy_jit import from the module initialization to streamline the API. * Enhancement: Improve layout annotation handling in T.copy and parallel loops Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify layout hints for the outermost parallel loop. Enhanced layout annotation handling in CopyNode and AtomicAddNode classes to ensure compatibility with SIMT operations. Added tests to validate the functionality of loop layout annotations, improving robustness and clarity in layout management for parallel loops. * Add out_idx parameter to @jit and validate it's only used in lazy mode * lintfix * Refactor: centralize mode inference in JITImpl.__call__ before parse_args * test fix * Enhancement: Introduce JITNoBuilderError for eager mode checks in T.Kernel and T.const functions. Ensure proper error handling when Builder is unavailable, improving robustness in JIT context management. * lint fix * lint fix * Enhancement: Update example_triton_sparse_gqa_decode_varlen_indice.py to improve performance metrics output by adding average FLOPS calculations. Modify layout_inference.cc to refine buffer layout inference checks by ensuring fragment buffers are considered. * lint fix * Enhancement: Refactor JIT implementation to centralize mode initialization and improve error handling for eager mode. Update matmul tests to utilize both Cython and TVM FFI JIT kernels, ensuring consistency in output validation. * Refactor: Remove unused simplify_prim_func decorator from example_gemm_intrinsics.py and related test files. Enhance layout.cc to simplify integer set analysis for multi-threaded compilation, ensuring constant integer outputs. Update lower_tile_op.cc to enforce constant integer checks in layout output shapes. Improve error handling in builder.py for JIT context management. * remove claude.md * Refactor: Clean up code formatting in layout.cc, lower_tile_op.cc, and proxy.py for improved readability. Add missing import in test_tilelang_kernel_int4_gemm_mma.py and ensure consistent spacing in builder.py. * Refactor: Replace logical conditions with tir_all in complex expressions within test_arith_hard.py for improved clarity and consistency. Update main execution flow to call test_hard_prove instead of tilelang.testing.main. * Enhancement: Add @tilelang.testing.requires_cuda decorator to test functions in test_tilelang_issue_1549.py and test_tilelang_issue_1601.py to ensure CUDA requirements are explicitly stated for relevant tests. * lint fix * Refactor: Simplify output shape calculation in layout.cc by directly setting max value from ist.max(). Update test_valid_loop in test_tilelang_fragment_loop_checker.py to disable certain tests and ensure proper execution flow. Enhance error handling in builder.py to correctly manage PrimFunc return types. * Refactor: Update boolean condition in test_boolop to use Or and Not for improved clarity and consistency in logical expressions. * Refactor: Update test execution flow in test_tilelang_fragment_loop_checker.py to call tilelang.testing.main. Modify JIT implementation to handle callable functions more generically. Enhance TirTemplate class in builder.py to return an empty tuple if matcher is None, improving robustness. * Refactor JIT function handling: Replace LazyJITFunc with JITFunc across the codebase, updating related function signatures and logic to support eager JIT compilation mode. This change enhances the clarity and functionality of JIT operations. * Update JITImpl documentation: Clarify output tensor index attribute description and remove outdated comment regarding implicit TIR construction. * Refactor TileLang to transition from v2 to eager language module. Update imports across multiple files to replace references to the deprecated v2 module with the new eager module, ensuring compatibility with the latest JIT compilation features. Remove obsolete v2 module files to streamline the codebase. * Enhance JIT compilation support: Update JIT function handling to improve compatibility with eager execution. Refactor related function signatures and logic, ensuring seamless integration with the latest compilation features. This change aims to streamline JIT operations and enhance overall code clarity.
409 lines
13 KiB
Python
409 lines
13 KiB
Python
import torch
|
|
import tilelang
|
|
from tilelang import tvm as tvm
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
from tilelang.intrinsics import (
|
|
make_mma_swizzle_layout as make_swizzle_layout,
|
|
)
|
|
from tilelang.transform import simplify_prim_func
|
|
from tilelang.intrinsics.mma_macro_generator import (
|
|
INT4TensorCoreIntrinEmitter,
|
|
INT4TensorCoreIntrinEmitterWithLadderTransform,
|
|
)
|
|
|
|
tilelang.testing.set_random_seed(42)
|
|
|
|
|
|
def tl_matmul(
|
|
M,
|
|
N,
|
|
K,
|
|
in_dtype,
|
|
out_dtype,
|
|
accum_dtype,
|
|
):
|
|
assert in_dtype in [
|
|
T.float16,
|
|
T.int8,
|
|
], "Currently only float16 and int8 are supported"
|
|
assert out_dtype in [
|
|
T.float16,
|
|
T.float32,
|
|
T.int32,
|
|
], "Currently only float16, float32 and int32 are supported"
|
|
|
|
K = K // 2
|
|
|
|
micro_size_x = micro_size_y = micro_size_k = 16
|
|
|
|
if accum_dtype == T.int32:
|
|
micro_size_k = 32
|
|
|
|
# This is a debug config
|
|
block_row_warps = 2
|
|
block_col_warps = 2
|
|
warp_row_tiles = 64
|
|
warp_col_tiles = 64
|
|
chunk = 32 if in_dtype == T.float16 else 64
|
|
shared_scope = "shared.dyn"
|
|
|
|
# Pipeline Stage
|
|
stage = 2
|
|
|
|
block_M = block_row_warps * warp_row_tiles
|
|
block_N = block_col_warps * warp_col_tiles
|
|
block_K = chunk
|
|
|
|
A_shape = (M, K) # int8 storage represents int4*2
|
|
B_shape = (N, K) # int8 storage represents int4*2
|
|
A_shared_shape = (block_M, block_K)
|
|
B_shared_shape = (block_N, block_K)
|
|
C_shared_shape = (
|
|
block_M // micro_size_x,
|
|
block_N // micro_size_y,
|
|
micro_size_x,
|
|
micro_size_y,
|
|
)
|
|
|
|
warp_size = 32
|
|
threads = warp_size * (block_row_warps * block_col_warps)
|
|
local_size_a = (micro_size_x * micro_size_k) // warp_size
|
|
local_size_b = (micro_size_y * micro_size_k) // warp_size
|
|
local_size_c = (micro_size_x * micro_size_y) // warp_size
|
|
warp_rows = warp_row_tiles // micro_size_x
|
|
warp_cols = warp_col_tiles // micro_size_y
|
|
|
|
# MMA Wrapper to Auto Generate Code for MMA
|
|
mma_emitter = INT4TensorCoreIntrinEmitter(
|
|
a_dtype=in_dtype,
|
|
b_dtype=in_dtype,
|
|
accum_dtype=accum_dtype,
|
|
a_transposed=False,
|
|
b_transposed=True,
|
|
block_row_warps=block_row_warps,
|
|
block_col_warps=block_col_warps,
|
|
warp_row_tiles=warp_row_tiles,
|
|
warp_col_tiles=warp_col_tiles,
|
|
chunk=chunk,
|
|
)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor(A_shape, in_dtype),
|
|
B: T.Tensor(B_shape, in_dtype),
|
|
C: T.Tensor((M, N), out_dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
|
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope)
|
|
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope)
|
|
C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope)
|
|
A_local = T.alloc_local((warp_rows * local_size_a), in_dtype)
|
|
B_local = T.alloc_local((warp_cols * local_size_b), in_dtype)
|
|
C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)
|
|
|
|
T.annotate_layout(
|
|
{
|
|
A_shared: make_swizzle_layout(A_shared),
|
|
B_shared: make_swizzle_layout(B_shared),
|
|
}
|
|
)
|
|
|
|
# Improve L2 Cache
|
|
T.use_swizzle(panel_size=10)
|
|
|
|
T.clear(C_local)
|
|
|
|
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=stage):
|
|
# Load A into shared memory
|
|
for i, k in T.Parallel(block_M, block_K):
|
|
A_shared[i, k] = A[by * block_M + i, ko * block_K + k]
|
|
|
|
# Load B into shared memory
|
|
for j, k in T.Parallel(block_N, block_K):
|
|
B_shared[j, k] = B[bx * block_N + j, ko * block_K + k]
|
|
|
|
for ki in T.serial(0, (block_K // micro_size_k)):
|
|
# Load A into fragment
|
|
mma_emitter.ldmatrix_a(
|
|
A_local,
|
|
A_shared,
|
|
ki,
|
|
)
|
|
|
|
# Load B into fragment
|
|
mma_emitter.ldmatrix_b(
|
|
B_local,
|
|
B_shared,
|
|
ki,
|
|
)
|
|
|
|
# Perform Matrix Multiplication
|
|
mma_emitter.mma(A_local, B_local, C_local)
|
|
|
|
# Perform STMatrix
|
|
mma_emitter.stmatrix(
|
|
C_local,
|
|
C_shared,
|
|
)
|
|
|
|
# Store shared into global
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
C[by * block_M + i, bx * block_N + j] = C_shared[
|
|
i // micro_size_x,
|
|
j // micro_size_y,
|
|
i % micro_size_x,
|
|
j % micro_size_y,
|
|
]
|
|
|
|
return main
|
|
|
|
|
|
def assert_tl_matmul_correctness(M, N, K, in_dtype, out_dtype, accum_dtype):
|
|
matmul = tl_matmul(M, N, K, in_dtype, out_dtype, accum_dtype)
|
|
kernel = tilelang.compile(
|
|
matmul,
|
|
out_idx=[2],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DEBUG_MERGE_SHARED_MEMORY_ALLOCATIONS: True,
|
|
},
|
|
)
|
|
print(kernel.get_kernel_source())
|
|
profiler = kernel.get_profiler()
|
|
|
|
src_code = kernel.get_kernel_source()
|
|
# src_code is the generated cuda source
|
|
assert src_code is not None
|
|
|
|
A = torch.randint(0, 4, (M, K), device="cuda", dtype=getattr(torch, in_dtype))
|
|
B = torch.randint(0, 4, (N, K), device="cuda", dtype=getattr(torch, in_dtype))
|
|
|
|
compressed_A = (A[:, ::2] & 0x0F) + ((A[:, 1::2] & 0x0F) << 4)
|
|
compressed_B = (B[:, ::2] & 0x0F) + ((B[:, 1::2] & 0x0F) << 4)
|
|
C = kernel(compressed_A, compressed_B)
|
|
print(C)
|
|
latency = profiler.do_bench()
|
|
print(latency)
|
|
# Ensure that the latency is not None
|
|
assert latency is not None
|
|
|
|
# Get Reference Result
|
|
ref_c = torch.matmul(A.to(torch.float32), B.T.to(torch.float32)).to(getattr(torch, accum_dtype))
|
|
|
|
print(ref_c)
|
|
torch.testing.assert_close(C, ref_c, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_assert_tl_matmul_correctness():
|
|
assert_tl_matmul_correctness(128, 128, 128, T.int8, T.int32, T.int32)
|
|
assert_tl_matmul_correctness(128, 128, 64, T.int8, T.int32, T.int32)
|
|
|
|
|
|
@simplify_prim_func
|
|
def tl_matmul_weight_only_transform(
|
|
M,
|
|
N,
|
|
K,
|
|
in_dtype,
|
|
out_dtype,
|
|
accum_dtype,
|
|
):
|
|
K = K // 2
|
|
assert in_dtype in [
|
|
T.float16,
|
|
T.int8,
|
|
], "Currently only float16 and int8 are supported"
|
|
assert out_dtype in [
|
|
T.float16,
|
|
T.float32,
|
|
T.int32,
|
|
], "Currently only float16, float32 and int32 are supported"
|
|
|
|
micro_size_x = micro_size_y = micro_size_k = 16
|
|
|
|
if out_dtype == T.int32:
|
|
micro_size_k = 32
|
|
|
|
transform_b = 3
|
|
|
|
# This is a debug config
|
|
block_row_warps = 2
|
|
block_col_warps = 2
|
|
warp_row_tiles = 64
|
|
warp_col_tiles = 64
|
|
chunk = 32 if in_dtype == T.float16 else 64
|
|
shared_scope = "shared.dyn"
|
|
|
|
# Pipeline Stage
|
|
stage = 2
|
|
|
|
block_M = block_row_warps * warp_row_tiles
|
|
block_N = block_col_warps * warp_col_tiles
|
|
block_K = chunk
|
|
|
|
A_shape = (M, K)
|
|
B_shape = (N // micro_size_y, K // micro_size_k, micro_size_y, micro_size_k)
|
|
A_shared_shape = (
|
|
block_M,
|
|
block_K,
|
|
)
|
|
B_shared_shape = (
|
|
block_N // micro_size_y,
|
|
block_K // micro_size_k,
|
|
micro_size_y,
|
|
micro_size_k,
|
|
)
|
|
C_shared_shape = (
|
|
block_M // micro_size_x,
|
|
block_N // micro_size_y,
|
|
micro_size_x,
|
|
micro_size_y,
|
|
)
|
|
warp_size = 32
|
|
threads = warp_size * (block_row_warps * block_col_warps)
|
|
local_size_a = (micro_size_x * micro_size_k) // warp_size
|
|
local_size_b = (micro_size_y * micro_size_k) // warp_size
|
|
local_size_c = (micro_size_x * micro_size_y) // warp_size
|
|
warp_rows = warp_row_tiles // micro_size_x
|
|
warp_cols = warp_col_tiles // micro_size_y
|
|
|
|
# MMA Wrapper to Auto Generate Code for MMA
|
|
mma_emitter = INT4TensorCoreIntrinEmitterWithLadderTransform(
|
|
a_dtype=in_dtype,
|
|
b_dtype=in_dtype,
|
|
accum_dtype=accum_dtype,
|
|
a_transposed=False,
|
|
b_transposed=True,
|
|
block_row_warps=block_row_warps,
|
|
block_col_warps=block_col_warps,
|
|
warp_row_tiles=warp_row_tiles,
|
|
warp_col_tiles=warp_col_tiles,
|
|
chunk=chunk,
|
|
transform_kind_b=transform_b,
|
|
)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor(A_shape, in_dtype),
|
|
B: T.Tensor(B_shape, in_dtype),
|
|
C: T.Tensor((M, N), out_dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
|
|
A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope)
|
|
B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope)
|
|
C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope)
|
|
A_local = T.alloc_local((warp_rows * local_size_a), in_dtype)
|
|
B_local = T.alloc_local((warp_cols * local_size_b), in_dtype)
|
|
C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)
|
|
|
|
T.annotate_layout(
|
|
{
|
|
A_shared: make_swizzle_layout(A_shared),
|
|
B_shared: make_swizzle_layout(B_shared),
|
|
}
|
|
)
|
|
|
|
# Improve L2 Cache
|
|
T.use_swizzle(panel_size=10)
|
|
|
|
T.clear(C_local)
|
|
|
|
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=stage):
|
|
# Load A into shared memory
|
|
for i, k in T.Parallel(block_M, block_K):
|
|
A_shared[i, k] = A[by * block_M + i, ko * block_K + k]
|
|
|
|
# Load B into shared memory
|
|
for j, k, jj, kk in T.Parallel(block_N // micro_size_y, block_K // micro_size_k, micro_size_y, micro_size_k):
|
|
B_shared[j, k, jj, kk] = B[bx * (block_N // micro_size_y) + j, ko * (block_K // micro_size_k) + k, jj, kk]
|
|
|
|
for ki in T.serial(0, (block_K // micro_size_k)):
|
|
# Load A into fragment
|
|
mma_emitter.ldmatrix_a(
|
|
A_local,
|
|
A_shared,
|
|
ki,
|
|
)
|
|
|
|
# Load B into fragment
|
|
mma_emitter.ldmatrix_b(
|
|
B_local,
|
|
B_shared,
|
|
ki,
|
|
)
|
|
|
|
# Perform Matrix Multiplication
|
|
mma_emitter.mma(A_local, B_local, C_local)
|
|
|
|
# Perform STMatrix
|
|
mma_emitter.stmatrix(
|
|
C_local,
|
|
C_shared,
|
|
)
|
|
|
|
# Store shared into global
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
C[by * block_M + i, bx * block_N + j] = C_shared[
|
|
i // micro_size_x,
|
|
j // micro_size_y,
|
|
i % micro_size_x,
|
|
j % micro_size_y,
|
|
]
|
|
|
|
return main
|
|
|
|
|
|
def assert_tl_matmul_weight_only_transform_correctness(M, N, K, in_dtype, out_dtype, accum_dtype):
|
|
import bitblas
|
|
|
|
matmul = tl_matmul_weight_only_transform(M, N, K, in_dtype, out_dtype, accum_dtype)
|
|
kernel = tilelang.compile(matmul, out_idx=[2])
|
|
profiler = kernel.get_profiler()
|
|
|
|
src_code = kernel.get_kernel_source()
|
|
# src_code is the generated cuda source
|
|
assert src_code is not None
|
|
transform_b = 3
|
|
|
|
A = torch.randint(0, 4, (M, K), device="cuda", dtype=getattr(torch, in_dtype))
|
|
B = torch.randint(0, 4, (N, K), device="cuda", dtype=getattr(torch, in_dtype))
|
|
compressed_A = (A[:, ::2] & 0x0F) + ((A[:, 1::2] & 0x0F) << 4)
|
|
compressed_B = (B[:, ::2] & 0x0F) + ((B[:, 1::2] & 0x0F) << 4)
|
|
|
|
ladder_permutate_config = bitblas.ops.LadderPermutateConfig(
|
|
M=N,
|
|
N=(K // 2),
|
|
datatype=T.int8,
|
|
storage_dtype=T.int8,
|
|
transform_kind=transform_b,
|
|
transpose_matrix=True,
|
|
)
|
|
|
|
ladder_permutate = bitblas.ops.LadderPermutate(ladder_permutate_config)
|
|
LB = ladder_permutate(compressed_B.cpu()).cuda()
|
|
C = kernel(compressed_A, LB)
|
|
|
|
latency = profiler.do_bench()
|
|
print(f"Latency: {latency}")
|
|
# Ensure that the latency is not None
|
|
assert latency is not None
|
|
|
|
# Get Reference Result
|
|
ref_c = torch.matmul(A.to(torch.float32), B.T.to(torch.float32)).to(getattr(torch, accum_dtype))
|
|
print(C)
|
|
print(ref_c)
|
|
torch.testing.assert_close(C, ref_c, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
@tilelang.testing.requires_package("bitblas")
|
|
@tilelang.testing.requires_llvm
|
|
@tilelang.testing.requires_cuda
|
|
def test_assert_tl_matmul_weight_only_transform():
|
|
assert_tl_matmul_weight_only_transform_correctness(128, 128, 128, T.int8, T.int32, T.int32)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# tilelang.testing.main()
|
|
assert_tl_matmul_correctness(128, 128, 128, T.int8, T.int32, T.int32)
|