gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[BACKEND] Add isRow attribute for DotOp tensors whose parent is mmav1 (#970)

Browse Source 
Co-authored-by: Yan Chunwei <yanchunwei@outlook.com>
This commit is contained in:
Philippe Tillet
2022-12-11 19:01:57 -08:00
committed by GitHub
parent 4fb048873a
commit 52accd4c2b
7 changed files with 186 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/tests/test_gemm.py
View File

@@ -32,7 +32,7 @@ def matmul_no_scf_kernel(
(shape, num_warps, trans_a, trans_b)
for shape in [
[128, 256, 32],
[256, 128, 16],
# [256, 128, 16],
[128, 16, 32],
[32, 128, 64],
[128, 128, 64],
@@ -43,8 +43,6 @@ def matmul_no_scf_kernel(
for trans_b in [False, True]
])
def test_gemm_no_scf(SHAPE, NUM_WARPS, TRANS_A, TRANS_B):
guard_for_volta(NUM_WARPS, TRANS_A, TRANS_B)
SIZE_M, SIZE_N, SIZE_K = SHAPE
if (TRANS_A):
a = torch.randn((SIZE_K, SIZE_M), device='cuda', dtype=torch.float16).T
@@ -83,7 +81,7 @@ def test_gemm_no_scf(SHAPE, NUM_WARPS, TRANS_A, TRANS_B):
for trans_b in [False, True]
])
def test_gemm_no_scf_int8(SHAPE, NUM_WARPS, TRANS_A, TRANS_B):
guard_for_volta(NUM_WARPS, TRANS_A, TRANS_B, is_int8=True)
guard_for_volta(is_int8=True)
SIZE_M, SIZE_N, SIZE_K = SHAPE
@@ -199,7 +197,6 @@ def get_proper_err(a, b, golden):
[128, 64, 128, 4, 128, 64, 32, False, True],
])
def test_gemm(SIZE_M, SIZE_N, SIZE_K, NUM_WARPS, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, TRANS_A, TRANS_B):
guard_for_volta(NUM_WARPS, TRANS_A, TRANS_B)
if (TRANS_A):
a = torch.randn((SIZE_K, SIZE_M), device='cuda', dtype=torch.float16).T
@@ -276,7 +273,7 @@ def test_gemm_fp32(M, N, K, num_warps, block_M, block_N, block_K, allow_tf32):
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, accumulator, c_mask)
guard_for_volta(num_warps, trans_a=False, trans_b=False, is_tf32=allow_tf32)
guard_for_volta(is_tf32=allow_tf32)
# Configure the pytorch counterpart
torch.backends.cuda.matmul.allow_tf32 = allow_tf32
@@ -302,7 +299,7 @@ def test_gemm_fp32(M, N, K, num_warps, block_M, block_N, block_K, allow_tf32):
torch.testing.assert_close(c, golden, rtol=max(1e-4, 1.5 * golden_rel_err), atol=max(1e-4, 1.5 * golden_abs_err))
def guard_for_volta(num_warps, trans_a, trans_b, is_int8=False, is_tf32=False):
def guard_for_volta(is_int8=False, is_tf32=False):
'''
Tell whether the test case is valid on Volta GPU.
Some features are WIP, so the corresponding support are missing.
@@ -311,8 +308,7 @@ def guard_for_volta(num_warps, trans_a, trans_b, is_int8=False, is_tf32=False):
is_on_Volta = capability[0] < 8
# TODO[Superjomn]: Remove the constraints below when features are ready
is_feature_supported = not (is_int8 or is_tf32)
is_feature_ready = not (trans_a or trans_b)
if is_on_Volta:
if (not is_feature_supported) or (not is_feature_ready):
pytest.skip("Not valid on Volta")
if (not is_feature_supported):
pytest.skip("Not valid on Volta")