diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 024bd9f22..1c7f49de9 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -50,6 +50,48 @@ K_PAD = 8 +def _sched_plan(reg_m, reg_n, reg_k, g2s_chunks): + """The order the k-tile body is meant to issue in, as (group, count) pairs. + + Left alone, LLVM hoists the whole ds_store block into the middle of the WMMA + stream, so each store's ``s_waitcnt vmcnt`` stalls the wave while the WMMAs + behind it wait. This spends the WMMA stream as cover instead: global loads + go out first, the next k-step's LDS reads hide under this k-step's math, and + the stores drain under the tail, which also leaves the ``lgkmcnt(0)`` in + front of the barrier with almost nothing left to wait for. + + Built here, from Python ints, because the kernel body is traced and cannot + branch on values. The counts are hints: a group the dependences cannot + satisfy is dropped rather than honoured. + """ + per_rk_wmma = reg_m * reg_n + per_rk_dsrd = 2 * (reg_m + reg_n) # each v16 operand is two 128-bit reads + chunk = max(1, reg_n) + + plan = [("vmem", g2s_chunks), ("dsrd", per_rk_dsrd)] + dsrd_left = per_rk_dsrd * (reg_k - 1) + dsrd_step = chunk + dswr_left = g2s_chunks + dswr_chunk = max(1, g2s_chunks // 2) + for _ in range(reg_k): + issued = 0 + while issued < per_rk_wmma: + n = min(chunk, per_rk_wmma - issued) + plan.append(("mfma", n)) + issued += n + if dsrd_left: + take = min(dsrd_step, dsrd_left) + plan.append(("dsrd", take)) + dsrd_left -= take + elif dswr_left: + take = min(dswr_chunk, dswr_left) + plan.append(("dswr", take)) + dswr_left -= take + if dswr_left: + plan.append(("dswr", dswr_left)) + return tuple(plan) + + def _group_width(grid_m, group_m): """Largest grouping width <= group_m that divides grid_m. @@ -95,6 +137,8 @@ def create_wmma_gemm_module( group_m=8, a_k_pad=K_PAD, b_k_pad=K_PAD, + lds_layout="pad", + sched_hint=False, ): gpu_arch = str(get_rocm_arch() or "") if not gpu_arch.startswith("gfx11"): @@ -109,7 +153,10 @@ def create_wmma_gemm_module( NUM_WAVES = waves_m * waves_n # 4 THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE # 128 - assert reg_k >= 2 and reg_k % 2 == 0 + # One WMMA K-step is 16 elements and the v16 operand reads it as two 8-wide + # chunks, so BLOCK_K only has to be a multiple of 16. reg_k=1 (BLOCK_K=16) is + # what keeps the LDS tile small enough to fit more than one workgroup per CU. + assert reg_k >= 1 assert rounding in ("rn", "rs"), f"rounding must be 'rn' or 'rs', got {rounding!r}" if rounding == "rs": assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only" @@ -121,13 +168,37 @@ def create_wmma_gemm_module( # ``row = tid // THRS_K, col = (tid % THRS_K) * LOAD_VEC``. THRS_K = BLOCK_K // LOAD_VEC THRS_M = THREADS_PER_BLOCK // THRS_K + # 128-bit chunks of A and B one thread moves per k-tile; equals both the + # global-load and the ds_store count in the loop body. + G2S_CHUNKS = (BLOCK_M + BLOCK_N) * BLOCK_K // THREADS_PER_BLOCK // LOAD_VEC + SCHED_PLAN = _sched_plan(reg_m, reg_n, reg_k, G2S_CHUNKS) if sched_hint else () 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 - LDS_A_SIZE = BLOCK_M * BLOCK_K_PAD_A - LDS_B_SIZE = BLOCK_N * BLOCK_K_PAD_B + # Two ways to lay a tile out in LDS, and the choice is what decides which + # macro tiles are reachable at all: + # + # "pad" row-major, BLOCK_K + 8 elements per row. The pad is what keeps + # the 128-bit reads off each other's banks, and it is the only + # pad that does: the row start has to stay 16-byte aligned, which + # leaves 0, 8, 16 and 24, and of those only 8 and 24 spread the + # 16 rows a read touches over all 32 banks. + # + # "kblock" k-major in groups of 8, so element (row, k) sits at + # ((k // 8) * rows + row) * 8 + k % 8. Consecutive rows are then + # 16 bytes apart, so the 8 lanes the LDS services per cycle cover + # 128 contiguous bytes -- every bank, once -- with no pad at all. + # + # Dropping the pad is worth 20% of the LDS budget, which is what brings a + # 256x256x32 tile inside the 64 KB a workgroup may allocate. + assert lds_layout in ("pad", "kblock") + if lds_layout == "kblock": + assert BLOCK_K % LOAD_VEC == 0 + a_k_pad = b_k_pad = 0 + ROW_STRIDE_A = BLOCK_K + a_k_pad + ROW_STRIDE_B = BLOCK_K + b_k_pad + LDS_A_SIZE = BLOCK_M * ROW_STRIDE_A + LDS_B_SIZE = BLOCK_N * ROW_STRIDE_B LDS_ONE_BUF = LDS_A_SIZE + LDS_B_SIZE LDS_TOTAL = 2 * LDS_ONE_BUF @@ -248,14 +319,30 @@ 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))) + if const_expr(lds_layout == "kblock"): + # (row, k) with k split as (k % 8, k // 8): the low part is + # contiguous so the 128-bit copy atom still sees eight adjacent + # elements, the high part strides a whole plane of rows. + layout = fx.make_layout( + (rows, (LOAD_VEC, BLOCK_K // LOAD_VEC)), + (LOAD_VEC, (1, rows * LOAD_VEC)), + ) + else: + layout = fx.make_layout((rows, BLOCK_K), (row_stride, 1)) + view = fx.make_view(fx.recast_iter(elem_dtype, ptr), layout) 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) + return _lds_dst(buf_offset, 0, BLOCK_M, ROW_STRIDE_A) def _pB_s(buf_offset): - return _lds_dst(buf_offset, LDS_A_SIZE, BLOCK_N, BLOCK_K_PAD_B) + return _lds_dst(buf_offset, LDS_A_SIZE, BLOCK_N, ROW_STRIDE_B) + + def _lds_elem(rows, row_stride, row, col): + """Element index of (row, col) inside one A- or B-tile of the buffer.""" + if const_expr(lds_layout == "kblock"): + return (col // LOAD_VEC * rows + row) * LOAD_VEC + col % LOAD_VEC + return row * row_stride + col frag_copy_A = fx.make_fragment_like(_pA_s(0)) frag_copy_B = fx.make_fragment_like(_pB_s(0)) @@ -283,8 +370,8 @@ def _load_b_from_lds(rk, buf_offset): col_hi = 16 * rk + 8 for rn in range_constexpr(reg_n): row = wave_n * (reg_n * WMMA_N) + 16 * rn + lane16 - lds_idx_lo = buf_offset + LDS_A_SIZE + row * BLOCK_K_PAD_B + col_lo - lds_idx_hi = buf_offset + LDS_A_SIZE + row * BLOCK_K_PAD_B + col_hi + lds_idx_lo = buf_offset + LDS_A_SIZE + _lds_elem(BLOCK_N, ROW_STRIDE_B, row, col_lo) + lds_idx_hi = buf_offset + LDS_A_SIZE + _lds_elem(BLOCK_N, ROW_STRIDE_B, row, col_hi) v_lo = _v8_load(lds_idx_lo // 8) v_hi = _v8_load(lds_idx_hi // 8) vecs.append(v_lo.shuffle(v_hi, _concat16_mask)) @@ -294,8 +381,8 @@ def _load_a_single_from_lds(rk, rm_val, buf_offset): col_lo = 16 * rk col_hi = 16 * rk + 8 row = wave_m * (reg_m * WMMA_M) + 16 * rm_val + lane16 - lds_idx_lo = buf_offset + row * BLOCK_K_PAD_A + col_lo - lds_idx_hi = buf_offset + row * BLOCK_K_PAD_A + col_hi + lds_idx_lo = buf_offset + _lds_elem(BLOCK_M, ROW_STRIDE_A, row, col_lo) + lds_idx_hi = buf_offset + _lds_elem(BLOCK_M, ROW_STRIDE_A, row, col_hi) v_lo = _v8_load(lds_idx_lo // 8) v_hi = _v8_load(lds_idx_hi // 8) return v_lo.shuffle(v_hi, _concat16_mask) @@ -310,11 +397,30 @@ def _barrier(): has_side_effects=True, ) - def _do_compute_rk(accs_in, rk, buf_offset): + def _do_compute_rk(accs_in, rk, buf_offset, b_vecs): new_accs = list(accs_in) - b_vecs = _load_b_from_lds(rk, buf_offset) + # Each A fragment feeds reg_n back-to-back WMMAs. Emitting its read + # immediately before its first consumer lets the register allocator + # give every rm the same quad, and then the read for rm+1 cannot + # issue until the WMMAs on rm have retired -- the body ends up with + # a full ``lgkmcnt(0)`` drain in front of each group of reg_n WMMAs + # instead of a partial wait. Issuing one fragment ahead is what + # breaks that: the read in flight and the one being consumed need + # different registers, so the wait has something to overlap. + # + # Only survives together with sched_hint. On its own the machine + # scheduler sinks the read back down onto its consumer and the ISA + # comes out unchanged at any prefetch depth; the group barriers are + # what hold the reads up front for the allocator to see. Measured as + # a pair at 256x256x32, worth 5% (109.1 -> 103.9 ms at 16384 cubed) + # and 9 lgkmcnt(0) drains per k-tile down to 4. Neither shows up in + # an instruction histogram -- the counts and the register total are + # the same, only the order changes. + a_next = _load_a_single_from_lds(rk, 0, buf_offset) for rm in range_constexpr(reg_m): - a_vec = _load_a_single_from_lds(rk, rm, buf_offset) + a_vec = a_next + if const_expr(rm + 1 < reg_m): + a_next = _load_a_single_from_lds(rk, rm + 1, buf_offset) for rn in range_constexpr(reg_n): idx = rm * reg_n + rn new_accs[idx] = _wmma_op( @@ -324,6 +430,29 @@ def _do_compute_rk(accs_in, rk, buf_offset): ) return new_accs + def _compute_k_tile(accs_in, buf_offset): + """All reg_k WMMA steps over one LDS buffer. + + Step rk reads B into the registers step rk-1 is still using, so its + reads cannot issue until that step's WMMAs retire, and the wait in + front of them is a full lgkmcnt(0) drain rather than a partial one. + Reading every step's B up front instead does remove that dependence, + and it loses: 2*reg_n*(reg_k-1) more live registers pushes the + allocator into recycling the A fragments harder, and the drains go + from 4 per k-tile to 8. Measured at 256x256x32, 104.0 -> 118.3 ms. + """ + new_accs = list(accs_in) + for rk in range_constexpr(reg_k): + new_accs = _do_compute_rk(new_accs, rk, buf_offset, + _load_b_from_lds(rk, buf_offset)) + return new_accs + + def _sched_k_tile(): + emit = {"vmem": rocdl.sched_vmem, "mfma": rocdl.sched_mfma, + "dsrd": rocdl.sched_dsrd, "dswr": rocdl.sched_dswr} + for group, count in SCHED_PLAN: + emit[group](count) + zero_acc = fx.full(8, 0.0, fx.Float32) accs = [zero_acc for _ in range_constexpr(reg_m * reg_n)] @@ -337,27 +466,33 @@ def _do_compute_rk(accs_in, rk, buf_offset): n_acc = reg_m * reg_n init_state = list(accs) - for iv, state in range(0, num_k_tiles - 1, 1, init=init_state): - s_accs = list(state[:n_acc]) - - read_off = iv % 2 * c_lds_buf_stride - write_off = (1 - iv % 2) * c_lds_buf_stride - - _gmem_load(iv + 1) - - for rk in range_constexpr(reg_k): - s_accs = _do_compute_rk(s_accs, rk, read_off) - + def _one_k_tile(s_accs, read_off, write_off, load_tile): + """Prefetch the next k-tile, consume this one, hand over, barrier.""" + _gmem_load(load_tile) + s_accs = _compute_k_tile(s_accs, read_off) _lds_store(write_off) + if const_expr(sched_hint): + _sched_k_tile() _barrier() - + return s_accs + + # The read buffer alternates with the trip counter, so both offsets are + # values derived from ``iv`` and the body spends ~10 VALU and SALU ops + # per trip on them. Stepping the loop by two fixes each half's parity at + # trace time and folds the arithmetic into ds_load immediates: overhead + # instructions drop from 0.93 to 0.31 per WMMA and VGPRs from 212 to + # 204. It is 1.7% slower (104.0 -> 105.8 ms at 16384 cubed). The loop is + # not issue-bound, so paying more instructions is not what it costs. + for iv, state in range(0, num_k_tiles - 1, 1, init=init_state): + s_accs = list(state[:n_acc]) + s_accs = _one_k_tile(s_accs, iv % 2 * c_lds_buf_stride, + (1 - iv % 2) * c_lds_buf_stride, iv + 1) results = yield list(s_accs) accs = list(results[:n_acc]) last_read_off = ((num_k_tiles - 1) % 2) * c_lds_buf_stride - for rk in range_constexpr(reg_k): - accs = _do_compute_rk(accs, rk, last_read_off) + accs = _compute_k_tile(accs, last_read_off) # ============================================================ # Store results to GMEM through the tiled copy diff --git a/kernels/gemm/rdna3_f16_gemm_autotune.py b/kernels/gemm/rdna3_f16_gemm_autotune.py new file mode 100644 index 000000000..48a706a89 --- /dev/null +++ b/kernels/gemm/rdna3_f16_gemm_autotune.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Tile selection for the RDNA3 WMMA GEMM. + +``rdna3_f16_gemm`` builds whatever tile it is handed and defaults to 128x128x32. +That tile is right once the problem fills the grid, but it cuts only 4 +workgroups at 256x256 and 16 at 512x512, so on a 96-CU part most CUs idle no +matter how good the inner loop is. Choosing the tile from the shape is worth up +to 3.0x there, and this module owns that decision in two layers: + + * ``pick_tile`` — a heuristic fitted to a sweep of every feasible tile on 27 + shapes. It needs no GPU and no measurement, and it is what a call resolves + to with nothing configured, so the wrapper benchmarks nothing by default. + * the shared autotuner — ``FLYDSL_AUTOTUNE=1`` sweeps ``feasible_tiles`` for + real on the GPU in hand, and the result can be frozen into an offline + artifact so other machines with the same device fingerprint skip the search. + +The second layer earns its keep mainly where the first cannot reach: ``NUM_CU`` +is hard-coded for gfx1100, so the thresholds do not transfer to a gfx11 part +with a different CU count, and shapes outside the fitted set are extrapolation. +It is also the least settled part of this — see ``_graph_bench`` for why its +verdict on the shortest kernels should not be taken at face value. The default +path does not benchmark and is unaffected. + +The tile is not a Constexpr argument of one compiled entry point: it decides the +block shape, the wave grid and the LDS budget, so each candidate is a separate +module. The dispatcher below therefore takes the tile as ordinary keyword +arguments and looks the built module up in a cache, the same shape of +indirection ``conv3d_implicit_autotune`` uses. +""" + +import functools + +import torch + +from flydsl.autotune import Config, autotune, do_bench +from kernels.gemm.rdna3_f16_gemm import K_PAD, WAVE_SIZE, WMMA_K, WMMA_M, WMMA_N, create_wmma_gemm_module + +# gfx1100 (W7900) exposes 96 CUs. Used only to decide when a tile is too coarse +# to fill the machine; being off by a little just shifts one ladder step. +NUM_CU = 96 + +# LDS budget per workgroup. K_PAD comes from the kernel so the feasibility check +# below cannot drift from the allocation it is predicting. +LDS_BYTES = 64 * 1024 + +# Named by the block tile they produce: (reg_m, reg_n, reg_k, waves_m, waves_n). +TILE_256x256x32 = (4, 4, 2, 4, 4) +TILE_128x256x32 = (4, 4, 2, 2, 4) +TILE_128x128x32 = (4, 4, 2, 2, 2) +TILE_128x64x32 = (4, 2, 2, 2, 2) +TILE_64x64x64 = (2, 2, 4, 2, 2) +TILE_32x64x64 = (2, 2, 4, 1, 2) +TILE_32x32x64 = (2, 2, 4, 1, 1) + +# Options a tile needs beyond its shape. Kept next to the ladder rather than in +# the ladder tuples so a Config stays the five tile fields the autotuner sweeps. +# +# 256x256x32 is only reachable unpadded: with K_PAD it wants 80 KB and the pad is +# what the LDS budget cannot afford at that width. Unpadded it is exactly 64 KB. +# The hints are not a tuning knob there but a requirement — without the immediate +# offsets the pad used to give it, LLVM assigns every A fragment the same register +# quad, so each ds_load waits on a full lgkmcnt(0) drain instead of overlapping +# the WMMA stream. Measured at 128x256x32: 53 TFLOP/s without the hints, 84 with. +# +# group_m is the width of the L2 grouping. The wide tiles read a whole 256-row +# band of A per workgroup, so the default 8 walks off the reuse the swizzle is +# there to get; 16 was the best of 1/4/8/16/32 for both, worth 1-2%. +_DEFAULT_OPTS = {"lds_layout": "pad", "sched_hint": False, "group_m": 8} +_TILE_OPTS = { + TILE_256x256x32: {"lds_layout": "kblock", "sched_hint": True, "group_m": 16}, + TILE_128x256x32: {"lds_layout": "pad", "sched_hint": False, "group_m": 16}, +} + + +def tile_opts(tile): + return _TILE_OPTS.get(tuple(tile), _DEFAULT_OPTS) + + +# Tile ladder, widest first. +# +# 128x64x32 exists only on the small-K ladder: with large K a workgroup runs long +# enough that the deeper k-tile (fewer barriers, twice as long for the gmem +# prefetch to land) pays off, while with small K the per-workgroup prologue and +# epilogue dominate and the wider tile amortizes them. +# +# The ladder is the feasibility-ordered search space; pick_tile does not walk it +# in order. 32x64x64 is here only as a fallback for shapes the others cannot +# divide -- it was not the fastest tile on any of the 27 shapes swept. +_LADDER_LARGE_K = [TILE_256x256x32, TILE_128x256x32, TILE_128x128x32, TILE_64x64x64, + TILE_32x64x64, TILE_32x32x64] +_LADDER_SMALL_K = [TILE_128x128x32, TILE_128x64x32, TILE_64x64x64, TILE_32x64x64, TILE_32x32x64] + + +def _tile_workgroups(M, N, K, cfg): + """Workgroup count for this tile, or None if the shape cannot use it.""" + reg_m, reg_n, reg_k, waves_m, waves_n = cfg + block_m = WMMA_M * reg_m * waves_m + block_n = WMMA_N * reg_n * waves_n + block_k = WMMA_K * reg_k + threads = waves_m * waves_n * WAVE_SIZE + if M % block_m or N % block_n or K % block_k: + return None + if K // block_k < 2: # the prefetch pipeline needs at least two k-tiles + return None + # Every thread must carry a whole 8-element vector of both tiles. + if (block_m * block_k) % (threads * 8) or (block_n * block_k) % (threads * 8): + return None + pad = 0 if tile_opts(cfg)["lds_layout"] == "kblock" else K_PAD + if 2 * (block_m + block_n) * (block_k + pad) * 2 > LDS_BYTES: # 2 buffers, 2 bytes/elem + return None + return (M // block_m) * (N // block_n) + + +def _ladder_for(K): + return _LADDER_SMALL_K if K <= 1024 else _LADDER_LARGE_K + + +def feasible_tiles(M, N, K): + """``(tile, workgroup count)`` for every ladder tile this shape can run, widest first. + + Also the search space the autotuner sweeps: anything excluded here does not + divide the shape, cannot fill the prefetch pipeline, or does not fit in LDS, + so benchmarking it would only measure a build failure. + """ + return [(cfg, wgs) for cfg in _ladder_for(K) if (wgs := _tile_workgroups(M, N, K, cfg)) is not None] + + +def pick_tile(M, N, K): + """Tile for this shape, fitted to a sweep of every ladder tile on 27 shapes. + + 64x64x64 is the default rather than the widest tile that covers the machine. + Measured on gfx1100 it is the fastest tile on 16 of the 27 shapes and holds + 50-59 TFLOP/s across the whole range, where 128x128x32 swings between 40 and + 72 depending on how its much coarser grid happens to land. Taking the widest + covering tile cost up to 37% (1664x1664x1024) and averaged 6.5%. + + Five exceptions, in order. The first two are the wide tiles, and both are + about operand traffic rather than about filling the machine: a tile reads + ``(BM + BN) / (BM * BN)`` of A and B per output, so 128x256 moves a quarter + less than 128x128 and 256x256 half as much again. Measured against rocBLAS TN + in one thermal window, as a fraction of its time: + + square 4096 8192 12288 16384 20480 + 128x128x32 0.805 0.903 0.919 0.875 0.829 + 128x256x32 0.826 0.933 0.990 0.988 0.928 + 256x256x32 0.797 0.913 0.979 1.012 0.999 + + So the widest tile is not simply best: 256x256 needs roughly 32 workgroups per + CU before its traffic saving outweighs how coarsely its grid lands, and below + that 128x256 leads. Above it the order reverses and stays reversed, because + 128x256 is the one that starts falling off as the footprint grows. + + * 256x256x32 once the grid is worth ~32 workgroups per CU. Only reachable + unpadded; see _TILE_OPTS for why that drags the scheduling hints with it. + * 128x256x32 from ~4 workgroups per CU up. Never behind 128x128x32 anywhere + it applies, by 2-11%. + * 128x128x32 once the grid is worth at least ~2.5 workgroups per CU. Its + compute intensity wins outright there, by 5-11% over 64x64x64. + * 128x64x32 (small-K ladder only) when it lands near one workgroup per CU. + That is the narrow band around 1024x1024, where it leads by 13-28%. + * 32x32x64 when 64x64x64 cannot fill a quarter of the machine and K is + long enough for the idle CUs to dominate: 256x256x4096, worth 23%. + + Against the per-shape fastest tile this averages 0.6%, worst case 8.1% at + 1152x1152x1024, where 128x128x32's grid happens to land well. + """ + feasible = dict(feasible_tiles(M, N, K)) + if not feasible: + return _ladder_for(K)[0] + + if feasible.get(TILE_256x256x32, 0) >= 32 * NUM_CU: + return TILE_256x256x32 + if feasible.get(TILE_128x256x32, 0) >= 4 * NUM_CU: + return TILE_128x256x32 + if feasible.get(TILE_128x128x32, 0) >= 2.5 * NUM_CU: + return TILE_128x128x32 + if NUM_CU <= feasible.get(TILE_128x64x32, 0) <= 1.5 * NUM_CU: + return TILE_128x64x32 + if TILE_64x64x64 in feasible: + starved = feasible[TILE_64x64x64] < NUM_CU / 4 and K >= 2048 + if starved and TILE_32x32x64 in feasible: + return TILE_32x32x64 + return TILE_64x64x64 + return list(feasible)[-1] + + +# Launches to capture per graph. Enough that replay is dominated by the kernel +# rather than by graph launch, small enough to keep capture cheap. +_GRAPH_LAUNCHES = 20 +_REPLAYS_PER_ROUND = 8 + +_TILE_FIELDS = ("reg_m", "reg_n", "reg_k", "waves_m", "waves_n") + +# Launcher for a call signature whose tile the autotuner has already resolved. +# The tuner re-derives its cache key from scratch on every call — fingerprinting +# the environment, toolchain and device — which costs more host time than a +# small GEMM takes on the GPU, and under FLYDSL_AUTOTUNE=1 it re-runs the whole +# search per call. Consulting it once per signature keeps steady-state dispatch +# as cheap as calling the built module directly. +_resolved = {} + + +def _tile_config(tile): + return Config(**dict(zip(_TILE_FIELDS, tile))) + + +@functools.lru_cache(maxsize=None) +def _build(M, N, K, in_dtype, out_dtype, rounding, reg_m, reg_n, reg_k, waves_m, waves_n): + launch_fn, _, _, _ = create_wmma_gemm_module( + M, + N, + K, + in_dtype=in_dtype, + out_dtype=out_dtype, + rounding=rounding, + reg_m=reg_m, + reg_n=reg_n, + reg_k=reg_k, + waves_m=waves_m, + waves_n=waves_n, + **tile_opts((reg_m, reg_n, reg_k, waves_m, waves_n)), + ) + return launch_fn + + +def rdna3_gemm_dispatch( + C, + A, + B_T, + M, + N, + K, + in_dtype="bf16", + out_dtype="bf16", + rounding="rn", + reg_m=None, + reg_n=None, + reg_k=None, + waves_m=None, + waves_n=None, + stream=None, + sr_seed=0, +): + """Run the GEMM on one tile. Unset tile fields fall back to ``pick_tile``. + + The stream is resolved here rather than by the caller so that this stays + capturable: under ``torch.cuda.graph`` the current stream is the capture + stream, and enqueueing onto a stream captured before then aborts the capture. + """ + if stream is None: + stream = torch.cuda.current_stream() + # Resolve before the cache key so a partially specified tile and the fully + # spelled-out one it means share a single built module. + tile = tuple( + auto if given is None else given + for auto, given in zip(pick_tile(M, N, K), (reg_m, reg_n, reg_k, waves_m, waves_n)) + ) + launch_fn = _build(M, N, K, in_dtype, out_dtype, rounding, *tile) + # A search calls this once per candidate and then once more on the winner, + # so the last write is the config the tuner settled on. + _resolved[(M, N, K, in_dtype, out_dtype, rounding)] = launch_fn + return launch_fn(C, A, B_T, stream, sr_seed) + + +def _default_config( + C=None, + A=None, + B_T=None, + M=None, + N=None, + K=None, + in_dtype="bf16", + out_dtype="bf16", + rounding="rn", + **_kwargs, +): + return _tile_config(pick_tile(M, N, K)) + + +def _search_configs( + C=None, + A=None, + B_T=None, + M=None, + N=None, + K=None, + in_dtype="bf16", + out_dtype="bf16", + rounding="rn", + **_kwargs, +): + candidates = [_tile_config(tile) for tile, _wgs in feasible_tiles(M, N, K)] + return candidates or [_default_config(M=M, N=N, K=K)] + + +def _graph_bench(fn, warmup=5, rep=25): + """Fastest observed ms per launch, timed by replaying a captured graph. + + The stock ``do_bench`` pays one launch plus one full sync per measurement, + about 90us of host time on this kernel. That is longer than the kernel runs + on any shape small enough for the tile to be worth choosing, so all the + candidates measure alike and the search ends up ranking dispatch noise. + Capturing the launches amortises that overhead away and makes a sweep + reproducible to within a percent. + + It is still not trustworthy on the shortest kernels. Measured in isolation, + 512x512x2048 runs at 28.6us on 64x64x64 and 31.0us on 32x32x64; measured + here the multi-wave tiles read about 5us high and the ranking inverts, so a + forced sweep of that shape emits an artifact for the slower tile. Treat a + tuned result for a sub-50us kernel as a hypothesis to confirm, not a fact. + + Falls back to the stock timer if the kernel turns out not to be capturable. + """ + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(graph): + for _ in range(_GRAPH_LAUNCHES): + fn() + except Exception: + return do_bench(fn, warmup=warmup, rep=rep) + + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + + # Time a batch of replays per round so the graph launch is amortised too, + # and keep the fastest round. This runs on shared nodes where a neighbour + # can inflate a whole round two- or threefold; taking the median carries + # those rounds into the result, taking the minimum drops them. + rounds = max(3, rep // _REPLAYS_PER_ROUND) + best = float("inf") + for _ in range(rounds): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(_REPLAYS_PER_ROUND): + graph.replay() + end.record() + torch.cuda.synchronize() + best = min(best, start.elapsed_time(end) / (_REPLAYS_PER_ROUND * _GRAPH_LAUNCHES)) + return best + + +_gemm_tuner = autotune( + configs=_search_configs, + key=["M", "N", "K", "in_dtype", "out_dtype", "rounding"], + default=_default_config, + do_bench=_graph_bench, + artifact_name="rdna3_f16_gemm", +)(rdna3_gemm_dispatch) + + +def rdna3_gemm_autotuned( + C, + A, + B_T, + in_dtype="bf16", + out_dtype="bf16", + rounding="rn", + stream=None, + sr_seed=0, +): + """``C = A @ B_T.T`` on the tile chosen for this shape.""" + M, K = A.shape + N = B_T.shape[0] + M, N, K = int(M), int(N), int(K) + + launch_fn = _resolved.get((M, N, K, in_dtype, out_dtype, rounding)) + if launch_fn is not None: + return launch_fn(C, A, B_T, torch.cuda.current_stream() if stream is None else stream, sr_seed) + + # Make the caller's stream current instead of passing it down, so the + # dispatcher picks it up while a benchmark capture still gets its own. + launch_stream = torch.cuda.current_stream() if stream is None else stream + with torch.cuda.device(A.device), torch.cuda.stream(launch_stream): + return _gemm_tuner( + C, + A, + B_T, + M=M, + N=N, + K=K, + in_dtype=in_dtype, + out_dtype=out_dtype, + rounding=rounding, + stream=None, + sr_seed=sr_seed, + ) diff --git a/tests/kernels/benchmark_common.py b/tests/kernels/benchmark_common.py index 8770f3915..2c02230f2 100644 --- a/tests/kernels/benchmark_common.py +++ b/tests/kernels/benchmark_common.py @@ -239,22 +239,31 @@ def _bench_flydsl_torch(*, op: str, M: int, N: int, dtype: str, warmup: int, ite # arch. from flydsl.runtime.device import get_rocm_arch as _get_arch - if str(_get_arch() or "").startswith("gfx11"): - from kernels.gemm.rdna3_f16_gemm import create_wmma_gemm_module - else: - from kernels.gemm.rdna_f16_gemm import create_wmma_gemm_module - K = N # square by default; caller can override via config torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16 - launch, *_ = create_wmma_gemm_module(M, N, K, in_dtype=dtype, out_dtype="bf16") A = torch.randn(M, K, dtype=torch_dtype, device="cuda") B_T = torch.randn(N, K, dtype=torch_dtype, device="cuda") C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda") - return bench_gpu_us_torch( - lambda: launch(C, A, B_T, torch.cuda.current_stream()), - warmup=warmup, - iters=iters, - ) + + if str(_get_arch() or "").startswith("gfx11"): + # Through the autotune wrapper, so the number reflects the tile + # chosen for the shape rather than the kernel's 128x128x32 default. + # With nothing configured that resolution still benchmarks nothing, + # and warmup absorbs the one-off dispatch through the tuner. + from kernels.gemm.rdna3_f16_gemm_autotune import rdna3_gemm_autotuned + + def run(): + rdna3_gemm_autotuned(C, A, B_T, in_dtype=dtype, out_dtype="bf16") + + else: + from kernels.gemm.rdna_f16_gemm import create_wmma_gemm_module + + launch, *_ = create_wmma_gemm_module(M, N, K, in_dtype=dtype, out_dtype="bf16") + + def run(): + launch(C, A, B_T, torch.cuda.current_stream()) + + return bench_gpu_us_torch(run, warmup=warmup, iters=iters) if op == "wmma_fp8_gemm": from kernels.gemm.rdna_fp8_preshuffle_gemm import ( diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index 77aab2beb..c48c9ba40 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/tests/kernels/test_rdna_gemm.py @@ -21,6 +21,7 @@ from flydsl.runtime.device import get_rocm_arch # noqa: E402 from kernels.gemm.rdna3_f16_gemm import create_wmma_gemm_module as _create_wmma_gemm_module_gfx11 # noqa: E402 +from kernels.gemm.rdna3_f16_gemm_autotune import pick_tile # noqa: E402 from kernels.gemm.rdna_f16_gemm import create_wmma_gemm_module as _create_wmma_gemm_module_gfx12 # noqa: E402 from kernels.gemm.rdna_fp8_preshuffle_gemm import ( # noqa: E402 compile_fp8_gemm, @@ -209,6 +210,47 @@ def test_f16_gemm_grid_m_not_a_multiple_of_the_group_width(M, N, K): assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05) +@pytest.mark.parametrize( + "M, N, K", + [ + pytest.param(256, 256, 4096, id="256x256x4096"), + pytest.param(1024, 1024, 1024, id="1024x1024x1024"), + ], +) +def test_f16_gemm_autotuned_matches_the_heuristic_path(M, N, K): + """The autotune wrapper, left unconfigured, is a pass-through. + + It resolves through the tuner once and then calls the built module directly: + the tuner re-derives its cache key on every call, which costs more host time + than these shapes take on the GPU. So this checks both halves — the same + tile as ``pick_tile``, and that the second call does not build again. + """ + _requires_rdna3() + from kernels.gemm import rdna3_f16_gemm_autotune as gemm_autotune + + torch.manual_seed(42) + 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") + + signature = (M, N, K, "bf16", "bf16", "rn") + gemm_autotune._resolved.pop(signature, None) + + gemm_autotune.rdna3_gemm_autotuned(C, A, B_T) + torch.cuda.synchronize() + assert verify_output(C.float(), A.float() @ B_T.float().T, atol=0.05, rtol=0.05) + + resolved = gemm_autotune._resolved[signature] + expected_tile = pick_tile(M, N, K) + assert resolved is gemm_autotune._build(M, N, K, "bf16", "bf16", "rn", *expected_tile) + + C.zero_() + gemm_autotune.rdna3_gemm_autotuned(C, A, B_T) + torch.cuda.synchronize() + assert verify_output(C.float(), A.float() @ B_T.float().T, atol=0.05, rtol=0.05) + assert gemm_autotune._resolved[signature] is resolved + + @pytest.mark.parametrize( "M, N, K", [ diff --git a/tests/unit/test_rdna3_gemm_autotune.py b/tests/unit/test_rdna3_gemm_autotune.py new file mode 100644 index 000000000..9d1ef3b5d --- /dev/null +++ b/tests/unit/test_rdna3_gemm_autotune.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors +"""GPU-free tests for the RDNA3 GEMM autotune wrapper. + +The wrapper's job is to expose tile selection through the shared autotuner +without a call having to opt into benchmarking, so the properties worth pinning +are about agreement rather than performance: + + * with nothing configured, the tuner's ``default`` is exactly ``pick_tile`` — + a call that never benchmarks must land on the heuristic's tile; + * the search space only contains tiles the shape can actually build, and it + always contains the default, so a search can confirm the heuristic and can + never be forced to pick something that fails at build time; + * ``feasible_tiles`` reports the same tiles, in the same order, that + ``pick_tile`` chooses from. + +Device-side coverage is ``tests/kernels/test_rdna_gemm.py``. +""" + +import pytest + +from kernels.gemm.rdna3_f16_gemm_autotune import ( + _TILE_FIELDS, + _default_config, + _ladder_for, + _search_configs, + _tile_workgroups, + feasible_tiles, + pick_tile, +) + +pytestmark = pytest.mark.l0_backend_agnostic + +SHAPES = [ + (M, N, K) + for M in (256, 512, 768, 1024, 1536, 2048, 4096) + for N in (256, 512, 1024, 2048) + for K in (512, 1024, 2048, 4096) +] + +DEEPSEEK_V4_FLASH_DECODE_SHAPES = [ + (M, N, K) + for M in (1, 5, 16) + for N, K in ( + (4096, 1024), + (1024, 4096), + (4096, 512), + (4096, 2048), + (2048, 4096), + ) +] + +_shape_id = lambda shape: "x".join(map(str, shape)) # noqa: E731 + + +def _tile_of(config): + return tuple(config.kwargs[field] for field in _TILE_FIELDS) + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_default_config_is_the_heuristic_tile(shape): + """An untuned call must be indistinguishable from calling the kernel directly.""" + M, N, K = shape + assert _tile_of(_default_config(M=M, N=N, K=K)) == pick_tile(*shape) + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_search_space_only_holds_buildable_tiles(shape): + """Every candidate must survive ``_tile_workgroups``. + + A candidate the shape cannot use does not merely lose the benchmark: it + raises out of ``create_wmma_gemm_module``, so the search would spend its + time measuring build failures. + """ + M, N, K = shape + if not feasible_tiles(*shape): + pytest.skip("no tile in the ladder fits this shape") + for config in _search_configs(M=M, N=N, K=K): + assert _tile_workgroups(*shape, _tile_of(config)) is not None + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_default_is_reachable_by_the_search(shape): + """The search must be able to return the heuristic's own answer.""" + M, N, K = shape + if not feasible_tiles(*shape): + pytest.skip("no tile in the ladder fits this shape") + tiles = [_tile_of(config) for config in _search_configs(M=M, N=N, K=K)] + assert pick_tile(*shape) in tiles + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_feasible_tiles_is_the_ladder_filtered_in_order(shape): + """``pick_tile`` walks this list, so it has to stay a subsequence of the ladder.""" + ladder = _ladder_for(shape[2]) + reported = feasible_tiles(*shape) + + assert [tile for tile, _ in reported] == [cfg for cfg in ladder if _tile_workgroups(*shape, cfg) is not None] + for tile, workgroups in reported: + assert workgroups == _tile_workgroups(*shape, tile) + + +@pytest.mark.parametrize("shape", DEEPSEEK_V4_FLASH_DECODE_SHAPES, ids=_shape_id) +def test_deepseek_v4_flash_decode_shapes_remain_out_of_scope(shape): + """Formal decode rows are small-M tail cases, not tune candidates yet. + + This pins the current scope boundary: these inference shapes have no legal + tile in the existing ladder and would reject at the kernel's ``M % BLOCK_M`` + assertion. Small-M/tail handling is a separate follow-on. + """ + assert feasible_tiles(*shape) == [] + assert all(_tile_workgroups(*shape, tile) is None for tile in _ladder_for(shape[2])) diff --git a/tests/unit/test_rdna3_tile_selection.py b/tests/unit/test_rdna3_tile_selection.py new file mode 100644 index 000000000..cfca1acac --- /dev/null +++ b/tests/unit/test_rdna3_tile_selection.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors +"""GPU-free tests for the RDNA3 GEMM host-side tile selection. + +``pick_tile`` decides which tile a shape gets. Beyond the measured expectations, +two properties make it safe to enable by default: it never returns a tile the +shape cannot build, and it never moves a shape off 128x128x32 while that tile +still has a deep enough grid to win. + +All plain integer arithmetic that runs before a kernel is built, so no GPU is +needed. The swizzle the chosen tile feeds into is covered by +``test_rdna3_grid_swizzle.py``, and the device-side counterpart is +``tests/kernels/test_rdna_gemm.py``. +""" + +import pytest + +from kernels.gemm.rdna3_f16_gemm import _group_width, _swizzle_tile_id +from kernels.gemm.rdna3_f16_gemm_autotune import ( + NUM_CU, + TILE_32x32x64, + TILE_32x64x64, + TILE_64x64x64, + TILE_128x64x32, + TILE_128x128x32, + _ladder_for, + _tile_workgroups, + feasible_tiles, + pick_tile, +) + +pytestmark = pytest.mark.l0_backend_agnostic + +DEFAULT_TILE = TILE_128x128x32 # the tile used before selection existed +DEFAULT_GROUP_M = 8 # create_wmma_gemm_module's default grouping cap + +# Every shape the tile heuristic was measured on, with the tile it is expected to +# pick, so that a heuristic edit has to state which shapes it moves. Each of +# these was timed against the whole ladder; the expectation is the fastest tile +# except where noted. +MEASURED_SHAPES = [ + pytest.param((256, 256, 4096), TILE_32x32x64, id="256x256x4096-32x32x64"), + pytest.param((256, 256, 1024), TILE_64x64x64, id="256x256x1024-64x64x64"), + pytest.param((384, 384, 2048), TILE_64x64x64, id="384x384x2048-64x64x64"), + pytest.param((512, 512, 1024), TILE_64x64x64, id="512x512x1024-64x64x64"), + pytest.param((512, 512, 4096), TILE_64x64x64, id="512x512x4096-64x64x64"), + pytest.param((256, 1024, 4096), TILE_64x64x64, id="256x1024x4096-64x64x64"), + pytest.param((768, 768, 2048), TILE_64x64x64, id="768x768x2048-64x64x64"), + pytest.param((1024, 1024, 512), TILE_128x64x32, id="1024x1024x512-128x64x32"), + pytest.param((1024, 1024, 1024), TILE_128x64x32, id="1024x1024x1024-128x64x32"), + pytest.param((1024, 1024, 4096), TILE_64x64x64, id="1024x1024x4096-64x64x64"), + # 128x128x32 is 8.1% faster here, the worst case the heuristic accepts. + pytest.param((1152, 1152, 1024), TILE_64x64x64, id="1152x1152x1024-64x64x64"), + pytest.param((1536, 1536, 1024), TILE_64x64x64, id="1536x1536x1024-64x64x64"), + pytest.param((1792, 1792, 1024), TILE_64x64x64, id="1792x1792x1024-64x64x64"), + pytest.param((2048, 2048, 512), DEFAULT_TILE, id="2048x2048x512-128x128x32"), + pytest.param((2048, 2048, 2048), DEFAULT_TILE, id="2048x2048x2048-128x128x32"), + pytest.param((3072, 3072, 1024), DEFAULT_TILE, id="3072x3072x1024-128x128x32"), + pytest.param((4096, 4096, 4096), DEFAULT_TILE, id="4096x4096x4096-128x128x32"), +] + +MEASURED_ONLY_SHAPES = [p.values[0] for p in MEASURED_SHAPES] + +# A spread for the property tests: square and skewed, K on both sides of the +# ladder split, sizes from "cannot fill the machine" up to "fills it easily". +PROPERTY_SHAPES = [ + (M, N, K) + for M in (256, 512, 768, 1024, 1536, 2048, 4096) + for N in (256, 512, 1024, 2048) + for K in (512, 1024, 2048, 4096) +] + +_shape_id = lambda shape: "x".join(map(str, shape)) # noqa: E731 + + +def _block_shape(cfg): + reg_m, reg_n, reg_k, waves_m, waves_n = cfg + return 16 * reg_m * waves_m, 16 * reg_n * waves_n, 16 * reg_k + + +def _tile_grid(shape): + """(grid_m, grid_n) the kernel would launch for this shape.""" + M, N, _ = shape + block_m, block_n, _ = _block_shape(pick_tile(*shape)) + return M // block_m, N // block_n + + +@pytest.mark.parametrize("shape, expected", MEASURED_SHAPES) +def test_pick_tile_matches_measured_shapes(shape, expected): + """The tile picked for each benchmarked shape is the one that was measured.""" + assert pick_tile(*shape) == expected + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_picked_tile_is_usable_by_the_shape(shape): + """Whenever the ladder has a usable tile, the picked one is usable. + + ``create_wmma_gemm_module`` asserts the divisibility and vectorization rules + that ``_tile_workgroups`` screens for, so returning an unusable tile turns + into an assertion failure at build time. + """ + if all(_tile_workgroups(*shape, cfg) is None for cfg in _ladder_for(shape[2])): + pytest.skip("no tile in the ladder fits this shape") + assert _tile_workgroups(*shape, pick_tile(*shape)) is not None + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_deep_grids_keep_the_original_tile(shape): + """A shape whose 128x128x32 grid is deep enough must not be moved off it. + + This is what makes the heuristic safe to turn on by default: the large + shapes take the same path as before selection existed, so they cannot + regress. Merely covering the machine is not enough — 128x128x32 lost by 19% + at 1792x1792x1024 and 37% at 1664x1664x1024, both well past one workgroup + per CU — so the bar is the measured one. + """ + wgs = _tile_workgroups(*shape, DEFAULT_TILE) + if wgs is None or wgs < 2.5 * NUM_CU: + pytest.skip("128x128x32's grid is not deep enough for this shape") + assert pick_tile(*shape) == DEFAULT_TILE + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_the_narrowest_tile_needs_a_starved_grid(shape): + """32x32x64 trades compute intensity for parallelism, and loses when it does not need to. + + It gave up 12-23% against 64x64x64 wherever 64x64x64 had the workgroups to + keep the machine busy, so it must be unreachable unless 64x64x64 either does + not fit or cannot fill the machine. + """ + if pick_tile(*shape) != TILE_32x32x64: + pytest.skip("this shape does not take the narrowest tile") + feasible = dict(feasible_tiles(*shape)) + assert TILE_64x64x64 not in feasible or feasible[TILE_64x64x64] < NUM_CU / 4 + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_32x64x64_is_only_ever_a_fallback(shape): + """32x64x64 was not the fastest tile on any of the 27 shapes swept. + + The old heuristic chose it for five of them, each time losing to 64x64x64. + It stays on the ladder only to cover shapes that nothing else divides. + """ + if pick_tile(*shape) != TILE_32x64x64: + pytest.skip("this shape does not take 32x64x64") + assert dict(feasible_tiles(*shape)).keys() == {TILE_32x64x64} + + +@pytest.mark.parametrize("shape", MEASURED_ONLY_SHAPES, ids=_shape_id) +def test_picked_tile_and_default_grouping_stay_inside_the_grid(shape): + """Selection and swizzle have to agree, not just work in isolation. + + The fault needs both halves to line up: the tile decides grid_m, and only then + does the grouping width become a divisor question. + """ + grid_m, grid_n = _tile_grid(shape) + width = _group_width(grid_m, DEFAULT_GROUP_M) + mapped = [_swizzle_tile_id(pid, grid_n, width) for pid in range(grid_m * grid_n)] + + assert set(mapped) == {(m, n) for m in range(grid_m) for n in range(grid_n)}