trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-24 07:04:00 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
6873e19da2b267d3cb7d0305561ba0bed63feae3
tilelang/testing/python/tilelibrary/test_tilelang_tilelibrary_gemm_sp.py
T

762 lines
24 KiB
Python
Raw Blame History

import pytest
from examples.gemm_sp.sparse_utils import compress, randn_semi_sparse, randint_semi_sparse, get_e_factor
from tilelang.utils.tensor import torch_assert_close
import tilelang.testing
import torch
import tilelang.language as T
def matmul(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
accum_dtype,
metadata_dtype,
E_factor,
num_stages,
threads,
):
A_sparse_shape = (M, K // 2) if not trans_A else (K // 2, M)
B_shape = (N, K) if trans_B else (K, N)
A_shared_shape = (block_M, block_K // 2) if not trans_A else (block_K // 2, block_M)
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
E_shape = (M, K // E_factor) if not trans_A else (K // E_factor, M)
E_shared_shape = (block_M, block_K // E_factor) if not trans_A else (block_K // E_factor, block_M)
@T.prim_func
def main(
A_sparse: T.Tensor(A_sparse_shape, in_dtype),
E: T.Tensor(E_shape, metadata_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)
E_shared = T.alloc_shared(E_shared_shape, metadata_dtype)
C_frag = T.alloc_fragment((block_M, block_N), accum_dtype)
T.clear(C_frag)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
if trans_A:
T.copy(E[k * block_K // E_factor, by * block_M], E_shared)
T.copy(A_sparse[k * block_K // 2, by * block_M], A_shared)
else:
T.copy(E[by * block_M, k * block_K // E_factor], E_shared)
T.copy(A_sparse[by * block_M, k * block_K // 2], 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_sp(A_shared, E_shared, B_shared, C_frag, trans_A, trans_A, trans_B)
T.copy(C_frag, 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,
num_threads,
meta_dtype,
):
metadata_dtype = meta_dtype
program = matmul(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
metadata_dtype,
get_e_factor(in_dtype, metadata_dtype),
num_stages,
num_threads,
)
kernel = tilelang.compile(
program,
out_idx=[3],
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
)
A, B = generate_dense_input(M, N, K, trans_A, trans_B, in_dtype)
A_sparse, E = compress(A.t().contiguous() if trans_A else A, meta_dtype=meta_dtype.as_torch())
if trans_A:
A_sparse = A_sparse.t().contiguous()
E = E.t().contiguous()
C_sp = kernel(A_sparse, E, B)
def _matmul(A, B):
if trans_A:
A = A.T
if trans_B:
B = B.T
A = A.to(torch.float32)
B = B.to(torch.float32)
return torch.matmul(A, B)
C = _matmul(A, B)
torch_assert_close(
C_sp.to(out_dtype.as_torch()).to(torch.float32),
C.to(out_dtype.as_torch()).to(torch.float32),
rtol=1e-3,
atol=1e-3,
base_name="tilelang_sp",
ref_name="ref_dense",
)
def generate_dense_input(M, N, K, trans_A, trans_B, in_dtype, seed=0):
torch.manual_seed(seed)
is_8bit = "8" in in_dtype
is_unsigned = "uint" in in_dtype
is_int = "int" in in_dtype
if is_int:
if is_8bit:
low, high = (0, 4) if is_unsigned else (-2, 2)
else:
low, high = (0, 128) if is_unsigned else (-64, 64)
A = randint_semi_sparse(M, K, low=low, high=high, dtype=in_dtype.as_torch(), device="cuda", transposed=trans_A)
B = torch.randint(size=(N, K) if trans_B else (K, N), low=low, high=high, dtype=in_dtype.as_torch(), device="cuda")
else:
A = randn_semi_sparse(M, K, dtype=in_dtype.as_torch(), device="cuda", transposed=trans_A)
B = torch.randn((N, K) if trans_B else (K, N), device="cuda", dtype=torch.float32).to(in_dtype.as_torch())
return A, B
@tilelang.testing.requires_cuda
@pytest.mark.parametrize(
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype",
[
(128, 128, 32, False, True, T.float16, T.float16, T.float, 128, 128, 32, 2, 128, T.int16),
(128, 128, 64, False, True, T.int8, T.int8, T.int32, 128, 128, 64, 2, 128, T.int32),
(128, 128, 32, False, False, T.float16, T.float16, T.float, 128, 128, 32, 2, 128, T.int16),
(64, 128, 32, True, False, T.float16, T.float16, T.float, 64, 128, 32, 2, 128, T.int16),
(64, 128, 32, True, True, T.float16, T.float16, T.float, 64, 128, 32, 2, 128, T.int16),
(128, 8, 64, False, True, T.float16, T.float16, T.float, 128, 8, 32, 0, 128, T.int16),
(128, 128, 32, False, True, T.bfloat16, T.bfloat16, T.float32, 128, 128, 32, 2, 128, T.int16),
(64, 128, 128, True, True, T.int8, T.int8, T.int32, 64, 128, 128, 2, 128, T.int32),
(128, 128, 64, False, True, T.int8, T.int8, T.int32, 128, 128, 64, 2, 128, T.int16),
(128, 128, 64, False, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 64, 2, 128, T.int32),
],
)
def test_gemm_ss(
M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype
):
run_gemm_ss(
M,
N,
K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
block_M,
block_N,
block_K,
num_stages,
num_threads,
meta_dtype=meta_dtype,
)
def matmul_rs(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
accum_dtype,
metadata_dtype,
E_factor,
num_stages,
threads,
):
A_sparse_shape = (M, K // 2) if not trans_A else (K // 2, M)
B_shape = (N, K) if trans_B else (K, N)
A_shared_shape = (block_M, block_K // 2) if not trans_A else (block_K // 2, block_M)
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
A_frag_shape = A_shared_shape
E_shape = (M, K // E_factor) if not trans_A else (K // E_factor, M)
E_shared_shape = (block_M, block_K // E_factor) if not trans_A else (block_K // E_factor, block_M)
import tilelang.language as T
@T.prim_func
def main(
A_sparse: T.Tensor(A_sparse_shape, in_dtype),
E: T.Tensor(E_shape, metadata_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)
E_shared = T.alloc_shared(E_shared_shape, metadata_dtype)
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
C_frag = T.alloc_fragment((block_M, block_N), accum_dtype)
T.annotate_layout(
{
A_shared: tilelang.layout.make_swizzled_layout(A_shared),
}
)
T.clear(C_frag)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
if trans_A:
T.copy(E[k * block_K // E_factor, by * block_M], E_shared)
T.copy(A_sparse[k * block_K // 2, by * block_M], A_shared)
else:
T.copy(E[by * block_M, k * block_K // E_factor], E_shared)
T.copy(A_sparse[by * block_M, k * block_K // 2], 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_sp(A_frag, E_shared, B_shared, C_frag, trans_A, trans_A, trans_B)
T.copy(C_frag, 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,
num_threads,
meta_dtype,
):
metadata_dtype = meta_dtype
program = matmul_rs(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
metadata_dtype,
get_e_factor(in_dtype, metadata_dtype),
num_stages,
num_threads,
)
kernel = tilelang.compile(
program,
out_idx=[3],
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
)
A, B = generate_dense_input(M, N, K, trans_A, trans_B, in_dtype)
A_sparse, E = compress(A.t().contiguous() if trans_A else A, meta_dtype=meta_dtype.as_torch())
if trans_A:
A_sparse = A_sparse.t().contiguous()
E = E.t().contiguous()
C_sp = kernel(A_sparse, E, B)
def _matmul(A, B):
if trans_A:
A = A.T
if trans_B:
B = B.T
A = A.to(torch.float32)
B = B.to(torch.float32)
return torch.matmul(A, B)
C = _matmul(A, B)
torch_assert_close(
C_sp.to(out_dtype.as_torch()).to(torch.float32),
C.to(out_dtype.as_torch()).to(torch.float32),
rtol=1e-3,
atol=1e-3,
base_name="tilelang_sp",
ref_name="ref_dense",
)
@tilelang.testing.requires_cuda
@pytest.mark.parametrize(
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype",
[
(128, 256, 32, False, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128, T.int16),
(128, 128, 64, False, True, T.int8, T.int8, T.int32, 128, 128, 64, 2, 128, T.int32),
(128, 256, 32, False, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128, T.int16),
(64, 256, 32, True, False, T.float16, T.float16, T.float32, 64, 256, 32, 2, 128, T.int16),
(64, 256, 32, True, True, T.float16, T.float16, T.float32, 64, 256, 32, 2, 128, T.int16),
(128, 8, 64, False, True, T.float16, T.float16, T.float32, 128, 8, 32, 0, 128, T.int16),
(128, 256, 32, False, True, T.bfloat16, T.bfloat16, T.float32, 128, 256, 32, 2, 128, T.int16),
(64, 128, 128, True, True, T.int8, T.int8, T.int32, 64, 128, 128, 2, 128, T.int32),
(128, 128, 64, False, True, T.int8, T.int8, T.int32, 128, 128, 64, 2, 128, T.int16),
(128, 128, 64, False, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 64, 2, 128, T.int32),
],
)
def test_gemm_rs(
M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype
):
run_gemm_rs(
M,
N,
K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
block_M,
block_N,
block_K,
num_stages,
num_threads,
meta_dtype=meta_dtype,
)
def matmul_sr(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
accum_dtype,
metadata_dtype,
E_factor,
num_stages,
threads,
):
A_sparse_shape = (M, K // 2) if not trans_A else (K // 2, M)
B_shape = (N, K) if trans_B else (K, N)
A_shared_shape = (block_M, block_K // 2) if not trans_A else (block_K // 2, block_M)
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_N)
B_frag_shape = B_shared_shape
E_shape = (M, K // E_factor) if not trans_A else (K // E_factor, M)
E_shared_shape = (block_M, block_K // E_factor) if not trans_A else (block_K // E_factor, block_M)
import tilelang.language as T
@T.prim_func
def main(
A_sparse: T.Tensor(A_sparse_shape, in_dtype),
E: T.Tensor(E_shape, metadata_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)
E_shared = T.alloc_shared(E_shared_shape, metadata_dtype)
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
C_frag = T.alloc_fragment((block_M, block_N), accum_dtype)
T.annotate_layout(
{
B_shared: tilelang.layout.make_swizzled_layout(B_shared),
}
)
T.clear(C_frag)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
if trans_A:
T.copy(E[k * block_K // E_factor, by * block_M], E_shared)
T.copy(A_sparse[k * block_K // 2, by * block_M], A_shared)
else:
T.copy(E[by * block_M, k * block_K // E_factor], E_shared)
T.copy(A_sparse[by * block_M, k * block_K // 2], 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_sp(A_shared, E_shared, B_frag, C_frag, trans_A, trans_A, trans_B)
T.copy(C_frag, 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,
num_threads,
meta_dtype,
):
metadata_dtype = meta_dtype
program = matmul_sr(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
metadata_dtype,
get_e_factor(in_dtype, metadata_dtype),
num_stages,
num_threads,
)
kernel = tilelang.compile(
program,
out_idx=[3],
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
)
A, B = generate_dense_input(M, N, K, trans_A, trans_B, in_dtype)
A_sparse, E = compress(A.t().contiguous() if trans_A else A, meta_dtype=meta_dtype.as_torch())
if trans_A:
A_sparse = A_sparse.t().contiguous()
E = E.t().contiguous()
C_sp = kernel(A_sparse, E, B)
def _matmul(A, B):
if trans_A:
A = A.T
if trans_B:
B = B.T
A = A.to(torch.float32)
B = B.to(torch.float32)
return torch.matmul(A, B)
C = _matmul(A, B)
torch_assert_close(
C_sp.to(out_dtype.as_torch()).to(torch.float32),
C.to(out_dtype.as_torch()).to(torch.float32),
rtol=1e-3,
atol=1e-3,
base_name="tilelang_sp",
ref_name="ref_dense",
)
@tilelang.testing.requires_cuda
@pytest.mark.parametrize(
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype",
[
(128, 256, 32, False, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128, T.int16),
(128, 128, 128, False, True, T.int8, T.int8, T.int32, 128, 128, 128, 2, 128, T.int32),
(128, 256, 32, False, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128, T.int16),
(64, 256, 32, True, False, T.float16, T.float16, T.float32, 64, 256, 32, 2, 128, T.int16),
(64, 256, 32, True, True, T.float16, T.float16, T.float32, 64, 256, 32, 2, 128, T.int16),
(128, 8, 64, False, True, T.float16, T.float16, T.float32, 128, 8, 32, 0, 128, T.int16),
(128, 256, 32, False, True, T.bfloat16, T.bfloat16, T.float32, 128, 256, 32, 2, 128, T.int16),
(64, 128, 128, True, True, T.int8, T.int8, T.int32, 64, 128, 128, 2, 128, T.int32),
(128, 128, 128, False, True, T.int8, T.int8, T.int32, 128, 128, 128, 2, 128, T.int16),
(128, 128, 64, False, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 64, 2, 128, T.int32),
],
)
def test_gemm_sr(
M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype
):
run_gemm_sr(
M,
N,
K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
block_M,
block_N,
block_K,
num_stages,
num_threads,
meta_dtype=meta_dtype,
)
def matmul_rr(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
accum_dtype,
metadata_dtype,
E_factor,
num_stages,
threads,
):
A_sparse_shape = (M, K // 2) if not trans_A else (K // 2, M)
B_shape = (N, K) if trans_B else (K, N)
A_shared_shape = (block_M, block_K // 2) if not trans_A else (block_K // 2, block_M)
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
E_shape = (M, K // E_factor) if not trans_A else (K // E_factor, M)
E_shared_shape = (block_M, block_K // E_factor) if not trans_A else (block_K // E_factor, block_M)
import tilelang.language as T
@T.prim_func
def main(
A_sparse: T.Tensor(A_sparse_shape, in_dtype),
E: T.Tensor(E_shape, metadata_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)
E_shared = T.alloc_shared(E_shared_shape, metadata_dtype)
A_frag = T.alloc_fragment(A_frag_shape, in_dtype)
B_frag = T.alloc_fragment(B_frag_shape, in_dtype)
C_frag = T.alloc_fragment((block_M, block_N), accum_dtype)
T.annotate_layout(
{
A_shared: tilelang.layout.make_swizzled_layout(A_shared),
B_shared: tilelang.layout.make_swizzled_layout(B_shared),
}
)
T.clear(C_frag)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
if trans_A:
T.copy(E[k * block_K // E_factor, by * block_M], E_shared)
T.copy(A_sparse[k * block_K // 2, by * block_M], A_shared)
else:
T.copy(E[by * block_M, k * block_K // E_factor], E_shared)
T.copy(A_sparse[by * block_M, k * block_K // 2], 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_sp(A_frag, E_shared, B_frag, C_frag, trans_A, trans_A, trans_B)
T.copy(C_frag, 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,
meta_dtype=T.int16,
):
metadata_dtype = meta_dtype
program = matmul_rr(
M,
N,
K,
block_M,
block_N,
block_K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
metadata_dtype,
get_e_factor(in_dtype, metadata_dtype),
num_stages,
num_threads,
)
kernel = tilelang.compile(
program,
out_idx=[3],
pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
)
A, B = generate_dense_input(M, N, K, trans_A, trans_B, in_dtype)
A_sparse, E = compress(A.t().contiguous() if trans_A else A, meta_dtype=meta_dtype.as_torch())
if trans_A:
A_sparse = A_sparse.t().contiguous()
E = E.t().contiguous()
C_sp = kernel(A_sparse, E, B)
def _matmul(A, B):
if trans_A:
A = A.T
if trans_B:
B = B.T
A = A.to(torch.float32)
B = B.to(torch.float32)
return torch.matmul(A, B)
C = _matmul(A, B)
torch_assert_close(
C_sp.to(out_dtype.as_torch()).to(torch.float32),
C.to(out_dtype.as_torch()).to(torch.float32),
rtol=1e-3,
atol=1e-3,
base_name="tilelang_sp",
ref_name="ref_dense",
)
@tilelang.testing.requires_cuda
@pytest.mark.parametrize(
"M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype",
[
(128, 256, 32, False, True, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128, T.int16),
(128, 128, 64, False, True, T.int8, T.int8, T.int32, 128, 128, 64, 2, 128, T.int32),
(128, 256, 32, False, False, T.float16, T.float16, T.float32, 128, 256, 32, 2, 128, T.int16),
(64, 256, 32, True, False, T.float16, T.float16, T.float32, 64, 256, 32, 2, 128, T.int16),
(64, 256, 32, True, True, T.float16, T.float16, T.float32, 64, 256, 32, 2, 128, T.int16),
(128, 256, 32, False, True, T.bfloat16, T.bfloat16, T.float32, 128, 256, 32, 2, 128, T.int16),
(128, 8, 32, False, True, T.float16, T.float16, T.float32, 128, 8, 32, 2, 128, T.int16),
(128, 8, 64, False, True, T.int8, T.int8, T.int32, 128, 8, 64, 2, 128, T.int32),
(64, 128, 128, True, True, T.int8, T.int8, T.int32, 64, 128, 128, 2, 128, T.int32),
(128, 128, 64, False, True, T.int8, T.int8, T.int32, 128, 128, 64, 2, 128, T.int16),
(128, 128, 64, False, True, T.float8_e5m2, T.float8_e5m2, T.float32, 128, 128, 64, 2, 128, T.int32),
],
)
def test_gemm_rr(
M, N, K, trans_A, trans_B, in_dtype, out_dtype, dtypeAccum, block_M, block_N, block_K, num_stages, num_threads, meta_dtype
):
run_gemm_rr(
M,
N,
K,
trans_A,
trans_B,
in_dtype,
out_dtype,
dtypeAccum,
block_M,
block_N,
block_K,
num_stages,
num_threads,
meta_dtype=meta_dtype,
)
@tilelang.testing.requires_cuda
@pytest.mark.parametrize(
"in_dtype, out_dtype, dtypeAccum, meta_dtype",
[
(T.float16, T.float16, T.float32, T.int16),
(T.int8, T.int8, T.int32, T.int32),
(
T.float16,
T.float16,
T.float32,
T.int8,
),
(
T.bfloat16,
T.bfloat16,
T.float32,
T.int8,
),
(
T.bfloat16,
T.bfloat16,
T.float32,
T.int16,
),
(
T.int8,
T.int8,
T.int32,
T.int8,
),
(
T.int8,
T.int8,
T.int32,
T.int16,
),
(
T.float8_e5m2,
T.float8_e5m2,
T.float32,
T.int8,
),
(
T.float8_e5m2,
T.float8_e5m2,
T.float32,
T.int16,
),
(
T.float8_e5m2,
T.float8_e5m2,
T.float32,
T.int32,
),
],
)
def test_compress_dtype_combinations(in_dtype, out_dtype, dtypeAccum, meta_dtype):
run_gemm_ss(128, 128, 128, False, True, in_dtype, out_dtype, dtypeAccum, 128, 128, 64, 2, 128, meta_dtype=meta_dtype)
if __name__ == "__main__":
tilelang.testing.main()
Reference in New Issue View Git Blame Copy Permalink