From 77359de61c0f28d155b80ea27940ea24121aa2d0 Mon Sep 17 00:00:00 2001 From: vlluvia Date: Fri, 7 Aug 2026 11:12:51 +0800 Subject: [PATCH 1/3] [Kernel] Migrate RDNA3 WMMA GEMM to the tiled copy/MMA API Replaces the hand-rolled index arithmetic in rdna3_f16_gemm with the layout-based API, and fixes a latent grid-swizzle fault found on the way. The epilogue now partitions C through a tiled_mma and make_tiled_copy_C instead of an explicit ``g_row = base + 2*si + klane`` store loop. The tiled_mma carries a permutation that reproduces the kernel's wave banding: the default stamping interleaves the repeats instead, which measured 62% slower at 3072x3072x1024. The GMEM to LDS path now goes through flat_divide plus a tiled copy with BufferCopy128b, replacing the precomputed offset tables and the flat _v8_store. The thread geometry is the same assignment the tables computed. The grid swizzle moves out of the kernel body into _group_width and _swizzle_tile_id. It derived bid_m from a group width that need not divide grid_m, so the last group addressed tiles past the end of the grid and the kernel wrote past C. Every shape in use happened to divide evenly, which is why it survived; it becomes reachable as soon as a caller asks for a tile narrower than 128x128. The block tile stays a parameter defaulting to 128x128x32, so this is a refactor: the generated ISA is byte-identical for rn and rs output across 128x128x32, 64x64x64 and 128x64x32 on gfx1100. Adds test_rdna3_wmma_atom for the gfx11 tiled copy atoms, which have to handle lanes 16-31 mirroring lanes 0-15 under the v16 operand ABI, and test_rdna3_grid_swizzle for the swizzle bijection. Co-authored-by: Cursor --- kernels/gemm/rdna3_f16_gemm.py | 312 ++++++++++++++++---------- tests/kernels/test_rdna3_wmma_atom.py | 210 +++++++++++++++++ tests/kernels/test_rdna_gemm.py | 43 ++++ tests/unit/test_rdna3_grid_swizzle.py | 74 ++++++ 4 files changed, 520 insertions(+), 119 deletions(-) create mode 100644 tests/kernels/test_rdna3_wmma_atom.py create mode 100644 tests/unit/test_rdna3_grid_swizzle.py diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 513d52b20..217d72fce 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 and fault. At the default 128x128 tile grid_m divides evenly for + every shape in use, which is why this only surfaces once a caller asks for a + narrower tile. + """ + 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,80 @@ 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 +334,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 +347,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 +364,41 @@ 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 +408,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..451041c6c 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,43 @@ 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 and the kernel faulted writing C. Every shape in use when + the swizzle was written happened to divide evenly, which is why it survived. + + 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..8ff287fc5 --- /dev/null +++ b/tests/unit/test_rdna3_grid_swizzle.py @@ -0,0 +1,74 @@ +#!/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 + (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, the kernel writes past C, and the launch faults. + """ + 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)} From bcabe9ec81a66ca471ecbf6d96dcb47c9a1e65a4 Mon Sep 17 00:00:00 2001 From: vlluvia Date: Fri, 7 Aug 2026 12:30:05 +0800 Subject: [PATCH 2/3] [Kernel] Correct how the grid-swizzle fault is described Benchmarking the migration against the unmodified kernel turned up that the swizzle bug is reachable at the default 128x128 tile, not only at narrower ones as the docstrings claimed, and that it has two symptoms rather than one. Measured on gfx1100 at 128x128x32: M of 1152, 1280 and 1664 return a wrong C, off by roughly 400x the bf16 rounding floor, while 1536 and 2560 fault the GPU. Which one you get depends on whether the address past the grid happens to be mapped, so the silent wrong answer is the common case and the fault is the lucky one. No behaviour change -- the fix itself already landed with the swizzle rewrite. This corrects the three places that understated it and adds grid_m of 9, 10 and 13 to the bijection test so the wrong-answer grids are covered alongside the two that faulted. Co-authored-by: Cursor --- kernels/gemm/rdna3_f16_gemm.py | 6 +++--- tests/kernels/test_rdna_gemm.py | 7 +++++-- tests/unit/test_rdna3_grid_swizzle.py | 7 ++++++- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 217d72fce..c50afb321 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -55,9 +55,9 @@ def _group_width(grid_m, group_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 and fault. At the default 128x128 tile grid_m divides evenly for - every shape in use, which is why this only surfaces once a caller asks for a - narrower tile. + 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) diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index 451041c6c..77aab2beb 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/tests/kernels/test_rdna_gemm.py @@ -182,8 +182,11 @@ def test_f16_gemm_grid_m_not_a_multiple_of_the_group_width(M, N, K): 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 and the kernel faulted writing C. Every shape in use when - the swizzle was written happened to divide evenly, which is why it survived. + 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. diff --git a/tests/unit/test_rdna3_grid_swizzle.py b/tests/unit/test_rdna3_grid_swizzle.py index 8ff287fc5..fedd73a26 100644 --- a/tests/unit/test_rdna3_grid_swizzle.py +++ b/tests/unit/test_rdna3_grid_swizzle.py @@ -59,6 +59,9 @@ def test_group_width_always_divides_the_grid(grid_m, group_m): (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), ], ) @@ -66,7 +69,9 @@ 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, the kernel writes past C, and the launch faults. + 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) From 55ec22458654914befad0173d7175ae53f81a8c4 Mon Sep 17 00:00:00 2001 From: vlluvia Date: Fri, 7 Aug 2026 15:00:17 +0800 Subject: [PATCH 3/3] [Kernel] Reformat the RDNA3 GEMM changes with black The repo formats at 120 columns; these files were wrapped nearer 100, so black wanted to rejoin several call arguments and expand the parametrize lists to one tuple per line. Formatting only, no behaviour change. Co-authored-by: Cursor --- kernels/gemm/rdna3_f16_gemm.py | 12 +++--------- tests/unit/test_rdna3_grid_swizzle.py | 25 ++++++++++++++++--------- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index c50afb321..024bd9f22 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -230,9 +230,7 @@ def _v8_load(v8_idx): # ============================================================ # GMEM -> registers -> LDS, through the tiled copy # ============================================================ - tA = fx.flat_divide(fx.rocdl.make_buffer_tensor(arg_a), fx.make_tile(BLOCK_M, BLOCK_K))[ - None, None, bid_m, None - ] + 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 ] @@ -250,9 +248,7 @@ def _v8_load(v8_idx): # 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)) - ) + 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): @@ -373,9 +369,7 @@ def _do_compute_rk(accs_in, rk, buf_offset): # # 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) - ] + 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 diff --git a/tests/unit/test_rdna3_grid_swizzle.py b/tests/unit/test_rdna3_grid_swizzle.py index fedd73a26..6a2c8b51d 100644 --- a/tests/unit/test_rdna3_grid_swizzle.py +++ b/tests/unit/test_rdna3_grid_swizzle.py @@ -35,9 +35,9 @@ def _swizzled_tiles(grid_m, grid_n, group_m=DEFAULT_GROUP_M): (3, 3), (5, 5), (8, 8), - (12, 6), # 1536 / 128 — used to take 8 and address tiles 12..15 + (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 + (20, 5), # 2560 / 128 — used to take 8 and address tiles 20..23 (32, 8), ], ) @@ -56,13 +56,20 @@ def test_group_width_always_divides_the_grid(grid_m, group_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), + (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):