mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +08:00
Lei Wang
91a7bb2be6
* Refactor GEMM and GEMM-SP operations to enhance clarity and maintainability - Removed deprecated prime factorization functions from `gemm.cc` and `gemm_sp.cc`. - Introduced a new `GemmWarpPolicy` class to manage warp policy attributes and methods, improving encapsulation. - Updated reflection methods to include the new policy structure, ensuring proper registration and introspection capabilities. - Enhanced `GetArchInt` function in `utils.cc` for better readability and type safety. - Added new `gemm_v2` function in `gemm.py` for improved GEMM operation with additional parameters and checks. * Refactor GEMM and frontend legalize operations for improved clarity and functionality - Updated `gemm_py.h` to include the correct header for GEMM operations. - Renamed `FrontendLegalizer` class to `LetInliner` and updated related methods to reflect this change, enhancing code clarity. - Modified the pass function from `FrontendLegalize` to `LetInline` for better alignment with its purpose. - Updated test cases to utilize the new `gemm_v2` function and adjusted the testing framework for improved output and clarity. - Removed obsolete test file `test_tilelang_transform_frontend_legalize.py` to streamline the test suite. - Enhanced the `LowerAndLegalize` function to utilize the new `LetInline` pass, improving the overall transformation process. * Enhance CUDA code generation and testing for GEMM operations - Added indentation printing in `codegen_cuda.cc` for improved assembly code formatting. - Updated `test_tilelang_tilelibrary_gemm.py` to include additional GEMM test cases and shared memory allocation with specified scope. - Introduced new `matmul_sr` and `run_gemm_sr` functions for GEMM operations with shared and fragment memory layouts. - Refactored layout inference in `mma_macro_generator.py` to improve clarity and correctness in shared memory handling. - Enhanced `gemm/__init__.py` to support new GEMM operation combinations and layout inference logic. These changes improve the clarity, functionality, and testing coverage of GEMM operations in the TileLang framework. * Refactor GEMM layout and testing for improved clarity and functionality - Updated `gemm_layouts.cc` to enhance the layout generation logic for transposed and non-transposed GEMM operations. - Renamed and modified functions in `test_tilelang_tilelibrary_gemm.py` to reflect changes in GEMM function signatures and improve test coverage. - Introduced new GEMM operation combinations in `gemm/__init__.py` to support additional layouts and configurations. - Enhanced layout inference in `mma_layout.py` and `mma_macro_generator.py` for better handling of shared memory layouts. These changes improve the clarity, functionality, and testing coverage of GEMM operations in the TileLang framework. * Refactor GEMM layout and Python integration for improved functionality - Updated `gemm_layouts.cc` to correct the order of layout replication and repetition for transposed and non-transposed GEMM operations. - Enhanced `gemm_py.cc` to handle block realization more robustly, ensuring correct assignment of global symbols and block attributes. - Refactored `inject_pipeline.cc` to streamline buffer read/write region handling, improving clarity and maintainability. - Cleaned up test cases in `test_tilelang_tilelibrary_gemm.py` by removing unnecessary print statements and adjusting function calls for better test execution flow. These changes enhance the clarity, functionality, and robustness of GEMM operations and their testing in the TileLang framework. * Refactor GEMM layout and testing for improved clarity and functionality - Updated `gemm_layouts.cc` to enhance layout generation logic for transposed and non-transposed GEMM operations. - Improved block realization handling in `gemm_py.cc` for better assignment of global symbols. - Streamlined buffer read/write region handling in `inject_pipeline.cc` for clarity. - Enhanced test cases in `test_tilelang_tilelibrary_gemm.py` by adjusting function calls and adding new GEMM operation combinations. These changes improve the clarity, functionality, and robustness of GEMM operations and their testing in the TileLang framework. * tfloat32 support. * lint fix * lint fix * Refactor shared memory allocation in GEMM tests - Removed unnecessary scope specification in shared memory allocation for matrices A and B in `test_tilelang_tilelibrary_gemm.py`. - This change simplifies the allocation process and aligns with the updated GEMM function signatures.
544 lines
17 KiB
Python
544 lines
17 KiB
Python
from tilelang import tvm as tvm
|
|
import tilelang.testing
|
|
|
|
|
|
def matmul(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
accum_dtype,
|
|
num_stages,
|
|
threads,
|
|
):
|
|
A_shape = (K, M) if trans_A else (M, K)
|
|
B_shape = (N, K) if trans_B else (K, N)
|
|
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
|
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
|
|
|
import tilelang.language as T
|
|
|
|
@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)
|
|
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
T.clear(C_local)
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
if trans_A:
|
|
T.copy(A[k * block_K, by * block_M], A_shared)
|
|
else:
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
if trans_B:
|
|
T.copy(B[bx * block_N, k * block_K], B_shared)
|
|
else:
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.gemm_v2(A_shared, B_shared, C_local, trans_A, trans_B)
|
|
# T.gemm(A_shared, B_shared, C_local, trans_A, trans_B)
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_gemm_ss(
|
|
M,
|
|
N,
|
|
K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages=3,
|
|
num_threads=128,
|
|
):
|
|
program = matmul(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
num_stages,
|
|
num_threads,
|
|
)
|
|
|
|
kernel = tilelang.compile(
|
|
program,
|
|
out_idx=[2],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
})
|
|
|
|
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
|
|
|
def ref_program(A, B):
|
|
import torch
|
|
|
|
if trans_A:
|
|
A = A.T
|
|
if trans_B:
|
|
B = B.T
|
|
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
|
C = C.to(torch.__getattribute__(out_dtype))
|
|
return C
|
|
|
|
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
def test_gemm_ss():
|
|
# More test case can be found in kernel/test_tilelang_kernel_gemm.py
|
|
# GEMM tests for float16
|
|
run_gemm_ss(512, 1024, 768, False, True, "float16", "float16", "float16", 128, 128, 32, 2)
|
|
run_gemm_ss(512, 1024, 768, False, False, "float16", "float16", "float16", 128, 128, 32, 2)
|
|
run_gemm_ss(512, 1024, 768, True, False, "float16", "float16", "float16", 128, 128, 32, 2)
|
|
run_gemm_ss(512, 1024, 768, True, True, "float16", "float16", "float16", 128, 128, 32, 2)
|
|
# n8 test
|
|
run_gemm_ss(128, 8, 32, False, True, "float16", "float16", "float16", 128, 8, 32, 0, 128)
|
|
|
|
# int8 test
|
|
run_gemm_ss(128, 128, 128, False, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_ss(128, 128, 128, False, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_ss(128, 128, 128, True, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_ss(128, 128, 128, True, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
|
|
# float8 tests
|
|
run_gemm_ss(128, 128, 128, True, True, "float8_e5m2", "float8_e5m2", "float32", 128, 128, 32, 2)
|
|
|
|
# tfloat32 test
|
|
run_gemm_ss(128, 128, 128, False, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_ss(128, 128, 128, False, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_ss(128, 128, 128, True, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_ss(128, 128, 128, True, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
|
|
|
|
def matmul_rs(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
accum_dtype,
|
|
num_stages,
|
|
threads,
|
|
):
|
|
A_shape = (K, M) if trans_A else (M, K)
|
|
B_shape = (N, K) if trans_B else (K, N)
|
|
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
|
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
|
A_frag_shape = A_shared_shape
|
|
|
|
import tilelang.language as T
|
|
|
|
@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)
|
|
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
|
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
T.clear(C_local)
|
|
T.annotate_layout({
|
|
A_shared: tilelang.layout.make_swizzled_layout(A_shared),
|
|
})
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
if trans_A:
|
|
T.copy(A[k * block_K, by * block_M], A_shared)
|
|
else:
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
if trans_B:
|
|
T.copy(B[bx * block_N, k * block_K], B_shared)
|
|
else:
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.copy(A_shared, A_frag)
|
|
T.gemm_v2(A_frag, B_shared, C_local, trans_A, trans_B)
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_gemm_rs(
|
|
M,
|
|
N,
|
|
K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages=3,
|
|
num_threads=128,
|
|
):
|
|
program = matmul_rs(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
num_stages,
|
|
num_threads,
|
|
)
|
|
|
|
kernel = tilelang.compile(
|
|
program,
|
|
out_idx=[2],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
})
|
|
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
|
|
|
def ref_program(A, B):
|
|
import torch
|
|
|
|
if trans_A:
|
|
A = A.T
|
|
if trans_B:
|
|
B = B.T
|
|
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
|
C = C.to(torch.__getattribute__(out_dtype))
|
|
return C
|
|
|
|
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
def test_gemm_rs():
|
|
# GEMM tests for float16
|
|
run_gemm_rs(512, 1024, 768, False, False, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rs(512, 1024, 768, False, True, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rs(512, 1024, 768, True, False, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rs(512, 1024, 768, True, True, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
|
|
# n8 tests
|
|
run_gemm_rs(128, 8, 32, False, True, "float16", "float16", "float16", 128, 8, 32, 0, 128)
|
|
|
|
# int8 tests
|
|
run_gemm_rs(128, 128, 128, False, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_rs(128, 128, 128, False, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_rs(128, 128, 128, True, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_rs(128, 128, 128, True, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
|
|
# float8 tests
|
|
run_gemm_rs(128, 128, 128, True, True, "float8_e5m2", "float8_e5m2", "float32", 128, 128, 32, 2)
|
|
|
|
# float32 tests
|
|
run_gemm_rs(128, 128, 128, False, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_rs(128, 128, 128, False, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_rs(128, 128, 128, True, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_rs(128, 128, 128, True, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
|
|
|
|
def matmul_sr(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
accum_dtype,
|
|
num_stages,
|
|
threads,
|
|
):
|
|
A_shape = (K, M) if trans_A else (M, K)
|
|
B_shape = (N, K) if trans_B else (K, N)
|
|
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
|
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
|
B_frag_shape = B_shared_shape
|
|
|
|
import tilelang.language as T
|
|
|
|
@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)
|
|
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
|
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
T.clear(C_local)
|
|
T.annotate_layout({
|
|
B_shared: tilelang.layout.make_swizzled_layout(B_shared),
|
|
})
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
if trans_A:
|
|
T.copy(A[k * block_K, by * block_M], A_shared)
|
|
else:
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
if trans_B:
|
|
T.copy(B[bx * block_N, k * block_K], B_shared)
|
|
else:
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.copy(B_shared, B_frag)
|
|
T.gemm_v2(A_shared, B_frag, C_local, trans_A, trans_B)
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_gemm_sr(
|
|
M,
|
|
N,
|
|
K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages=3,
|
|
num_threads=128,
|
|
):
|
|
program = matmul_sr(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
num_stages,
|
|
num_threads,
|
|
)
|
|
|
|
kernel = tilelang.compile(
|
|
program,
|
|
out_idx=[2],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
})
|
|
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
|
|
|
def ref_program(A, B):
|
|
import torch
|
|
|
|
if trans_A:
|
|
A = A.T
|
|
if trans_B:
|
|
B = B.T
|
|
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
|
C = C.to(torch.__getattribute__(out_dtype))
|
|
return C
|
|
|
|
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
def test_gemm_sr():
|
|
# GEMM tests for float16
|
|
run_gemm_sr(512, 1024, 768, False, False, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_sr(512, 1024, 768, False, True, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_sr(512, 1024, 768, True, False, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_sr(512, 1024, 768, True, True, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
|
|
# n8 tests
|
|
run_gemm_sr(128, 8, 32, False, True, "float16", "float16", "float16", 128, 8, 32, 0, 128)
|
|
|
|
# int8 tests
|
|
run_gemm_sr(128, 128, 32, False, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_sr(128, 128, 32, False, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_sr(128, 128, 32, True, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_sr(128, 128, 32, True, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
|
|
# float8 tests
|
|
run_gemm_sr(128, 128, 128, True, True, "float8_e5m2", "float8_e5m2", "float32", 128, 128, 32, 2)
|
|
|
|
# float32 tests
|
|
run_gemm_sr(128, 128, 128, False, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_sr(128, 128, 128, False, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_sr(128, 128, 128, True, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_sr(128, 128, 128, True, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
|
|
|
|
def matmul_rr(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
accum_dtype,
|
|
num_stages,
|
|
threads,
|
|
):
|
|
A_shape = (K, M) if trans_A else (M, K)
|
|
B_shape = (N, K) if trans_B else (K, N)
|
|
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
|
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
|
|
A_frag_shape = A_shared_shape
|
|
B_frag_shape = B_shared_shape
|
|
|
|
import tilelang.language as T
|
|
|
|
@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)
|
|
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
|
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
|
|
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
T.clear(C_local)
|
|
T.annotate_layout({
|
|
A_shared: tilelang.layout.make_swizzled_layout(A_shared),
|
|
B_shared: tilelang.layout.make_swizzled_layout(B_shared),
|
|
})
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
if trans_A:
|
|
T.copy(A[k * block_K, by * block_M], A_shared)
|
|
else:
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
if trans_B:
|
|
T.copy(B[bx * block_N, k * block_K], B_shared)
|
|
else:
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.copy(A_shared, A_frag)
|
|
T.copy(B_shared, B_frag)
|
|
T.gemm_v2(A_frag, B_frag, C_local, trans_A, trans_B)
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_gemm_rr(
|
|
M,
|
|
N,
|
|
K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages=3,
|
|
num_threads=128,
|
|
):
|
|
program = matmul_rr(
|
|
M,
|
|
N,
|
|
K,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
trans_A,
|
|
trans_B,
|
|
in_dtype,
|
|
out_dtype,
|
|
dtypeAccum,
|
|
num_stages,
|
|
num_threads,
|
|
)
|
|
|
|
kernel = tilelang.compile(
|
|
program,
|
|
out_idx=[2],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
})
|
|
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
|
|
|
|
def ref_program(A, B):
|
|
import torch
|
|
|
|
if trans_A:
|
|
A = A.T
|
|
if trans_B:
|
|
B = B.T
|
|
C = torch.matmul(A.to(torch.float), B.to(torch.float))
|
|
C = C.to(torch.__getattribute__(out_dtype))
|
|
return C
|
|
|
|
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
def test_gemm_rr():
|
|
# GEMM tests for float16
|
|
run_gemm_rr(512, 1024, 768, False, False, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rr(512, 1024, 768, False, True, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rr(512, 1024, 768, True, False, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rr(512, 1024, 768, True, True, "float16", "float16", "float16", 128, 256, 32, 2)
|
|
run_gemm_rr(512, 1024, 768, False, True, "bfloat16", "bfloat16", "float", 128, 256, 32, 2)
|
|
# n8 tests
|
|
run_gemm_rr(128, 8, 128, False, True, "float16", "float16", "float16", 128, 8, 32, 2)
|
|
run_gemm_rr(128, 8, 128, False, True, "int8", "int8", "int32", 128, 8, 32, 2)
|
|
|
|
# int8 tests
|
|
run_gemm_rr(128, 128, 128, False, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_rr(128, 128, 128, False, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_rr(128, 128, 128, True, False, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
run_gemm_rr(128, 128, 128, True, True, "int8", "int8", "int32", 128, 128, 32, 2)
|
|
|
|
# float8 tests
|
|
run_gemm_rr(128, 128, 128, True, True, "float8_e5m2", "float8_e5m2", "float32", 128, 128, 32, 2)
|
|
|
|
# float32 tests
|
|
run_gemm_rr(128, 128, 128, False, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_rr(128, 128, 128, False, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_rr(128, 128, 128, True, False, "float", "float", "float32", 128, 128, 32, 2)
|
|
run_gemm_rr(128, 128, 128, True, True, "float", "float", "float32", 128, 128, 32, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|