diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 513d52b20..024bd9f22 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -24,19 +24,55 @@ Computes C[M,N] = A[M,K] @ B_T[N,K]^T (same interface as ``rdna_f16_gemm.create_wmma_gemm_module``). + +The block tile is a parameter and defaults to 128x128x32. Deciding it from the +shape -- which is worth up to 3.0x on shapes too small to fill the grid -- is +the job of ``rdna3_f16_gemm_autotune``; this module only builds what it is told. """ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import llvm as _llvm -from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl._mlir.dialects import vector +from flydsl.expr import as_ir_value, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch -from kernels.common import buffer_ops from kernels.common.kernels_common import cvt_sr_f32_to_bf16 WMMA_M = 16 WMMA_N = 16 WMMA_K = 16 +WAVE_SIZE = 32 + +# The k-padding both operands pay to break LDS bank conflicts. Also the default +# for a_k_pad/b_k_pad below, so a caller sizing a tile against the LDS budget can +# assume it without the two drifting apart. +K_PAD = 8 + + +def _group_width(grid_m, group_m): + """Largest grouping width <= group_m that divides grid_m. + + ``_swizzle_tile_id`` derives bid_m from a fixed group width, so a grid_m that + is not a multiple of it makes the final group address tiles past the end of + the grid. That is reachable at the default 128x128 tile: on gfx1100 it writes + a wrong C at M of 1152, 1280 and 1664, and faults outright at 1536 and 2560, + depending on whether the address past the grid happens to be mapped. + """ + return max(d for d in range(1, min(group_m, grid_m) + 1) if grid_m % d == 0) + + +def _swizzle_tile_id(pid, grid_n, group_width): + """Linear workgroup id -> (bid_m, bid_n). + + Walks group_width tiles down M before stepping in N, so workgroups that run + concurrently share B tiles in L2. Plain integer arithmetic, so it evaluates + the same on a host int as on the kernel's block id. + """ + num_pid_in_group = group_width * grid_n + group_id = pid // num_pid_in_group + pid_in_group = pid % num_pid_in_group + return group_id * group_width + (pid_in_group % group_width), pid_in_group // group_width def create_wmma_gemm_module( @@ -47,14 +83,18 @@ def create_wmma_gemm_module( out_dtype="bf16", *, rounding="rn", # "rn" (round to nearest) or "rs" (stochastic rounding) + # 128x128x32. That tile is right once the problem is large enough to fill the + # grid, but it cuts only 4 workgroups at 256x256 on a 96-CU part, so most CUs + # idle. Choosing a tile from the shape is worth up to 3.0x there and lives in + # rdna3_f16_gemm_autotune, which drives these arguments. reg_m=4, reg_n=4, reg_k=2, waves_m=2, waves_n=2, group_m=8, - a_k_pad=8, - b_k_pad=8, + a_k_pad=K_PAD, + b_k_pad=K_PAD, ): gpu_arch = str(get_rocm_arch() or "") if not gpu_arch.startswith("gfx11"): @@ -67,7 +107,6 @@ def create_wmma_gemm_module( BLOCK_N = WMMA_N * reg_n * waves_n # 128 BLOCK_K = WMMA_K * reg_k # 32 NUM_WAVES = waves_m * waves_n # 4 - WAVE_SIZE = 32 THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE # 128 assert reg_k >= 2 and reg_k % 2 == 0 @@ -76,10 +115,14 @@ def create_wmma_gemm_module( assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only" LOAD_VEC = 8 # 8 bf16 = 128-bit GMEM/LDS load - A_TILE_ELEMS = BLOCK_M * BLOCK_K - NUM_A_LOADS = A_TILE_ELEMS // (THREADS_PER_BLOCK * LOAD_VEC) - B_TILE_ELEMS = BLOCK_N * BLOCK_K - NUM_B_LOADS = B_TILE_ELEMS // (THREADS_PER_BLOCK * LOAD_VEC) + # G2S thread geometry: thread (tk, tm) moves the 128-bit chunk at columns + # [tk*LOAD_VEC, +LOAD_VEC) of row tm, and the tiled copy repeats over M to + # cover the tile. Same assignment the hand-rolled offset tables computed as + # ``row = tid // THRS_K, col = (tid % THRS_K) * LOAD_VEC``. + THRS_K = BLOCK_K // LOAD_VEC + THRS_M = THREADS_PER_BLOCK // THRS_K + assert THRS_K * THRS_M == THREADS_PER_BLOCK + assert BLOCK_M % THRS_M == 0 and BLOCK_N % THRS_M == 0 BLOCK_K_PAD_A = BLOCK_K + a_k_pad # 40 BLOCK_K_PAD_B = BLOCK_K + b_k_pad # 40 @@ -98,6 +141,9 @@ def create_wmma_gemm_module( grid_m = M // BLOCK_M grid_n = N // BLOCK_N + + group_width = _group_width(grid_m, group_m) + is_bf16 = in_dtype == "bf16" def _wmma_op(a_vec, b_vec, acc): @@ -109,6 +155,8 @@ def _wmma_op(a_vec, b_vec, acc): return rocdl.wmma_f32_16x16x16_f16(acc.type, a_vec, b_vec, acc).result elem_dtype = fx.BFloat16 if is_bf16 else fx.Float16 + out_elem_cls = {"bf16": fx.BFloat16, "f16": fx.Float16, "f32": fx.Float32}[out_dtype] + acc_size = 8 * reg_m * reg_n # accumulator f32 VGPRs per thread # ── Shared-memory storage for double-buffered A+B LDS tiles ────────── # One flat bf16/f16 array; v8 chunks are addressed by byte_offset // 2 @@ -124,6 +172,8 @@ def wmma_gemm_kernel( arg_c: fx.Tensor, arg_a: fx.Tensor, arg_bt: fx.Tensor, + tiled_mma: fx.TiledMma, + tiled_copy_g2s: fx.TiledCopy, sr_seed: fx.Int32, # runtime seed; only read on the stochastic-rounding path ): lds_storage = fx.SharedAllocator().allocate(_SharedStorage).peek() @@ -138,12 +188,6 @@ def _v8_load(v8_idx): typed_ptr = fx.recast_iter(elem_dtype, ptr_off) return fx.make_view(typed_ptr, fx.make_layout(8, 1)).load() - def _v8_store(v8_idx, value): - elem_off = fx.Int32(v8_idx * 8) - ptr_off = fx.add_offset(lds_ptr, fx.make_int_tuple(elem_off)) - typed_ptr = fx.recast_iter(elem_dtype, ptr_off) - fx.make_view(typed_ptr, fx.make_layout(8, 1)).store(value) - tid = gpu.thread_id("x") pid = gpu.block_id("x") @@ -153,82 +197,76 @@ def _v8_store(v8_idx, value): # M (or N) row is selected by ``lane % 16`` only. No klane shift # in the K dimension — each lane carries all 16 K-elements. lane16 = lane % 16 - klane = lane // 16 # used only for the gfx11 accumulator store-back - - # Swizzle workgroup mapping for L2 locality - effective_group_m = min(group_m, grid_m) - num_pid_in_group = effective_group_m * grid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * effective_group_m - group_size_m = effective_group_m - pid_in_group = pid % num_pid_in_group - bid_m = first_pid_m + (pid_in_group % group_size_m) - bid_n = pid_in_group // group_size_m + bid_m, bid_n = _swizzle_tile_id(pid, grid_n, group_width) wave_m = wave_id // waves_n wave_n = wave_id % waves_n - tile_m0 = bid_m * BLOCK_M - tile_n0 = bid_n * BLOCK_N - - a_rsrc = buffer_ops.create_buffer_resource(arg_a, max_size=True) - bt_rsrc = buffer_ops.create_buffer_resource(arg_bt, max_size=True) - c_rsrc = buffer_ops.create_buffer_resource(arg_c, max_size=True) + # Wave wm owns the contiguous row band [wm*reg_m*16, +reg_m*16). A + # tiled_mma stamps its wave grid across the tile instead, putting repeat + # rm at row (rm*waves_m + wm)*16; measured on gfx1100 that interleaving + # costs 62% at 3072x3072x1024 (269 -> 436 us) while gaining 3-8% on the + # medium shapes, so a tiled_mma standing in for this loop has to carry a + # permutation that restores the banding rather than adopt the default. + + # Result partition. The tiled_mma carries a permutation that reproduces + # the wave banding above, so the accumulators land where the hand-rolled + # ``g_row = base + 2*si + klane`` store used to put them. + tC = fx.flat_divide(fx.rocdl.make_buffer_tensor(arg_c), fx.make_tile(BLOCK_M, BLOCK_N))[ + None, None, bid_m, bid_n + ] + thr_mma = tiled_mma.thr_slice(tid) + frag_C = thr_mma.make_fragment_C(tC) + copy_out = fx.make_copy_atom(fx.rocdl.BufferCopy(out_elem_cls.width), out_elem_cls) + thr_r2g_C = fx.make_tiled_copy_C(copy_out, tiled_mma).get_slice(tid) + pC_g = thr_r2g_C.partition_S(tC) + if const_expr(out_elem_cls is fx.Float32): + frag_C_out = frag_C + else: + frag_C_out = fx.make_fragment_like(frag_C, out_elem_cls.ir_type) + frag_C_retile = thr_r2g_C.retile(frag_C_out) # ============================================================ - # Pre-compute GMEM offsets and LDS addresses (same as gfx12) + # GMEM -> registers -> LDS, through the tiled copy # ============================================================ - a_lds_info = [] - for al in range_constexpr(NUM_A_LOADS): - a_lin = tid * LOAD_VEC + (al * THREADS_PER_BLOCK * LOAD_VEC) - a_load_row = a_lin // BLOCK_K - a_load_col = a_lin % BLOCK_K - lds_rel = a_load_row * BLOCK_K_PAD_A + a_load_col - g_row = tile_m0 + a_load_row - a_lds_info.append((g_row, a_load_col, lds_rel)) - - b_lds_info = [] - for bl in range_constexpr(NUM_B_LOADS): - b_lin = tid * LOAD_VEC + (bl * THREADS_PER_BLOCK * LOAD_VEC) - b_load_row = b_lin // BLOCK_K - b_load_col = b_lin % BLOCK_K - lds_rel = LDS_A_SIZE + b_load_row * BLOCK_K_PAD_B + b_load_col - g_row = tile_n0 + b_load_row - b_lds_info.append((g_row, b_load_col, lds_rel)) - - def _gmem_load(k_base): - raw_data = [] - for al in range_constexpr(NUM_A_LOADS): - g_row, a_load_col, _ = a_lds_info[al] - g_col = k_base + a_load_col - elem_off = g_row * K + g_col - f32_off = elem_off // 2 - a_raw = buffer_ops.buffer_load(a_rsrc, f32_off, vec_width=4, dtype=fx.Float32) - raw_data.append(a_raw) - - for bl in range_constexpr(NUM_B_LOADS): - g_row, b_load_col, _ = b_lds_info[bl] - g_col = k_base + b_load_col - elem_off = g_row * K + g_col - f32_off = elem_off // 2 - b_raw = buffer_ops.buffer_load(bt_rsrc, f32_off, vec_width=4, dtype=fx.Float32) - raw_data.append(b_raw) - - return raw_data - - def _lds_store(raw_data, buf_offset): - for al in range_constexpr(NUM_A_LOADS): - _, _, lds_rel = a_lds_info[al] - a_vec = raw_data[al].bitcast(fx.BFloat16 if is_bf16 else fx.Float16) - lds_idx = buf_offset + lds_rel - _v8_store(lds_idx // 8, a_vec) - - for bl in range_constexpr(NUM_B_LOADS): - _, _, lds_rel = b_lds_info[bl] - b_vec = raw_data[NUM_A_LOADS + bl].bitcast(fx.BFloat16 if is_bf16 else fx.Float16) - lds_idx = buf_offset + lds_rel - _v8_store(lds_idx // 8, b_vec) + tA = fx.flat_divide(fx.rocdl.make_buffer_tensor(arg_a), fx.make_tile(BLOCK_M, BLOCK_K))[None, None, bid_m, None] + tB = fx.flat_divide(fx.rocdl.make_buffer_tensor(arg_bt), fx.make_tile(BLOCK_N, BLOCK_K))[ + None, None, bid_n, None + ] + + thr_g2s = tiled_copy_g2s.get_slice(tid) + pA_g = thr_g2s.partition_S(tA) + pB_g = thr_g2s.partition_S(tB) + + buf_copy = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_dtype) + uni_copy = fx.make_copy_atom(fx.UniversalCopy128b(), elem_dtype) + + # The destination buffer alternates on ``iv % 2``, a loop-carried value, + # so it cannot be selected from a list of per-stage views at trace time. + # The view is built from the running element offset into the single LDS + # allocation instead, which is what the flat _v8_store did by hand. + def _lds_dst(buf_offset, base, rows, row_stride): + ptr = fx.add_offset(lds_ptr, fx.make_int_tuple(buf_offset + base)) + view = fx.make_view(fx.recast_iter(elem_dtype, ptr), fx.make_layout((rows, BLOCK_K), (row_stride, 1))) + return thr_g2s.partition_D(view)[None, None, None] + + def _pA_s(buf_offset): + return _lds_dst(buf_offset, 0, BLOCK_M, BLOCK_K_PAD_A) + + def _pB_s(buf_offset): + return _lds_dst(buf_offset, LDS_A_SIZE, BLOCK_N, BLOCK_K_PAD_B) + + frag_copy_A = fx.make_fragment_like(_pA_s(0)) + frag_copy_B = fx.make_fragment_like(_pB_s(0)) + + def _gmem_load(k_tile): + fx.copy(buf_copy, pA_g[None, None, None, k_tile], frag_copy_A) + fx.copy(buf_copy, pB_g[None, None, None, k_tile], frag_copy_B) + + def _lds_store(buf_offset): + fx.copy(uni_copy, frag_copy_A, _pA_s(buf_offset)) + fx.copy(uni_copy, frag_copy_B, _pB_s(buf_offset)) # ============================================================ # LDS read helpers — v16 by concatenating two v8 loads @@ -292,8 +330,8 @@ def _do_compute_rk(accs_in, rk, buf_offset): c_lds_buf_stride = LDS_ONE_BUF # --- PROLOGUE --- - prologue_data = _gmem_load(0) - _lds_store(prologue_data, 0) + _gmem_load(fx.Int32(0)) + _lds_store(0) _barrier() n_acc = reg_m * reg_n @@ -305,13 +343,12 @@ def _do_compute_rk(accs_in, rk, buf_offset): read_off = iv % 2 * c_lds_buf_stride write_off = (1 - iv % 2) * c_lds_buf_stride - next_k = (iv + 1) * BLOCK_K - next_data = _gmem_load(next_k) + _gmem_load(iv + 1) for rk in range_constexpr(reg_k): s_accs = _do_compute_rk(s_accs, rk, read_off) - _lds_store(next_data, write_off) + _lds_store(write_off) _barrier() results = yield list(s_accs) @@ -323,38 +360,39 @@ def _do_compute_rk(accs_in, rk, buf_offset): accs = _do_compute_rk(accs, rk, last_read_off) # ============================================================ - # Store results to GMEM (gfx11 layout: stride-2 rows) + # Store results to GMEM through the tiled copy # ============================================================ - # gfx11 v8f32 acc layout: lane L holds D[2*si + (L/16)][L%16] - # for si in 0..7 — i.e. lanes 0-15 carry even rows, lanes 16-31 - # carry odd rows of the same 16 columns. - for rm in range_constexpr(reg_m): - for rn in range_constexpr(reg_n): - idx = rm * reg_n + rn - wmma_m_off = wave_m * (reg_m * WMMA_M) + 16 * rm - wmma_n_off = wave_n * (reg_n * WMMA_N) + 16 * rn - if const_expr(rounding == "rs"): - # One Philox draw per 8-element block, keyed on the block's - # base output element: the 4 random words cover all 8 - # stores, each taking a distinct 16-bit slice (low/high of a - # word), so the f32 -> bf16 store is unbiased in expectation - # without a per-element draw. - base_off = (tile_m0 + wmma_m_off + klane) * N + (tile_n0 + wmma_n_off + lane16) - rand_words = fx.random.randint4x(fx.Uint32(sr_seed), fx.Uint32(base_off)) + # The gfx11 v8f32 accumulator (lane L holds D[2*si + L/16][L%16]) and the + # wave banding are both encoded in the tiled_mma, so the row arithmetic + # that used to live here is gone. What remains is the value transform, + # which no copy atom can express. + # + # frag_C flattens as si + 8*(rm + reg_m*rn), so each run of 8 elements is + # exactly one atom's accumulator, and one Philox draw still covers one run. + ordered_accs = [accs[rm * reg_n + rn] for rn in range_constexpr(reg_n) for rm in range_constexpr(reg_m)] + if const_expr(rounding == "rs"): + # The 4 random words cover all 8 values, each taking a distinct + # 16-bit slice (low/high of a word), so the f32 -> bf16 store is + # unbiased in expectation without a per-element draw. Keying on the + # thread's slot rather than the output coordinate keeps the draw + # independent of where the tiled copy lands the fragment. + out_elems = [] + for g, acc in enumerate(ordered_accs): + base_off = (pid * THREADS_PER_BLOCK + tid) * acc_size + 8 * g + words = fx.random.randint4x(fx.Uint32(sr_seed), fx.Uint32(base_off)) for si in range_constexpr(8): - g_row = tile_m0 + wmma_m_off + 2 * si + klane - g_col = tile_n0 + wmma_n_off + lane16 - val = accs[idx][si] - elem_off = g_row * N + g_col - if const_expr(rounding == "rs"): - word = rand_words[si // 2] - rbits = word if si % 2 == 0 else (word >> fx.Uint32(16)) - val = cvt_sr_f32_to_bf16(val, rbits) - elif const_expr(out_dtype == "bf16"): - val = val.to(fx.BFloat16) - elif const_expr(out_dtype == "f16"): - val = val.to(fx.Float16) - buffer_ops.buffer_store(val, c_rsrc, elem_off) + word = words[si // 2] + rbits = word if si % 2 == 0 else (word >> fx.Uint32(16)) + out_elems.append(cvt_sr_f32_to_bf16(acc[si], rbits)) + elif const_expr(out_elem_cls is fx.Float32): + out_elems = [acc[si] for acc in ordered_accs for si in range_constexpr(8)] + else: + out_elems = [acc[si].to(out_elem_cls) for acc in ordered_accs for si in range_constexpr(8)] + + frag_C_out.store( + vector.from_elements(T.vec(acc_size, out_elem_cls.ir_type), [as_ir_value(e) for e in out_elems]) + ) + fx.copy(copy_out, frag_C_retile, pC_g) @flyc.jit def launch_gemm( @@ -364,11 +402,41 @@ def launch_gemm( stream: fx.Stream, sr_seed: fx.Int32 = 0, ): + # 16x16x16 v16 WMMA atom (gfx11.wmma) over a waves_m x waves_n wave grid. + # The permutation spans the whole block tile and remaps the natural + # (atom, wave, repeat) coordinate so wave wm keeps the contiguous band + # [wm*reg_m*16, +reg_m*16); the default stamping interleaves the repeats + # instead, which measured 62% slower at 3072x3072x1024. + mma_atom = fx.make_mma_atom(fx.rocdl.WMMA(WMMA_M, WMMA_N, WMMA_K, elem_dtype, fx.Float32)) + tiled_mma = fx.make_tiled_mma( + mma_atom, + fx.make_layout((waves_m, waves_n, 1), (waves_n, 1, 0)), + permutation=( + fx.make_layout((WMMA_M, waves_m, reg_m), (1, WMMA_M * reg_m, WMMA_M)), + fx.make_layout((WMMA_N, waves_n, reg_n), (1, WMMA_N * reg_n, WMMA_N)), + WMMA_K, + ), + ) + # G2S tiled copy: thread (tk, tm) moves one 128-bit contiguous chunk of + # row tm, repeated over M to cover the block tile. + tiled_copy_g2s = fx.make_tiled_copy( + fx.make_copy_atom(fx.UniversalCopy128b(), elem_dtype), + fx.make_layout( + ((THRS_K, THRS_M), (1, LOAD_VEC)), + ((THRS_M * LOAD_VEC, 1), (1, THRS_M)), + ), + fx.make_tile(THRS_M, BLOCK_K), + ) + + arg_a_2d = fx.make_view(fx.get_iter(arg_a), fx.make_layout((M, K), (K, 1))) + arg_bt_2d = fx.make_view(fx.get_iter(arg_bt), fx.make_layout((N, K), (K, 1))) + arg_c_2d = fx.make_view(fx.get_iter(arg_c), fx.make_layout((M, N), (N, 1))) + c1 = 1 total_blocks = grid_m * grid_n bk = THREADS_PER_BLOCK - launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt, sr_seed) + launcher = wmma_gemm_kernel(arg_c_2d, arg_a_2d, arg_bt_2d, tiled_mma, tiled_copy_g2s, sr_seed) launcher.launch( grid=(total_blocks, c1, c1), block=(bk, c1, c1), diff --git a/tests/kernels/test_rdna3_wmma_atom.py b/tests/kernels/test_rdna3_wmma_atom.py new file mode 100644 index 000000000..df2790450 --- /dev/null +++ b/tests/kernels/test_rdna3_wmma_atom.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Device correctness tests for the RDNA3 (gfx11*) WMMA MMA atom. + +The gfx11 counterpart of ``test_rdna4_wmma_atom.py``. RDNA3 keeps the same +16x16x16 shapes as RDNA4 but uses the legacy v16-operand register ABI, modelled +by the ``gfx11.wmma`` atom (``lib/Dialect/FlyROCDL/GFX11/MmaAtom.cpp``): + + * A/B carry 16 elements per lane, and lanes 16-31 hold a *replica* of what + lanes 0-15 hold. The atom expresses that as stride 0 on the ``lane/16`` axis + of the thread layout, so the A/B partition is deliberately not injective. + * C/D interleaves rows between the two lane halves (``row = 2*val + lane/16``) + rather than giving each half a contiguous run of rows. + +Both of those are exactly the places where a hand-written kernel gets the lane +math wrong, and both are what ``make_tiled_copy_{A,B,C}`` has to reproduce for +gfx11 kernels to be written against the layout API instead of raw intrinsics. +``kernels/gemm/rdna3_f16_gemm.py`` still spells the math out by hand; these +tests are what has to pass before it can stop doing that. + +The integer atom is covered separately in ``test_rdna3_integer_wmma_atom.py``, +which builds its fragments by hand and so does not exercise the tiled copies. +""" + +import os +import sys + +import pytest + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +import flydsl # noqa: E402,F401 -- preload comgr before torch/HIP loads LLVM +import flydsl.compiler as flyc # noqa: E402 +import flydsl.expr as fx # noqa: E402 +from flydsl.expr import range_constexpr # noqa: E402 +from flydsl.runtime.device import get_rocm_arch # noqa: E402 + +try: + import torch +except ImportError: + torch = None + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +if torch is None or not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + +_ARCH = str(get_rocm_arch() or "") +if not _ARCH.startswith("gfx11"): + pytest.skip(f"RDNA3 WMMA atom requires gfx11*, got {_ARCH}", allow_module_level=True) + +WAVE_SIZE = 32 +WMMA_M = WMMA_N = WMMA_K = 16 + +# Unlike RDNA4, gfx11 WMMA does not return a bit-exact f32 accumulator even for +# integer-valued fp16/bf16 inputs: results land within about one f32 ULP of the +# exact sum. The residue is in the instruction, not in the fragment layouts — +# assembling the fragments by hand and going through make_tiled_copy produce +# bit-identical output. So this tolerance has to be loose enough to absorb one +# ULP but far tighter than any layout error, which mixes in whole other rows or +# columns and is wrong by O(1) or more. +_ATOL = 1e-3 +_RTOL = 1e-6 + + +def _compile_single_wmma(elem_cls): + """One wave, one atom: C[16,16] = A[16,16] @ B[16,16].T.""" + f32 = fx.Float32 + + @flyc.kernel + def wmma_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor): + tid = fx.thread_idx.x + + bA = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(A)), fx.make_layout((WMMA_M, WMMA_K), (WMMA_K, 1))) + bB = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(B)), fx.make_layout((WMMA_N, WMMA_K), (WMMA_K, 1))) + bC = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(C)), fx.make_layout((WMMA_M, WMMA_N), (WMMA_N, 1))) + + mma_atom = fx.make_mma_atom(fx.rocdl.WMMA(WMMA_M, WMMA_N, WMMA_K, elem_cls, f32)) + tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((1, 1, 1), (0, 0, 0))) + thr_mma = tiled_mma.thr_slice(tid) + + frag_A = thr_mma.make_fragment_A(bA) + frag_B = thr_mma.make_fragment_B(bB) + frag_C = thr_mma.make_fragment_C(bC) + + copy_ab = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_cls.width), elem_cls) + copy_c = fx.make_copy_atom(fx.rocdl.BufferCopy(f32.width), f32) + thr_copy_A = fx.make_tiled_copy_A(copy_ab, tiled_mma).get_slice(tid) + thr_copy_B = fx.make_tiled_copy_B(copy_ab, tiled_mma).get_slice(tid) + thr_copy_C = fx.make_tiled_copy_C(copy_c, tiled_mma).get_slice(tid) + + fx.copy(copy_ab, thr_copy_A.partition_S(bA), thr_copy_A.retile(frag_A)) + fx.copy(copy_ab, thr_copy_B.partition_S(bB), thr_copy_B.retile(frag_B)) + + frag_C.fill(0) + fx.gemm(mma_atom, frag_C, frag_A, frag_B, frag_C) + fx.copy(copy_c, thr_copy_C.retile(frag_C), thr_copy_C.partition_S(bC)) + + @flyc.jit + def launch(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + wmma_kernel(A, B, C).launch(grid=(1, 1, 1), block=(WAVE_SIZE, 1, 1), stream=stream) + + return launch + + +def _compile_tiled_wmma(elem_cls, tile_m, tile_n, tile_k, waves_m, waves_n): + """A wave grid with atom repeats, the geometry a real GEMM block tile uses.""" + f32 = fx.Float32 + threads = waves_m * waves_n * WAVE_SIZE + k_iters = tile_k // WMMA_K + + @flyc.kernel + def wmma_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor): + tid = fx.thread_idx.x + + bA = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(A)), fx.make_layout((tile_m, tile_k), (tile_k, 1))) + bB = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(B)), fx.make_layout((tile_n, tile_k), (tile_k, 1))) + bC = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(C)), fx.make_layout((tile_m, tile_n), (tile_n, 1))) + + mma_atom = fx.make_mma_atom(fx.rocdl.WMMA(WMMA_M, WMMA_N, WMMA_K, elem_cls, f32)) + tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((waves_m, waves_n, 1), (waves_n, 1, 0))) + thr_mma = tiled_mma.thr_slice(tid) + + frag_A = thr_mma.make_fragment_A(bA) + frag_B = thr_mma.make_fragment_B(bB) + frag_C = thr_mma.make_fragment_C(bC) + + copy_ab = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_cls.width), elem_cls) + copy_c = fx.make_copy_atom(fx.rocdl.BufferCopy(f32.width), f32) + thr_copy_A = fx.make_tiled_copy_A(copy_ab, tiled_mma).get_slice(tid) + thr_copy_B = fx.make_tiled_copy_B(copy_ab, tiled_mma).get_slice(tid) + thr_copy_C = fx.make_tiled_copy_C(copy_c, tiled_mma).get_slice(tid) + + fx.copy(copy_ab, thr_copy_A.partition_S(bA), thr_copy_A.retile(frag_A)) + fx.copy(copy_ab, thr_copy_B.partition_S(bB), thr_copy_B.retile(frag_B)) + + frag_C.fill(0) + for ki in range_constexpr(k_iters): + fx.gemm(tiled_mma, frag_C, frag_A[None, None, ki], frag_B[None, None, ki], frag_C) + fx.copy(copy_c, thr_copy_C.retile(frag_C), thr_copy_C.partition_S(bC)) + + @flyc.jit + def launch(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + wmma_kernel(A, B, C).launch(grid=(1, 1, 1), block=(threads, 1, 1), stream=stream) + + return launch + + +def _exact_operands(m, n, k, torch_dtype): + """Small integer inputs, so the f32 result is exact and any mismatch is layout.""" + torch.manual_seed(0) + a = (torch.randn(m, k, device="cuda") * 4).round().to(torch_dtype) + b = (torch.randn(n, k, device="cuda") * 4).round().to(torch_dtype) + c = torch.zeros(m, n, dtype=torch.float32, device="cuda") + return a, b, c + + +@pytest.mark.parametrize( + "elem_cls, torch_dtype", + [ + (fx.BFloat16, torch.bfloat16), + (fx.Float16, torch.float16), + ], + ids=["bf16", "f16"], +) +def test_single_wmma_atom(elem_cls, torch_dtype): + """A single gfx11.wmma atom call must match A @ B.T exactly. + + This is the narrowest check that ``make_tiled_copy_A`` can partition an + operand whose thread layout replicates across the two lane halves. + """ + a, b, c = _exact_operands(WMMA_M, WMMA_N, WMMA_K, torch_dtype) + + launch = _compile_single_wmma(elem_cls) + launch(a, b, c, stream=torch.cuda.current_stream()) + torch.cuda.synchronize() + + torch.testing.assert_close(c, a.float() @ b.float().T, atol=_ATOL, rtol=_RTOL) + + +@pytest.mark.parametrize( + "tile_m, tile_n, tile_k, waves_m, waves_n", + [ + pytest.param(32, 32, 16, 2, 2, id="32x32x16-2x2waves"), + pytest.param(64, 64, 32, 2, 2, id="64x64x32-2x2waves-2x2x2repeats"), + pytest.param(32, 32, 32, 1, 1, id="32x32x32-1wave-2x2x2repeats"), + ], +) +def test_tiled_wmma_gemm(tile_m, tile_n, tile_k, waves_m, waves_n): + """Atom repeats across a wave grid, the geometry a block tile actually uses. + + A single atom can be right while the tiled partition is wrong, because the + repeat and wave axes are what stack on top of the replicated lane axis. + """ + a, b, c = _exact_operands(tile_m, tile_n, tile_k, torch.bfloat16) + + launch = _compile_tiled_wmma(fx.BFloat16, tile_m, tile_n, tile_k, waves_m, waves_n) + launch(a, b, c, stream=torch.cuda.current_stream()) + torch.cuda.synchronize() + + torch.testing.assert_close(c, a.float() @ b.float().T, atol=_ATOL, rtol=_RTOL) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index b6bbf5db4..77aab2beb 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/tests/kernels/test_rdna_gemm.py @@ -49,6 +49,12 @@ def _requires_rdna_wmma(): pytest.skip(f"RDNA WMMA GEMM requires gfx11* or gfx120*, got {ARCH}") +def _requires_rdna3(): + """Only rdna3_f16_gemm picks its tile from the shape; gfx12 still uses a fixed one.""" + if not ARCH.startswith("gfx11"): + pytest.skip(f"gfx11-only behaviour, got {ARCH}") + + def create_wmma_gemm_module(*args, **kwargs): """Pick the kernel variant matching the current arch. @@ -163,6 +169,46 @@ def test_f16_gemm_f32_output(M, N, K): assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05) +@pytest.mark.parametrize( + "M, N, K", + [ + pytest.param(384, 384, 2048, id="384x384x2048"), + pytest.param(1152, 1152, 1024, id="1152x1152x1024"), + pytest.param(2560, 2560, 1024, id="2560x2560x1024", marks=pytest.mark.large_shape), + ], +) +def test_f16_gemm_grid_m_not_a_multiple_of_the_group_width(M, N, K): + """Shapes whose M-tile count is not a multiple of the L2 grouping cap. + + The grid swizzle derives bid_m from a fixed group width, so before the width + was snapped down to a divisor of grid_m the last group addressed tiles past + the end of the grid. This is reachable at the default 128x128 tile, not only + at narrower ones: measured on gfx1100, 1152, 1280 and 1664 square came back + wrong by roughly 400x the bf16 rounding floor, while 1536 and 2560 square + faulted the GPU. Which of the two you get depends on whether the address past + the grid happens to be mapped, so the silent wrong answer is the common case. + + At the default 128x128 tile these give grid_m of 3, 9 and 20, none of them a + multiple of the group width of 8. + """ + _requires_rdna3() + torch.manual_seed(42) + + launch_fn, BLOCK_M, _, _ = _create_wmma_gemm_module_gfx11(M, N, K, in_dtype="bf16", out_dtype="bf16") + grid_m = M // BLOCK_M + assert grid_m % 8, f"grid_m={grid_m} divides the default group width; shape no longer covers the fault" + + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 0.1 + B_T = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda") + + launch_fn(C, A, B_T, torch.cuda.current_stream()) + torch.cuda.synchronize() + + C_ref = A.float() @ B_T.float().T + assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05) + + @pytest.mark.parametrize( "M, N, K", [ diff --git a/tests/unit/test_rdna3_grid_swizzle.py b/tests/unit/test_rdna3_grid_swizzle.py new file mode 100644 index 000000000..6a2c8b51d --- /dev/null +++ b/tests/unit/test_rdna3_grid_swizzle.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors +"""GPU-free tests for the RDNA3 GEMM workgroup swizzle. + +``_group_width`` / ``_swizzle_tile_id`` map a linear workgroup id onto the tile +grid, walking down M before stepping in N so concurrent workgroups share B tiles +in L2. The mapping must be a bijection onto the grid. It was not, for any grid_m +that the grouping cap did not divide: the last group ran off the end, bid_m +exceeded grid_m, and the kernel wrote past C. + +All plain integer arithmetic that runs before a kernel is built, so no GPU is +needed. The device-side counterpart is +``tests/kernels/test_rdna_gemm.py::test_f16_gemm_grid_m_not_a_multiple_of_the_group_width``. +""" + +import pytest + +from kernels.gemm.rdna3_f16_gemm import _group_width, _swizzle_tile_id + +pytestmark = pytest.mark.l0_backend_agnostic + +DEFAULT_GROUP_M = 8 # create_wmma_gemm_module's default grouping cap + + +def _swizzled_tiles(grid_m, grid_n, group_m=DEFAULT_GROUP_M): + width = _group_width(grid_m, group_m) + return [_swizzle_tile_id(pid, grid_n, width) for pid in range(grid_m * grid_n)] + + +@pytest.mark.parametrize( + "grid_m, expected", + [ + (1, 1), + (3, 3), + (5, 5), + (8, 8), + (12, 6), # 1536 / 128 — used to take 8 and address tiles 12..15 + (16, 8), + (20, 5), # 2560 / 128 — used to take 8 and address tiles 20..23 + (32, 8), + ], +) +def test_group_width_is_the_largest_divisor_within_the_cap(grid_m, expected): + assert _group_width(grid_m, DEFAULT_GROUP_M) == expected + + +@pytest.mark.parametrize("group_m", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("grid_m", range(1, 65)) +def test_group_width_always_divides_the_grid(grid_m, group_m): + width = _group_width(grid_m, group_m) + assert grid_m % width == 0 + assert 1 <= width <= min(group_m, grid_m) + + +@pytest.mark.parametrize( + "grid_m, grid_n", + [ + (1, 1), + (1, 8), + (8, 1), + (12, 12), # 1536x1536 at 128x128: the grid that faulted + (20, 20), # 2560x2560 at 128x128: likewise + (9, 9), # 1152x1152: returned a wrong C instead of faulting + (10, 10), # 1280x1280: likewise + (13, 13), # 1664x1664: likewise + (3, 7), + (5, 4), + (7, 3), + (16, 16), + (12, 5), + (20, 3), + ], +) +def test_swizzle_covers_every_tile_exactly_once(grid_m, grid_n): + """The swizzle must be a bijection onto the grid. + + With a grouping width that does not divide grid_m the last group runs off the + end: bid_m exceeds grid_m and the kernel addresses past C. Measured on gfx1100 + that is a wrong result at grid_m 9, 10 and 13 and a hard fault at 12 and 20, + so a bijection here is a memory-safety property, not just a tidy mapping. + """ + mapped = _swizzled_tiles(grid_m, grid_n) + + assert all(0 <= m < grid_m and 0 <= n < grid_n for m, n in mapped) + assert set(mapped) == {(m, n) for m in range(grid_m) for n in range(grid_n)}