mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-28 07:43:01 +08:00
Kurisandcoderabbitai[bot]
0b6336b583
* Refactor GEMM tests to use parameterized pytest fixtures - Converted multiple test cases for GEMM operations in `test_tilelang_tilelibrary_gemm_sp.py` to use `pytest.mark.parametrize` for better maintainability and readability. - Similar refactoring applied to `test_tilelang_tilelibrary_gemm_sp_v2.py`, consolidating test cases for `run_gemm_ss`, `run_gemm_rs`, `run_gemm_sr`, and `run_gemm_rr` into parameterized tests. - This change reduces code duplication and enhances the clarity of test configurations. * Update testing/python/amd/test_tilelang_gemm_mfma_preshuffle.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
108 lines
3.6 KiB
Python
108 lines
3.6 KiB
Python
import pytest
|
|
import torch
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
str2dtype = {
|
|
"float32": torch.float32,
|
|
"float16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"float8_e4m3": torch.float8_e4m3fn,
|
|
"float8_e5m2": torch.float8_e5m2,
|
|
}
|
|
|
|
|
|
@tilelang.jit
|
|
def vectorized_cast_kernel(M: int, dtype_A: str, dtype_B: str):
|
|
assert M % 256 == 0
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(M,), dtype_A], # noqa: F821
|
|
B: T.Tensor[(M,), dtype_B], # noqa: F821
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
T.copy(A, B)
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.jit
|
|
def parallel_vectorized_cast_kernel(M: int, dtype_A: str, dtype_B: str):
|
|
assert M % 256 == 0
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(M,), dtype_A], # noqa: F821
|
|
B: T.Tensor[(M,), dtype_B], # noqa: F821
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_local = T.alloc_fragment((M,), dtype_A)
|
|
B_local = T.alloc_fragment((M,), dtype_B)
|
|
|
|
T.copy(A, A_local)
|
|
for i in T.Parallel(M):
|
|
B_local[i] = A_local[i]
|
|
T.copy(B_local, B)
|
|
|
|
return main
|
|
|
|
|
|
def run_vectorized_cast(src_dtype_str: str, dst_dtype_str: str, check_str: str, lanes: int = 2):
|
|
"""Run the vectorized cast kernel and check the correctness.
|
|
Args:
|
|
src_dtype_str: The source data type string.
|
|
dst_dtype_str: The destination data type string.
|
|
check_str: Used to ensure vectorized cast is used.
|
|
lanes: The number of lanes of the source and destination data types.
|
|
"""
|
|
|
|
M = 128 * lanes
|
|
kernel = vectorized_cast_kernel(M, src_dtype_str, dst_dtype_str)
|
|
kernel_parallel = parallel_vectorized_cast_kernel(M, src_dtype_str, dst_dtype_str)
|
|
|
|
A_float = torch.randn(M, dtype=torch.float32, device="cuda")
|
|
A = A_float.to(str2dtype[src_dtype_str])
|
|
B = torch.zeros(M, dtype=str2dtype[dst_dtype_str], device="cuda")
|
|
C = torch.zeros(M, dtype=str2dtype[dst_dtype_str], device="cuda")
|
|
|
|
kernel(A, B)
|
|
kernel_parallel(A, C)
|
|
|
|
torch.testing.assert_close(A.to(str2dtype[dst_dtype_str]), B)
|
|
torch.testing.assert_close(A.to(str2dtype[dst_dtype_str]), C)
|
|
|
|
code = kernel.get_kernel_source()
|
|
code_parallel = kernel_parallel.get_kernel_source()
|
|
|
|
assert check_str in code and check_str in code_parallel, f"Cast {src_dtype_str} to {dst_dtype_str} with {lanes=} is not vectorized!"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"src_dtype, dst_dtype, check_str, lanes",
|
|
[
|
|
("float32", "float16", "__float22half2_rn", 2),
|
|
("float32", "float16", "__float22half2_rn", 4),
|
|
("float16", "float32", "__half22float2", 2),
|
|
("float16", "float32", "__half22float2", 4),
|
|
("float32", "float8_e4m3", "__nv_cvt_float2_to_fp8x2", 2),
|
|
("float32", "float8_e4m3", "__nv_cvt_float2_to_fp8x2", 4),
|
|
("float32", "float8_e5m2", "__nv_cvt_float2_to_fp8x2", 2),
|
|
("float32", "float8_e5m2", "__nv_cvt_float2_to_fp8x2", 4),
|
|
("float32", "bfloat16", "__float22bfloat162_rn", 2),
|
|
("float32", "bfloat16", "__float22bfloat162_rn", 4),
|
|
("bfloat16", "float32", "__bfloat1622float2", 2),
|
|
("bfloat16", "float32", "__bfloat1622float2", 4),
|
|
("float8_e4m3", "float32", "__tl_cvt_fp8x2_to_float2", 2),
|
|
("float8_e4m3", "float32", "__tl_cvt_fp8x2_to_float2", 4),
|
|
("float8_e5m2", "float32", "__tl_cvt_fp8x2_to_float2", 2),
|
|
("float8_e5m2", "float32", "__tl_cvt_fp8x2_to_float2", 4),
|
|
],
|
|
)
|
|
def test_vectorized_cast(src_dtype, dst_dtype, check_str, lanes):
|
|
run_vectorized_cast(src_dtype, dst_dtype, check_str, lanes)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|