mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-27 07:32:58 +08:00
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
import pytest
|
|
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def _make_gemm(M, N, K):
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), "float16"),
|
|
B: T.Tensor((K, N), "float16"),
|
|
C: T.Tensor((M, N), "float16"),
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_shared = T.alloc_shared((M, K), "float16")
|
|
B_shared = T.alloc_shared((K, N), "float16")
|
|
C_local = T.alloc_fragment((M, N), "float16")
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
return main
|
|
|
|
|
|
def _make_gemm_sp(M, N, Kc):
|
|
@T.prim_func
|
|
def main(
|
|
A_sparse: T.Tensor((M, Kc), "float16"),
|
|
E: T.Tensor((M, Kc // 4), "uint8"),
|
|
B: T.Tensor((2 * Kc, N), "float16"),
|
|
C: T.Tensor((M, N), "float16"),
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_shared = T.alloc_shared((M, Kc), "float16")
|
|
E_shared = T.alloc_shared((M, Kc // 4), "uint8")
|
|
B_shared = T.alloc_shared((2 * Kc, N), "float16")
|
|
C_local = T.alloc_fragment((M, N), "float16")
|
|
T.gemm_sp(A_shared, E_shared, B_shared, C_local)
|
|
|
|
return main
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("M", "N", "K"),
|
|
[(T.dynamic("M"), 128, 32), (128, T.dynamic("N"), 32), (128, 128, T.dynamic("K"))],
|
|
)
|
|
def test_gemm_rejects_symbolic_dimension(M, N, K):
|
|
with pytest.raises(ValueError):
|
|
_make_gemm(M, N, K)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("M", "N", "K"),
|
|
[(T.dynamic("M"), 128, 32), (128, T.dynamic("N"), 32), (128, 128, T.dynamic("K"))],
|
|
)
|
|
def test_gemm_sp_rejects_symbolic_dimension(M, N, K):
|
|
with pytest.raises(ValueError):
|
|
_make_gemm_sp(M, N, K)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|