chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,690 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""Tests for the CUDA synchronous ``gemm`` (mma.sync) tensor-core dispatch.
|
||||
|
||||
The dispatch lowers ``tirx.tile.gemm`` over pure-register fragments to warp-level
|
||||
``mma.sync.aligned.m16n8k16/k8`` for bf16/f16 inputs with f32 accumulation.
|
||||
|
||||
The fragment layouts below are the standard m16n8 register maps (PTX ISA
|
||||
§9.7.13; see ``tests/python/tirx-base/test_tir_ptx_mma.py``):
|
||||
|
||||
lane = 4*g + t (g = lane >> 2 in [0, 8), t = lane & 3 in [0, 4))
|
||||
D/C[M, N]: M = g + 8*rM, N = 2*t + rN, c_id = 2*rM + rN
|
||||
A[M, K]: M = g + 8*rM, K = 2*t + p + 8*kHi, ma = p + 2*rM + 4*kHi
|
||||
B[K, N]: N = g, K = 2*t + p + 8*kHi, mb = p + 2*kHi
|
||||
|
||||
Most assertions run the CPU-only ``LowerTIRx`` transform; the numerical check
|
||||
is guarded by ``requires_cuda`` since it needs a real device.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.tirx import tile as Tx
|
||||
from tvm.testing import env
|
||||
from tvm.tirx.layout import S, TileLayout, laneid
|
||||
from tvm.tirx.operator.tile_primitive import list_registered_schedules
|
||||
|
||||
# Single-tile m16n8k8 fragment layouts -- the smallest unit everything else is
|
||||
# built from. A is 16x8, B is 8x8 as [K, N], D/C is 16x8 (the accumulator does
|
||||
# not depend on K). Every other layout (the k16 single tile, all tilings, and
|
||||
# the transposed orientations) is derived from these three via ``tile_to`` /
|
||||
# ``group`` + ``permute_by_groups``.
|
||||
D_FRAG = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
|
||||
A_FRAG_K8 = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
|
||||
B_FRAG_K8 = TileLayout(S[(4, 2, 8) : (1 @ laneid, 1, 4 @ laneid)])
|
||||
# m16n8k16 single tile = two k8 tiles stacked along K.
|
||||
A_FRAG = A_FRAG_K8.tile_to([16, 16], [16, 8])
|
||||
B_FRAG = B_FRAG_K8.tile_to([16, 8], [8, 8])
|
||||
|
||||
|
||||
def _transpose_frag(layout, shape):
|
||||
"""Swap the two logical axes of a 2D fragment layout.
|
||||
|
||||
The transposed input orientations (A as [K, M], B as [N, K]) hold the exact
|
||||
same per-lane/per-register element distribution as the K-major fragments --
|
||||
only the buffer's logical axes are swapped. So instead of writing them out
|
||||
by hand, derive them: ``group`` the shard into the logical dims, then
|
||||
``permute_by_groups`` to exchange the two groups.
|
||||
"""
|
||||
grouped, seps = layout.group(shape)
|
||||
return grouped.permute_by_groups(seps, [1, 0])
|
||||
|
||||
|
||||
# Transposed input orientations of the same single tile: A as [K, M], B as
|
||||
# [N, K]. The dispatch swaps axes per the transpose flags; the .row.col mma is
|
||||
# unchanged.
|
||||
A_KM_FRAG = _transpose_frag(A_FRAG, [16, 16])
|
||||
B_NK_FRAG = _transpose_frag(B_FRAG, [16, 8])
|
||||
|
||||
|
||||
def _frag(Mt, Nt, Kt, kinst):
|
||||
"""Fragment layouts for an Mt x Nt x Kt tiling of m16n8k{8,16}.
|
||||
|
||||
Logical shapes: A = (16*Mt, kinst*Kt), B = (kinst*Kt, 8*Nt), D/C = (16*Mt, 8*Nt).
|
||||
Each operand's tiled layout is the single-tile base ``tile_to`` the full
|
||||
logical shape -- ``tile_to`` repeats the base's per-lane/per-register element
|
||||
map over the tile grid, so a tiling is just a grid of the single-tile
|
||||
fragments (the base is the single source of truth, k8 and k16 alike).
|
||||
"""
|
||||
A_base = A_FRAG if kinst == 16 else A_FRAG_K8
|
||||
B_base = B_FRAG if kinst == 16 else B_FRAG_K8
|
||||
D = D_FRAG.tile_to([16 * Mt, 8 * Nt], [16, 8])
|
||||
A = A_base.tile_to([16 * Mt, kinst * Kt], [16, kinst])
|
||||
B = B_base.tile_to([kinst * Kt, 8 * Nt], [kinst, 8])
|
||||
return D, A, B
|
||||
|
||||
|
||||
def _build_tiled(Mt, Nt, Kt, kinst, *, beta=0.0, dtype="float16", store=False):
|
||||
"""A single-warp kernel issuing one ``T.gemm`` over an Mt x Nt x Kt tiling.
|
||||
|
||||
With ``store=True`` the result is written back to a global buffer (a full
|
||||
kernel for codegen); otherwise only the ``T.gemm`` is emitted (for
|
||||
``LowerTIRx`` dispatch checks).
|
||||
"""
|
||||
Dl, Al, Bl = _frag(Mt, Nt, Kt, kinst)
|
||||
M, N, K = 16 * Mt, 8 * Nt, kinst * Kt
|
||||
|
||||
if not store:
|
||||
|
||||
@T.prim_func
|
||||
def gemm():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
_lane = T.lane_id([32])
|
||||
A = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
|
||||
B = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
D = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
|
||||
|
||||
return gemm
|
||||
|
||||
@T.prim_func
|
||||
def gemm(D_ptr: T.handle):
|
||||
D_g = T.match_buffer(D_ptr, (M, N), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
|
||||
B = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
D = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
|
||||
# Decode D's per-thread registers (c = ((mt*Nt + nt)*2 + rM)*2 + rN)
|
||||
# back to logical (M, N) and store, exercising the whole tiling.
|
||||
D_reg = D.local(Mt * Nt * 4)
|
||||
for c in T.unroll(Mt * Nt * 4):
|
||||
rN = c % 2
|
||||
rM = (c // 2) % 2
|
||||
nt = (c // 4) % Nt
|
||||
mt = c // (4 * Nt)
|
||||
D_g[mt * 16 + lane // 4 + rM * 8, nt * 8 + (lane % 4) * 2 + rN] = D_reg[c]
|
||||
|
||||
return gemm
|
||||
|
||||
|
||||
def _build_gemm(alpha=1.0, beta=0.0, dtype="bfloat16"):
|
||||
"""A single-warp kernel issuing one ``T.gemm`` over register fragments."""
|
||||
|
||||
@T.prim_func
|
||||
def gemm_min():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_tid = T.thread_id([32])
|
||||
D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
A = T.alloc_buffer((16, 16), dtype, scope="local", layout=A_FRAG)
|
||||
B = T.alloc_buffer((16, 8), dtype, scope="local", layout=B_FRAG)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=alpha, beta=beta)
|
||||
|
||||
return gemm_min
|
||||
|
||||
|
||||
def _build_transpose(transpose_A, transpose_B, *, store=False):
|
||||
"""Single m16n8k16 tile with the requested A/B input orientations."""
|
||||
Al = A_KM_FRAG if transpose_A else A_FRAG
|
||||
Bl = B_NK_FRAG if transpose_B else B_FRAG
|
||||
A_shape = (16, 16) # [K, M] or [M, K] -- both 16x16 for one tile
|
||||
B_shape = (8, 16) if transpose_B else (16, 8)
|
||||
|
||||
if not store:
|
||||
|
||||
@T.prim_func
|
||||
def gemm():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
_lane = T.lane_id([32])
|
||||
A = T.alloc_buffer(A_shape, "float16", scope="local", layout=Al)
|
||||
B = T.alloc_buffer(B_shape, "float16", scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
Tx.warp.gemm(
|
||||
D,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A=transpose_A,
|
||||
transpose_B=transpose_B,
|
||||
alpha=1.0,
|
||||
beta=0.0,
|
||||
)
|
||||
|
||||
return gemm
|
||||
|
||||
@T.prim_func
|
||||
def gemm(D_ptr: T.handle):
|
||||
D_g = T.match_buffer(D_ptr, (16, 8), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A = T.alloc_buffer(A_shape, "float16", scope="local", layout=Al)
|
||||
B = T.alloc_buffer(B_shape, "float16", scope="local", layout=Bl)
|
||||
C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
Tx.warp.gemm(
|
||||
D,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
transpose_A=transpose_A,
|
||||
transpose_B=transpose_B,
|
||||
alpha=1.0,
|
||||
beta=0.0,
|
||||
)
|
||||
D_reg = D.local(4)
|
||||
for c in T.unroll(4):
|
||||
D_g[lane // 4 + (c // 2) * 8, (lane % 4) * 2 + c % 2] = D_reg[c]
|
||||
|
||||
return gemm
|
||||
|
||||
|
||||
def _build_dtypes(a_dtype, b_dtype, c_dtype, d_dtype):
|
||||
"""Single tile with explicit per-operand dtypes (for decline checks)."""
|
||||
|
||||
@T.prim_func
|
||||
def gemm_min():
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_tid = T.thread_id([32])
|
||||
D = T.alloc_buffer((16, 8), d_dtype, scope="local", layout=D_FRAG)
|
||||
C = T.alloc_buffer((16, 8), c_dtype, scope="local", layout=D_FRAG)
|
||||
A = T.alloc_buffer((16, 16), a_dtype, scope="local", layout=A_FRAG)
|
||||
B = T.alloc_buffer((16, 8), b_dtype, scope="local", layout=B_FRAG)
|
||||
Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0)
|
||||
|
||||
return gemm_min
|
||||
|
||||
|
||||
def _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype):
|
||||
"""End-to-end ``T.gemm`` over an Mt x Nt x Kt tiling, with the A/B inputs
|
||||
loaded and the D output stored register-by-register.
|
||||
|
||||
Fragments are indexed through their per-register multi-dim ``.local()`` views
|
||||
(the shard's non-lane dims, in shard order): A = [Mt, rM(2), Kt, kHi, kp],
|
||||
B = [Kt, kHi, kp, Nt], D/C = [Mt, rM(2), Nt, rN(2)]. The lane owns g = lane>>2
|
||||
and t = lane&3; within a tile M = mt*16 + rM*8 + g, N = nt*8 + t*2 + rN,
|
||||
K = kt*kinst + kHi*8 + t*2 + kp.
|
||||
"""
|
||||
Dl, Al, Bl = _frag(Mt, Nt, Kt, kinst)
|
||||
M, N, K = 16 * Mt, 8 * Nt, kinst * Kt
|
||||
KP = 2
|
||||
kHi_n = kinst // (4 * KP)
|
||||
|
||||
@T.prim_func
|
||||
def gemm(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, D_ptr: T.handle):
|
||||
A_g = T.match_buffer(A_ptr, (M, K), dtype)
|
||||
B_g = T.match_buffer(B_ptr, (K, N), dtype)
|
||||
C_g = T.match_buffer(C_ptr, (M, N), "float32")
|
||||
D_g = T.match_buffer(D_ptr, (M, N), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A_f = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
|
||||
B_f = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
|
||||
C_f = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
D_f = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
|
||||
A_reg = A_f.local(Mt, 2, Kt, kHi_n, KP)
|
||||
for mt, rM, kt, kHi, kp in T.grid(Mt, 2, Kt, kHi_n, KP):
|
||||
A_reg[mt, rM, kt, kHi, kp] = A_g[
|
||||
mt * 16 + lane // 4 + 8 * rM,
|
||||
kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
|
||||
]
|
||||
B_reg = B_f.local(Kt, kHi_n, KP, Nt)
|
||||
for kt, kHi, kp, nt in T.grid(Kt, kHi_n, KP, Nt):
|
||||
B_reg[kt, kHi, kp, nt] = B_g[
|
||||
kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
|
||||
nt * 8 + lane // 4,
|
||||
]
|
||||
if beta == 1.0:
|
||||
C_reg = C_f.local(Mt, 2, Nt, 2)
|
||||
for mt, rM, nt, rN in T.grid(Mt, 2, Nt, 2):
|
||||
C_reg[mt, rM, nt, rN] = C_g[
|
||||
mt * 16 + lane // 4 + 8 * rM, nt * 8 + 2 * (lane % 4) + rN
|
||||
]
|
||||
Tx.warp.gemm(D_f, A_f, B_f, C_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
|
||||
D_reg = D_f.local(Mt, 2, Nt, 2)
|
||||
for mt, rM, nt, rN in T.grid(Mt, 2, Nt, 2):
|
||||
D_g[mt * 16 + lane // 4 + 8 * rM, nt * 8 + 2 * (lane % 4) + rN] = D_reg[mt, rM, nt, rN]
|
||||
|
||||
return gemm, M, N, K
|
||||
|
||||
|
||||
def _build_transpose_numeric(transpose_A, transpose_B, dtype="float16"):
|
||||
"""End-to-end single-tile ``T.gemm`` for one A/B input orientation.
|
||||
|
||||
The transposed A fragment (``A_KM_FRAG``) carries its registers in the
|
||||
[kHi, kp, rM] shard order (vs [rM, kHi, kp] for the K-major ``A_FRAG``); B's
|
||||
register order ([kHi, kp]) is the same for both orientations. The buffer
|
||||
index axes swap with the orientation, but each register still holds the same
|
||||
logical (M, K) / (K, N) element.
|
||||
"""
|
||||
Al = A_KM_FRAG if transpose_A else A_FRAG
|
||||
Bl = B_NK_FRAG if transpose_B else B_FRAG
|
||||
A_shape = (16, 16)
|
||||
B_shape = (8, 16) if transpose_B else (16, 8)
|
||||
|
||||
@T.prim_func
|
||||
def gemm(A_ptr: T.handle, B_ptr: T.handle, D_ptr: T.handle):
|
||||
A_g = T.match_buffer(A_ptr, A_shape, dtype)
|
||||
B_g = T.match_buffer(B_ptr, B_shape, dtype)
|
||||
D_g = T.match_buffer(D_ptr, (16, 8), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A_f = T.alloc_buffer(A_shape, dtype, scope="local", layout=Al)
|
||||
B_f = T.alloc_buffer(B_shape, dtype, scope="local", layout=Bl)
|
||||
D_f = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
A_reg = A_f.local(2, 2, 2)
|
||||
if transpose_A:
|
||||
# A_KM_FRAG register order is [kHi, kp, rM]; buffer is [K, M].
|
||||
for kHi, kp, rM in T.grid(2, 2, 2):
|
||||
A_reg[kHi, kp, rM] = A_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4 + 8 * rM]
|
||||
else:
|
||||
# A_FRAG register order is [rM, kHi, kp]; buffer is [M, K].
|
||||
for rM, kHi, kp in T.grid(2, 2, 2):
|
||||
A_reg[rM, kHi, kp] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
|
||||
B_reg = B_f.local(2, 2)
|
||||
if transpose_B:
|
||||
# B_NK_FRAG buffer is [N, K].
|
||||
for kHi, kp in T.grid(2, 2):
|
||||
B_reg[kHi, kp] = B_g[lane // 4, 2 * (lane % 4) + kp + 8 * kHi]
|
||||
else:
|
||||
for kHi, kp in T.grid(2, 2):
|
||||
B_reg[kHi, kp] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
|
||||
Tx.warp.gemm(
|
||||
D_f,
|
||||
A_f,
|
||||
B_f,
|
||||
D_f,
|
||||
transpose_A=transpose_A,
|
||||
transpose_B=transpose_B,
|
||||
alpha=1.0,
|
||||
beta=0.0,
|
||||
)
|
||||
D_reg = D_f.local(2, 2)
|
||||
for rM, rN in T.grid(2, 2):
|
||||
D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[rM, rN]
|
||||
|
||||
return gemm
|
||||
|
||||
|
||||
def _lower(func):
|
||||
with tvm.target.Target("cuda"):
|
||||
return tvm.tirx.transform.LowerTIRx()(tvm.IRModule({"main": func}))
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_variant_is_registered():
|
||||
# Importing tvm.tirx registers all per-target schedule variants. The new
|
||||
# synchronous CUDA mma path must show up for ("gemm", "cuda"). The registry
|
||||
# keys ops by their full name (``op.name`` == "tirx.tile.gemm").
|
||||
schedules = list_registered_schedules()
|
||||
cuda_gemm = schedules.get("tirx.tile.gemm", {}).get("cuda", [])
|
||||
assert "mma.m16n8k*" in cuda_gemm, (
|
||||
f"mma.m16n8k* not registered; tirx.tile.gemm schedules = {schedules.get('tirx.tile.gemm')}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16", "float16"])
|
||||
def test_cuda_gemm_mma_lowers_to_mma_sync(dtype):
|
||||
"""beta=0: the dispatch clears D, then issues a single accumulating mma with
|
||||
the registers laid out in the fixed PTX fragment order."""
|
||||
script = _lower(_build_gemm(alpha=1.0, beta=0.0, dtype=dtype))["main"].script()
|
||||
|
||||
assert "T.ptx.mma(" in script
|
||||
assert "m16n8k16" in script
|
||||
# beta == 0 clears the accumulator before the K loop.
|
||||
assert "T.float32(0" in script
|
||||
# D accumulator: c_id = 2*rM + rN -> regs 0..3.
|
||||
for r in range(4):
|
||||
assert f"d_local[{r}]" in script
|
||||
# A multiplicand: b32 = rM + 2*kHi (kHi outer) -> ma in {0, 2, 4, 6}.
|
||||
for r in (0, 2, 4, 6):
|
||||
assert f"a_local[{r}]" in script
|
||||
# B multiplicand: b32 = kHi -> mb in {0, 2}.
|
||||
for r in (0, 2):
|
||||
assert f"b_local[{r}]" in script
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_accumulates_c_when_beta_one():
|
||||
"""beta=1: the accumulator is initialized by copying C instead of zeroing."""
|
||||
script = _lower(_build_gemm(alpha=1.0, beta=1.0))["main"].script()
|
||||
|
||||
assert "T.ptx.mma(" in script
|
||||
assert "m16n8k16" in script
|
||||
# The init reads C into D; nothing is zeroed.
|
||||
assert "c_local[" in script
|
||||
assert "T.float32(0" not in script
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_rejects_nonunit_alpha():
|
||||
"""alpha != 1 is unsupported (ptx mma has no scale); dispatch must fail."""
|
||||
with pytest.raises(RuntimeError, match="dispatch failed"):
|
||||
_lower(_build_gemm(alpha=2.0, beta=0.0))
|
||||
|
||||
|
||||
def test_cuda_gemm_mma_rejects_fractional_beta():
|
||||
"""beta must be 0 or 1 (mma only accumulates 1*C); other values must fail."""
|
||||
with pytest.raises(RuntimeError, match="dispatch failed"):
|
||||
_lower(_build_gemm(alpha=1.0, beta=0.5))
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
|
||||
def test_cuda_gemm_mma_numerical(dtype):
|
||||
"""End-to-end D = A @ B on a single m16n8k16 tile (one warp).
|
||||
|
||||
A is [M, K] = [16, 16], B is [K, N] = [16, 8], D is [M, N] = [16, 8].
|
||||
|
||||
The lane-distributed register fragments cannot be filled with a whole-tile
|
||||
``T.copy`` (the per-thread axis can't be matched coordinate-wise), so each
|
||||
of a lane's registers is loaded/stored by decoding the m16n8k16 register map
|
||||
with ``g = lane >> 2`` and ``t = lane & 3``. The per-register *slot* order
|
||||
matches the dispatch's fragment register layout:
|
||||
|
||||
A reg slot = 4*rM + 2*kHi + kp -> M = g + 8*rM, K = 2*t + kp + 8*kHi
|
||||
B reg slot = 2*kHi + kp -> K = 2*t + kp + 8*kHi, N = g
|
||||
D reg slot = 2*rM + rN -> M = g + 8*rM, N = 2*t + rN
|
||||
"""
|
||||
if dtype == "bfloat16":
|
||||
ml_dtypes = pytest.importorskip("ml_dtypes")
|
||||
np_dtype = ml_dtypes.bfloat16
|
||||
else:
|
||||
np_dtype = np.float16
|
||||
|
||||
@T.prim_func
|
||||
def gemm(A_ptr: T.handle, B_ptr: T.handle, D_ptr: T.handle):
|
||||
A_g = T.match_buffer(A_ptr, (16, 16), dtype)
|
||||
B_g = T.match_buffer(B_ptr, (16, 8), dtype)
|
||||
D_g = T.match_buffer(D_ptr, (16, 8), "float32")
|
||||
T.device_entry()
|
||||
_cta = T.cta_id([1])
|
||||
_warp = T.warp_id([1])
|
||||
lane = T.lane_id([32])
|
||||
A_f = T.alloc_buffer((16, 16), dtype, scope="local", layout=A_FRAG)
|
||||
B_f = T.alloc_buffer((16, 8), dtype, scope="local", layout=B_FRAG)
|
||||
D_f = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
|
||||
A_reg = A_f.local(8)
|
||||
for s in T.unroll(8):
|
||||
kp = s % 2
|
||||
kHi = (s // 2) % 2
|
||||
rM = s // 4
|
||||
A_reg[s] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
|
||||
B_reg = B_f.local(4)
|
||||
for s in T.unroll(4):
|
||||
kp = s % 2
|
||||
kHi = s // 2
|
||||
B_reg[s] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
|
||||
Tx.warp.gemm(D_f, A_f, B_f, D_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0)
|
||||
D_reg = D_f.local(4)
|
||||
for s in T.unroll(4):
|
||||
rN = s % 2
|
||||
rM = s // 2
|
||||
D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[s]
|
||||
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.compile(tvm.IRModule({"main": gemm}), target="cuda", tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = np.random.uniform(-1, 1, (16, 16)).astype(np.float32)
|
||||
B_np = np.random.uniform(-1, 1, (16, 8)).astype(np.float32)
|
||||
golden = A_np @ B_np
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np.astype(np_dtype), dev)
|
||||
B_dev = tvm.runtime.tensor(B_np.astype(np_dtype), dev)
|
||||
D_dev = tvm.runtime.tensor(np.zeros((16, 8), np.float32), dev)
|
||||
mod(A_dev, B_dev, D_dev)
|
||||
tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=1e-2, rtol=1e-2)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
# (Mt, Nt, Kt, kinst) tilings: single tile, each dim multi-tiled, fully tiled,
|
||||
# M = 64, and the m16n8k8 (kHi == 1) variants including a non-16-divisible K.
|
||||
_TILED_SHAPES = [
|
||||
(1, 1, 1, 16), # single m16n8k16 tile
|
||||
(2, 1, 1, 16), # two M-tiles
|
||||
(1, 2, 1, 16), # two N-tiles
|
||||
(1, 1, 2, 16), # two K-tiles (accumulated in place)
|
||||
(2, 2, 2, 16), # every dim tiled
|
||||
(4, 1, 1, 16), # M = 64
|
||||
(1, 1, 1, 8), # m16n8k8 single tile (kHi == 1)
|
||||
(1, 1, 3, 8), # K = 24 -> three k8 tiles
|
||||
(2, 2, 3, 8), # k8, every dim tiled
|
||||
]
|
||||
# (dtype, beta) input modes crossed against every shape: f16/bf16 inputs, with
|
||||
# beta = 0 (D = A @ B) and beta = 1 (D = A @ B + C, accumulating C in place).
|
||||
_TILED_MODES = [
|
||||
("float16", 0.0),
|
||||
("bfloat16", 0.0),
|
||||
("float16", 1.0),
|
||||
("bfloat16", 1.0),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize("Mt, Nt, Kt, kinst", _TILED_SHAPES)
|
||||
@pytest.mark.parametrize("dtype, beta", _TILED_MODES)
|
||||
def test_cuda_gemm_mma_numerical_tiled(dtype, beta, Mt, Nt, Kt, kinst):
|
||||
"""End-to-end D = A @ B (+ C when beta==1) over an Mt x Nt x Kt tiling.
|
||||
|
||||
The two stacked ``parametrize`` decorators form the cartesian product of
|
||||
every tiling shape with every (dtype, beta) input mode, so each combination
|
||||
is an independent pytest item (pytest-xdist runs them in parallel)."""
|
||||
if dtype == "bfloat16":
|
||||
ml_dtypes = pytest.importorskip("ml_dtypes")
|
||||
np_dtype = ml_dtypes.bfloat16
|
||||
else:
|
||||
np_dtype = np.float16
|
||||
|
||||
func, M, N, K = _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype)
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.compile(tvm.IRModule({"main": func}), target="cuda", tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_np = np.random.uniform(-1, 1, (M, K)).astype(np.float32)
|
||||
B_np = np.random.uniform(-1, 1, (K, N)).astype(np.float32)
|
||||
C_np = np.random.uniform(-1, 1, (M, N)).astype(np.float32)
|
||||
golden = A_np @ B_np + (C_np if beta == 1.0 else 0.0)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_np.astype(np_dtype), dev)
|
||||
B_dev = tvm.runtime.tensor(B_np.astype(np_dtype), dev)
|
||||
C_dev = tvm.runtime.tensor(C_np, dev)
|
||||
D_dev = tvm.runtime.tensor(np.zeros((M, N), np.float32), dev)
|
||||
mod(A_dev, B_dev, C_dev, D_dev)
|
||||
tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=2e-2, rtol=2e-2)
|
||||
|
||||
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("dtype", ["float16", "bfloat16"])
|
||||
@pytest.mark.parametrize(
|
||||
"transpose_A, transpose_B",
|
||||
[(False, False), (True, False), (False, True), (True, True)],
|
||||
)
|
||||
def test_cuda_gemm_mma_numerical_transpose(transpose_A, transpose_B, dtype):
|
||||
"""End-to-end D = A @ B for every A/B input orientation, crossed with dtype.
|
||||
|
||||
The orientation and dtype decorators form a cartesian product, so each
|
||||
(transpose_A, transpose_B, dtype) is an independent pytest item."""
|
||||
if dtype == "bfloat16":
|
||||
ml_dtypes = pytest.importorskip("ml_dtypes")
|
||||
np_dtype = ml_dtypes.bfloat16
|
||||
else:
|
||||
np_dtype = np.float16
|
||||
|
||||
func = _build_transpose_numeric(transpose_A, transpose_B, dtype)
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.compile(tvm.IRModule({"main": func}), target="cuda", tir_pipeline="tirx")
|
||||
|
||||
np.random.seed(0)
|
||||
A_log = np.random.uniform(-1, 1, (16, 16)).astype(np.float32) # logical A[M, K]
|
||||
B_log = np.random.uniform(-1, 1, (16, 8)).astype(np.float32) # logical B[K, N]
|
||||
A_buf = (A_log.T if transpose_A else A_log).astype(np_dtype)
|
||||
B_buf = (B_log.T if transpose_B else B_log).astype(np_dtype)
|
||||
golden = A_log @ B_log
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
A_dev = tvm.runtime.tensor(A_buf, dev)
|
||||
B_dev = tvm.runtime.tensor(B_buf, dev)
|
||||
D_dev = tvm.runtime.tensor(np.zeros((16, 8), np.float32), dev)
|
||||
mod(A_dev, B_dev, D_dev)
|
||||
tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=2e-2, rtol=2e-2)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"Mt, Nt, Kt, kinst",
|
||||
[
|
||||
(2, 1, 1, 16), # M = 32 (two M-tiles)
|
||||
(1, 2, 1, 16), # N = 16 (two N-tiles)
|
||||
(1, 1, 2, 16), # K = 32 (two K-tiles, accumulated in place)
|
||||
(2, 2, 2, 16), # 8 tiles
|
||||
(4, 1, 1, 16), # M = 64
|
||||
(1, 1, 1, 8), # m16n8k8 single tile (kHi == 1)
|
||||
(1, 1, 3, 8), # K = 24 -> three k8 tiles (16 does not divide 24)
|
||||
(2, 2, 3, 8), # k8, every dim tiled
|
||||
],
|
||||
)
|
||||
def test_cuda_gemm_mma_lowers_tiled(Mt, Nt, Kt, kinst):
|
||||
"""Every tiling we expect to dispatch must lower, selecting the right mma.
|
||||
|
||||
The k8 cases are the regression guard for the kHi == 1 fragment grouping
|
||||
(an extent-1 high-K register group must not be rejected as a thread axis).
|
||||
"""
|
||||
script = _lower(_build_tiled(Mt, Nt, Kt, kinst))["main"].script()
|
||||
assert "T.ptx.mma(" in script
|
||||
assert f"m16n8k{kinst}" in script
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize(
|
||||
"Mt, Nt, Kt, kinst",
|
||||
[
|
||||
(1, 1, 1, 16),
|
||||
(2, 2, 2, 16),
|
||||
(4, 1, 1, 16),
|
||||
(1, 1, 1, 8),
|
||||
(1, 1, 3, 8),
|
||||
(2, 2, 3, 8),
|
||||
],
|
||||
)
|
||||
def test_cuda_gemm_mma_codegen_issue_count(Mt, Nt, Kt, kinst):
|
||||
"""Full pipeline (UnrollLoop + CUDA codegen) emits one mma per (Mt, Nt, Kt)
|
||||
tile; K-tiles accumulate in place, so D is cleared once per output tile."""
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
|
||||
with target:
|
||||
mod = tvm.compile(
|
||||
tvm.IRModule({"main": _build_tiled(Mt, Nt, Kt, kinst, store=True)}),
|
||||
target=target,
|
||||
tir_pipeline="tirx",
|
||||
)
|
||||
src = mod.mod.imports[0].inspect_source()
|
||||
assert f"mma.sync.aligned.m16n8k{kinst}" in src
|
||||
# mma is emitted as one __device__ helper, invoked once per tile.
|
||||
helper = f"ptx_mma_m16n8k{kinst}_row_col"
|
||||
assert src.count(helper) - 1 == Mt * Nt * Kt
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"transpose_A, transpose_B",
|
||||
[(False, False), (True, False), (False, True), (True, True)],
|
||||
)
|
||||
def test_cuda_gemm_mma_lowers_transpose(transpose_A, transpose_B):
|
||||
"""All four A/B orientations dispatch to the same m16n8k16. transpose only
|
||||
describes the input's logical orientation; the .row.col mma is unchanged."""
|
||||
script = _lower(_build_transpose(transpose_A, transpose_B))["main"].script()
|
||||
assert "T.ptx.mma(" in script
|
||||
assert "m16n8k16" in script
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize(
|
||||
"transpose_A, transpose_B",
|
||||
[(False, False), (True, False), (False, True), (True, True)],
|
||||
)
|
||||
def test_cuda_gemm_mma_codegen_transpose(transpose_A, transpose_B):
|
||||
"""Every orientation codegens to a valid m16n8k16 kernel."""
|
||||
target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
|
||||
with target:
|
||||
mod = tvm.compile(
|
||||
tvm.IRModule({"main": _build_transpose(transpose_A, transpose_B, store=True)}),
|
||||
target=target,
|
||||
tir_pipeline="tirx",
|
||||
)
|
||||
assert "mma.sync.aligned.m16n8k16" in mod.mod.imports[0].inspect_source()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a, b, c, d",
|
||||
[
|
||||
("float16", "float16", "float16", "float16"), # f16 accumulate
|
||||
("bfloat16", "float16", "float32", "float32"), # mixed A/B inputs
|
||||
("float32", "float32", "float32", "float32"), # f32 (tf32) inputs
|
||||
("int8", "int8", "int32", "int32"), # integer
|
||||
],
|
||||
)
|
||||
def test_cuda_gemm_mma_rejects_unsupported_dtype(a, b, c, d):
|
||||
"""The table holds only (bf16|f16, same, f32, f32); any other dtype
|
||||
signature must decline rather than emit a wrong mma."""
|
||||
with pytest.raises(RuntimeError, match="dispatch failed"):
|
||||
_lower(_build_dtypes(a, b, c, d))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user