Skip to content

[CI][TIRx] Compile tests before GPU execution gates #19985

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tlopex wants to merge 1 commit into
base: main
Choose a base branch
from
+1,777 −870
Open

[CI][TIRx] Compile tests before GPU execution gates #19985

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 34 additions & 8 deletions tests/python/tirx/codegen/test_codegen_ampere.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
D/C: 4 f32 accumulator registers (0,1,2,3)
"""

import functools

import numpy as np
import pytest

Expand All @@ -39,13 +41,27 @@


def _get_source(func: tvm.tirx.PrimFunc):
target = tvm.target.Target("cuda")
target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
mod = tvm.IRModule({"main": func})
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
with target:
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
src = mod.mod.imports[0].inspect_source()
return src, mod


def _xfail_pointer_list_parser(func):
@functools.wraps(func)
def wrapped(*args, **kwargs):
try:
return func(*args, **kwargs)
except tvm.error.DiagnosticError as err:
if "Cannot automatically inference the type. value=ir.Call(" in str(err):
pytest.xfail("TIRx parser cannot yet bind PTX pointer-list calls")
raise

return wrapped


def _np_in(dtype):
if dtype == "bfloat16":
return __import__("ml_dtypes").bfloat16
Expand All @@ -54,6 +70,8 @@ def _np_in(dtype):

def _run_mma(mod, K, no_c_ptr, np_in):
"""Run an m16n8kK mma kernel and check D == A @ B (+ C) against numpy."""
if not env.has_cuda_compute(8, exact=True):
pytest.skip("requires a CUDA compute capability 8.0 device to execute")
np.random.seed(0)
A_np = np.random.randn(16, K).astype(np_in)
B_np = np.random.randn(K, 8).astype(np_in)
Expand All @@ -74,11 +92,13 @@ def run_and_check():
tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
@pytest.mark.parametrize(
"execute", [False, pytest.param(True, marks=pytest.mark.gpu)], ids=["compile", "run"]
)
@pytest.mark.parametrize("a_type", ["float16", "bfloat16"])
@pytest.mark.parametrize("no_c_ptr", [False, True])
def test_ptx_mma_m16n8k16(a_type, no_c_ptr):
@_xfail_pointer_list_parser
def test_ptx_mma_m16n8k16(execute, a_type, no_c_ptr):
"""m16n8k16 row.col mma, f32 accumulate: A is 16x16 (4 b32/lane), B is 16x8
as [K, N] (2 b32/lane), D/C is 16x8 (4 f32/lane)."""
if a_type == "bfloat16":
Expand Down Expand Up @@ -142,14 +162,18 @@ def G2L(buf_local, buf_global, block_8x8, mode="row"):

src, mod = _get_source(main)
assert "mma.sync.aligned.m16n8k16.row.col" in src
if not execute:
return
_run_mma(mod, 16, no_c_ptr, _np_in(a_type))


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
@pytest.mark.parametrize(
"execute", [False, pytest.param(True, marks=pytest.mark.gpu)], ids=["compile", "run"]
)
@pytest.mark.parametrize("a_type", ["float16", "bfloat16"])
@pytest.mark.parametrize("no_c_ptr", [False, True])
def test_ptx_mma_m16n8k8(a_type, no_c_ptr):
@_xfail_pointer_list_parser
def test_ptx_mma_m16n8k8(execute, a_type, no_c_ptr):
"""m16n8k8 row.col mma, f32 accumulate: A is 16x8 (2 b32/lane), B is 8x8
as [K, N] (1 b32/lane), D/C is 16x8 (4 f32/lane)."""
if a_type == "bfloat16":
Expand Down Expand Up @@ -213,6 +237,8 @@ def G2L(buf_local, buf_global, block_8x8, mode="row"):

src, mod = _get_source(main)
assert "mma.sync.aligned.m16n8k8.row.col" in src
if not execute:
return
_run_mma(mod, 8, no_c_ptr, _np_in(a_type))


Expand Down
107 changes: 55 additions & 52 deletions tests/python/tirx/codegen/test_codegen_blackwell.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,17 +24,27 @@
from tvm.script.tirx import tile as Tx
from tvm.testing import env

compile_and_run = pytest.mark.parametrize(
"run",
[False, pytest.param(True, marks=pytest.mark.gpu)],
ids=["compile", "run"],
)

def _get_source(func: tvm.tirx.PrimFunc) -> str:
target = tvm.target.Target("cuda")

def _require_cuda_compute_10x():
if not env.has_cuda_compute(10, exact=True):
pytest.skip("requires a CUDA compute capability 10.0 device to execute")


def _get_source(func: tvm.tirx.PrimFunc) -> tuple[str, tvm.runtime.Module]:
target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
mod = tvm.IRModule({"main": func})
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
with target:
mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
src = mod.mod.imports[0].inspect_source()
return src, mod


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_tmem_alloc_dealloc_relinquish():
N_COLS = 512
cta_group = 1
Expand All @@ -61,16 +71,12 @@ def test_tmem(A: T.Buffer((16, 16), "float16")):
T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
# fmt: on

target = tvm.target.Target("cuda")
with target:
src, _ = _get_source(test_tmem)
assert f"tcgen05.alloc.cta_group::{cta_group}.sync.aligned.shared::cta.b32" in src
assert f"tcgen05.dealloc.cta_group::{cta_group}.sync.aligned.b32" in src
assert f"tcgen05.relinquish_alloc_permit.cta_group::{cta_group}.sync.aligned" in src
src, _ = _get_source(test_tmem)
assert f"tcgen05.alloc.cta_group::{cta_group}.sync.aligned.shared::cta.b32" in src
assert f"tcgen05.dealloc.cta_group::{cta_group}.sync.aligned.b32" in src
assert f"tcgen05.relinquish_alloc_permit.cta_group::{cta_group}.sync.aligned" in src


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_mbarrier_try_wait_once_codegen():
# fmt: off
@T.prim_func
Expand All @@ -82,15 +88,11 @@ def test_try_wait_once(A: T.Buffer((16, 16), "float16")):
T.evaluate(T.ptx.mbarrier.try_wait_once(T.address_of(bar), 0, 0))
# fmt: on

target = tvm.target.Target("cuda")
with target:
src, _ = _get_source(test_try_wait_once)
assert "mbarrier.try_wait.parity.shared::cta.b64" in src
assert "selp.u32" in src
src, _ = _get_source(test_try_wait_once)
assert "mbarrier.try_wait.parity.shared::cta.b64" in src
assert "selp.u32" in src


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_fence_before_after_thread_sync():
# fmt: off
@T.prim_func
Expand All @@ -105,16 +107,13 @@ def test_fence(A: T.Buffer((16, 16), "float16")):
T.ptx.tcgen05.fence.after_thread_sync()
# fmt: on

target = tvm.target.Target("cuda")
with target:
src, _ = _get_source(test_fence)
assert "tcgen05.fence::after_thread_sync" in src
assert "tcgen05.fence::before_thread_sync" in src
src, _ = _get_source(test_fence)
assert "tcgen05.fence::after_thread_sync" in src
assert "tcgen05.fence::before_thread_sync" in src


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_tcgen05_ld_st_roundtrip():
@compile_and_run
def test_tcgen05_ld_st_roundtrip(run):
HEIGHT = 128
WIDTH = 256
N_COLS = 512
Expand Down Expand Up @@ -164,11 +163,13 @@ def test_ld_st(A: T.Buffer((HEIGHT, WIDTH), "float32"), B: T.Buffer((HEIGHT, WID
T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
# fmt: on

target = tvm.target.Target("cuda")
with target:
src, mod = _get_source(test_ld_st)
assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
assert "tcgen05.st.sync.aligned.32x32b.x1.b32" in src
src, mod = _get_source(test_ld_st)
assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
assert "tcgen05.st.sync.aligned.32x32b.x1.b32" in src

if not run:
return
_require_cuda_compute_10x()

def run_and_check():
dev = tvm.cuda(0)
Expand All @@ -182,9 +183,8 @@ def run_and_check():
tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_tcgen05_cp_ld_roundtrip():
@compile_and_run
def test_tcgen05_cp_ld_roundtrip(run):
dtype = "float32"
dtype_bits = tvm.DataType(dtype).bits
HEIGHT = 128
Expand Down Expand Up @@ -249,11 +249,13 @@ def test_cp_ld(A: T.Buffer((HEIGHT, WIDTH), dtype, layout=T.TileLayout(T.S[(HEIG
T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
# fmt: on

target = tvm.target.Target("cuda")
with target:
src, mod = _get_source(test_cp_ld)
assert "tcgen05.cp.cta_group::1.128x256b" in src
assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
src, mod = _get_source(test_cp_ld)
assert "tcgen05.cp.cta_group::1.128x256b" in src
assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src

if not run:
return
_require_cuda_compute_10x()

def run_and_check():
dev = tvm.cuda(0)
Expand All @@ -268,9 +270,8 @@ def run_and_check():


@pytest.mark.parametrize("swizzle", [0, 1, 2, 3])
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_tcgen05_mma_ss_no_tma(swizzle):
@compile_and_run
def test_tcgen05_mma_ss_no_tma(swizzle, run):
d_type, a_type, b_type = "float32", "float16", "float16"
M, N, K = 128, 128, 64
MMA_K = 16
Expand Down Expand Up @@ -385,14 +386,16 @@ def test_mma_ss_no_tma(A: T.Buffer((M, K), a_type, layout=T.TileLayout(T.S[M, K]
import torch

torch.manual_seed(42)
target = tvm.target.Target("cuda")
with target:
src, mod = _get_source(test_mma_ss_no_tma)
print(src)
assert "tcgen05.mma.cta_group::1.kind::f16" in src
assert "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" in src
assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
assert "tcgen05.wait::ld.sync.aligned" in src
src, mod = _get_source(test_mma_ss_no_tma)
print(src)
assert "tcgen05.mma.cta_group::1.kind::f16" in src
assert "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" in src
assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
assert "tcgen05.wait::ld.sync.aligned" in src

if not run:
return
_require_cuda_compute_10x()

def run_and_check():
dev = tvm.cuda(0)
Expand Down
Loading
Loading