diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py index 4cf2188bb323..93d906a91d93 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py @@ -98,7 +98,10 @@ DIRECT_HALO_VARIANTS as _HALO_VARIANTS, DIRECT_MFMA_DGRAD as _MFMA_DGRAD, DIRECT_SPLIT_REDUCE as _SPLIT_REDUCE, + DIRECT_WGRAD_FOLD as _WGRAD_FOLD, + DIRECT_WGRAD_FOLDED as _WGRAD_FOLDED, direct_dgrad_weight_transforms, + direct_wgrad_fold_knobs, direct_dgrad_workspaces, direct_spec_for_problem, make_direct_spec as make_spec, @@ -416,6 +419,8 @@ def _conv_step(identity, hsaco_path, meta, spec, problem, direction, buffers): if isinstance(spec, DirectNongroupedConvSpec): grid, block = spec.grid(), (spec.threads_per_block, 1, 1) + elif hasattr(spec, "launch_geometry"): + grid, block = spec.launch_geometry() else: grid, block = direct_launch_geometry(spec) return DirectStep( @@ -518,6 +523,45 @@ def plan_for( identity, hsaco_path, meta, spec, problem, "fwd", ("x", "w", "y") ) return DirectPlan(identity, (step,)), "ok" + if identity.direction == "direct_wgrad" and variant in _WGRAD_FOLDED: + # The pass stores one partial slot per partition into the scratch + # ``ws``; the fold sums them into dW (written whole, no zeroing). + from kernels.common.conv_direct_dw_wgrad import ( + DW_WGRAD_FOLD_ARGS, + dw_wgrad_fold_signature, + ) + + caps = _caps_of(identity) + fold_knobs = direct_wgrad_fold_knobs(variant, caps.KH, caps.KW, caps.cpg) + fold_id = _helper_identity( + identity.arch, + identity.wave_size, + _WGRAD_FOLD, + caps, + fold_knobs, + identity.dtype_a, + ) + entry = cache.get(fold_id) + if entry is None: + return None, "missing cached wgrad fold (rebuild the cache)" + fold_spec = make_spec(_WGRAD_FOLD, problem, fold_knobs) + channels = ( + spec.fold_channels() if hasattr(spec, "fold_channels") else problem.groups + ) + grid, block = fold_spec.launch_geometry(channels) + main = _conv_step( + identity, hsaco_path, meta, spec, problem, "wgrad", ("dy", "x", "ws") + ) + fold = DirectStep( + kernel_name=entry[1]["kernel_name"], + hsaco_path=cache.hsaco_path(fold_id), + signature=dw_wgrad_fold_signature(fold_spec), + grid=grid, + block=block, + buffers=("ws", None, "dw"), + scalars=tuple(zip(DW_WGRAD_FOLD_ARGS, (spec.n_partitions(), channels))), + ) + return DirectPlan(identity, (main, fold), {"ws": spec.workspace_bytes()}), "ok" if identity.direction == "direct_wgrad": # dY rides A, X rides B, and D is the fp32 dW the kernel accumulates # into atomically (the caller zeroes it). diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_dw_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_dw_wgrad.py new file mode 100644 index 000000000000..789d4584b8cc --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_dw_wgrad.py @@ -0,0 +1,768 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Depthwise wgrad kernels on the VALU, and the partial-slot fold. + +A depthwise wgrad does ``2 * KH * KW`` FLOPs per dY element, so it is bound by +the bytes of X and dY: the kernels here read both once and keep the MFMA out +of it. Every block reduces over the spatial work it owns and stores one +partial slot ``ws[part][tap][c]`` with a plain store; :class:`DwWgradFoldSpec` +sums the partitions into ``dW``. (fp32 atomics on a depthwise ``dW`` -- a few +thousand addresses -- serialize: same-address atomics cost ~100 ns each.) + +* :class:`DwWgradLdsSpec` -- ``C % 8 == 0``: a block owns 8 channels per lane + and a column tile, streams output rows, and DMAs the S new X rows and one + dY row per step into LDS rings. +* :class:`DwWgradColsSpec` -- small ``C`` (stem layers, ``C = 3``): lanes own a + filter row and a strip of output columns instead of channels. + +dW is ``[C, KH, KW, 1]``; the fold writes it whole, so it needs no zeroing. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List + +from rocke.core.ir import F32, IRBuilder, KernelDef, Value +from rocke.helpers.spec import ceil_div + +from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + _io_type, + dilation_reason, + emit_direct_params, +) + +_WAVE = 64 + + +# --------------------------------------------------------------------------- +# Fold: dW[c, tap] = sum_part ws[part][tap][c] +# --------------------------------------------------------------------------- + +_FOLD_SLICES = 16 +_FOLD_COLS = 64 +# Partition loads each slice keeps in flight. +_FOLD_UNROLL = 4 + + +@dataclass(frozen=True) +class DwWgradFoldSpec: + """Fold of ``taps`` partial gradients per output channel into an fp32 dW.""" + + taps: int + name: str = "dw_wgrad_fold" + + def kernel_name(self) -> str: + from rocke.helpers.spec import kernel_name_join + + return kernel_name_join(self.name, f"t{self.taps}") + + def launch_geometry(self, channels: int): + return (ceil_div(channels * self.taps, _FOLD_COLS), 1, 1), ( + _FOLD_SLICES * _WAVE, + 1, + 1, + ) + + +# Kernargs after the A/D pointers and byte sizes, in order. +DW_WGRAD_FOLD_ARGS = ("parts", "channels") + + +def dw_wgrad_fold_signature(spec: DwWgradFoldSpec) -> list: + """A = the fp32 partial slots, D = dW, as the AOT plans bind helper steps.""" + from rocke.helpers.spec import SignatureBuilder + + sb = ( + SignatureBuilder() + .ptr("A", "fp32") + .ptr("D", "fp32") + .scalar("A_bytes", "i32") + .scalar("D_bytes", "i32") + ) + for name in DW_WGRAD_FOLD_ARGS: + sb = sb.scalar(name, "i32") + return sb.build() + + +def build_dw_wgrad_fold(spec: DwWgradFoldSpec, arch: str = "gfx950") -> KernelDef: + """Each block owns 64 outputs; 16 waves split the partition axis, meet in LDS.""" + from rocke.core.ir import I32, PtrType + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = _FOLD_SLICES * _WAVE + ws = b.param("A", PtrType(F32, "global"), noalias=True, readonly=True, align=16) + dw = b.param( + "D", + PtrType(F32, "global"), + noalias=True, + writeonly=True, + align=4, + ) + b.param("A_bytes", I32) + b.param("D_bytes", I32) + parts = b.param("parts", I32) + chans = b.param("channels", I32) + + def ci(v): + return b.const_i32(v) + + tid = b.thread_id_x() + lane = b.mod(tid, ci(_WAVE)) + slc = b.div(tid, ci(_WAVE)) + n_out = b.mul(chans, ci(spec.taps)) + o = b.add(b.mul(b.block_id_x(), ci(_FOLD_COLS)), lane) + o_ok = b.cmp_lt(o, n_out) + o_safe = b.select(o_ok, o, ci(0)) + zero = b.const_f32(0.0) + U = _FOLD_UNROLL + # Partition p = slc + SLICES * (U * it + j). + n_it = b.div(b.add(parts, ci(_FOLD_SLICES * U - 1)), ci(_FOLD_SLICES * U)) + loop = b.scf_for_iter( + ci(0), + n_it, + ci(1), + [(f"s{j}", zero) for j in range(U)], + iv_name="it", + elide_trailing_barrier=False, + ) + with loop as (it, acc): + acc = list(acc) + for j in range(U): + pidx = b.add(slc, b.mul(b.add(b.mul(it, ci(U)), ci(j)), ci(_FOLD_SLICES))) + ok = b.cmp_lt(pidx, parts) + v = b.global_load_f32( + ws, b.add(b.mul(b.select(ok, pidx, ci(0)), n_out), o_safe) + ) + acc[j] = b.fadd(acc[j], b.select(ok, v, zero)) + b.scf_yield(*acc) + tot = loop.results[0] + for v in loop.results[1:]: + tot = b.fadd(tot, v) + red = b.smem_alloc(F32, [1, _FOLD_SLICES * _WAVE], name_hint="fold") + b.smem_store_vN_f32(red, [ci(0), tid], tot, 1) + b.sync() + with b.scf_if(b.land(b.cmp_eq(slc, ci(0)), o_ok)): + t = None + for q in range(_FOLD_SLICES): + v = b.vec_extract( + b.smem_load_vN_f32(red, ci(0), b.add(lane, ci(q * _WAVE)), n=1), 0 + ) + t = v if t is None else b.fadd(t, v) + tap = b.div(o, chans) + c = b.mod(o, chans) + off = b.add(b.mul(c, ci(spec.taps)), tap) + b.global_store(dw, off, t, align=4) + return b.kernel + + +# --------------------------------------------------------------------------- +# Small-C depthwise wgrad: lanes along output columns +# --------------------------------------------------------------------------- +# +# At C = 3 a channel-vectorised lane would idle most of its width, so here a +# lane owns one filter row r and a strip of `strip` consecutive output +# columns, keeps KW x C accumulators, and per output row loads the dY strip +# and the X row segment it touches as whole 16-byte vectors (the pixels are +# C * 2 bytes, so the vectors straddle pixels; edge columns are masked with +# per-lane masks that do not change along the rows). Block x is the filter +# row; block y is a partition that walks a contiguous run of (image, output +# row, column group) units and stores one partial slot ws[part][r][s][c], +# which DwWgradFoldSpec folds. + + +@dataclass(frozen=True) +class DwWgradColsSpec: + problem: DirectConvProblem + name: str = "dw_wgrad_cols" + strip: int = 8 + threads: int = 64 + target_blocks: int = 2048 + + def kernel_name(self) -> str: + from rocke.helpers.spec import kernel_name_join + + p = self.problem + flags = {"bf16": True} if p.dtype == "bf16" else {} + return kernel_name_join( + self.name, + f"f{p.KH}x{p.KW}p{p.PAD}s{p.stride}c{p.groups}", + f"st{self.strip}", + f"t{self.threads}", + f"tb{self.target_blocks}", + flags=flags, + ) + + def n_units(self) -> int: + p = self.problem + return ceil_div(p.N * p.Ho * (p.Wo // self.strip), self.threads) + + def n_partitions(self) -> int: + return min(self.n_units(), ceil_div(self.target_blocks, self.problem.KH)) + + def workspace_bytes(self) -> int: + p = self.problem + return self.n_partitions() * p.KH * p.KW * p.groups * 4 + + def launch_geometry(self): + return (self.problem.KH, self.n_partitions(), 1), (self.threads, 1, 1) + + def lds_bytes(self) -> int: + """The block reduction buffer: KW x C fp32 accumulators per thread.""" + p = self.problem + return self.threads * p.KW * p.groups * 4 + + +def is_valid_dw_wgrad_cols_spec(spec: DwWgradColsSpec, arch: str = "gfx950"): + p = spec.problem + why = dilation_reason(p) + if why is not None: + return False, why + if p.cpg != 1 or p.kpg != 1: + return False, "depthwise only" + if p.dtype not in ("fp16", "bf16"): + return False, "fp16/bf16 only" + if p.groups > 8: + return False, "small-C kernel: C <= 8" + if p.KW * p.groups > 96: + return False, "too many accumulators" + if spec.threads < _WAVE or spec.threads % _WAVE or spec.threads > 1024: + return False, "threads must be a multiple of 64" + C = p.groups + # Whole 16-byte vectors at a build-time phase: the strip and the rows + # must keep every vector start on an 8-element boundary. + if (spec.strip * C) % 8 or (spec.strip * p.stride * C) % 8: + return False, "strip * C (and strip * stride * C) must be multiples of 8" + if (p.W * C) % 8 or (p.Wo * C) % 8 or spec.strip < 1 or p.Wo % spec.strip: + return False, "W*C and Wo*C must be multiples of 8 and Wo of the strip" + if spec.target_blocks < 1: + return False, "target_blocks >= 1" + if spec.lds_bytes() > 160 * 1024: + return False, f"LDS {spec.lds_bytes()} B > 160 KiB" + return True, "ok" + + +def build_dw_wgrad_cols(spec: DwWgradColsSpec, arch: str = "gfx950") -> KernelDef: + ok, why = is_valid_dw_wgrad_cols_spec(spec, arch) + if not ok: + raise ValueError(f"invalid DwWgradColsSpec: {why}") + p = spec.problem + KH, KW, PAD, S, C = p.KH, p.KW, p.PAD, p.stride, p.groups + WS = spec.strip + T = spec.threads + io_type = _io_type(p.dtype) + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = T + params = emit_direct_params(b, direction="wgrad", io_type=io_type) + + def ci(v: int) -> Value: + return b.const_i32(v) + + c0, c1, c2 = ci(0), ci(1), ci(2) + oob = ci((1 << 31) - 1) + p_H, p_W = params["p_Hi"], params["p_Wi"] + p_Ho, p_Wo = params["p_Ho"], params["p_Wo"] + tid = b.thread_id_x() + r_tap = b.block_id_x() # filter row + part = b.block_id_y() + strips_per_row = b.div(p_Wo, ci(WS)) + n_strips = b.mul(b.mul(params["p_N"], p_Ho), strips_per_row) + n_units = b.div(b.add(n_strips, ci(T - 1)), ci(T)) + n_part = b.smin(n_units, ci(ceil_div(spec.target_blocks, KH))) + u_lo = b.div(b.mul(part, n_units), n_part) + u_hi = b.div(b.mul(b.add(part, c1), n_units), n_part) + + a_rsrc = b.buffer_rsrc(params["A"], params["A_bytes"]) # dY + b_rsrc = b.buffer_rsrc(params["B"], params["B_bytes"]) # X + X_COLS = (WS - 1) * S + KW + # Window element e (column e // C of the window) sits at vector element + # e + PH: the window starts PAD columns left of an 8-aligned strip. + PH = (-PAD * C) % 8 + n_dy_vec = WS * C // 8 + n_x_vec = ceil_div(PH + X_COLS * C, 8) + + def load8(rsrc, elem_off: Value, ok: Value) -> Value: + voff = b.select(ok, b.mul(elem_off, c2), oob) + if p.dtype == "bf16": + return b.buffer_load_vN_bf16(rsrc, voff, c0, 4) + return b.buffer_load_vN_f16(rsrc, voff, c0, 4) + + def elem(vecs, e): + return b.cast_to_f32(b.vec_extract(vecs[e // 8], e % 8)) + + zero = b.const_f32(0.0) + n_acc = KW * C + loop = b.scf_for_iter( + u_lo, + u_hi, + c1, + [(f"acc{i}", zero) for i in range(n_acc)], + iv_name="unit", + elide_trailing_barrier=False, + ) + with loop as (unit, carried): + accs = list(carried) + strip = b.add(b.mul(unit, ci(T)), tid) + live = b.cmp_lt(strip, n_strips) + cs = b.mod(strip, strips_per_row) + rest = b.div(strip, strips_per_row) + ho = b.mod(rest, p_Ho) + n_img = b.div(rest, p_Ho) + wo0 = b.mul(cs, ci(WS)) + hi = b.add(b.mul(ho, ci(S)), b.sub(r_tap, ci(PAD))) + hi_ok = b.land(live, b.land(b.cmp_ge(hi, c0), b.cmp_lt(hi, p_H))) + wi0 = b.sub(b.mul(wo0, ci(S)), ci(PAD)) + dy_base = b.add( + b.add( + b.mul(n_img, params["p_dY_stride_n"]), + b.mul(ho, params["p_dY_stride_ho"]), + ), + b.mul(wo0, params["p_dY_stride_wo"]), + ) + x_base = b.add( + b.add( + b.mul(n_img, params["p_X_stride_n"]), b.mul(hi, params["p_X_stride_hi"]) + ), + b.mul(wi0, params["p_X_stride_wi"]), + ) + dy_vecs = [ + load8(a_rsrc, b.add(dy_base, ci(8 * i)), live) for i in range(n_dy_vec) + ] + x_vecs = [ + load8(b_rsrc, b.add(x_base, ci(8 * i - PH)), hi_ok) for i in range(n_x_vec) + ] + dy = [[elem(dy_vecs, j * C + c) for c in range(C)] for j in range(WS)] + xs = [] + for col in range(X_COLS): + wi = b.add(wi0, ci(col)) + okc = b.land(b.cmp_ge(wi, c0), b.cmp_lt(wi, p_W)) + xs.append( + [b.select(okc, elem(x_vecs, PH + col * C + c), zero) for c in range(C)] + ) + for j in range(WS): + for s in range(KW): + for c in range(C): + i = s * C + c + accs[i] = b.fma(dy[j][c], xs[j * S + s][c], accs[i]) + b.scf_yield(*accs) + accs = list(loop.results) + + # Block reduction: every lane holds this block's r row; sum over lanes. + red = b.smem_alloc(F32, [1, T * n_acc], name_hint="red") + for i in range(n_acc): + b.smem_store_vN_f32(red, [c0, b.add(b.mul(ci(i), ci(T)), tid)], accs[i], 1) + b.sync() + # Thread tid sums outputs tid, tid + T, ... (KW * C may exceed T). + for k in range(ceil_div(n_acc, T)): + o = tid if k == 0 else b.add(tid, ci(k * T)) + with b.scf_if(b.cmp_lt(o, ci(n_acc))): + tot = None + for q in range(T): + v = b.vec_extract( + b.smem_load_vN_f32(red, c0, b.add(b.mul(o, ci(T)), ci(q)), n=1), 0 + ) + tot = v if tot is None else b.fadd(tot, v) + # ws[part][tap][c], tap = r*KW + s; output o is (s, c) = divmod(o, C). + s_i = b.div(o, ci(C)) + c_i = b.mod(o, ci(C)) + tap = b.add(b.mul(r_tap, ci(KW)), s_i) + off = b.add(b.mul(b.add(b.mul(part, ci(KH * KW)), tap), ci(C)), c_i) + b.global_store(params["D"], off, tot, align=4) + return b.kernel + + +# --------------------------------------------------------------------------- +# Depthwise wgrad, LDS row rings +# --------------------------------------------------------------------------- +# +# DwWgradSpec re-reads its KH input rows per output row from L2, and with a +# few MB of rows in flight per XCD the L2 thrashes (HBM traffic ~1.75x the +# tensors). Here the block owns a channel tile and a column tile and streams +# output rows: per step it DMAs S new X rows and one dY row into LDS rings +# (prefetch steps ahead), and every lane reads its window out of LDS. Each +# byte of X and dY leaves HBM once per (channel tile, column tile). +# +# Ring slots fold to constants by unrolling U = prefetch + ceil(KH / S) steps: +# the X ring has S*U row slots and the dY ring U. + + +@dataclass(frozen=True) +class DwWgradLdsSpec: + problem: DirectConvProblem + name: str = "dw_wgrad_lds" + lanes_c: int = 8 # lanes along channels (8 channels each) + lanes_w: int = 32 # lanes along columns + cols_per_lane: int = 1 + prefetch: int = 2 + rows_per_seg: int = 32 # output rows per segment (rounded up to the unroll) + target_blocks: int = 1024 + + @property + def block_c(self) -> int: + return 8 * self.lanes_c + + @property + def block_wo(self) -> int: + return self.lanes_w * self.cols_per_lane + + @property + def threads(self) -> int: + return self.lanes_c * self.lanes_w + + @property + def unroll(self) -> int: + return self.prefetch + ceil_div(self.problem.KH, self.problem.stride) + + def x_cols(self) -> int: + p = self.problem + return (self.block_wo - 1) * p.stride + p.KW + + def rps(self) -> int: + return ceil_div(self.rows_per_seg, self.unroll) * self.unroll + + def kernel_name(self) -> str: + from rocke.helpers.spec import kernel_name_join + + p = self.problem + flags = {"bf16": True} if p.dtype == "bf16" else {} + return kernel_name_join( + self.name, + f"f{p.KH}x{p.KW}p{p.PAD}s{p.stride}", + f"lc{self.lanes_c}", + f"lw{self.lanes_w}", + f"cpl{self.cols_per_lane}", + f"pf{self.prefetch}", + f"rps{self.rows_per_seg}", + f"tb{self.target_blocks}", + flags=flags, + ) + + def n_channel_tiles(self) -> int: + return ceil_div(self.problem.groups, self.block_c) + + def n_units(self) -> int: + p = self.problem + return p.N * ceil_div(p.Wo, self.block_wo) * ceil_div(p.Ho, self.rps()) + + def n_partitions(self) -> int: + return min(self.n_units(), ceil_div(self.target_blocks, self.n_channel_tiles())) + + def workspace_bytes(self) -> int: + p = self.problem + return self.n_partitions() * p.KH * p.KW * p.groups * 4 + + def launch_geometry(self): + return (self.n_channel_tiles(), self.n_partitions(), 1), (self.threads, 1, 1) + + def dma_instrs(self, cols: int) -> int: + """16-byte DMA instructions per wave for one LDS row of ``cols`` pixels.""" + n_waves = self.threads // _WAVE + return ceil_div(ceil_div(cols * (self.block_c // 8), _WAVE), n_waves) + + def x_row_elems(self) -> int: + """Elements of one X ring slot: whole DMA instructions of every wave.""" + return self.threads * self.dma_instrs(self.x_cols()) * 8 + + def dy_row_elems(self) -> int: + return self.threads * self.dma_instrs(self.block_wo) * 8 + + def lds_bytes(self) -> int: + """The X and dY rings as the builder allocates them, plus the reduction.""" + U = self.unroll + x = self.problem.stride * U * self.x_row_elems() * 2 + dy = U * self.dy_row_elems() * 2 + red = self.lanes_w * self.block_c * 4 + return x + dy + red + + +def is_valid_dw_wgrad_lds_spec(spec: DwWgradLdsSpec, arch: str = "gfx950"): + from rocke.core.arch import ArchTarget + + try: + target = ArchTarget.from_gfx(arch) + except KeyError as e: + return False, str(e) + p = spec.problem + why = dilation_reason(p) + if why is not None: + return False, why + if target.wave_size != 64 or target.async_lds_max_dwords < 4: + return False, f"needs wave64 + 16-byte buffer_load_lds (absent on {arch})" + if p.cpg != 1 or p.kpg != 1: + return False, "depthwise only" + if p.dtype not in ("fp16", "bf16"): + return False, "fp16/bf16 only" + if p.groups % 8: + return False, "C must be a multiple of 8" + if spec.threads < _WAVE or spec.threads % _WAVE or spec.threads > 1024: + return False, "threads must be a multiple of 64, at most 1024" + if p.KH * p.KW * 8 > 128: + return False, "too many accumulators" + if spec.prefetch < 1: + return False, "prefetch >= 1" + if spec.lds_bytes() > 160 * 1024: + return False, f"LDS {spec.lds_bytes()} B > 160 KiB" + if spec.rows_per_seg < 1: + return False, "rows_per_seg >= 1" + if spec.target_blocks < 1: + return False, "target_blocks >= 1" + return True, "ok" + + +def build_dw_wgrad_lds(spec: DwWgradLdsSpec, arch: str = "gfx950") -> KernelDef: + from rocke.core.ir import I64 + + ok, why = is_valid_dw_wgrad_lds_spec(spec, arch) + if not ok: + raise ValueError(f"invalid DwWgradLdsSpec: {why}") + p = spec.problem + KH, KW, PAD, S = p.KH, p.KW, p.PAD, p.stride + TAPS = KH * KW + LC, LW, WL = spec.lanes_c, spec.lanes_w, spec.cols_per_lane + BC, BWO = spec.block_c, spec.block_wo + CH = BC // 8 # 16-byte chunks per pixel + XC = spec.x_cols() + T = spec.threads + D = spec.prefetch + U = spec.unroll + NX, NY = S * U, U + RPS = spec.rps() + PRE = max(KH - S, 0) # X rows a segment needs before its first step's own + io_type = _io_type(p.dtype) + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = T + b.kernel.attrs["lds_alias_scopes"] = True + params = emit_direct_params(b, direction="wgrad", io_type=io_type) + + def ci(v: int) -> Value: + return b.const_i32(v) + + c0, c1, c2 = ci(0), ci(1), ci(2) + oob = ci((1 << 31) - 1) + p_C = params["p_groups"] + p_H, p_W = params["p_Hi"], params["p_Wi"] + p_Ho, p_Wo = params["p_Ho"], params["p_Wo"] + tid = b.thread_id_x() + lane = b.mod(tid, ci(_WAVE)) + wave = b.readfirstlane(b.div(tid, ci(_WAVE))) + lc = b.mod(tid, ci(LC)) + lw = b.div(tid, ci(LC)) + ctile, part = b.block_id_x(), b.block_id_y() + c_blk0 = b.mul(ctile, ci(BC)) + n_ct = b.div(b.add(p_Wo, ci(BWO - 1)), ci(BWO)) + n_hb = b.div(b.add(p_Ho, ci(RPS - 1)), ci(RPS)) + n_units = b.mul(b.mul(params["p_N"], n_ct), n_hb) + n_tiles = b.div(b.add(p_C, ci(BC - 1)), ci(BC)) + n_part = b.smin( + n_units, b.div(b.add(ci(spec.target_blocks), b.sub(n_tiles, c1)), n_tiles) + ) + + a_rsrc = b.buffer_rsrc(params["A"], params["A_bytes"]) # dY + b_rsrc = b.buffer_rsrc(params["B"], params["B_bytes"]) # X + sdy_n, sdy_h, sdy_w = ( + params["p_dY_stride_n"], + params["p_dY_stride_ho"], + params["p_dY_stride_wo"], + ) + sx_n, sx_h, sx_w = ( + params["p_X_stride_n"], + params["p_X_stride_hi"], + params["p_X_stride_wi"], + ) + + # ---- LDS rings: one allocation per slot so the waitcnt pass can tell them apart + IX = spec.dma_instrs(XC) # DMA instrs per wave per X row + IY = spec.dma_instrs(BWO) + X_ELEMS = spec.x_row_elems() + Y_ELEMS = spec.dy_row_elems() + x_bufs = [ + b.smem_alloc(io_type, [1, X_ELEMS], name_hint=f"xr{i}") for i in range(NX) + ] + y_bufs = [ + b.smem_alloc(io_type, [1, Y_ELEMS], name_hint=f"yr{i}") for i in range(NY) + ] + x_addr = [b.smem_addr_of(s) for s in x_bufs] + y_addr = [b.smem_addr_of(s) for s in y_bufs] + + # DMA lane geometry: slot L = (wave * I + i) * 64 + lane -> (col, chunk). + def dma_lanes(n_instr: int, cols: int): + out = [] + for i in range(n_instr): + slot = b.add(b.mul(b.add(b.mul(wave, ci(n_instr)), ci(i)), ci(_WAVE)), lane) + col = b.div(slot, ci(CH)) + ch = b.mul(b.mod(slot, ci(CH)), ci(8)) + out.append((col, ch, b.cmp_lt(col, ci(cols)))) + return out + + x_lanes = dma_lanes(IX, XC) + y_lanes = dma_lanes(IY, BWO) + + def dma(rsrc, addr, n_instr, i, voff): + lds = b.smem_ptr_add( + addr, + b.zext(b.mul(b.add(b.mul(wave, ci(n_instr)), ci(i)), ci(_WAVE * 16)), I64), + ) + b.async_buffer_load_lds(rsrc, lds, voff, c0, 4) + + # Fragment read element indices (lane: chunk lc of a column). + def col_index(col: Value) -> Value: + return b.add(b.mul(col, ci(BC)), b.mul(lc, ci(8))) + + def cvt8(v: Value) -> List[Value]: + return [b.cast_to_f32(b.vec_extract(v, e)) for e in range(8)] + + zero = b.const_f32(0.0) + n_acc = TAPS * 8 + zero_accs = [(f"acc{i}", zero) for i in range(n_acc)] + + seg_loop = b.scf_for_iter( + part, n_units, n_part, zero_accs, iv_name="unit", elide_trailing_barrier=False + ) + with seg_loop as (unit, carried): + accs = list(carried) + hb = b.mod(unit, n_hb) + rest = b.div(unit, n_hb) + ct = b.mod(rest, n_ct) + n_img = b.div(rest, n_ct) + h0 = b.mul(hb, ci(RPS)) + q0 = b.mul(ct, ci(BWO)) + xi0 = b.sub(b.mul(q0, ci(S)), ci(PAD)) # first X column of the tile + x_row0 = b.sub(b.mul(h0, ci(S)), ci(PAD)) # X row of relative index 0 + n_steps = b.smin(ci(RPS), b.sub(p_Ho, h0)) + + # Per-lane DMA offsets of this unit (rows added per issue). + x_offs = [] + for col, ch, okl in x_lanes: + wi = b.add(xi0, col) + okc = b.land(okl, b.land(b.cmp_ge(wi, c0), b.cmp_lt(wi, p_W))) + cc = b.add(c_blk0, ch) + okc = b.land(okc, b.cmp_lt(cc, p_C)) + x_offs.append((b.add(b.add(b.mul(n_img, sx_n), b.mul(wi, sx_w)), cc), okc)) + y_offs = [] + for col, ch, okl in y_lanes: + wo = b.add(q0, col) + cc = b.add(c_blk0, ch) + oky = b.land(okl, b.land(b.cmp_lt(wo, p_Wo), b.cmp_lt(cc, p_C))) + y_offs.append( + (b.add(b.add(b.mul(n_img, sdy_n), b.mul(wo, sdy_w)), cc), oky) + ) + + def issue_x(xr: Value, slot: int): + hi = b.add(x_row0, xr) + # Rows of later steps than the segment's own stay masked. + ok_r = b.land( + b.land(b.cmp_ge(hi, c0), b.cmp_lt(hi, p_H)), + b.cmp_lt(xr, b.add(b.mul(n_steps, ci(S)), ci(PRE))), + ) + for i, (off, okl) in enumerate(x_offs): + voff = b.select( + b.land(okl, ok_r), b.mul(b.add(off, b.mul(hi, sx_h)), c2), oob + ) + dma(b_rsrc, x_addr[slot], IX, i, voff) + + def issue_y(t: Value, slot: int): + ho = b.add(h0, t) + ok_r = b.cmp_lt(t, n_steps) + for i, (off, okl) in enumerate(y_offs): + voff = b.select( + b.land(okl, ok_r), b.mul(b.add(off, b.mul(ho, sdy_h)), c2), oob + ) + dma(a_rsrc, y_addr[slot], IY, i, voff) + + def issue_step(t: Value, tm: int): + """DMA the S X rows that step t adds and dY row t (tm = t mod U).""" + for j in range(S): + xr = b.add(b.mul(t, ci(S)), ci(PRE + j)) + issue_x(xr, (S * tm + PRE + j) % NX) + issue_y(t, tm % NY) + + PER_STEP = S * IX + IY + # Prologue: the PRE leading X rows (waited for in full), then steps 0..D-1. + for j in range(PRE): + issue_x(ci(j), j % NX) + for t in range(D): + issue_step(ci(t), t) + + def compute(tm: int, accs): + dyv = [] + for jw in range(WL): + col = b.add(b.mul(lw, ci(WL)), ci(jw)) + dyv.append( + cvt8( + b.smem_load_vN( + y_bufs[tm % NY], c0, col_index(col), dtype=io_type, n=8 + ) + ) + ) + accs = list(accs) + for r in range(KH): + slot = (S * tm + r) % NX + win = {} + for jw in range(WL): + for s in range(KW): + xc = jw * S + s + if xc not in win: + col = b.add(b.mul(lw, ci(WL * S)), ci(xc)) + win[xc] = cvt8( + b.smem_load_vN( + x_bufs[slot], c0, col_index(col), dtype=io_type, n=8 + ) + ) + xv = win[xc] + for e in range(8): + i = (r * KW + s) * 8 + e + accs[i] = b.fma(dyv[jw][e], xv[e], accs[i]) + return accs + + def step(t: Value, tm: int, accs): + b.s_waitcnt(vmcnt=(D - 1) * PER_STEP, lgkmcnt=0) + b.s_barrier_bare() + issue_step(b.add(t, ci(D)), (tm + D) % U) + return compute(tm, accs) + + trips = b.smax(b.div(b.add(n_steps, ci(U - 1)), ci(U)), c1) + loop = b.scf_for_iter( + c0, + trips, + c1, + [(f"a{i}", v) for i, v in enumerate(accs)], + iv_name="rb", + elide_trailing_barrier=False, + ) + with loop as (rb, la): + la = list(la) + for u in range(U): + la = step(b.add(b.mul(rb, ci(U)), ci(u)), u, la) + b.scf_yield(*la) + accs_out = list(loop.results) + b.s_waitcnt(vmcnt=0, lgkmcnt=0) + b.s_barrier_bare() + b.scf_yield(*accs_out) + accs = list(seg_loop.results) + + # ---- reduce over the LW column lanes, store ws[part][tap][c] ---- + red = b.smem_alloc(F32, [1, LW * BC], name_hint="red") + for t in range(TAPS): + base = b.add(b.mul(lw, ci(BC)), b.mul(lc, ci(8))) + for e0 in (0, 4): + vec = b.vec_pack([accs[t * 8 + e0 + e] for e in range(4)], F32) + b.smem_store_vN_f32(red, [c0, b.add(base, ci(e0))], vec, 4) + b.sync() + for k in range(ceil_div(BC, T)): + o = b.add(tid, ci(k * T)) + o_ok = b.cmp_lt(o, ci(BC)) + tot = None + for q in range(LW): + idx = b.select(o_ok, b.add(ci(q * BC), o), c0) + v = b.vec_extract(b.smem_load_vN_f32(red, c0, idx, n=1), 0) + tot = v if tot is None else b.fadd(tot, v) + c_glb = b.add(c_blk0, o) + okw = b.land(o_ok, b.cmp_lt(c_glb, p_C)) + off = b.add(b.mul(b.add(b.mul(part, ci(TAPS)), ci(t)), p_C), c_glb) + with b.scf_if(okw): + b.global_store(params["D"], off, tot, align=4) + b.sync() + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_wgrad_gather.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_wgrad_gather.py new file mode 100644 index 000000000000..54a730468308 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_wgrad_gather.py @@ -0,0 +1,526 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Direct wgrad for dense layers with few input channels, any stride. + +Computes, for ``groups == 1``:: + + dW[k, r, s, c] = sum_{n, ho, wo} dY[n, ho, wo, k] * X[n, ho*S + r - PAD, wo*S + s - PAD, c] + +as one GEMM per filter row ``r``: ``M = k``, ``N = (s, c)`` (``KW * C`` columns, +padded to 16), reduced over the output columns of a 32-wide strip by +``mfma_f32_16x16x32``. That is the shape for a stem layer (``C = 3``) or a +narrow one (``C = 24``), where the per-tap ``(k, c)`` tile of +``conv_direct_wgrad_tiled`` would leave most of the MFMA's ``N`` empty. + +The block streams OUTPUT rows of one strip. Per step it DMAs ``S`` new X rows +and one dY row into LDS rings (``prefetch`` steps ahead). The dY fragment is a +transpose read of the ``[32 px][K]`` row, as in the tiled kernel. The X +operand is gathered: lane ``(n = (s, c), 8 columns)`` reads its 8 elements out +of the raw ``[column][C]`` X row at a stride of ``S * C``. + +X rows are DMAed as whole 16-byte chunks from an 8-element-aligned start, so a +strip at the image edge stages neighbouring-row data in its padding columns; +those strips zero them in LDS before the row is read. + +The epilogue stores one partial slot per block (``ws[part][(r, s, c)][k]``), +folded into ``dW`` by ``conv_direct_dw_wgrad.DwWgradFoldSpec`` with +``taps = KH * KW * C`` and ``channels = K``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Tuple + +from rocke.core.ir import I64, IRBuilder, KernelDef, Value +from rocke.helpers.spec import ceil_div + +from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + _io_type, + _mfma, + dilation_reason, + emit_direct_params, +) +from kernels.common.conv_direct_wgrad_tiled import pick_swizzle + +_ATOM = 16 +_WO = 32 +_VEC = 8 +_WAVE = 64 + + +@dataclass(frozen=True) +class WgradGatherSpec: + problem: DirectConvProblem + name: str = "direct_wgrad_gather" + wave_k16: int = 2 # k atoms per wave + waves_k: int = 1 + prefetch: int = 2 + rows_per_seg: int = 16 # output rows per segment (rounded up to the unroll) + target_blocks: int = 1024 + + @property + def block_k(self) -> int: + return _ATOM * self.wave_k16 * self.waves_k + + @property + def threads(self) -> int: + return self.waves_k * _WAVE + + @property + def n_atoms(self) -> int: + p = self.problem + return ceil_div(p.KW * p.cpg, _ATOM) + + @property + def unroll(self) -> int: + p = self.problem + return self.prefetch + ceil_div(p.KH, p.stride) + + def rps(self) -> int: + return ceil_div(self.rows_per_seg, self.unroll) * self.unroll + + def x_phase(self) -> int: + p = self.problem + return (-p.PAD * p.cpg) % _VEC + + def x_elems(self) -> int: + """Elements of one staged X row: phase + the strip's columns.""" + p = self.problem + return self.x_phase() + ((_WO - 1) * p.stride + p.KW) * p.cpg + + def kernel_name(self) -> str: + from rocke.helpers.spec import kernel_name_join + + p = self.problem + flags = {"bf16": True} if p.dtype == "bf16" else {} + return kernel_name_join( + self.name, + f"f{p.KH}x{p.KW}p{p.PAD}s{p.stride}c{p.cpg}", + f"wk{self.wave_k16}x{self.waves_k}", + f"pf{self.prefetch}", + f"rps{self.rows_per_seg}", + f"tb{self.target_blocks}", + flags=flags, + ) + + def n_k_tiles(self) -> int: + return ceil_div(self.problem.kpg, self.block_k) + + def n_units(self) -> int: + p = self.problem + return p.N * ceil_div(p.Wo, _WO) * ceil_div(p.Ho, self.rps()) + + def n_partitions(self) -> int: + return min(self.n_units(), ceil_div(self.target_blocks, self.n_k_tiles())) + + def fold_taps(self) -> int: + p = self.problem + return p.KH * p.KW * p.cpg + + def fold_channels(self) -> int: + return self.problem.kpg + + def workspace_bytes(self) -> int: + return self.n_partitions() * self.fold_taps() * self.fold_channels() * 4 + + def launch_geometry(self): + return (self.n_k_tiles(), self.n_partitions(), 1), (self.threads, 1, 1) + + def dma_instrs(self) -> Tuple[int, int]: + """16-byte DMA instructions per wave and row: (X, dY).""" + ix = ceil_div(ceil_div(ceil_div(self.x_elems(), _VEC), _WAVE), self.waves_k) + iy = ceil_div(ceil_div(_WO * self.block_k // _VEC, _WAVE), self.waves_k) + return ix, iy + + def x_slot_elems(self) -> int: + """Elements of one X ring slot: whole DMA instructions of every wave, + +16 because a transpose read of the last column atom runs past it.""" + return self.threads * self.dma_instrs()[0] * _VEC + 16 + + def dy_slot_elems(self) -> int: + return self.threads * self.dma_instrs()[1] * _VEC + + def lds_bytes(self) -> int: + """The X and dY rings as the builder allocates them.""" + U = self.unroll + return ( + self.problem.stride * U * self.x_slot_elems() + U * self.dy_slot_elems() + ) * 2 + + +def is_valid_wgrad_gather_spec( + spec: WgradGatherSpec, arch: str = "gfx950" +) -> Tuple[bool, str]: + from rocke.core.arch import ArchTarget + + try: + target = ArchTarget.from_gfx(arch) + except KeyError as e: + return False, str(e) + p = spec.problem + why = dilation_reason(p) + if why is not None: + return False, why + if p.groups != 1: + return False, "dense layers only" + if p.dtype not in ("fp16", "bf16"): + return False, "fp16/bf16 only" + if not target.memory.has_ds_read_tr: + return False, f"needs ds_read_b64_tr_b16 (absent on {arch})" + if target.wave_size != 64 or target.async_lds_max_dwords < 4: + return False, f"needs wave64 + 16-byte buffer_load_lds (absent on {arch})" + if p.kpg % _VEC: + return False, "K must be a multiple of 8" + if p.KW * p.cpg > 96: + return False, "KW * C must be at most 96" + # Whole 16-byte chunks from an aligned start at a build-time phase. + if (p.W * p.cpg) % _VEC or (_WO * p.stride * p.cpg) % _VEC: + return False, "W * C and 32 * stride * C must be multiples of 8" + acc = p.KH * spec.wave_k16 * spec.n_atoms * 4 + if acc > 192: + return False, f"{acc} accumulator VGPRs (> 192)" + if spec.wave_k16 < 1 or spec.waves_k < 1 or spec.threads > 1024: + return False, "wave_k16 >= 1 and waves_k in 1..16" + if spec.prefetch < 1: + return False, "prefetch >= 1" + if spec.lds_bytes() > 160 * 1024: + return False, f"LDS {spec.lds_bytes()} B > 160 KiB" + if spec.rows_per_seg < 1: + return False, "rows_per_seg >= 1" + if spec.target_blocks < 1: + return False, "target_blocks >= 1" + return True, "ok" + + +def build_wgrad_gather(spec: WgradGatherSpec, arch: str = "gfx950") -> KernelDef: + ok, why = is_valid_wgrad_gather_spec(spec, arch) + if not ok: + raise ValueError(f"invalid WgradGatherSpec for {arch}: {why}") + p = spec.problem + KH, KW, PAD, S, C = p.KH, p.KW, p.PAD, p.stride, p.cpg + io_dtype = p.dtype + io_type = _io_type(io_dtype) + WK16 = spec.wave_k16 + NA = spec.n_atoms + BK = spec.block_k + CHK = BK // _VEC + PH = spec.x_phase() + XCH = ceil_div(spec.x_elems(), _VEC) # 16-byte chunks of a staged X row + D = spec.prefetch + U = spec.unroll + NX, NY = S * U, U + RPS = spec.rps() + PRE = max(KH - S, 0) + TAPS = KH * KW * C + + IX, IY = spec.dma_instrs() + X_ELEMS = spec.x_slot_elems() + # When a pixel row of the strip (S * C elements) and the phase keep every + # lane's 8-byte read aligned, the raw row is itself a [pixel][(s, c)] + # tile with contiguous columns, so the X operand is a transpose read too. + X_TR = (S * C) % 4 == 0 and PH % 4 == 0 + Y_ELEMS = spec.dy_slot_elems() + PER_STEP = S * IX + IY + + swz_k = pick_swizzle(CHK, (0,), _WO, _WO, tuple(range(0, BK, _ATOM))) + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = spec.threads + b.kernel.attrs["lds_alias_scopes"] = True + params = emit_direct_params(b, direction="wgrad", io_type=io_type) + + def ci(v: int) -> Value: + return b.const_i32(v) + + c0, c1, c2 = ci(0), ci(1), ci(2) + oob = ci((1 << 31) - 1) + p_N = params["p_N"] + p_H, p_W = params["p_Hi"], params["p_Wi"] + p_Ho, p_Wo = params["p_Ho"], params["p_Wo"] + p_K = params["p_total_k"] + tid = b.thread_id_x() + lane = b.mod(tid, ci(_WAVE)) + wave = b.readfirstlane(b.div(tid, ci(_WAVE))) + k_tile, part = b.block_id_x(), b.block_id_y() + k_blk0 = b.mul(k_tile, ci(BK)) + n_q = b.div(b.add(p_Wo, ci(_WO - 1)), ci(_WO)) + n_hb = b.div(b.add(p_Ho, ci(RPS - 1)), ci(RPS)) + n_units = b.mul(b.mul(p_N, n_q), n_hb) + n_kt = b.div(b.add(p_K, ci(BK - 1)), ci(BK)) + n_part = b.smin( + n_units, b.div(b.add(ci(spec.target_blocks), b.sub(n_kt, c1)), n_kt) + ) + + a_rsrc = b.buffer_rsrc(params["A"], params["A_bytes"]) # dY + b_rsrc = b.buffer_rsrc(params["B"], params["B_bytes"]) # X + sdy_n, sdy_h, sdy_w = ( + params["p_dY_stride_n"], + params["p_dY_stride_ho"], + params["p_dY_stride_wo"], + ) + sx_n, sx_h = params["p_X_stride_n"], params["p_X_stride_hi"] + + x_bufs = [ + b.smem_alloc(io_type, [1, X_ELEMS], name_hint=f"xg{i}") for i in range(NX) + ] + y_bufs = [ + b.smem_alloc(io_type, [1, Y_ELEMS], name_hint=f"yg{i}") for i in range(NY) + ] + x_addr = [b.smem_addr_of(s) for s in x_bufs] + y_addr = [b.smem_addr_of(s) for s in y_bufs] + + def swz_val(row: Value, bits): + v = None + for j, bit in enumerate(bits): + term = b.shl(b.land(b.lshr(row, ci(bit)), c1), ci(j + 1)) + v = term if v is None else b.lor(v, term) + return v + + def dma(rsrc, addr, n_instr, i, voff): + lds = b.smem_ptr_add( + addr, + b.zext(b.mul(b.add(b.mul(wave, ci(n_instr)), ci(i)), ci(_WAVE * 16)), I64), + ) + b.async_buffer_load_lds(rsrc, lds, voff, c0, 4) + + # X: chunk L of a row is elements [8 L, 8 L + 8) from the aligned start. + x_chunks = [] + for i in range(IX): + slot = b.add(b.mul(b.add(b.mul(wave, ci(IX)), ci(i)), ci(_WAVE)), lane) + x_chunks.append((b.mul(slot, ci(_VEC)), b.cmp_lt(slot, ci(XCH)))) + # dY: slot -> (pixel row, swizzled k chunk). + y_chunks = [] + for i in range(IY): + slot = b.add(b.mul(b.add(b.mul(wave, ci(IY)), ci(i)), ci(_WAVE)), lane) + row = b.div(slot, ci(CHK)) + phys = b.mod(slot, ci(CHK)) + sw = swz_val(row, swz_k) + log = phys if sw is None else b.xor(phys, sw) + y_chunks.append((row, b.mul(log, ci(_VEC)), b.cmp_lt(row, ci(_WO)))) + + # ---- fragment addresses ------------------------------------------------- + lane_hi = b.div(lane, ci(16)) + lane_q = b.mod(b.div(lane, ci(4)), ci(4)) + lane_c4 = b.mod(lane, ci(4)) + lane_n = b.mod(lane, ci(16)) + wave_k0 = b.mul(wave, ci(WK16 * _ATOM)) + + def dy_index(rd: int, col0: Value) -> Value: + row = b.add(b.add(b.mul(lane_hi, ci(8)), lane_q), ci(rd * 4)) + col = b.add(col0, b.mul(lane_c4, ci(4))) + chunk = b.div(col, ci(_VEC)) + sw = swz_val(row, swz_k) + phys = chunk if sw is None else b.xor(chunk, sw) + return b.add( + b.mul(b.add(b.mul(row, ci(CHK)), phys), ci(_VEC)), b.mod(col, ci(_VEC)) + ) + + dy_idx = [ + [dy_index(rd, b.add(wave_k0, ci(mi * _ATOM))) for rd in (0, 1)] + for mi in range(WK16) + ] + # X gather base per atom: element PH + (S * px0 + s) * C + c, px0 = 8 * (lane / 16). + x_base = [] + for nj in range(NA): + n = b.add(lane_n, ci(nj * _ATOM)) + s_i = b.smin(b.div(n, ci(C)), ci(KW - 1)) # n >= KW*C: any in-range column + c_i = b.mod(n, ci(C)) + x_base.append( + b.add( + ci(PH), b.add(b.mul(b.add(b.mul(lane_hi, ci(8 * S)), s_i), ci(C)), c_i) + ) + ) + + def read_dy(buf: int) -> List[Value]: + out = [] + for mi in range(WK16): + lo = b.ds_read_tr16_b64(y_bufs[buf], c0, dy_idx[mi][0], dtype=io_type) + hi = b.ds_read_tr16_b64(y_bufs[buf], c0, dy_idx[mi][1], dtype=io_type) + out.append(b.vec_concat(lo, hi)) + return out + + x_tr_idx = [ + [ + b.add( + ci(PH + nj * _ATOM), + b.add( + b.mul( + b.add(b.add(b.mul(lane_hi, ci(8)), lane_q), ci(rd * 4)), + ci(S * C), + ), + b.mul(lane_c4, ci(4)), + ), + ) + for rd in (0, 1) + ] + for nj in range(NA) + ] + + def read_x(buf: int, nj: int) -> Value: + if X_TR: + lo = b.ds_read_tr16_b64(x_bufs[buf], c0, x_tr_idx[nj][0], dtype=io_type) + hi = b.ds_read_tr16_b64(x_bufs[buf], c0, x_tr_idx[nj][1], dtype=io_type) + return b.vec_concat(lo, hi) + vals = [ + b.vec_extract( + b.smem_load_vN( + x_bufs[buf], + c0, + b.add(x_base[nj], ci(j * S * C)), + dtype=io_type, + n=1, + ), + 0, + ) + for j in range(8) + ] + return b.vec_pack(vals, io_type) + + zero_acc = b.zero_vec_f32(4) + zero_elem = b.vec_extract(b.zero_vec(io_type, 2), 0) + n_acc = KH * WK16 * NA + + def acc_index(r, mi, nj): + return (r * WK16 + mi) * NA + nj + + seg_loop = b.scf_for_iter( + part, + n_units, + n_part, + [(f"acc{i}", zero_acc) for i in range(n_acc)], + iv_name="unit", + elide_trailing_barrier=False, + ) + with seg_loop as (unit, carried): + accs = list(carried) + hb = b.mod(unit, n_hb) + rest = b.div(unit, n_hb) + q = b.mod(rest, n_q) + n_img = b.div(rest, n_q) + h0 = b.mul(hb, ci(RPS)) + q0 = b.mul(q, ci(_WO)) + col0 = b.sub(b.mul(q0, ci(S)), ci(PAD)) # first X column of the strip + x_row0 = b.sub(b.mul(h0, ci(S)), ci(PAD)) + n_steps = b.smin(ci(RPS), b.sub(p_Ho, h0)) + # Strip X start, 8-aligned: element col0 * C sits at PH. + x_start = b.add(b.mul(n_img, sx_n), b.sub(b.mul(col0, ci(C)), ci(PH))) + y_offs = [] + for row, ch, okr in y_chunks: + wo = b.add(q0, row) + kk = b.add(k_blk0, ch) + ok = b.land(okr, b.land(b.cmp_lt(wo, p_Wo), b.cmp_lt(kk, p_K))) + y_offs.append((b.add(b.add(b.mul(n_img, sdy_n), b.mul(wo, sdy_w)), kk), ok)) + # Padding columns staged from neighbouring rows: [lo, hi) element + # ranges of the row to zero, for the edge strips. + pad_lo = b.smax(b.mul(b.sub(c0, col0), ci(C)), c0) # cols < 0 + cols_in = b.sub(p_W, col0) # cols >= W start here (relative) + pad_hi0 = b.add(ci(PH), b.mul(cols_in, ci(C))) + edge = b.lor(b.cmp_lt(col0, c0), b.cmp_lt(pad_hi0, ci(spec.x_elems()))) + + def issue_x(xr: Value, slot: int): + hi = b.add(x_row0, xr) + ok_r = b.land( + b.land(b.cmp_ge(hi, c0), b.cmp_lt(hi, p_H)), + b.cmp_lt(xr, b.add(b.mul(n_steps, ci(S)), ci(PRE))), + ) + row_off = b.add(x_start, b.mul(hi, sx_h)) + for i, (e, okl) in enumerate(x_chunks): + voff = b.select(b.land(okl, ok_r), b.mul(b.add(row_off, e), c2), oob) + dma(b_rsrc, x_addr[slot], IX, i, voff) + + def issue_y(t: Value, slot: int): + ho = b.add(h0, t) + ok_r = b.cmp_lt(t, n_steps) + for i, (off, okl) in enumerate(y_offs): + voff = b.select( + b.land(okl, ok_r), b.mul(b.add(off, b.mul(ho, sdy_h)), c2), oob + ) + dma(a_rsrc, y_addr[slot], IY, i, voff) + + def issue_step(t: Value, tm: int): + for j in range(S): + issue_x(b.add(b.mul(t, ci(S)), ci(PRE + j)), (S * tm + PRE + j) % NX) + issue_y(t, tm % NY) + + def fix_padding(slots: List[int]): + """Zero the padding columns of freshly landed X rows (edge strips).""" + with b.scf_if(edge): + for slot in slots: + for k in range(ceil_div(spec.x_elems(), spec.threads)): + e = b.add(tid, ci(k * spec.threads)) + rel = b.sub(e, ci(PH)) + is_pad = b.land( + b.cmp_lt(e, ci(spec.x_elems())), + b.lor(b.cmp_lt(rel, pad_lo), b.cmp_ge(e, pad_hi0)), + ) + with b.scf_if(is_pad): + b.smem_store_vN(x_bufs[slot], [c0, e], zero_elem, 1) + b.s_waitcnt(lgkmcnt=0) + b.s_barrier_bare() + + for j in range(PRE): + issue_x(ci(j), j % NX) + for t in range(D): + issue_step(ci(t), t) + + def step(t: Value, tm: int, accs): + b.s_waitcnt(vmcnt=(D - 1) * PER_STEP, lgkmcnt=0) + b.s_barrier_bare() + new_slots = [(S * tm + PRE + j) % NX for j in range(S)] + if tm == 0: + # Step 0 of a segment also consumes the PRE leading rows. + new_slots = [j % NX for j in range(PRE)] + new_slots + fix_padding(new_slots) + issue_step(b.add(t, ci(D)), (tm + D) % U) + dyf = read_dy(tm % NY) + accs = list(accs) + for r in range(KH): + slot = (S * tm + r) % NX + for nj in range(NA): + xf = read_x(slot, nj) + for mi in range(WK16): + i = acc_index(r, mi, nj) + accs[i] = _mfma(b, io_dtype, "16x16x32", dyf[mi], xf, accs[i]) + return accs + + trips = b.smax(b.div(b.add(n_steps, ci(U - 1)), ci(U)), c1) + loop = b.scf_for_iter( + c0, + trips, + c1, + [(f"a{i}", v) for i, v in enumerate(accs)], + iv_name="rb", + elide_trailing_barrier=False, + ) + with loop as (rb, la): + la = list(la) + for u in range(U): + la = step(b.add(b.mul(rb, ci(U)), ci(u)), u, la) + b.scf_yield(*la) + accs_out = list(loop.results) + b.s_waitcnt(vmcnt=0, lgkmcnt=0) + b.s_barrier_bare() + b.scf_yield(*accs_out) + accs = list(seg_loop.results) + + # ---- epilogue: partial slot ws[part][(r, s, c)][k] ---------------------- + # acc lane layout: k = (lane / 16) * 4 + slot of atom mi, n = lane % 16 of atom nj. + D_ptr = params["D"] + k_lane0 = b.add(b.add(k_blk0, wave_k0), b.mul(lane_hi, ci(4))) + for r in range(KH): + for nj in range(NA): + n = b.add(lane_n, ci(nj * _ATOM)) + n_ok = b.cmp_lt(n, ci(KW * C)) + tap = b.add(ci(r * KW * C), n) # (r, s, c) flattened + row_base = b.mul(b.add(b.mul(part, ci(TAPS)), tap), p_K) + for mi in range(WK16): + for slot in range(4): + kk = b.add(k_lane0, ci(mi * _ATOM + slot)) + ok = b.land(n_ok, b.cmp_lt(kk, p_K)) + val = b.vec_extract(accs[acc_index(r, mi, nj)], slot) + with b.scf_if(ok): + b.global_store(D_ptr, b.add(row_base, kk), val, align=4) + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_wgrad_tiled.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_wgrad_tiled.py new file mode 100644 index 000000000000..724abefc46b3 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_wgrad_tiled.py @@ -0,0 +1,853 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Direct wgrad with multi-atom wave tiles and an LDS-DMA row pipeline. + +Computes the weight gradient of a stride-1 NHWC convolution:: + + dW[k, r, s, c] = sum_{n, ho, wo} dY[n, ho, wo, k] * X[n, ho + r - PAD, wo + s - PAD, c] + +The row loop walks INPUT rows ``hi`` of one 32-column strip. Input row ``hi`` +pairs with output rows ``hi + PAD - r`` for every filter row ``r``, so each dY +row is read out of LDS once and kept in a ``KH``-deep register ring, and one X +row strip of ``32 + KW - 1`` columns serves every filter column ``s`` at a +one-column shift. + +What sets it apart from :func:`build_direct_conv_wgrad` (16x16 per wave): + +* A wave owns ``wave_k16*16 x wave_c16*16`` channels for all ``KH*KW`` taps + (``32 x 32 x 9`` = 144 fp32 accumulator registers at 3x3), so one dY + fragment and one X fragment feed ``wave_c16`` resp. ``wave_k16`` MFMAs. +* The rows are workgroup-shared: every wave DMAs one 1 KiB slice of each row + straight into LDS (``buffer_load ... lds``), with ``prefetch`` rows in + flight, so no VGPR staging and one barrier per row. +* A 16-byte channel chunk lands at ``chunk ^ swizzle(row)``; the swizzle is + chosen at build time so the transpose reads are bank-conflict free. +* A persistent grid: ``ceil(target_blocks / channel_tiles)`` blocks walk the + (image, column strip, row segment) list with a stride, accumulate in + registers across segments, and atomically add once at the end. + +dW is fp32; the caller zero-initialises it. +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass +from functools import lru_cache +from typing import List, Optional, Tuple + +from rocke.core.ir import I64, IRBuilder, KernelDef, Value +from rocke.helpers.spec import ceil_div + +from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + _io_type, + _mfma, + dilation_reason, + emit_direct_params, +) + +# MFMA 16x16x32: M = output channels (k), N = input channels (c), K = columns. +_ATOM = 16 +_WO = 32 +_VEC = 8 # 16-bit elements per 16-byte chunk +_WAVE = 64 +# Bank model used to pick the swizzle: 64 four-byte banks, a 32-lane half of +# a ds_read_b64 serviced per pass. +_BANKS = 64 +_HALF = 32 + + +@dataclass(frozen=True) +class DirectWgradTiledSpec: + problem: DirectConvProblem + name: str = "direct_wgrad_tiled" + wave_k16: int = 2 + wave_c16: int = 2 + waves_k: int = 4 + waves_c: int = 2 + rows_per_seg: int = 30 + prefetch: int = 2 + target_blocks: int = 256 + # Main-loop unroll override (a multiple of prefetch + 1); 0 = default. + row_unroll: int = 0 + # Images packed into one 32-column strip: the MFMA reduction runs over + # unfold_n images x 32 / unfold_n columns, so a narrow image (W = 40, say) + # does not leave most of a 32-wide strip empty. + unfold_n: int = 1 + # Two wave groups per block, each streaming its own segments of the same + # channel tile; their partial gradients meet in LDS before the atomics, + # which halves the atomic count. With ``pingpong`` the second group runs + # half a row behind the first (one extra barrier at the start), so one + # group's MFMA phase covers the other's DMA issue and fragment reads. + q_groups: int = 1 + pingpong: bool = False + + @property + def wave_k(self) -> int: + return _ATOM * self.wave_k16 + + @property + def wave_c(self) -> int: + return _ATOM * self.wave_c16 + + @property + def block_k(self) -> int: + return self.wave_k * self.waves_k + + @property + def block_c(self) -> int: + return self.wave_c * self.waves_c + + @property + def n_waves(self) -> int: + return self.waves_k * self.waves_c + + @property + def threads_per_block(self) -> int: + return self.n_waves * self.q_groups * _WAVE + + @property + def strip_w(self) -> int: + """Output columns of one image in a strip.""" + return _WO // self.unfold_n + + @property + def n_buffers(self) -> int: + return self.prefetch + 1 + + @property + def unroll(self) -> int: + # The LDS buffer index has to fold to a constant, so the unroll is a + # multiple of n_buffers. The dY ring is SSA renaming and would work at + # any unroll, but one that is not a multiple of KH rotates the ring + # through the loop-carried values, and the copies cost registers. + if self.row_unroll: + return self.row_unroll + return math.lcm(self.n_buffers, self.problem.KH) + + def acc_vgprs(self) -> int: + p = self.problem + return p.KH * p.KW * self.wave_k16 * self.wave_c16 * 4 + + def dma_instrs(self) -> Tuple[int, int]: + """1 KiB DMA instructions per wave and row: (dY, X).""" + seg = self.strip_w + self.problem.KW - 1 + ik = ceil_div(ceil_div(_WO * self.block_k // _VEC, _WAVE), self.n_waves) + ic = ceil_div( + ceil_div(self.unfold_n * seg * self.block_c // _VEC, _WAVE), self.n_waves + ) + return ik, ic + + def lds_bytes(self) -> int: + """The padded dY / X ring buffers, plus the q-group reduction buffers.""" + ik, ic = self.dma_instrs() + slot = self.q_groups * self.n_waves * (ik + ic) * _WAVE * _VEC * 2 + red = 0 + if self.q_groups > 1: + red = 2 * self.n_waves * self.wave_k16 * self.wave_c16 * _WAVE * 4 * 4 + return self.n_buffers * slot + red + + def kernel_name(self) -> str: + from rocke.helpers.spec import kernel_name_join + + p = self.problem + flags = {"bf16": True} if p.dtype == "bf16" else {} + if self.row_unroll: + flags[f"u{self.row_unroll}"] = True + if self.unfold_n > 1: + flags[f"un{self.unfold_n}"] = True + if self.q_groups > 1: + flags[f"qg{self.q_groups}"] = True + if self.pingpong: + flags["pp"] = True + return kernel_name_join( + self.name, + f"f{p.KH}x{p.KW}p{p.PAD}", + f"wt{self.wave_k}x{self.wave_c}", + f"w{self.waves_k}x{self.waves_c}", + f"rps{self.rows_per_seg}", + f"pf{self.prefetch}", + f"tb{self.target_blocks}", + flags=flags, + ) + + # ---- host-side launch geometry (mirrors the in-kernel decode) ---- + def n_channel_tiles(self) -> int: + p = self.problem + return p.groups * ceil_div(p.kpg, self.block_k) * ceil_div(p.cpg, self.block_c) + + def n_segments(self) -> int: + p = self.problem + return ( + ceil_div(p.N, self.unfold_n) + * ceil_div(p.Wo, self.strip_w) + * ceil_div(p.H, self.rows_per_seg) + ) + + def n_splits(self) -> int: + return min( + ceil_div(self.n_segments(), self.q_groups), + ceil_div(self.target_blocks, self.n_channel_tiles()), + ) + + def launch_geometry(self): + return (self.n_channel_tiles() * self.n_splits(), 1, 1), ( + self.threads_per_block, + 1, + 1, + ) + + +def is_valid_wgrad_tiled_spec(spec: DirectWgradTiledSpec, arch: str = "gfx950"): + from rocke.core.arch import ArchTarget + + try: + target = ArchTarget.from_gfx(arch) + except KeyError as e: + return False, str(e) + p = spec.problem + why = dilation_reason(p) + if why is not None: + return False, why + if p.dtype not in ("fp16", "bf16"): + return False, f"unsupported dtype {p.dtype!r}" + if p.stride != 1: + return False, f"stride-1 algorithm (got stride={p.stride})" + if target.wave_size != 64 or not target.memory.has_ds_read_tr: + return False, f"needs wave64 + ds_read_b64_tr_b16 (absent on {arch})" + if target.async_lds_max_dwords < 4: + return False, f"needs 16-byte buffer_load_lds (absent on {arch})" + ab = "bf16" if p.dtype == "bf16" else "f16" + if not target.mma.has_shape( + a_dtype=ab, b_dtype=ab, c_dtype="fp32", m=16, n=16, k=32 + ): + return False, f"missing mfma_f32_16x16x32_{ab} on {arch}" + if spec.wave_k16 not in (1, 2, 4) or spec.wave_c16 not in (1, 2, 4): + return False, "wave_k16 / wave_c16 must be 1, 2 or 4" + for n, v in (("waves_k", spec.waves_k), ("waves_c", spec.waves_c)): + if v < 1 or v & (v - 1): + return False, f"{n} must be a power of two" + if spec.unfold_n not in (1, 2, 4): + return False, "unfold_n must be 1, 2 or 4" + if spec.q_groups not in (1, 2): + return False, "q_groups must be 1 or 2" + if spec.pingpong and spec.q_groups != 2: + return False, "pingpong needs q_groups == 2" + if spec.threads_per_block > target.max_threads_per_block: + return False, "too many threads per block" + if p.cpg % _VEC or p.kpg % _VEC: + return False, "cpg and kpg must be multiples of 8 (16-byte chunks)" + if p.kpg < spec.wave_k or p.cpg < spec.wave_c: + return False, "wave tile wider than the group's channels" + if spec.acc_vgprs() > 192: + return False, f"accumulators need {spec.acc_vgprs()} VGPRs (> 192)" + if spec.row_unroll < 0: + return False, "row_unroll must be >= 0" + if spec.target_blocks < 1: + return False, "target_blocks must be >= 1" + if spec.row_unroll % spec.n_buffers: + return False, "row_unroll must be a multiple of prefetch + 1" + if spec.prefetch < 1: + return False, "prefetch must be >= 1" + if spec.rows_per_seg <= 0 or spec.rows_per_seg % spec.unroll: + return False, f"rows_per_seg must be a positive multiple of {spec.unroll}" + if p.KH > 8 or p.KW > 8: + return False, "filter extent above 8" + if spec.lds_bytes() > 160 * 1024: + return False, f"LDS {spec.lds_bytes()} B > 160 KiB" + return True, "ok" + + +# --------------------------------------------------------------------------- +# LDS swizzle +# --------------------------------------------------------------------------- + + +def _frag_row(lane: int, rd: int, shift: int, wu: int, seg: int) -> int: + """LDS row a ``ds_read_b64_tr_b16`` lane reads for reduction index block rd. + + Lane ``l`` supplies MFMA reduction elements ``(l // 16) * 8 + 0..7``; with + ``unfold_n`` images in a strip those are columns of image ``kk // wu``, + whose rows start at ``image * seg`` in the tile. + """ + kk0 = (lane // 16) * 8 + return (kk0 // wu) * seg + kk0 % wu + rd * 4 + (lane // 4) % 4 + shift + + +def _swizzle_ok(swz, chunks: int, rows_of, col_origins) -> bool: + """No two lanes of a 32-lane half hit one bank, for every read the kernel makes.""" + for row_of in rows_of: + for origin in col_origins: + for rd in (0, 1): + for half in (0, 1): + banks = set() + for lane in range(half * _HALF, (half + 1) * _HALF): + row = row_of(lane, rd) + col = origin + (lane % 4) * 4 + phys = (col // _VEC) ^ swz(row) + byte = (row * chunks + phys) * 16 + (col % _VEC) * 2 + for w in range(2): + banks.add((byte // 4 + w) % _BANKS) + if len(banks) != 2 * _HALF: + return False + return True + + +@lru_cache(maxsize=None) +def pick_swizzle(chunks: int, shifts: Tuple[int, ...], wu: int, seg: int, col_origins): + """``bits``: chunk XOR ``2 * sum(((row >> bits[j]) & 1) << j)``. + + The XOR moves 32-byte chunk PAIRS so a lane's 8-byte half-chunk read stays + inside its chunk. Returns the first bit selection that makes every read + conflict-free, or no swizzle when none does. + """ + rows_of = [ + (lambda lane, rd, sh=sh: _frag_row(lane, rd, sh, wu, seg)) for sh in shifts + ] + n_pairs = chunks // 2 + width = int(math.log2(n_pairs)) if n_pairs > 1 else 0 + + def make(bits): + def swz(row): + v = 0 + for j, bit in enumerate(bits): + v |= ((row >> bit) & 1) << j + return 2 * v + + return swz + + if _swizzle_ok(make(()), chunks, rows_of, col_origins): + return () + for nbits in range(1, width + 1): + for bits in itertools.permutations(range(6), nbits): + if _swizzle_ok(make(bits), chunks, rows_of, col_origins): + return bits + return () + + +# --------------------------------------------------------------------------- +# Builder +# --------------------------------------------------------------------------- + + +def build_direct_wgrad_tiled( + spec: DirectWgradTiledSpec, arch: str = "gfx950" +) -> KernelDef: + """Build the tiled direct wgrad kernel (see the module docstring). + + Grid: ``n_channel_tiles * n_splits`` blocks (``launch_geometry``); block + ``split`` walks segments ``split, split + n_splits, ...`` of the + (image group, column strip, row block) list. A segment streams + ``rows_per_seg`` input rows of one strip through the row pipeline: + + step t (input row hi = h0 - (KH - 1) + t): + wait for row t's DMA, barrier, DMA row t + prefetch into the free + buffer, transpose-read the dY fragment into the ring and the KW X + fragments, then KH * KW * wave_k16 * wave_c16 MFMAs. + + The first KH - 1 steps of a segment only prime the dY ring. + """ + ok, why = is_valid_wgrad_tiled_spec(spec, arch) + if not ok: + raise ValueError(f"invalid DirectWgradTiledSpec for {arch}: {why}") + + p = spec.problem + KH, KW, PAD = p.KH, p.KW, p.PAD + io_dtype = p.dtype + io_type = _io_type(io_dtype) + NW = spec.n_waves + WK16, WC16 = spec.wave_k16, spec.wave_c16 + BK, BC = spec.block_k, spec.block_c + CHK, CHC = BK // _VEC, BC // _VEC + UN = spec.unfold_n + WU = spec.strip_w + SEG = WU + KW - 1 # X tile rows (columns of the image) per packed image + STRIP = UN * SEG + D = spec.prefetch + NB = spec.n_buffers + HPB = spec.rows_per_seg + U = spec.unroll + T_PRO = KH - 1 + + # 1 KiB DMA instructions per wave and row, per operand. Slots past the + # tile load zero (OOB sentinel) into the buffer's pad. + IK, IC = spec.dma_instrs() + N_INSTR = IK + IC + QG = spec.q_groups + # A ring buffer holds one slice per wave group, group g's at g * G*_ELEMS. + GDY_ELEMS = NW * IK * _WAVE * _VEC + GX_ELEMS = NW * IC * _WAVE * _VEC + DY_ELEMS = QG * GDY_ELEMS + X_ELEMS = QG * GX_ELEMS + + k_origins = tuple( + wk * spec.wave_k + mi * _ATOM + for wk in range(spec.waves_k) + for mi in range(WK16) + ) + c_origins = tuple( + wc * spec.wave_c + ni * _ATOM + for wc in range(spec.waves_c) + for ni in range(WC16) + ) + swz_k = pick_swizzle(CHK, (0,), WU, WU, k_origins) + swz_c = pick_swizzle(CHC, tuple(range(KW)), WU, SEG, c_origins) + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = spec.threads_per_block + # Separate scopes per row buffer, or the waitcnt pass drains the DMA ring. + b.kernel.attrs["lds_alias_scopes"] = True + # Two waves per SIMD: the register budget the 32x32 tile is sized for. + b.kernel.attrs["waves_per_eu"] = (2, 8) + params = emit_direct_params(b, direction="wgrad", io_type=io_type) + + def ci(v: int) -> Value: + return b.const_i32(v) + + c0, c1, c_two = ci(0), ci(1), ci(2) + oob = ci((1 << 31) - 1) + + p_N = params["p_N"] + p_H = params["p_Hi"] + p_W = params["p_Wi"] + p_Ho = params["p_Ho"] + p_Wo = params["p_Wo"] + p_groups = params["p_groups"] + c_cpg = b.div(params["p_total_c"], p_groups) + c_kpg = b.div(params["p_total_k"], p_groups) + + tid = b.thread_id_x() + lane = b.mod(tid, ci(_WAVE)) + wave_all = b.readfirstlane(b.div(tid, ci(_WAVE))) + if QG > 1: + qg = b.div(wave_all, ci(NW)) + wave = b.mod(wave_all, ci(NW)) + else: + qg, wave = c0, wave_all + + # ---- grid decode -------------------------------------------------------- + n_k_tiles = b.div(b.add(c_kpg, ci(BK - 1)), ci(BK)) + n_c_tiles = b.div(b.add(c_cpg, ci(BC - 1)), ci(BC)) + n_ch_tiles = b.mul(b.mul(p_groups, n_k_tiles), n_c_tiles) + n_q = b.div(b.add(p_Wo, ci(WU - 1)), ci(WU)) + n_ng = b.div(b.add(p_N, ci(UN - 1)), ci(UN)) + n_hb = b.div(b.add(p_H, ci(HPB - 1)), ci(HPB)) + total_segs = b.mul(b.mul(n_ng, n_q), n_hb) + # Same formula as DirectWgradTiledSpec.n_splits, which sizes the grid. + n_splits = b.smin( + b.div(b.add(total_segs, ci(spec.q_groups - 1)), ci(spec.q_groups)), + b.div(b.add(ci(spec.target_blocks), b.sub(n_ch_tiles, c1)), n_ch_tiles), + ) + + # Dispatch deals block ids round-robin over 8 XCDs; renumber so each XCD + # gets a contiguous run, i.e. the channel tiles of the same segments, + # which then share that XCD's L2. + lin = b.block_id_x() + n_blocks = b.mul(n_ch_tiles, n_splits) + xcd = b.mod(lin, ci(8)) + pos = b.div(lin, ci(8)) + per = b.div(n_blocks, ci(8)) + rem = b.mod(n_blocks, ci(8)) + lin = b.add(b.add(b.mul(xcd, per), b.smin(xcd, rem)), pos) + ch_tile = b.mod(lin, n_ch_tiles) + split = b.div(lin, n_ch_tiles) + c_tile = b.mod(ch_tile, n_c_tiles) + gk = b.div(ch_tile, n_c_tiles) + k_tile = b.mod(gk, n_k_tiles) + group = b.div(gk, n_k_tiles) + k_blk0 = b.add(b.mul(group, c_kpg), b.mul(k_tile, ci(BK))) # global k + c_blk0 = b.mul(c_tile, ci(BC)) # in-group c + c_glb0 = b.add(b.mul(group, c_cpg), c_blk0) + + wave_k_id = b.mod(wave, ci(spec.waves_k)) + wave_c_id = b.div(wave, ci(spec.waves_k)) + + a_rsrc = b.buffer_rsrc(params["A"], params["A_bytes"]) + b_rsrc = b.buffer_rsrc(params["B"], params["B_bytes"]) + + # ---- LDS: one allocation per ring buffer (one alias scope each) --------- + dy_bufs = [ + b.smem_alloc(io_type, [1, DY_ELEMS], name_hint=f"dy{i}") for i in range(NB) + ] + x_bufs = [ + b.smem_alloc(io_type, [1, X_ELEMS], name_hint=f"xs{i}") for i in range(NB) + ] + dy_addr = [b.smem_addr_of(s) for s in dy_bufs] + x_addr = [b.smem_addr_of(s) for s in x_bufs] + + if QG > 1: + qg_dy_bytes = b.zext(b.mul(qg, ci(GDY_ELEMS * 2)), I64) + qg_x_bytes = b.zext(b.mul(qg, ci(GX_ELEMS * 2)), I64) + + def swz_val(row: Value, bits) -> Optional[Value]: + v = None + for j, bit in enumerate(bits): + term = b.shl(b.land(b.lshr(row, ci(bit)), c1), ci(j + 1)) + v = term if v is None else b.lor(v, term) + return v + + # ---- DMA lane geometry (segment invariant part) ------------------------- + # Slot L = (wave * I + i) * 64 + lane of a 16-byte chunk; tile row = L / CH, + # physical chunk = L % CH, logical chunk = physical ^ swizzle(row). A tile + # row is (packed image, column) = divmod(row, rows per image). + def dma_lanes(n_instr: int, chunks: int, rows: int, per_img: int, bits): + out = [] + for i in range(n_instr): + slot = b.add(b.mul(b.add(b.mul(wave, ci(n_instr)), ci(i)), ci(_WAVE)), lane) + row = b.div(slot, ci(chunks)) + phys = b.mod(slot, ci(chunks)) + sw = swz_val(row, bits) + log = phys if sw is None else b.xor(phys, sw) + img = b.div(row, ci(per_img)) if UN > 1 else c0 + col = b.mod(row, ci(per_img)) if UN > 1 else row + out.append((img, col, b.mul(log, ci(_VEC)), b.cmp_lt(row, ci(rows)))) + return out + + dy_lanes = dma_lanes(IK, CHK, _WO, WU, swz_k) + x_lanes = dma_lanes(IC, CHC, STRIP, SEG, swz_c) + + sdy_n = params["p_dY_stride_n"] + sdy_h = params["p_dY_stride_ho"] + sdy_w = params["p_dY_stride_wo"] + sx_n = params["p_X_stride_n"] + sx_h = params["p_X_stride_hi"] + sx_w = params["p_X_stride_wi"] + + def seg_lane_offsets(n0: Value, q0: Value): + """Per-instruction (lane offset in elements, lane valid) for one strip.""" + dy = [] + for img, col, ch, okr in dy_lanes: + n_i = b.add(n0, img) + wo = b.add(q0, col) + ok = b.land(okr, b.land(b.cmp_lt(wo, p_Wo), b.cmp_lt(n_i, p_N))) + off = b.add(b.add(b.mul(n_i, sdy_n), b.mul(wo, sdy_w)), b.add(k_blk0, ch)) + dy.append((off, ok)) + xs = [] + for img, col, ch, okr in x_lanes: + n_i = b.add(n0, img) + wi = b.add(q0, b.sub(col, ci(PAD))) + ok = b.land( + okr, + b.land(b.land(b.cmp_ge(wi, c0), b.cmp_lt(wi, p_W)), b.cmp_lt(n_i, p_N)), + ) + off = b.add(b.add(b.mul(n_i, sx_n), b.mul(wi, sx_w)), b.add(c_glb0, ch)) + xs.append((off, ok)) + return dy, xs + + def issue_row(buf: int, t: Value, h0: Value, h_end: Value, lane_offs, t_end: Value): + """DMA row ``t`` of the segment into buffer ``buf``.""" + dy_offs, x_offs = lane_offs + hi = b.add(h0, b.sub(t, ci(T_PRO))) + ho = b.add(hi, ci(PAD)) + live = b.cmp_lt(t, t_end) + dy_ok = b.land(live, b.land(b.cmp_ge(ho, c0), b.cmp_lt(ho, p_Ho))) + x_ok = b.land( + live, + b.land( + b.cmp_ge(t, ci(T_PRO)), b.land(b.cmp_ge(hi, c0), b.cmp_lt(hi, h_end)) + ), + ) + dy_row = b.mul(ho, sdy_h) + x_row = b.mul(hi, sx_h) + for i, (off, ok) in enumerate(dy_offs): + voff = b.select(b.land(ok, dy_ok), b.mul(b.add(off, dy_row), c_two), oob) + lds = b.smem_ptr_add( + dy_addr[buf] if QG == 1 else b.smem_ptr_add(dy_addr[buf], qg_dy_bytes), + b.zext(b.mul(b.add(b.mul(wave, ci(IK)), ci(i)), ci(_WAVE * 16)), I64), + ) + b.async_buffer_load_lds(a_rsrc, lds, voff, c0, 4) + for i, (off, ok) in enumerate(x_offs): + voff = b.select(b.land(ok, x_ok), b.mul(b.add(off, x_row), c_two), oob) + lds = b.smem_ptr_add( + x_addr[buf] if QG == 1 else b.smem_ptr_add(x_addr[buf], qg_x_bytes), + b.zext(b.mul(b.add(b.mul(wave, ci(IC)), ci(i)), ci(_WAVE * 16)), I64), + ) + b.async_buffer_load_lds(b_rsrc, lds, voff, c0, 4) + + # ---- fragment read addresses (element index into a buffer) ------------- + # See _frag_row: a lane's 8 reduction elements are 8 columns of one image. + lane_hi = b.div(lane, ci(16)) + lane_q = b.mod(b.div(lane, ci(4)), ci(4)) + lane_c4 = b.mod(lane, ci(4)) + kk0 = b.mul(lane_hi, ci(8)) + if UN > 1: + img_l = b.div(kk0, ci(WU)) + jl = b.mod(kk0, ci(WU)) + else: + img_l, jl = c0, kk0 + + def frag_index( + shift: int, rd: int, col0: Value, chunks: int, bits, seg: int, g_elems: int + ) -> Value: + row = b.add(b.add(b.mul(img_l, ci(seg)), b.add(jl, lane_q)), ci(shift + rd * 4)) + col = b.add(col0, b.mul(lane_c4, ci(4))) + chunk = b.div(col, ci(_VEC)) + sw = swz_val(row, bits) + phys = chunk if sw is None else b.xor(chunk, sw) + idx = b.add( + b.mul(b.add(b.mul(row, ci(chunks)), phys), ci(_VEC)), b.mod(col, ci(_VEC)) + ) + if QG > 1: + idx = b.add(idx, b.mul(qg, ci(g_elems))) + return idx + + k_wave0 = b.mul(wave_k_id, ci(spec.wave_k)) + c_wave0 = b.mul(wave_c_id, ci(spec.wave_c)) + dy_idx = [ + [ + frag_index(0, rd, b.add(k_wave0, ci(mi * _ATOM)), CHK, swz_k, WU, GDY_ELEMS) + for rd in (0, 1) + ] + for mi in range(WK16) + ] + x_idx = [ + [ + [ + frag_index( + s, rd, b.add(c_wave0, ci(ni * _ATOM)), CHC, swz_c, SEG, GX_ELEMS + ) + for rd in (0, 1) + ] + for ni in range(WC16) + ] + for s in range(KW) + ] + + def read_frag(smem: Value, idx: List[Value]) -> Value: + lo = b.ds_read_tr16_b64(smem, c0, idx[0], dtype=io_type) + hi = b.ds_read_tr16_b64(smem, c0, idx[1], dtype=io_type) + return b.vec_concat(lo, hi) + + zero_acc = b.zero_vec_f32(4) + zero_frag = b.zero_vec(io_type, _VEC) + n_acc = KH * KW * WK16 * WC16 + + def acc_index(r, s, mi, ni): + return ((r * KW + s) * WK16 + mi) * WC16 + ni + + def step(buf: int, t: Value, ring, accs, compute: bool, issue): + """One row: wait, barrier, DMA row t + D, read, MFMA.""" + b.s_waitcnt(vmcnt=(D - 1) * N_INSTR, lgkmcnt=0) + b.s_barrier_bare() + issue((buf + D) % NB, b.add(t, ci(D))) + newest = [read_frag(dy_bufs[buf], dy_idx[mi]) for mi in range(WK16)] + ring = ring[1:] + [newest] + if not compute: + return ring, accs + xf = [ + [read_frag(x_bufs[buf], x_idx[s][ni]) for ni in range(WC16)] + for s in range(KW) + ] + if spec.pingpong: + # The other group's half-step: its reads run under these MFMAs, + # which take issue priority while they do. + b.s_waitcnt(lgkmcnt=0) + b.s_barrier_bare() + b.s_setprio(1) + accs = list(accs) + for r in range(KH): + dyf = ring[KH - 1 - r] + for s in range(KW): + for mi in range(WK16): + for ni in range(WC16): + i = acc_index(r, s, mi, ni) + accs[i] = _mfma( + b, io_dtype, "16x16x32", dyf[mi], xf[s][ni], accs[i] + ) + if spec.pingpong: + b.s_setprio(0) + return ring, accs + + # ---- persistent segment loop ------------------------------------------- + def segment(accs, n0, q, h0, h_end, valid=None): + """Rows [h0, h_end) of strip q of images n0 .. n0 + unfold_n - 1. + + ``valid`` false runs it with every load masked (a wave group with no + segment left keeps step with the other). + """ + lane_offs = seg_lane_offsets(n0, b.mul(q, ci(WU))) + t_end = ci(T_PRO + HPB) + if valid is not None: + t_end = b.select(valid, t_end, c0) + + def issue(buf, t): + issue_row(buf, t, h0, h_end, lane_offs, t_end) + + for t in range(D): + issue(t % NB, ci(t)) + + ring = [[zero_frag] * WK16 for _ in range(KH)] + for t in range(T_PRO): + ring, accs = step(t % NB, ci(t), ring, accs, False, issue) + + n_iter_args = [(f"a{i}", v) for i, v in enumerate(accs)] + n_iter_args += [ + (f"r{j}_{mi}", ring[j][mi]) for j in range(1, KH) for mi in range(WK16) + ] + row_loop = b.scf_for_iter( + c0, + ci(HPB // U), + c1, + n_iter_args, + iv_name="rb", + elide_trailing_barrier=False, + ) + with row_loop as (rb, rc): + accs_i = list(rc[:n_acc]) + ring_i = [[zero_frag] * WK16] + [ + list(rc[n_acc + (j - 1) * WK16 : n_acc + j * WK16]) + for j in range(1, KH) + ] + t_base = b.add(b.mul(rb, ci(U)), ci(T_PRO)) + for u in range(U): + ring_i, accs_i = step( + (T_PRO + u) % NB, b.add(t_base, ci(u)), ring_i, accs_i, True, issue + ) + b.scf_yield(*accs_i, *[v for j in range(1, KH) for v in ring_i[j]]) + accs_out = list(row_loop.results[:n_acc]) + # Drain the dummy tail DMAs before the next segment reuses the buffers. + b.s_waitcnt(vmcnt=0, lgkmcnt=0) + b.s_barrier_bare() + return accs_out + + if spec.pingpong: + # Group 1 starts one barrier late: from here on its every barrier + # pairs with group 0's next one, half a row ahead of it. + with b.scf_if(b.cmp_eq(qg, c1)): + b.s_barrier_bare() + zero_accs = [(f"acc{i}", zero_acc) for i in range(n_acc)] + if QG == 1: + seg_loop = b.scf_for_iter( + split, + total_segs, + n_splits, + zero_accs, + iv_name="seg", + elide_trailing_barrier=False, + ) + with seg_loop as (seg, carried): + hb = b.mod(seg, n_hb) + rest = b.div(seg, n_hb) + h0 = b.mul(hb, ci(HPB)) + accs_out = segment( + list(carried), + b.mul(b.div(rest, n_q), ci(UN)), + b.mod(rest, n_q), + h0, + b.smin(b.add(h0, ci(HPB)), p_H), + ) + b.scf_yield(*accs_out) + else: + # Group g of block `split` takes virtual split split * QG + g; both + # groups run group 0's segment count (the larger) in lockstep. + n_vsplits = b.mul(n_splits, ci(QG)) + vsplit0 = b.mul(split, ci(QG)) + n_iter = b.div( + b.add(b.sub(total_segs, vsplit0), b.sub(n_vsplits, c1)), n_vsplits + ) + seg_loop = b.scf_for_iter( + c0, n_iter, c1, zero_accs, iv_name="segi", elide_trailing_barrier=False + ) + with seg_loop as (segi, carried): + seg = b.add(b.add(vsplit0, qg), b.mul(segi, n_vsplits)) + hb = b.mod(seg, n_hb) + rest = b.div(seg, n_hb) + h0 = b.mul(hb, ci(HPB)) + accs_out = segment( + list(carried), + b.mul(b.div(rest, n_q), ci(UN)), + b.mod(rest, n_q), + h0, + b.smin(b.add(h0, ci(HPB)), p_H), + valid=b.cmp_lt(seg, total_segs), + ) + b.scf_yield(*accs_out) + accs = list(seg_loop.results) + if spec.pingpong: + with b.scf_if(b.cmp_eq(qg, c0)): + b.s_barrier_bare() + + # ---- epilogue: atomic add into dW[total_k, KH, KW, cpg] ----------------- + D_ptr = params["D"] + rs_stride = c_cpg + k_stride = b.mul(c_cpg, ci(KH * KW)) + k_lane0 = b.add(b.add(k_blk0, k_wave0), b.mul(lane_hi, ci(4))) + k_in_grp0 = b.sub(k_lane0, b.mul(group, c_kpg)) + c_lane0 = b.add(b.add(c_blk0, c_wave0), b.mod(lane, ci(16))) + if QG > 1: + # Tap t is summed and written by group t % 2; the other group hands + # its accumulators over through LDS, double-buffered on the tap + # parity, one barrier per tap. + n_atoms = WK16 * WC16 + red = [ + b.smem_alloc_f32([1, NW * n_atoms * _WAVE * 4], name_hint=f"red{i}") + for i in range(2) + ] + + def red_index(atom: int) -> Value: + return b.mul( + b.add( + b.mul(b.add(b.mul(wave, ci(n_atoms)), ci(atom)), ci(_WAVE)), lane + ), + ci(4), + ) + + for r in range(KH): + for s in range(KW): + tap = r * KW + s + buf = red[tap % 2] + is_recv = b.cmp_eq(qg, ci(tap % 2)) + with b.scf_if(b.lnot(is_recv)): + for mi in range(WK16): + for ni in range(WC16): + b.smem_store_vN_f32( + buf, + [c0, red_index(mi * WC16 + ni)], + accs[acc_index(r, s, mi, ni)], + 4, + ) + b.s_waitcnt(lgkmcnt=0) + b.s_barrier_bare() + with b.scf_if(is_recv): + for ni in range(WC16): + c_in = b.add(c_lane0, ci(ni * _ATOM)) + c_ok = b.cmp_lt(c_in, c_cpg) + for mi in range(WK16): + tot = b.fadd( + accs[acc_index(r, s, mi, ni)], + b.smem_load_vN_f32( + buf, c0, red_index(mi * WC16 + ni), n=4 + ), + ) + for slot in range(4): + k_in = b.add(k_in_grp0, ci(mi * _ATOM + slot)) + k_glb = b.add(k_lane0, ci(mi * _ATOM + slot)) + ok = b.land(c_ok, b.cmp_lt(k_in, c_kpg)) + off = b.add( + b.add(b.mul(k_glb, k_stride), c_in), + b.mul(ci(tap), rs_stride), + ) + with b.scf_if(ok): + b.global_atomic_add( + D_ptr, off, b.vec_extract(tot, slot) + ) + return b.kernel + for ni in range(WC16): + c_in = b.add(c_lane0, ci(ni * _ATOM)) + c_ok = b.cmp_lt(c_in, c_cpg) + for mi in range(WK16): + for slot in range(4): + k_in = b.add(k_in_grp0, ci(mi * _ATOM + slot)) + k_glb = b.add(k_lane0, ci(mi * _ATOM + slot)) + ok = b.land(c_ok, b.cmp_lt(k_in, c_kpg)) + base = b.add(b.mul(k_glb, k_stride), c_in) + for r in range(KH): + for s in range(KW): + val = b.vec_extract(accs[acc_index(r, s, mi, ni)], slot) + off = b.add(base, b.mul(ci(r * KW + s), rs_stride)) + with b.scf_if(ok): + b.global_atomic_add(D_ptr, off, val) + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_aot_jobs.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_aot_jobs.py index 8834cac70de7..20b172860e6c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_aot_jobs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_aot_jobs.py @@ -25,6 +25,9 @@ DIRECT_HALO_VARIANTS as _HALO_VARIANTS, DIRECT_MFMA_DGRAD as _MFMA_DGRAD, DIRECT_SPLIT_REDUCE as _SPLIT_REDUCE, + DIRECT_WGRAD_FOLD as _WGRAD_FOLD, + DIRECT_WGRAD_FOLDED as _WGRAD_FOLDED, + direct_wgrad_fold_knobs, direct_dgrad_weight_transforms, direct_mfma_knobs_reason as _mfma_knobs_reason, make_direct_spec as make_spec, @@ -136,6 +139,39 @@ WGRAD_HPB = (30, 60, 120) WGRAD_MFMA_K = (32,) +# Tiled direct wgrad (DirectWgradTiledSpec): (wave_k16, wave_c16, waves_k, +# waves_c) arrangements, (prefetch, row_unroll) pipelines, rows per segment, +# workgroups launched, images packed per strip. Measured on gfx950 bf16 over +# 3x3 C = K = 128 layers: 32x32 wave tiles in four-wave blocks (two per CU) +# lead; eight-wave blocks are kept for the shapes that fill few channel tiles; +# unfold_n = 4 is what a 40-wide image needs (Wo % 32 != 0). +WGRAD_TILED_WAVES = ((2, 2, 2, 2), (2, 2, 4, 2), (2, 1, 2, 2)) +WGRAD_TILED_PIPE = ((2, 0), (5, 6)) +WGRAD_TILED_RPS = (15, 30) +WGRAD_TILED_BLOCKS = (256, 512) +WGRAD_TILED_UNFOLD = (1, 4) +# Eight-wave blocks of two ping-ponged wave groups, on +# the plain two-row prefetch of four-wave 32x32 tiles. +WGRAD_TILED_PINGPONG = ((2, 2, 2, 2), (2, 0), (30,), (256,), (1,)) +# Depthwise LDS-ring wgrad (DwWgradLdsSpec): lanes along columns, prefetch, +# output rows per segment, workgroups. +WGRAD_DW_LANES_W = (8, 16, 32) +WGRAD_DW_PREFETCH = (1, 2) +WGRAD_DW_RPS = (16, 32, 64) +WGRAD_DW_BLOCKS = (1024, 2048) +# Small-C depthwise (DwWgradColsSpec): channel counts it bakes, strip widths, +# threads, workgroups. +WGRAD_DW_COLS_CHANNELS = (1, 2, 3, 4, 8) +WGRAD_DW_COLS_STRIP = (8, 16) +WGRAD_DW_COLS_THREADS = (128, 256) +WGRAD_DW_COLS_BLOCKS = (2048, 4096) +# Dense small-C gather wgrad (WgradGatherSpec): (wave_k16, waves_k), prefetch, +# output rows per segment, workgroups. +WGRAD_GATHER_WAVES = ((2, 1), (1, 3), (2, 3), (1, 6), (2, 2)) +WGRAD_GATHER_PREFETCH = (1, 2) +WGRAD_GATHER_RPS = (8, 16) +WGRAD_GATHER_BLOCKS = (256, 1024, 2048) + # Runtime-channel MFMA dgrad (direct_dgrad_mfma_rtc). Its compute pass is the # forward kernel on the transposed problem, whose K reduction runs over the # dgrad's kpg (dY channels) and whose M-tiles cover its cpg (dX channels). @@ -225,6 +261,8 @@ def _grouped_channels() -> Tuple[int, ...]: # Wgrad filters the cache carries. The kernel takes any KH, KW <= 8; a 1x1 # wgrad is a plain GEMM with no halo to reuse, so it is left to implicit GEMM. _WGRAD_FILTERS = (3, 5, 7) +# Input channel counts the dense small-C gather wgrad bakes (KW * C <= 96). +_GATHER_CHANNELS = (3, 4, 8, 16, 24, 32) # Dilated 3x3 rows of the non-grouped kernel: (PAD, stride, dil_h, dil_w). The # kernel bounds its rows and columns by the runtime output extents, so any # padding works; the rows bake the "same" padding of the symmetric dilations @@ -390,6 +428,45 @@ def _col_pads(kh: int) -> Tuple[int, ...]: "dgrad", (DirectCaps(KH=3, PAD=1, stride=1, cpg=0, kpg=0),), ), + # Tiled wgrad: dense and grouped, channel counts runtime, stride 1. + "direct_wgrad_tiled": ( + "wgrad", + tuple( + DirectCaps(KH=k, PAD=pad, stride=1, cpg=0, kpg=0) + for k in (3,) + for pad in _col_pads(k) + ), + ), + # Depthwise wgrad through LDS rings: 3x3 (8 channels x 9 taps of fp32 + # partial sums per lane), stride 1 and 2, C runtime (a multiple of 8). + "direct_wgrad_dw": ( + "wgrad", + tuple( + DirectCaps(KH=3, PAD=pad, stride=s, cpg=1, kpg=1) + for s in _STRIDES + for pad in _col_pads(3) + ), + ), + # Small-C depthwise wgrad (stem layers): every depthwise filter; the + # channel count is a knob (``channels``). + "direct_wgrad_dw_cols": ( + "wgrad", + tuple( + DirectCaps(KH=k, PAD=pad, stride=s, cpg=1, kpg=1) + for k, s in itertools.product(_DEPTHWISE_FILTERS, _STRIDES) + for pad in _col_pads(k) + ), + ), + # Dense small-C wgrad, any stride: C is baked, K runtime. + "direct_wgrad_gather": ( + "wgrad", + tuple( + DirectCaps(KH=3, PAD=pad, stride=s, cpg=c, kpg=0) + for c in _GATHER_CHANNELS + for s in _STRIDES + for pad in _col_pads(3) + ), + ), # Wgrad (grouped and non-grouped), channel counts runtime. The kernel is a # stride-1 algorithm; it bounds its rows and columns against the runtime # extents, so it serves "valid" padding as well as "same". @@ -523,6 +600,69 @@ def _knob_grid( k_atoms=ka, m_tiles=mt, ) + elif variant == "direct_wgrad_tiled": + for (wk16, wc16, wk, wc), (pf, ru), rps, tb, un in itertools.product( + WGRAD_TILED_WAVES, + WGRAD_TILED_PIPE, + WGRAD_TILED_RPS, + WGRAD_TILED_BLOCKS, + WGRAD_TILED_UNFOLD, + ): + yield dict( + wave_k16=wk16, + wave_c16=wc16, + waves_k=wk, + waves_c=wc, + prefetch=pf, + row_unroll=ru, + rows_per_seg=rps, + target_blocks=tb, + unfold_n=un, + ) + (wk16, wc16, wk, wc), (pf, ru), rpss, tbs, uns = WGRAD_TILED_PINGPONG + for rps, tb, un in itertools.product(rpss, tbs, uns): + yield dict( + wave_k16=wk16, + wave_c16=wc16, + waves_k=wk, + waves_c=wc, + prefetch=pf, + row_unroll=ru, + rows_per_seg=rps, + target_blocks=tb, + unfold_n=un, + q_groups=2, + pingpong=True, + ) + elif variant == "direct_wgrad_dw": + for lw, pf, rps, tb in itertools.product( + WGRAD_DW_LANES_W, WGRAD_DW_PREFETCH, WGRAD_DW_RPS, WGRAD_DW_BLOCKS + ): + yield dict( + lanes_c=8, lanes_w=lw, prefetch=pf, rows_per_seg=rps, target_blocks=tb + ) + elif variant == "direct_wgrad_dw_cols": + for c, st, th, tb in itertools.product( + WGRAD_DW_COLS_CHANNELS, + WGRAD_DW_COLS_STRIP, + WGRAD_DW_COLS_THREADS, + WGRAD_DW_COLS_BLOCKS, + ): + yield dict(channels=c, strip=st, threads=th, target_blocks=tb) + elif variant == "direct_wgrad_gather": + for (wk16, wk), pf, rps, tb in itertools.product( + WGRAD_GATHER_WAVES, + WGRAD_GATHER_PREFETCH, + WGRAD_GATHER_RPS, + WGRAD_GATHER_BLOCKS, + ): + yield dict( + wave_k16=wk16, + waves_k=wk, + prefetch=pf, + rows_per_seg=rps, + target_blocks=tb, + ) elif variant == "direct_wgrad": for (wk, wc, wq), hpb, mk in itertools.product( wgrad_waves(), WGRAD_HPB, WGRAD_MFMA_K @@ -612,7 +752,18 @@ def _helper_identity( ) -> KernelIdentity: # The weight transforms depend on the filter and channel counts only; # stride and padding are zeroed so every pipeline entry that shares a - # filter shares the helper binary. + # filter shares the helper binary. The wgrad fold depends on its knobs + # alone (the dW extent it folds), so it zeroes the channel counts too. + if variant == _WGRAD_FOLD: + return _identity( + arch, + wave_size, + variant, + "direct_wgrad_helper", + replace(caps, PAD=0, stride=0, cpg=0, kpg=0), + knobs, + dtype, + ) return _identity( arch, wave_size, @@ -664,6 +815,21 @@ def _caps_of(identity: KernelIdentity) -> DirectCaps: ) +def _dw_cols_probe_width(caps: DirectCaps, channels: int, strip: int) -> int: + """A width the small-C depthwise wgrad takes at ``caps``, or 64 if none. + + The kernel needs Wo to be a multiple of its strip and both W * C and + Wo * C to be whole 16-byte vectors. W is a runtime extent, so any width + that passes probes the same binary; a fixed 64 fails every one of them + at PAD 0 (Wo = 62 at 3x3), which left those capability rows unbuilt. + """ + for w in range(64, 64 + 16 * strip * caps.stride): + wo = (w + 2 * caps.PAD - caps.KH) // caps.stride + 1 + if (w * channels) % 8 == 0 and (wo * channels) % 8 == 0 and wo % strip == 0: + return w + return 64 + + def probe_problem( caps: DirectCaps, dtype: str = DIRECT_DTYPE, @@ -691,6 +857,36 @@ def probe_problem( """ from kernels.common.conv_direct_grouped import DirectConvProblem + if variant == "direct_wgrad_dw_cols": + # The binary bakes C (knob ``channels``): one group per channel. + return DirectConvProblem( + N=1, + H=64, + W=_dw_cols_probe_width(caps, knobs["channels"], knobs["strip"]), + groups=knobs["channels"], + cpg=1, + kpg=1, + KH=caps.KH, + KW=caps.KW, + PAD=caps.PAD, + stride=caps.stride, + dtype=dtype, + ) + if variant == "direct_wgrad_gather": + # Dense, C baked, K runtime: a K wider than every block_k. + return DirectConvProblem( + N=1, + H=64, + W=64, + groups=1, + cpg=caps.cpg, + kpg=96, + KH=caps.KH, + KW=caps.KW, + PAD=caps.PAD, + stride=caps.stride, + dtype=dtype, + ) if variant == "direct_dgrad_mfma_rtc": atom = 32 if knobs["fold_k32"] else 16 return DirectConvProblem( @@ -802,14 +998,16 @@ def _caps_jobs( out: List[Tuple[str, BuildJob]] = [] n_raw = 0 rtc_dgrad = variant == "direct_dgrad_mfma_rtc" - # The runtime-channel dgrad probe follows each knob's tile counts. - probe = None if rtc_dgrad else probe_problem(caps, dtype, variant) + # The runtime-channel dgrad probe follows each knob's tile counts, the + # small-C depthwise wgrad's its baked channel count. + knob_probe = rtc_dgrad or variant == "direct_wgrad_dw_cols" + probe = None if knob_probe else probe_problem(caps, dtype, variant) for knobs in _knob_grid(variant, arch, dtype, caps): n_raw += 1 - if rtc_dgrad: + if knob_probe: probe = probe_problem(caps, dtype, variant, knobs) - if dgrad_rtc_spills(knobs, caps.KH, caps.KW, arch): - continue + if rtc_dgrad and dgrad_rtc_spills(knobs, caps.KH, caps.KW, arch): + continue if variant in _MFMA_DGRAD and _mfma_knobs_reason(probe, knobs): continue spec = make_spec(variant, probe, knobs) @@ -828,6 +1026,17 @@ def _caps_jobs( knobs, ) ] + if variant in _WGRAD_FOLDED: + fold_knobs = direct_wgrad_fold_knobs(variant, caps.KH, caps.KW, caps.cpg) + cell.append( + _job( + _helper_identity( + arch, wave_size, _WGRAD_FOLD, caps, fold_knobs, dtype + ), + replace(caps, PAD=0, stride=0, cpg=0, kpg=0), + fold_knobs, + ) + ) if variant == "direct_nongrouped" and knobs.get("split_c", 1) > 1: cell.append(_split_reduce_job(arch, wave_size, knobs["split_c"], dtype)) if variant in _MFMA_DGRAD: @@ -921,7 +1130,7 @@ def direct_job_for_identity(identity: KernelIdentity) -> BuildJob: f"the direct grid cannot build this identity: {identity.label()}" ) return _job(identity, _NO_CAPS, {"split_c": split_c}) - if identity.direction == "direct_dgrad_helper": + if identity.direction in ("direct_dgrad_helper", "direct_wgrad_helper"): rebuilt = _helper_identity( identity.arch, identity.wave_size, @@ -945,7 +1154,7 @@ def direct_job_for_identity(identity: KernelIdentity) -> BuildJob: want, got = identity.to_dict(), rebuilt.to_dict() diff = {k: want[k] for k in want if want[k] != got[k]} raise LookupError(f"the direct grid cannot build this identity: {diff}") - if identity.direction != "direct_dgrad_helper": + if identity.direction not in ("direct_dgrad_helper", "direct_wgrad_helper"): # The fields can be consistent and still name a kernel no variant # computes (a dilated grouped kernel, say): ask the kernel's validator. variant = identity.algorithm diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_variants.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_variants.py index c28388939425..3aca472670d3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_variants.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_variants.py @@ -17,6 +17,8 @@ from dataclasses import replace __all__ = [ + "DIRECT_WGRAD_FOLD", + "DIRECT_WGRAD_FOLDED", "DIRECT_DGRAD_REORGANIZE", "DIRECT_DGRAD_REORGANIZE_RTC", "DIRECT_DGRAD_TRANSPOSE", @@ -30,6 +32,7 @@ "direct_dgrad_workspaces", "direct_mfma_knobs_reason", "direct_spec_for_problem", + "direct_wgrad_fold_knobs", "direct_uses_coalesced_weights", "make_direct_spec", "select_direct_kernel_builder", @@ -58,11 +61,21 @@ "direct_grouped_dgrad_mfma": (DIRECT_DGRAD_TRANSPOSE, DIRECT_DGRAD_REORGANIZE), "direct_dgrad_mfma_rtc": (DIRECT_DGRAD_TRANSPOSE_RTC, DIRECT_DGRAD_REORGANIZE_RTC), } +# Partial-slot fold of the wgrad kernels that store one slot per partition +# instead of atomically adding into dW; cached as its own entry, keyed by the +# number of dW elements per output channel it folds (``taps``). +DIRECT_WGRAD_FOLD = "direct_wgrad_fold" +DIRECT_WGRAD_FOLDED = ( + "direct_wgrad_dw", + "direct_wgrad_dw_cols", + "direct_wgrad_gather", +) _HELPERS = ( DIRECT_DGRAD_TRANSPOSE, DIRECT_DGRAD_REORGANIZE, DIRECT_DGRAD_TRANSPOSE_RTC, DIRECT_DGRAD_REORGANIZE_RTC, + DIRECT_WGRAD_FOLD, ) # Register budget for the MFMA dgrad fprop pass's preloaded weights and # accumulators (see direct_mfma_knobs_reason): three quarters of the 512-VGPR @@ -140,6 +153,36 @@ def make_direct_spec(variant: str, problem, knobs: dict): runtime_channels=True, **knobs, ) + if variant == "direct_wgrad_tiled": + from kernels.common.conv_direct_wgrad_tiled import DirectWgradTiledSpec + + return DirectWgradTiledSpec( + problem=problem, name="rocke_direct_wgrad_tiled", **knobs + ) + if variant == "direct_wgrad_dw": + from kernels.common.conv_direct_dw_wgrad import DwWgradLdsSpec + + return DwWgradLdsSpec(problem=problem, name="rocke_direct_wgrad_dw", **knobs) + if variant == "direct_wgrad_dw_cols": + from kernels.common.conv_direct_dw_wgrad import DwWgradColsSpec + + # ``channels`` names the C the binary bakes; the spec reads it off + # the problem, and direct_spec_for_problem holds the two equal. + return DwWgradColsSpec( + problem=problem, + name="rocke_direct_wgrad_dw_cols", + **{k: v for k, v in knobs.items() if k != "channels"}, + ) + if variant == "direct_wgrad_gather": + from kernels.common.conv_direct_wgrad_gather import WgradGatherSpec + + return WgradGatherSpec( + problem=problem, name="rocke_direct_wgrad_gather", **knobs + ) + if variant == DIRECT_WGRAD_FOLD: + from kernels.common.conv_direct_dw_wgrad import DwWgradFoldSpec + + return DwWgradFoldSpec(name="rocke_direct_wgrad_fold", **knobs) if variant == DIRECT_SPLIT_REDUCE: return dn.DirectSplitReduceSpec(dtype=problem.dtype, **knobs) if variant == DIRECT_DGRAD_TRANSPOSE: @@ -267,6 +310,24 @@ def validate_direct_spec(variant: str, spec, arch: str) -> tuple[bool, str]: if variant == "direct_wgrad": spec.validate() return dc.is_valid_wgrad_spec(spec, arch=arch) + if variant == "direct_wgrad_tiled": + from kernels.common.conv_direct_wgrad_tiled import is_valid_wgrad_tiled_spec + + return is_valid_wgrad_tiled_spec(spec, arch=arch) + if variant == "direct_wgrad_dw": + from kernels.common.conv_direct_dw_wgrad import is_valid_dw_wgrad_lds_spec + + return is_valid_dw_wgrad_lds_spec(spec, arch=arch) + if variant == "direct_wgrad_dw_cols": + from kernels.common.conv_direct_dw_wgrad import is_valid_dw_wgrad_cols_spec + + return is_valid_dw_wgrad_cols_spec(spec, arch=arch) + if variant == "direct_wgrad_gather": + from kernels.common.conv_direct_wgrad_gather import ( + is_valid_wgrad_gather_spec, + ) + + return is_valid_wgrad_gather_spec(spec, arch=arch) if variant in _HELPERS: return True, "ok" if variant == DIRECT_SPLIT_REDUCE: @@ -286,10 +347,18 @@ def select_direct_kernel_builder(variant: str, spec, arch: str): function for the ``(spec, arch)`` builder signature, and this one takes an extra ``variant``. """ + from kernels.common import conv_direct_dw_wgrad as dw from kernels.common import conv_direct_grouped as dc from kernels.common import conv_direct_nongrouped as dn + from kernels.common import conv_direct_wgrad_gather as wg + from kernels.common import conv_direct_wgrad_tiled as wt builders = { + "direct_wgrad_tiled": wt.build_direct_wgrad_tiled, + "direct_wgrad_dw": dw.build_dw_wgrad_lds, + "direct_wgrad_dw_cols": dw.build_dw_wgrad_cols, + "direct_wgrad_gather": wg.build_wgrad_gather, + DIRECT_WGRAD_FOLD: dw.build_dw_wgrad_fold, "direct_grouped": dc.build_direct_conv, "direct_nongrouped": dn.build_direct_conv_nongrouped, DIRECT_GROUPED_HALO: dn.build_direct_conv_nongrouped, @@ -329,6 +398,10 @@ def direct_spec_for_problem( why = direct_mfma_knobs_reason(problem, knobs) if why: return None, why + if variant == "direct_wgrad_dw_cols" and knobs["channels"] != problem.groups: + return None, ( + f"binary bakes C={knobs['channels']}, problem has C={problem.groups}" + ) spec = make_direct_spec(variant, problem, knobs) ok, why = validate_direct_spec(variant, spec, arch) if not ok: @@ -360,3 +433,15 @@ def direct_spec_for_problem( f"of a merged block" ) return spec, "ok" + + +def direct_wgrad_fold_knobs(variant: str, KH: int, KW: int, cpg: int) -> dict: + """Knobs of the fold a partial-slot wgrad ``variant`` runs after its pass. + + ``taps`` is the dW extent per output channel: ``KH * KW`` for depthwise, + ``KH * KW * C`` for the dense gather kernel (cpg baked). + """ + if variant not in DIRECT_WGRAD_FOLDED: + raise ValueError(f"{variant!r} does not fold partial slots") + taps = KH * KW * (cpg if variant == "direct_wgrad_gather" else 1) + return {"taps": taps} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c index 3dabf39f71b5..a56a659fde82 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c @@ -14,7 +14,9 @@ * parametric rocke_direct_conv_spec_t and the MFMA dgrad weight transforms, * baked and runtime-channel; from index 95 the split-C non-grouped spec and its * rocke_direct_conv_split_reduce_spec_t second pass; from index 103 the grouped / - * group-merge non-grouped spec) identically to the Python emitter + * group-merge non-grouped spec; from index 114 the direct wgrad kernels: + * rocke_direct_wgrad_tiled_spec_t, rocke_dw_wgrad_{fold,cols,lds}_spec_t and + * rocke_direct_wgrad_gather_spec_t) identically to the Python emitter * conv_direct_grouped_emit.py, builds the * kernel via the matching rocke_build_direct_conv_*_new function and lowers via * rocke_lower_kernel_to_llvm (per-config arch, flavor AUTO) and prints the .ll @@ -24,8 +26,11 @@ #include #include +#include "rocke/instance_conv_direct_dw_wgrad.h" #include "rocke/instance_conv_direct_grouped.h" #include "rocke/instance_conv_direct_nongrouped.h" +#include "rocke/instance_conv_direct_wgrad_gather.h" +#include "rocke/instance_conv_direct_wgrad_tiled.h" #include "rocke/ir.h" #include "rocke/ir_serialize.h" #include "rocke/lower_llvm.h" @@ -48,7 +53,12 @@ enum KIND_GENERIC = 12, KIND_TRANSPOSE = 13, KIND_REORG = 14, - KIND_SPLIT_REDUCE = 15 + KIND_SPLIT_REDUCE = 15, + KIND_WGRAD_TILED = 16, + KIND_DW_WGRAD_FOLD = 17, + KIND_DW_WGRAD_COLS = 18, + KIND_DW_WGRAD_LDS = 19, + KIND_WGRAD_GATHER = 20 }; #define NONGROUPED_CFG_BASE 42 @@ -68,6 +78,9 @@ enum /* Grouped (grouped = true, group_merge) non-grouped configs follow the 8 split-C * ones. */ #define NONGROUPED_GROUPED_CFG_BASE 103 +/* The direct wgrad kernels (tiled, depthwise fold / cols / lds, gather) follow + * the 11 grouped non-grouped configs. */ +#define WGRAD_DIRECT_CFG_BASE 114 static rocke_direct_conv_problem_t mk_dw_problem(int N, int H, int W, int groups, int KH, int PAD, int stride, const char* dtype) @@ -1535,6 +1548,314 @@ static int make_cfg(int idx, } } +/* ---- direct wgrad configs (see _wgrad_direct_spec in the Python emitter) ---- */ +typedef struct wgrad_direct_specs +{ + rocke_direct_wgrad_tiled_spec_t tiled; + rocke_dw_wgrad_fold_spec_t fold; + rocke_dw_wgrad_cols_spec_t cols; + rocke_dw_wgrad_lds_spec_t lds; + rocke_direct_wgrad_gather_spec_t gather; +} wgrad_direct_specs_t; + +/* _wp(): N=2 H=16 W=32 groups=1 cpg=kpg=64 3x3 PAD=1 stride=1 bf16. */ +static rocke_direct_conv_problem_t mk_wp(void) +{ + rocke_direct_conv_problem_t p = rocke_direct_conv_problem_default(); + p.N = 2; + p.H = 16; + p.W = 32; + p.groups = 1; + p.cpg = 64; + p.kpg = 64; + p.KH = 3; + p.KW = 3; + p.PAD = 1; + p.stride = 1; + p.dtype = "bf16"; + return p; +} + +/* _dwp(C): depthwise over C channels. */ +static rocke_direct_conv_problem_t mk_dwp(int C) +{ + rocke_direct_conv_problem_t p = mk_wp(); + p.groups = C; + p.cpg = 1; + p.kpg = 1; + return p; +} + +static void mk_filter(rocke_direct_conv_problem_t* p, int KH, int KW, int PAD, int stride) +{ + p->KH = KH; + p->KW = KW; + p->PAD = PAD; + p->stride = stride; +} + +/* Returns 0 on success, -1 if unknown. */ +static int make_wgrad_direct_cfg(int idx, int* kind, wgrad_direct_specs_t* w, const char** arch) +{ + rocke_direct_wgrad_tiled_spec_t* ts = &w->tiled; + rocke_dw_wgrad_cols_spec_t* cs = &w->cols; + rocke_dw_wgrad_lds_spec_t* ls = &w->lds; + rocke_direct_wgrad_gather_spec_t* gs = &w->gather; + *ts = rocke_direct_wgrad_tiled_spec_default(); + w->fold = rocke_dw_wgrad_fold_spec_default(); + *cs = rocke_dw_wgrad_cols_spec_default(); + *ls = rocke_dw_wgrad_lds_spec_default(); + *gs = rocke_direct_wgrad_gather_spec_default(); + *arch = "gfx950"; + if(idx <= 6 || idx == 20) + *kind = KIND_WGRAD_TILED; + else if(idx <= 8) + *kind = KIND_DW_WGRAD_FOLD; + else if(idx <= 11) + *kind = KIND_DW_WGRAD_COLS; + else if(idx <= 15) + *kind = KIND_DW_WGRAD_LDS; + else + *kind = KIND_WGRAD_GATHER; + ts->problem = mk_wp(); + gs->problem = mk_wp(); + switch(idx) + { + case 0: + ts->problem.cpg = 128; + ts->problem.kpg = 128; + ts->problem.H = 30; + return 0; + case 1: + ts->problem.groups = 2; + ts->problem.kpg = 32; + ts->problem.dtype = "fp16"; + ts->wave_k16 = 1; + ts->wave_c16 = 4; + ts->waves_k = 2; + ts->waves_c = 1; + ts->rows_per_seg = 12; + ts->prefetch = 1; + return 0; + case 2: + ts->problem.N = 4; + ts->problem.H = 15; + ts->problem.W = 40; + ts->unfold_n = 4; + ts->rows_per_seg = 15; + ts->row_unroll = 3; + return 0; + case 3: + ts->q_groups = 2; + ts->waves_k = 2; + ts->rows_per_seg = 6; + return 0; + case 4: + ts->q_groups = 2; + ts->pingpong = true; + ts->waves_k = 2; + ts->wave_k16 = 1; + ts->rows_per_seg = 12; + ts->prefetch = 3; + ts->target_blocks = 64; + return 0; + case 5: + mk_filter(&ts->problem, 1, 1, 0, 1); + ts->problem.cpg = 32; + ts->problem.kpg = 32; + ts->problem.dtype = "fp16"; + ts->waves_k = 1; + ts->waves_c = 1; + ts->rows_per_seg = 6; + ts->unfold_n = 2; + return 0; + case 6: + mk_filter(&ts->problem, 5, 5, 2, 1); + ts->problem.cpg = 48; + ts->problem.kpg = 40; + ts->wave_k16 = 1; + ts->wave_c16 = 1; + ts->waves_k = 2; + ts->waves_c = 2; + ts->rows_per_seg = 15; + ts->unfold_n = 2; + return 0; + case 7: + w->fold.taps = 9; + return 0; + case 8: + w->fold.taps = 27; + w->fold.name = "fold27"; + return 0; + case 9: + cs->problem = mk_dwp(3); + mk_filter(&cs->problem, 7, 7, 3, 2); + cs->problem.H = 20; + return 0; + case 10: + cs->problem = mk_dwp(3); + mk_filter(&cs->problem, 11, 11, 5, 1); + cs->problem.H = 12; + cs->problem.W = 48; + cs->problem.dtype = "fp16"; + cs->strip = 16; + cs->threads = 256; + cs->target_blocks = 4096; + return 0; + case 11: + cs->problem = mk_dwp(8); + cs->problem.H = 10; + cs->problem.W = 24; + cs->threads = 128; + return 0; + case 12: + ls->problem = mk_dwp(64); + return 0; + case 13: + ls->problem = mk_dwp(192); + ls->problem.stride = 2; + ls->problem.H = 30; + ls->problem.W = 40; + ls->problem.dtype = "fp16"; + ls->lanes_w = 8; + ls->prefetch = 1; + ls->target_blocks = 2048; + return 0; + case 14: + ls->problem = mk_dwp(40); + ls->problem.H = 20; + ls->problem.W = 20; + ls->lanes_c = 4; + ls->lanes_w = 16; + ls->cols_per_lane = 2; + ls->prefetch = 3; + ls->rows_per_seg = 8; + return 0; + case 15: + ls->problem = mk_dwp(16); + mk_filter(&ls->problem, 3, 1, 0, 1); + ls->problem.H = 9; + ls->problem.W = 12; + ls->lanes_c = 2; + ls->rows_per_seg = 5; + return 0; + case 16: /* gathered X (S * C = 6), the stem shape */ + gs->problem.cpg = 3; + gs->problem.kpg = 24; + gs->problem.stride = 2; + gs->problem.H = 32; + gs->problem.W = 64; + return 0; + case 17: /* transpose-read X, phase 0 */ + gs->problem.cpg = 24; + gs->problem.kpg = 96; + gs->problem.stride = 2; + gs->problem.H = 30; + gs->problem.W = 40; + gs->wave_k16 = 1; + gs->waves_k = 6; + gs->rows_per_seg = 8; + gs->target_blocks = 256; + return 0; + case 18: /* transpose-read X at phase 4 */ + gs->problem.cpg = 4; + gs->problem.kpg = 32; + gs->problem.dtype = "fp16"; + gs->prefetch = 1; + return 0; + case 19: /* 7x7 stride 2: five leading X rows per segment, two k-waves */ + gs->problem.cpg = 3; + gs->problem.kpg = 64; + mk_filter(&gs->problem, 7, 7, 3, 2); + gs->problem.H = 40; + gs->problem.W = 64; + gs->waves_k = 2; + gs->rows_per_seg = 8; + return 0; + case 20: /* must reject: the tiled kernel is stride 1 only */ + ts->problem.stride = 2; + return 0; + case 21: /* must reject: no ds_read_b64_tr_b16 on gfx942 */ + gs->problem.cpg = 3; + gs->problem.kpg = 24; + gs->problem.stride = 2; + gs->problem.W = 64; + *arch = "gfx942"; + return 0; + case 22: /* KW * C = 72 outputs over 64 threads: a second epilogue pass */ + *kind = KIND_DW_WGRAD_COLS; + cs->problem = mk_dwp(8); + mk_filter(&cs->problem, 9, 9, 4, 1); + return 0; + case 23: /* must reject: a 512 x 88 fp32 reduction buffer is 176 KiB of LDS */ + *kind = KIND_DW_WGRAD_COLS; + cs->problem = mk_dwp(8); + mk_filter(&cs->problem, 11, 11, 5, 1); + cs->threads = 512; + return 0; + case 24: /* must reject: the padded rings are 200,704 B (logical rows 145,152) */ + *kind = KIND_DW_WGRAD_LDS; + ls->problem = mk_dwp(64); + ls->problem.stride = 2; + ls->prefetch = 5; + return 0; + case 25: /* must reject: gfx942 DMAs one dword per lane into LDS */ + case 26: /* must reject: gfx1250 is wave32 */ + *kind = KIND_DW_WGRAD_LDS; + ls->problem = mk_dwp(64); + *arch = idx == 25 ? "gfx942" : "gfx1250"; + return 0; + case 27: /* must reject: sixteen waves pad each X slot to 16 KiB (196,864 B) */ + *kind = KIND_WGRAD_GATHER; + gs->problem.cpg = 3; + gs->problem.stride = 2; + gs->wave_k16 = 1; + gs->waves_k = 16; + return 0; + case 28: /* must reject: six 32 KiB ring buffers */ + *kind = KIND_WGRAD_TILED; + ts->problem.cpg = 256; + ts->problem.kpg = 256; + ts->waves_k = 8; + ts->waves_c = 2; + ts->prefetch = 5; + ts->row_unroll = 6; + return 0; + case 29: /* must reject: row_unroll -3 divides by 3 but runs no rows */ + *kind = KIND_WGRAD_TILED; + ts->row_unroll = -3; + return 0; + case 30: + *kind = KIND_WGRAD_TILED; + ts->target_blocks = 0; + return 0; + case 31: + *kind = KIND_WGRAD_TILED; + ts->problem.dtype = NULL; + return 0; + case 32: /* valid padding, three k-waves of one atom (two k tiles) */ + *kind = KIND_WGRAD_GATHER; + gs->problem.cpg = 8; + gs->problem.kpg = 96; + gs->problem.PAD = 0; + gs->wave_k16 = 1; + gs->waves_k = 3; + return 0; + case 33: /* valid padding at stride 2, three k-waves of two atoms (a partial tile) */ + *kind = KIND_WGRAD_GATHER; + gs->problem.cpg = 3; + gs->problem.kpg = 104; + gs->problem.PAD = 0; + gs->problem.stride = 2; + gs->problem.W = 64; + gs->waves_k = 3; + gs->prefetch = 1; + return 0; + default: + return -1; + } +} + int main(int argc, char** argv) { if(argc < 2) @@ -1561,9 +1882,18 @@ int main(int argc, char** argv) rocke_direct_conv_spec_t sgen; rocke_direct_transpose_weights_dgrad_spec_t stp; rocke_direct_reorganize_weights_spec_t sreorg; + wgrad_direct_specs_t swd; rocke_direct_conv_split_reduce_spec_t sred; const char* arch = "gfx950"; - if(idx >= NONGROUPED_GROUPED_CFG_BASE) + if(idx >= WGRAD_DIRECT_CFG_BASE) + { + if(make_wgrad_direct_cfg(idx - WGRAD_DIRECT_CFG_BASE, &kind, &swd, &arch) != 0) + { + fprintf(stderr, "unknown config index %d\n", idx); + return 2; + } + } + else if(idx >= NONGROUPED_GROUPED_CFG_BASE) { kind = KIND_NONGROUPED; if(make_nongrouped_grouped_cfg(idx - NONGROUPED_GROUPED_CFG_BASE, &snongrouped, &arch) != 0) @@ -1638,6 +1968,16 @@ int main(int argc, char** argv) kernel = rocke_build_direct_reorganize_weights_new(&b, &sreorg, arch); else if(kind == KIND_DW_TILED) kernel = rocke_build_direct_depthwise_tiled_new(&b, &stiled, arch); + else if(kind == KIND_WGRAD_TILED) + kernel = rocke_build_direct_wgrad_tiled_new(&b, &swd.tiled, arch); + else if(kind == KIND_DW_WGRAD_FOLD) + kernel = rocke_build_dw_wgrad_fold_new(&b, &swd.fold, arch); + else if(kind == KIND_DW_WGRAD_COLS) + kernel = rocke_build_dw_wgrad_cols_new(&b, &swd.cols, arch); + else if(kind == KIND_DW_WGRAD_LDS) + kernel = rocke_build_dw_wgrad_lds_new(&b, &swd.lds, arch); + else if(kind == KIND_WGRAD_GATHER) + kernel = rocke_build_direct_wgrad_gather_new(&b, &swd.gather, arch); else if(kind == KIND_SPLIT_REDUCE) kernel = rocke_build_direct_conv_split_reduce_new(&b, &sred, arch); else diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py index eed9933626ba..e5188d388efa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py @@ -14,8 +14,10 @@ # must reject, from index 71 the runtime-channel DirectConvWgradSpec and from # index 75 the parametric DirectConvSpec and the MFMA dgrad weight transforms, # baked and runtime-channel, from index 95 the split-C non-grouped -# DirectNongroupedConvSpec and its DirectSplitReduceSpec second pass, and from -# index 103 its grouped / group-merge mode), builds the kernel +# DirectNongroupedConvSpec and its DirectSplitReduceSpec second pass, from +# index 103 its grouped / group-merge mode, and from index 114 the direct +# wgrad kernels: DirectWgradTiledSpec, DwWgradFoldSpec, DwWgradColsSpec, +# DwWgradLdsSpec and WgradGatherSpec), builds the kernel # via the matching build_direct_conv_* function (arch=) and prints # _native_lower(arch=) to stdout so it can be byte-compared with # the C emitter conv_direct_grouped_emit.c. @@ -58,6 +60,22 @@ build_direct_conv_nongrouped, build_direct_conv_split_reduce, ) +from kernels.common.conv_direct_dw_wgrad import ( + DwWgradColsSpec, + DwWgradFoldSpec, + DwWgradLdsSpec, + build_dw_wgrad_cols, + build_dw_wgrad_fold, + build_dw_wgrad_lds, +) +from kernels.common.conv_direct_wgrad_gather import ( + WgradGatherSpec, + build_wgrad_gather, +) +from kernels.common.conv_direct_wgrad_tiled import ( + DirectWgradTiledSpec, + build_direct_wgrad_tiled, +) try: from rocke.core.lower_llvm import _lower_kernel_to_llvm_python as _native_lower @@ -102,6 +120,242 @@ _NONGROUPED_GROUPED_CFG_BASE = 103 +# --------------------------------------------------------------------------- +# The direct wgrad kernels, indices 114..147: DirectWgradTiledSpec (wave tiles, +# unfold_n, row_unroll, q_groups with and without pingpong, 1x1 and 5x5), the +# DwWgradFoldSpec partial-slot fold, DwWgradColsSpec (small C), DwWgradLdsSpec +# (row rings, strided, cols_per_lane, a partial channel tile), WgradGatherSpec +# (gathered and transpose-read X, with and without a phase), then two specs +# both engines must reject; a small-C depthwise spec with more outputs than +# threads; then eight more rejections (LDS footprints over budget, targets +# without 16-byte LDS DMA, a negative row_unroll, no blocks, no dtype); then +# WgradGatherSpec at valid padding over three k-waves. Follows the 11 grouped +# non-grouped configs. +# --------------------------------------------------------------------------- +_WGRAD_DIRECT_CFG_BASE = 114 + + +def _wp(**over): + kw = dict( + N=2, + H=16, + W=32, + groups=1, + cpg=64, + kpg=64, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + kw.update(over) + return DirectConvProblem(**kw) + + +def _dwp(C, **over): + return _wp(groups=C, cpg=1, kpg=1, **over) + + +def _wgrad_direct_spec(idx: int): + """Return (kind, spec, arch) for config ``_WGRAD_DIRECT_CFG_BASE + idx``.""" + ts, ls, cs, gs = ( + DirectWgradTiledSpec, + DwWgradLdsSpec, + DwWgradColsSpec, + WgradGatherSpec, + ) + table = { + 0: ("wgrad_tiled", ts(problem=_wp(cpg=128, kpg=128, H=30))), + 1: ( + "wgrad_tiled", + ts( + problem=_wp(groups=2, cpg=64, kpg=32, dtype="fp16"), + wave_k16=1, + wave_c16=4, + waves_k=2, + waves_c=1, + rows_per_seg=12, + prefetch=1, + ), + ), + 2: ( + "wgrad_tiled", + ts( + problem=_wp(N=4, H=15, W=40), + unfold_n=4, + rows_per_seg=15, + row_unroll=3, + ), + ), + 3: ("wgrad_tiled", ts(problem=_wp(), q_groups=2, waves_k=2, rows_per_seg=6)), + 4: ( + "wgrad_tiled", + ts( + problem=_wp(), + q_groups=2, + pingpong=True, + waves_k=2, + wave_k16=1, + rows_per_seg=12, + prefetch=3, + target_blocks=64, + ), + ), + 5: ( + "wgrad_tiled", + ts( + problem=_wp(KH=1, KW=1, PAD=0, cpg=32, kpg=32, dtype="fp16"), + waves_k=1, + waves_c=1, + rows_per_seg=6, + unfold_n=2, + ), + ), + 6: ( + "wgrad_tiled", + ts( + problem=_wp(KH=5, KW=5, PAD=2, cpg=48, kpg=40), + wave_k16=1, + wave_c16=1, + waves_k=2, + waves_c=2, + rows_per_seg=15, + unfold_n=2, + ), + ), + 7: ("dw_wgrad_fold", DwWgradFoldSpec(taps=9)), + 8: ("dw_wgrad_fold", DwWgradFoldSpec(taps=27, name="fold27")), + 9: ( + "dw_wgrad_cols", + cs(problem=_dwp(3, KH=7, KW=7, PAD=3, stride=2, H=20)), + ), + 10: ( + "dw_wgrad_cols", + cs( + problem=_dwp(3, KH=11, KW=11, PAD=5, H=12, W=48, dtype="fp16"), + strip=16, + threads=256, + target_blocks=4096, + ), + ), + 11: ( + "dw_wgrad_cols", + cs(problem=_dwp(8, H=10, W=24), threads=128), + ), + 12: ("dw_wgrad_lds", ls(problem=_dwp(64))), + 13: ( + "dw_wgrad_lds", + ls( + problem=_dwp(192, stride=2, H=30, W=40, dtype="fp16"), + lanes_w=8, + prefetch=1, + target_blocks=2048, + ), + ), + 14: ( + "dw_wgrad_lds", + ls( + problem=_dwp(40, H=20, W=20), + lanes_c=4, + lanes_w=16, + cols_per_lane=2, + prefetch=3, + rows_per_seg=8, + ), + ), + 15: ( + "dw_wgrad_lds", + ls( + problem=_dwp(16, KH=3, KW=1, PAD=0, H=9, W=12), + lanes_c=2, + rows_per_seg=5, + ), + ), + # Gathered X (S * C = 6), the stem shape. + 16: ("wgrad_gather", gs(problem=_wp(cpg=3, kpg=24, stride=2, H=32, W=64))), + # Transpose-read X, phase 0. + 17: ( + "wgrad_gather", + gs( + problem=_wp(cpg=24, kpg=96, stride=2, H=30, W=40), + wave_k16=1, + waves_k=6, + rows_per_seg=8, + target_blocks=256, + ), + ), + # Transpose-read X at phase 4. + 18: ( + "wgrad_gather", + gs(problem=_wp(cpg=4, kpg=32, dtype="fp16"), prefetch=1), + ), + # 7x7 stride 2: five leading X rows per segment, two k-waves. + 19: ( + "wgrad_gather", + gs( + problem=_wp(cpg=3, kpg=64, KH=7, KW=7, PAD=3, stride=2, H=40, W=64), + waves_k=2, + rows_per_seg=8, + ), + ), + # Must reject: the tiled kernel is stride 1 only. + 20: ("wgrad_tiled", ts(problem=_wp(stride=2))), + # Must reject: no ds_read_b64_tr_b16 on gfx942. + 21: ("wgrad_gather", gs(problem=_wp(cpg=3, kpg=24, stride=2, W=64)), "gfx942"), + # KW * C = 72 outputs over 64 threads: a second epilogue pass. + 22: ("dw_wgrad_cols", cs(problem=_dwp(8, KH=9, KW=9, PAD=4))), + # Must reject: a 512 x 88 fp32 reduction buffer is 176 KiB of LDS. + 23: ( + "dw_wgrad_cols", + cs(problem=_dwp(8, KH=11, KW=11, PAD=5), threads=512), + ), + # Must reject: the padded rings are 200,704 B (logical rows 145,152). + 24: ("dw_wgrad_lds", ls(problem=_dwp(64, stride=2), prefetch=5)), + # Must reject: gfx942 DMAs one dword per lane into LDS, gfx1250 is wave32. + 25: ("dw_wgrad_lds", ls(problem=_dwp(64)), "gfx942"), + 26: ("dw_wgrad_lds", ls(problem=_dwp(64)), "gfx1250"), + # Must reject: sixteen waves pad each X slot to 16 KiB (196,864 B). + 27: ( + "wgrad_gather", + gs(problem=_wp(cpg=3, kpg=64, stride=2), wave_k16=1, waves_k=16), + ), + # Must reject: six 32 KiB ring buffers. + 28: ( + "wgrad_tiled", + ts( + problem=_wp(cpg=256, kpg=256), + waves_k=8, + waves_c=2, + prefetch=5, + row_unroll=6, + ), + ), + # Must reject: row_unroll -3 divides by 3 but runs no rows. + 29: ("wgrad_tiled", ts(problem=_wp(), row_unroll=-3)), + 30: ("wgrad_tiled", ts(problem=_wp(), target_blocks=0)), + 31: ("wgrad_tiled", ts(problem=_wp(dtype=None))), + # Valid padding, three k-waves of one atom (two k tiles). + 32: ( + "wgrad_gather", + gs(problem=_wp(cpg=8, kpg=96, PAD=0), wave_k16=1, waves_k=3), + ), + # Valid padding at stride 2, three k-waves of two atoms (a partial tile). + 33: ( + "wgrad_gather", + gs( + problem=_wp(cpg=3, kpg=104, PAD=0, stride=2, W=64), + waves_k=3, + prefetch=1, + ), + ), + } + sel = table.get(idx) + if sel is None: + return None + return (sel[0], sel[1], sel[2] if len(sel) > 2 else "gfx950") + + def _gp(groups, cpg, kpg, *, N=2, H=20, W=24, KH=3, PAD=1, dtype="fp16"): return DirectConvProblem( N=N, @@ -1391,6 +1645,11 @@ def _spec(idx: int): DirectDepthwiseSpec(problem=p, block_w=8, block_waves=2), "gfx950", ) + if idx >= _WGRAD_DIRECT_CFG_BASE: + sel = _wgrad_direct_spec(idx - _WGRAD_DIRECT_CFG_BASE) + if sel is not None: + return sel + raise SystemExit(f"unknown config index {idx}") if idx >= _NONGROUPED_GROUPED_CFG_BASE: sel = _nongrouped_grouped_spec(idx - _NONGROUPED_GROUPED_CFG_BASE) if sel is not None: @@ -1540,6 +1799,16 @@ def main() -> int: kernel = build_direct_reorganize_weights(spec, arch=arch) elif kind == "dw_tiled": kernel = build_direct_depthwise_tiled(spec, arch=arch) + elif kind == "wgrad_tiled": + kernel = build_direct_wgrad_tiled(spec, arch=arch) + elif kind == "dw_wgrad_fold": + kernel = build_dw_wgrad_fold(spec, arch=arch) + elif kind == "dw_wgrad_cols": + kernel = build_dw_wgrad_cols(spec, arch=arch) + elif kind == "dw_wgrad_lds": + kernel = build_dw_wgrad_lds(spec, arch=arch) + elif kind == "wgrad_gather": + kernel = build_wgrad_gather(spec, arch=arch) elif kind == "split_reduce": kernel = build_direct_conv_split_reduce(spec, arch=arch) else: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py index f7d878b78cb3..3876542c0651 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py @@ -340,6 +340,12 @@ "kernels.common.conv_direct_grouped.DirectMfmaDgradSpec": ("problem",), "kernels.common.conv_direct_grouped.DirectConvWgradSpec": ("problem",), "kernels.common.conv_direct_nongrouped.DirectNongroupedConvSpec": ("problem",), + "kernels.common.conv_direct_wgrad_tiled.DirectWgradTiledSpec": ("problem",), + "kernels.common.conv_direct_dw_wgrad.DwWgradLdsSpec": ("problem",), + "kernels.common.conv_direct_dw_wgrad.DwWgradColsSpec": ("problem",), + # The dW extent per output channel the fold sums partial slots over. + "kernels.common.conv_direct_dw_wgrad.DwWgradFoldSpec": ("taps",), + "kernels.common.conv_direct_wgrad_gather.WgradGatherSpec": ("problem",), "kernels.common.conv_direct_nongrouped.DirectSplitReduceSpec": (), "kernels.common.deep_fused_conv_pool.FusedConvPoolProblem": ("conv",), "kernels.common.deep_fused_conv_pool.DeepFusedConvPoolSpec": ("problem",), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py index 4c170c39f136..6cbc09936dcc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py @@ -3077,5 +3077,304 @@ def test_mfma_dgrad_runtime_channels(self): self._check(label, knobs, fold_k32, shapes, dtype, True) +# --------------------------------------------------------------------------- +# The AOT sweep's direct wgrad kernels: conv_direct_wgrad_tiled (fp32 atomics), +# conv_direct_dw_wgrad (depthwise, partial slots) and conv_direct_wgrad_gather +# (dense small C, partial slots). +# --------------------------------------------------------------------------- +# +# Each case is built the way the kernel cache builds it -- on the probe problem +# of its capability row -- and launched on the problem listed here, so every +# extent the binary takes as a kernarg (and the host/kernel agreement on the +# partition and split counts) is exercised. The partial-slot kernels run the +# fold after the pass, as the sweep's two-step plan does. + +# (id, variant, DirectConvProblem fields, knobs) +_AOT_WGRAD_CASES = [ + ( + "tiled_c64", + "direct_wgrad_tiled", + dict(N=2, H=16, W=32, groups=1, cpg=64, kpg=64), + dict(rows_per_seg=30), + ), + # Valid padding, a partial wave tile in K and C, grouped. + ( + "tiled_g2_pad0", + "direct_wgrad_tiled", + dict(N=2, H=13, W=37, groups=2, cpg=48, kpg=40, PAD=0), + dict(wave_k16=1, wave_c16=1, waves_k=2, waves_c=2, rows_per_seg=12), + ), + # Four images packed into one 32-column strip of a narrow map. + ( + "tiled_unfold4", + "direct_wgrad_tiled", + dict(N=6, H=15, W=8, groups=1, cpg=32, kpg=32), + dict(waves_k=1, waves_c=1, rows_per_seg=15, unfold_n=4), + ), + # Two wave groups, the second half a row behind (pingpong); segment counts + # that leave the second group a dummy segment. + ( + "tiled_qg2_pp", + "direct_wgrad_tiled", + dict(N=3, H=9, W=40, groups=1, cpg=32, kpg=64), + dict( + wave_k16=1, + waves_k=2, + rows_per_seg=12, + prefetch=3, + q_groups=2, + pingpong=True, + target_blocks=8, + ), + ), + ( + "dw_c64", + "direct_wgrad_dw", + dict(N=2, H=17, W=33, groups=64, cpg=1, kpg=1), + dict(lanes_c=8, lanes_w=8, prefetch=2, rows_per_seg=16, target_blocks=1024), + ), + # Stride 2 and a channel count that leaves a partial channel tile. + ( + "dw_c40_s2", + "direct_wgrad_dw", + dict(N=2, H=20, W=24, groups=40, cpg=1, kpg=1, stride=2), + dict(lanes_c=8, lanes_w=16, prefetch=1, rows_per_seg=16, target_blocks=2048), + ), + # The stem: C = 3, 7x7, stride 2. + ( + "dw_cols_c3_7x7_s2", + "direct_wgrad_dw_cols", + dict(N=2, H=20, W=64, groups=3, cpg=1, kpg=1, KH=7, KW=7, PAD=3, stride=2), + dict(channels=3, strip=8, threads=128, target_blocks=2048), + ), + # Valid padding (the capability rows the probe width makes buildable). + ( + "dw_cols_c4_pad0", + "direct_wgrad_dw_cols", + dict(N=2, H=12, W=34, groups=4, cpg=1, kpg=1, PAD=0), + dict(channels=4, strip=16, threads=256, target_blocks=4096), + ), + # Gathered X operand (S * C = 6). + ( + "gather_c3_s2", + "direct_wgrad_gather", + dict(N=2, H=24, W=64, groups=1, cpg=3, kpg=24, stride=2), + dict(wave_k16=2, waves_k=1, prefetch=2, rows_per_seg=8, target_blocks=256), + ), + # Transpose-read X operand. + ( + "gather_c24_s2", + "direct_wgrad_gather", + dict(N=2, H=30, W=40, groups=1, cpg=24, kpg=96, stride=2), + dict(wave_k16=1, waves_k=6, prefetch=2, rows_per_seg=8, target_blocks=256), + ), + # Stride 1 with a K that leaves a partial k tile. + ( + "gather_c4_s1", + "direct_wgrad_gather", + dict(N=1, H=16, W=32, groups=1, cpg=4, kpg=40, stride=1), + dict(wave_k16=2, waves_k=1, prefetch=1, rows_per_seg=16, target_blocks=1024), + ), + # Valid padding, three k-waves of one atom each over two k tiles. + ( + "gather_c8_pad0_w1x3", + "direct_wgrad_gather", + dict(N=2, H=20, W=32, groups=1, cpg=8, kpg=96, PAD=0, stride=1), + dict(wave_k16=1, waves_k=3, prefetch=2, rows_per_seg=8, target_blocks=256), + ), + # Valid padding at stride 2, three k-waves of two atoms, a partial k tile. + ( + "gather_c3_pad0_w2x3_s2", + "direct_wgrad_gather", + dict(N=2, H=21, W=64, groups=1, cpg=3, kpg=104, PAD=0, stride=2), + dict(wave_k16=2, waves_k=3, prefetch=1, rows_per_seg=16, target_blocks=1024), + ), + # The cases above give every block one segment. Two blocks walking several + # segments each (ten, two): accumulation across segments and ring reuse + # after the end-of-segment barrier. + ( + "dw_c64_tb2", + "direct_wgrad_dw", + dict(N=2, H=40, W=40, groups=64, cpg=1, kpg=1), + dict(lanes_c=8, lanes_w=8, prefetch=2, rows_per_seg=16, target_blocks=2), + ), + ( + "gather_c3_s2_tb2", + "direct_wgrad_gather", + dict(N=2, H=24, W=64, groups=1, cpg=3, kpg=24, stride=2), + dict(wave_k16=2, waves_k=1, prefetch=2, rows_per_seg=8, target_blocks=2), + ), + # 128 partitions: the fold loops over its 64-partition stride twice. + ( + "dw_c64_parts128", + "direct_wgrad_dw", + dict(N=4, H=64, W=64, groups=64, cpg=1, kpg=1), + dict(lanes_c=8, lanes_w=8, prefetch=2, rows_per_seg=16, target_blocks=1024), + ), + # KW * C = 72 outputs over 64 threads (a second epilogue pass), one block + # per filter row walking two units. + ( + "dw_cols_c8_9x9_t64", + "direct_wgrad_dw_cols", + dict(N=2, H=16, W=32, groups=8, cpg=1, kpg=1, KH=9, KW=9, PAD=4), + dict(channels=8, strip=8, threads=64, target_blocks=9), + ), +] + + +def _run_aot_wgrad_one(arch: str, variant: str, problem, knobs: dict) -> float: + """Build ``variant`` on its cache probe, launch it on ``problem``, return + the max error relative to the largest reference dW entry.""" + import torch + from kernels.common.conv_abi import conv_direct_args_signature + from kernels.common.conv_args import ConvArgs + from kernels.common.conv_direct_dw_wgrad import ( + DW_WGRAD_FOLD_ARGS, + dw_wgrad_fold_signature, + ) + from kernels.conv.direct_aot_jobs import DirectCaps, probe_problem + from kernels.conv.direct_variants import ( + DIRECT_WGRAD_FOLD, + DIRECT_WGRAD_FOLDED, + direct_wgrad_fold_knobs, + make_direct_spec, + select_direct_kernel_builder, + validate_direct_spec, + ) + from rocke import compile_kernel + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + p = problem + folded = variant in DIRECT_WGRAD_FOLDED + if variant == "direct_wgrad_tiled": + caps = DirectCaps(KH=p.KH, PAD=p.PAD, stride=p.stride, cpg=0, kpg=0) + elif variant == "direct_wgrad_gather": + caps = DirectCaps(KH=p.KH, PAD=p.PAD, stride=p.stride, cpg=p.cpg, kpg=0) + else: + caps = DirectCaps(KH=p.KH, PAD=p.PAD, stride=p.stride, cpg=1, kpg=1) + probe = probe_problem(caps, p.dtype, variant, knobs) + build_spec = make_direct_spec(variant, probe, knobs) + spec = make_direct_spec(variant, p, knobs) + for s, what in ((build_spec, "probe"), (spec, "problem")): + ok, why = validate_direct_spec(variant, s, arch) + assert ok, f"{variant} rejects its {what}: {why}" + assert build_spec.kernel_name() == spec.kernel_name() + art = compile_kernel( + select_direct_kernel_builder(variant, build_spec, arch), arch=arch + ) + + torch.manual_seed(7) + td = torch.bfloat16 if p.dtype == "bf16" else torch.float16 + total_c, total_k = p.groups * p.cpg, p.groups * p.kpg + if folded and variant != "direct_wgrad_gather": + total_k = total_c = p.groups # depthwise + X = torch.empty(p.N, p.H, p.W, total_c, dtype=td).uniform_(-0.5, 0.5) + dY = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=td).uniform_(-0.5, 0.5) + dw_elems = total_k * p.KH * p.KW * p.cpg + + W_ref = torch.zeros( + total_k, p.cpg, p.KH, p.KW, dtype=torch.float32, device="cuda" + ).requires_grad_(True) + out = torch.nn.functional.conv2d( + X.float().cuda().permute(0, 3, 1, 2), + W_ref, + padding=p.PAD, + stride=p.stride, + groups=p.groups, + ) + out.backward(dY.float().cuda().permute(0, 3, 1, 2)) + ref = W_ref.grad.permute(0, 2, 3, 1).contiguous().cpu() # [K, KH, KW, cpg] + + rt = Runtime() + bufs = { + "x": (rt.alloc(X.nbytes), X.nbytes), + "dy": (rt.alloc(dY.nbytes), dY.nbytes), + "dw": (rt.alloc(dw_elems * 4), dw_elems * 4), + } + if folded: + ws_bytes = spec.workspace_bytes() + bufs["ws"] = (rt.alloc(ws_bytes), ws_bytes) + try: + rt.memcpy_h2d(bufs["x"][0], _u8(X), X.nbytes) + rt.memcpy_h2d(bufs["dy"][0], _u8(dY), dY.nbytes) + # The atomic kernel adds into a zeroed dW; the fold writes dW whole, + # so poison it to catch an element the fold misses. + rt.memset(bufs["dw"][0], 0 if not folded else 0xFF, bufs["dw"][1]) + out_buf = "ws" if folded else "dw" + grid, block = spec.launch_geometry() + launcher = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=conv_direct_args_signature(p.dtype, direction="wgrad"), + ) + values = ConvArgs.from_problem(p, direction="wgrad").to_launch_values( + int(bufs["dy"][0]), + int(bufs["x"][0]), + int(bufs[out_buf][0]), + bufs["dy"][1], + bufs["x"][1], + bufs[out_buf][1], + ) + launcher(values, config=LaunchConfig(grid=grid, block=block, fence=True)) + if folded: + fold_knobs = direct_wgrad_fold_knobs(variant, p.KH, p.KW, p.cpg) + fold_spec = make_direct_spec(DIRECT_WGRAD_FOLD, p, fold_knobs) + fold_art = compile_kernel( + select_direct_kernel_builder(DIRECT_WGRAD_FOLD, fold_spec, arch), + arch=arch, + ) + channels = ( + spec.fold_channels() if variant == "direct_wgrad_gather" else p.groups + ) + fgrid, fblock = fold_spec.launch_geometry(channels) + fold = KernelLauncher( + hsaco=fold_art.hsaco, + kernel_name=fold_art.kernel_name, + signature=dw_wgrad_fold_signature(fold_spec), + ) + fvalues = { + "A": int(bufs["ws"][0]), + "D": int(bufs["dw"][0]), + "A_bytes": bufs["ws"][1], + "D_bytes": bufs["dw"][1], + **dict(zip(DW_WGRAD_FOLD_ARGS, (spec.n_partitions(), channels))), + } + fold(fvalues, config=LaunchConfig(grid=fgrid, block=fblock, fence=True)) + dW = torch.empty(dw_elems, dtype=torch.float32) + rt.memcpy_d2h(_u8(dW), bufs["dw"][0], bufs["dw"][1]) + finally: + for ptr, _ in bufs.values(): + rt.free(ptr) + synchronize_and_release(0) + + got = dW.view(total_k, p.KH, p.KW, p.cpg) + scale = ref.abs().max().clamp(min=1.0) + return float(((got - ref).abs().max() / scale).nan_to_num(float("inf"))) + + +@unittest.skipUnless(not _WGRAD_SKIP_REASON, _WGRAD_SKIP_REASON or "no GPU") +class TestDirectWgradAotKernelsCorrectness(unittest.TestCase): + """The tiled, depthwise and gather direct wgrad kernels of the AOT sweep.""" + + def _check(self, dtype: str) -> None: + from kernels.common.conv_direct_grouped import DirectConvProblem + + tol = _TOL_BF16 if dtype == "bf16" else _TOL + for case_id, variant, shape, knobs in _AOT_WGRAD_CASES: + problem = DirectConvProblem(dtype=dtype, **shape) + with self.subTest(case=case_id, dtype=dtype): + rel_err = _run_aot_wgrad_one(GPU_ARCH, variant, problem, knobs) + print(f" {case_id:24s} {dtype} rel_err={rel_err:.2e}", flush=True) + self.assertLess(rel_err, tol, f"{case_id} {problem.short()} {dtype}") + + def test_aot_wgrad_fp16(self): + self._check("fp16") + + def test_aot_wgrad_bf16(self): + self._check("bf16") + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py index c113a458d52b..0bbe1b4825f4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py @@ -833,3 +833,76 @@ def test_dgrad_rtc_rejects_what_no_binary_fits(cache, groups, cpg, kpg, why): assert not _rtc(plans) reasons = [r for i, r in rejected if i.algorithm == "direct_dgrad_mfma_rtc"] assert reasons and any(why in r for r in reasons), reasons[:3] + + +@pytest.mark.parametrize( + "variant,problem", + [ + ( + "direct_wgrad_dw", + _problem(groups=192, cpg=1, kpg=1, stride=2, H=30, W=40), + ), + ( + "direct_wgrad_dw_cols", + _problem(groups=3, cpg=1, kpg=1, KH=7, KW=7, PAD=3, stride=2, H=40, W=64), + ), + # A PAD 0 row: built only since the probe picks a width the kernel takes. + ("direct_wgrad_dw_cols", _problem(groups=4, cpg=1, kpg=1, PAD=0, W=34)), + ( + "direct_wgrad_gather", + _problem(groups=1, cpg=3, kpg=24, stride=2, H=32, W=64), + ), + ( + "direct_wgrad_gather", + _problem(groups=1, cpg=24, kpg=96, stride=2, H=30, W=40), + ), + ], +) +def test_folded_wgrad_plans_run_the_pass_then_the_fold(cache, variant, problem): + """A partial-slot wgrad is a two-step plan: the pass writes one slot per + partition into the scratch ``ws`` and the fold sums them into dW. Both + steps, and the scratch size, must agree on the partition and channel + counts of the problem actually launched.""" + plans, _ = dks.direct_plans(cache, problem, "wgrad", _ARCH) + plans = [p for p in plans if p.identity.algorithm == variant] + assert plans + names = ("dy", "x", "dw", "ws") + ptrs = {n: 0x1000 * (i + 1) for i, n in enumerate(names)} + sizes = {n: 1 << 20 for n in names} + for plan in plans: + main, fold = plan.steps + assert main.buffers == ("dy", "x", "ws") + assert fold.buffers == ("ws", None, "dw") + spec = dks.make_spec(variant, problem, json.loads(plan.identity.knobs)) + assert (main.grid, main.block) == spec.launch_geometry() + if variant == "direct_wgrad_gather": + channels, taps = spec.fold_channels(), spec.fold_taps() + else: + channels, taps = problem.groups, problem.KH * problem.KW + assert dict(fold.scalars) == { + "parts": spec.n_partitions(), + "channels": channels, + } + assert ( + plan.workspaces + == {"ws": spec.n_partitions() * taps * channels * 4} + == {"ws": spec.workspace_bytes()} + ) + # The fold's 64-output blocks cover every element of dW. + assert fold.grid[0] * 64 >= channels * taps + for step in plan.steps: + compile_packer(step.signature)(dks.launch_values(step, ptrs, sizes)) + + +def test_small_c_depthwise_wgrad_binary_bakes_its_channel_count(cache): + """The small-C kernel bakes C, so a problem is only offered the binaries + built for its own channel count.""" + problem = _problem(groups=3, cpg=1, kpg=1, KH=7, KW=7, PAD=3, stride=2, W=64) + plans, rejected = dks.direct_plans(cache, problem, "wgrad", _ARCH) + cols = [p for p in plans if p.identity.algorithm == "direct_wgrad_dw_cols"] + assert cols + assert all(json.loads(p.identity.knobs)["channels"] == 3 for p in cols) + assert any( + i.algorithm == "direct_wgrad_dw_cols" and "bakes C=" in why + for i, why in rejected + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_wgrad_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_wgrad_specs.py new file mode 100644 index 000000000000..70c5e46468dd --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_wgrad_specs.py @@ -0,0 +1,232 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Validator gates of the direct wgrad kernels (CPU only, no GPU, no torch). + +Every gate has a twin in the C++ engine; parity configs 117..128 of +``tests/parity/conv_direct_grouped_emit.py`` pin that both engines build or +reject the same specs. +""" + +import pytest + +from kernels.common.conv_direct_dw_wgrad import ( + DwWgradColsSpec, + DwWgradLdsSpec, + build_dw_wgrad_cols, + build_dw_wgrad_lds, + is_valid_dw_wgrad_cols_spec, + is_valid_dw_wgrad_lds_spec, +) +from kernels.common.conv_direct_grouped import DirectConvProblem +from kernels.common.conv_direct_wgrad_gather import ( + WgradGatherSpec, + build_wgrad_gather, + is_valid_wgrad_gather_spec, +) +from kernels.common.conv_direct_wgrad_tiled import ( + DirectWgradTiledSpec, + build_direct_wgrad_tiled, + is_valid_wgrad_tiled_spec, +) +from rocke.core.ir import type_bits + +_LDS_BUDGET = 160 * 1024 + + +def _p(**over) -> DirectConvProblem: + kw = dict(N=2, H=16, W=32, groups=1, cpg=64, kpg=64, dtype="bf16") + kw.update(over) + return DirectConvProblem(**kw) + + +def _dw(C: int, **over) -> DirectConvProblem: + return _p(groups=C, cpg=1, kpg=1, **over) + + +def _walk(ops): + for op in ops: + yield op + for region in op.regions: + yield from _walk(region.ops) + + +def _lds_allocated(kernel) -> int: + total = 0 + for op in _walk(kernel.body.ops): + if op.name == "tile.smem_alloc": + t = op.results[0].type + n = 1 + for d in t.shape: + n *= d + total += n * type_bits(t.elem) // 8 + return total + + +_BUILT = [ + (build_dw_wgrad_cols, DwWgradColsSpec(problem=_dw(3, KH=7, KW=7, PAD=3, W=64))), + ( + build_dw_wgrad_cols, + DwWgradColsSpec(problem=_dw(8, KH=9, KW=9, PAD=4), threads=64), + ), + (build_dw_wgrad_lds, DwWgradLdsSpec(problem=_dw(64))), + ( + build_dw_wgrad_lds, + DwWgradLdsSpec(problem=_dw(64, stride=2, W=40), lanes_w=8, prefetch=1), + ), + ( + build_dw_wgrad_lds, + DwWgradLdsSpec(problem=_dw(40), lanes_c=4, lanes_w=16, cols_per_lane=2), + ), + (build_wgrad_gather, WgradGatherSpec(problem=_p(cpg=3, kpg=24, stride=2))), + ( + build_wgrad_gather, + WgradGatherSpec(problem=_p(cpg=24, kpg=96, stride=2), wave_k16=1, waves_k=6), + ), + (build_wgrad_gather, WgradGatherSpec(problem=_p(cpg=4, kpg=32), prefetch=1)), + (build_direct_wgrad_tiled, DirectWgradTiledSpec(problem=_p())), + ( + build_direct_wgrad_tiled, + DirectWgradTiledSpec(problem=_p(), q_groups=2, waves_k=2, rows_per_seg=6), + ), + ( + build_direct_wgrad_tiled, + DirectWgradTiledSpec( + problem=_p(N=4, W=8, cpg=32, kpg=32), + waves_k=1, + waves_c=1, + rows_per_seg=15, + unfold_n=4, + ), + ), +] + + +@pytest.mark.parametrize("build,spec", _BUILT, ids=lambda x: getattr(x, "name", "")) +def test_lds_bytes_is_what_the_builder_allocates(build, spec): + """The validators gate on lds_bytes(); it has to count the padded DMA + slots and the reduction buffers the builder actually allocates.""" + assert spec.lds_bytes() == _lds_allocated(build(spec)) + + +@pytest.mark.parametrize( + "spec,valid,rings", + [ + # 3x3 stride 2, 8 x 32 lanes, prefetch 5: the logical rows are 145,152 B, + # the whole-wave DMA slots 200,704 B (plus the 8 KiB reduction). + ( + DwWgradLdsSpec(problem=_dw(64, stride=2), prefetch=5), + is_valid_dw_wgrad_lds_spec, + 200_704, + ), + # C = 3 stride 2 over sixteen k-waves: 68,736 B logical, 196,864 B allocated. + ( + WgradGatherSpec( + problem=_p(cpg=3, kpg=64, stride=2), wave_k16=1, waves_k=16 + ), + is_valid_wgrad_gather_spec, + 196_864, + ), + # Sixteen waves of 256 x 256 channels, six ring buffers of 32 KiB. + ( + DirectWgradTiledSpec( + problem=_p(cpg=256, kpg=256), + waves_k=8, + waves_c=2, + prefetch=5, + row_unroll=6, + ), + is_valid_wgrad_tiled_spec, + 196_608, + ), + ], +) +def test_padded_lds_footprint_over_budget_is_rejected(spec, valid, rings): + assert spec.lds_bytes() >= rings > _LDS_BUDGET + ok, why = valid(spec) + assert not ok and "LDS" in why + + +def test_small_c_reduction_buffer_over_budget_is_rejected(): + spec = DwWgradColsSpec(problem=_dw(8, KH=11, KW=11, PAD=5), threads=512) + assert spec.lds_bytes() == 512 * 88 * 4 + ok, why = is_valid_dw_wgrad_cols_spec(spec) + assert not ok and "LDS" in why + + +def test_small_c_epilogue_covers_more_outputs_than_threads(): + """KW * C = 72 accumulators over 64 threads: two epilogue passes.""" + spec = DwWgradColsSpec(problem=_dw(8, KH=9, KW=9, PAD=4), threads=64) + assert is_valid_dw_wgrad_cols_spec(spec)[0] + stores = [ + op + for op in _walk(build_dw_wgrad_cols(spec).body.ops) + if op.name.startswith("memref.global_store") + ] + assert len(stores) == 2 + + +_LDS_DMA_SPECS = [ + (DwWgradLdsSpec(problem=_dw(64)), is_valid_dw_wgrad_lds_spec), + (WgradGatherSpec(problem=_p(cpg=3, kpg=24, stride=2)), is_valid_wgrad_gather_spec), + (DirectWgradTiledSpec(problem=_p()), is_valid_wgrad_tiled_spec), +] + + +@pytest.mark.parametrize("spec,valid", _LDS_DMA_SPECS) +@pytest.mark.parametrize("arch", ["gfx942", "gfx1151", "gfx1201", "gfx1250"]) +def test_lds_dma_kernels_need_wave64_and_16_byte_dma(spec, valid, arch): + """The rings are filled by 16-byte buffer_load ... lds from 64-lane waves: + gfx942 DMAs one dword per lane, the RDNA targets are wave32.""" + assert valid(spec, "gfx950")[0] + ok, why = valid(spec, arch) + assert not ok and arch in why + + +@pytest.mark.parametrize("spec,valid", _LDS_DMA_SPECS) +def test_unknown_target_is_rejected(spec, valid): + ok, why = valid(spec, "gfx000") + assert not ok and "unknown gfx target" in why + + +@pytest.mark.parametrize( + "spec,valid", + [ + ( + DwWgradColsSpec(problem=_dw(3, W=64), target_blocks=0), + is_valid_dw_wgrad_cols_spec, + ), + (DwWgradLdsSpec(problem=_dw(64), target_blocks=0), is_valid_dw_wgrad_lds_spec), + ( + WgradGatherSpec(problem=_p(cpg=3, kpg=24, stride=2), target_blocks=0), + is_valid_wgrad_gather_spec, + ), + ( + DirectWgradTiledSpec(problem=_p(), target_blocks=0), + is_valid_wgrad_tiled_spec, + ), + # -3 % 3 == 0 and 30 % -3 == 0: without the sign check the row loop + # runs no iterations and dW stays zero. + (DirectWgradTiledSpec(problem=_p(), row_unroll=-3), is_valid_wgrad_tiled_spec), + ], +) +def test_degenerate_knobs_are_rejected(spec, valid): + assert not valid(spec)[0] + + +@pytest.mark.parametrize( + "spec,valid", + [ + ( + DwWgradColsSpec(problem=_dw(3, W=64, dtype=None)), + is_valid_dw_wgrad_cols_spec, + ), + (DwWgradLdsSpec(problem=_dw(64, dtype=None)), is_valid_dw_wgrad_lds_spec), + ( + WgradGatherSpec(problem=_p(cpg=3, kpg=24, stride=2, dtype=None)), + is_valid_wgrad_gather_spec, + ), + (DirectWgradTiledSpec(problem=_p(dtype=None)), is_valid_wgrad_tiled_spec), + ], +) +def test_missing_dtype_is_rejected(spec, valid): + assert not valid(spec)[0] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp index 9a34fdf2459b..5ee3e0d761db 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp @@ -1477,6 +1477,10 @@ void rocke_ll_compute_smem_layout(rocke_lower_t* L) int* offsets = (int*)rocke_arena_calloc(&L->arena, n * sizeof(int)); if(!offsets) rocke_ll_fail(L, ROCKE_ERR_OOM, "smem_layout: offsets alloc"); + /* Bytes each placed allocation spans (Python _smem_sizes); 0 if dead. */ + int* sizes = (int*)rocke_arena_calloc(&L->arena, n * sizeof(int)); + if(!sizes) + rocke_ll_fail(L, ROCKE_ERR_OOM, "smem_layout: sizes alloc"); for(size_t si = 0; si < n; si++) { @@ -1492,6 +1496,7 @@ void rocke_ll_compute_smem_layout(rocke_lower_t* L) continue; const rocke_type_t* stype = L->smem_globals.data[gi].stype; int seg = ll_smem_seg_size(stype); + sizes[gi] = seg; int aln = ll_smem_align(stype); int first_seq = intervals[gi].first_seq; int last_seq = intervals[gi].last_seq; @@ -2335,7 +2340,6 @@ static void ll_lower_into(rocke_lower_t* L, { L->debug = rocke_ll_debug_create(L, kernel->name); } - /* Entry block (ll_make_block raises on OOM). */ ll_make_block(L, "entry"); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp index 2c7b20c18f04..b20aa841ce08 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp @@ -891,8 +891,11 @@ static void _op_tile_ds_read_tr16_b64(rocke_lower_t* L, const rocke_op_t* op) } rocke_ll_need(L, "ds.read.tr16.b64"); raw = rocke_ll_fresh(L, "tr.raw"); - rocke_ll_emitf( - L, " %s = call <4 x i16> @llvm.amdgcn.ds.read.tr16.b64(ptr addrspace(3) %s)", raw, base); + rocke_ll_emitf(L, + " %s = call <4 x i16> @llvm.amdgcn.ds.read.tr16.b64(ptr addrspace(3) %s)%s", + raw, + base, + rocke_ll_lds_aa_suffix(L, rocke_ll_smem_global_name(L, op->operands[0], NULL))); elem_ty = rocke_ll_llvm_type(L, op->results[0]->type->elem); rocke_ll_emitf(L, " %s = bitcast <4 x i16> %s to <4 x %s>", ll_result_name(op), raw, elem_ty); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index 126478af3787..7b4c303cd535 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -1499,13 +1499,14 @@ static void op_tile_async_buffer_load_lds(rocke_lower_t* L, const rocke_op_t* op rocke_ll_emitf(L, " call void @llvm.amdgcn.raw.ptr.buffer.load.lds(" "ptr addrspace(8) %s, ptr addrspace(3) %s, i32 %lld, i32 %s, i32 %s, " - "i32 0, i32 %lld)", + "i32 0, i32 %lld)%s", rocke_ll_operand(L, rsrc), lds, (long long)bytes_per_lane, rocke_ll_operand(L, voffset), rocke_ll_operand(L, soffset), - (long long)aux); + (long long)aux, + rocke_ll_lds_aa_suffix(L, rocke_ll_lds_prov_get(L, lds_ptr))); } static void op_tile_buffer_load_lds_async(rocke_lower_t* L, const rocke_op_t* op) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_dw_wgrad.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_dw_wgrad.h new file mode 100644 index 000000000000..969a1400e9c7 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_dw_wgrad.h @@ -0,0 +1,191 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_direct_dw_wgrad.h -- C++ engine port of the depthwise + * wgrad kernels and the partial-slot fold in + * library/kernels/common/conv_direct_dw_wgrad.py. + * + * Every block of a depthwise wgrad reduces over the spatial work it owns and + * stores one partial slot ws[part][tap][c] with a plain store (same-address + * fp32 atomics on a depthwise dW serialize); the fold sums the partitions. + * + * Python (conv_direct_dw_wgrad.py) C (this header) + * ------------------------------------ ----------------------------------- + * DwWgradFoldSpec / build_dw_wgrad_fold rocke_dw_wgrad_fold_spec_t / + * rocke_build_dw_wgrad_fold(_new) + * DwWgradColsSpec / is_valid_.._cols_spec rocke_dw_wgrad_cols_spec_t / + * / build_dw_wgrad_cols rocke_dw_wgrad_cols_is_valid_spec / + * rocke_build_dw_wgrad_cols(_new) + * DwWgradLdsSpec / is_valid_.._lds_spec rocke_dw_wgrad_lds_spec_t / + * / build_dw_wgrad_lds rocke_dw_wgrad_lds_is_valid_spec / + * rocke_build_dw_wgrad_lds(_new) + * .kernel_name() / launch geometry rocke_dw_wgrad_*_kernel_name / + * rocke_dw_wgrad_*_launch_geometry + * + * The cols and lds kernels declare the direct-conv kernarg block of + * rocke_conv_direct_arg_names("wgrad"), with D bound to the fp32 partial-slot + * workspace. The fold declares A (the fp32 slots), D (dW), A_bytes, D_bytes, + * parts, channels. + */ +#ifndef ROCKE_INSTANCE_CONV_DIRECT_DW_WGRAD_H +#define ROCKE_INSTANCE_CONV_DIRECT_DW_WGRAD_H + +#include +#include + +#include "rocke/instance_conv_direct_grouped.h" /* rocke_direct_conv_problem_t */ +#include "rocke/ir.h" +#include "rocke/lower_llvm.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* ===================================================================== * + * DwWgradFoldSpec: dW[c, tap] = sum_part ws[part][tap][c] + * + * @dataclass(frozen=True) + * class DwWgradFoldSpec: + * taps: int + * name: str = "dw_wgrad_fold" + * ===================================================================== */ +typedef struct rocke_dw_wgrad_fold_spec +{ + int taps; + const char* name; /* default "dw_wgrad_fold" */ +} rocke_dw_wgrad_fold_spec_t; + +rocke_dw_wgrad_fold_spec_t rocke_dw_wgrad_fold_spec_default(void); +/* kernel_name_join(name, t{taps}). */ +rocke_status_t rocke_dw_wgrad_fold_kernel_name(const rocke_dw_wgrad_fold_spec_t* spec, + char* out, + size_t out_cap); +void rocke_dw_wgrad_fold_launch_geometry(const rocke_dw_wgrad_fold_spec_t* spec, + int channels, + int grid[3], + int block[3]); +rocke_kernel_def_t* rocke_build_dw_wgrad_fold(rocke_ir_builder_t* b, + const rocke_dw_wgrad_fold_spec_t* spec, + const char* arch); +rocke_kernel_def_t* rocke_build_dw_wgrad_fold_new(rocke_ir_builder_t* b, + const rocke_dw_wgrad_fold_spec_t* spec, + const char* arch); +rocke_status_t rocke_dw_wgrad_fold_lower_to_llvm(const rocke_dw_wgrad_fold_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +/* ===================================================================== * + * DwWgradColsSpec: small-C depthwise wgrad, lanes along output columns + * + * @dataclass(frozen=True) + * class DwWgradColsSpec: + * problem: DirectConvProblem + * name: str = "dw_wgrad_cols" + * strip: int = 8 + * threads: int = 64 + * target_blocks: int = 2048 + * ===================================================================== */ +typedef struct rocke_dw_wgrad_cols_spec +{ + rocke_direct_conv_problem_t problem; + const char* name; /* default "dw_wgrad_cols" */ + int strip; /* default 8 */ + int threads; /* default 64 */ + int target_blocks; /* default 2048 */ +} rocke_dw_wgrad_cols_spec_t; + +rocke_dw_wgrad_cols_spec_t rocke_dw_wgrad_cols_spec_default(void); +int rocke_dw_wgrad_cols_n_partitions(const rocke_dw_wgrad_cols_spec_t* spec); +long rocke_dw_wgrad_cols_workspace_bytes(const rocke_dw_wgrad_cols_spec_t* spec); +void rocke_dw_wgrad_cols_launch_geometry(const rocke_dw_wgrad_cols_spec_t* spec, + int grid[3], + int block[3]); +/* DwWgradColsSpec.lds_bytes(): the block reduction buffer. */ +long rocke_dw_wgrad_cols_lds_bytes(const rocke_dw_wgrad_cols_spec_t* spec); +/* kernel_name_join(name, f{KH}x{KW}p{PAD}s{stride}c{groups}, st{strip}, + * t{threads}, tb{n}, flags [bf16]). */ +rocke_status_t rocke_dw_wgrad_cols_kernel_name(const rocke_dw_wgrad_cols_spec_t* spec, + char* out, + size_t out_cap); +bool rocke_dw_wgrad_cols_is_valid_spec(const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch, + char* reason, + size_t reason_cap); +rocke_kernel_def_t* rocke_build_dw_wgrad_cols(rocke_ir_builder_t* b, + const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch); +rocke_kernel_def_t* rocke_build_dw_wgrad_cols_new(rocke_ir_builder_t* b, + const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch); +rocke_status_t rocke_dw_wgrad_cols_lower_to_llvm(const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +/* ===================================================================== * + * DwWgradLdsSpec: depthwise wgrad, LDS row rings (C % 8 == 0) + * + * @dataclass(frozen=True) + * class DwWgradLdsSpec: + * problem: DirectConvProblem + * name: str = "dw_wgrad_lds" + * lanes_c: int = 8 + * lanes_w: int = 32 + * cols_per_lane: int = 1 + * prefetch: int = 2 + * rows_per_seg: int = 32 + * target_blocks: int = 1024 + * ===================================================================== */ +typedef struct rocke_dw_wgrad_lds_spec +{ + rocke_direct_conv_problem_t problem; + const char* name; /* default "dw_wgrad_lds" */ + int lanes_c; /* default 8 */ + int lanes_w; /* default 32 */ + int cols_per_lane; /* default 1 */ + int prefetch; /* default 2 */ + int rows_per_seg; /* default 32 */ + int target_blocks; /* default 1024 */ +} rocke_dw_wgrad_lds_spec_t; + +rocke_dw_wgrad_lds_spec_t rocke_dw_wgrad_lds_spec_default(void); +int rocke_dw_wgrad_lds_threads(const rocke_dw_wgrad_lds_spec_t* spec); +/* DwWgradLdsSpec.lds_bytes(): the padded X / dY rings plus the reduction. */ +long rocke_dw_wgrad_lds_lds_bytes(const rocke_dw_wgrad_lds_spec_t* spec); +int rocke_dw_wgrad_lds_n_partitions(const rocke_dw_wgrad_lds_spec_t* spec); +long rocke_dw_wgrad_lds_workspace_bytes(const rocke_dw_wgrad_lds_spec_t* spec); +void rocke_dw_wgrad_lds_launch_geometry(const rocke_dw_wgrad_lds_spec_t* spec, + int grid[3], + int block[3]); +/* kernel_name_join(name, f{KH}x{KW}p{PAD}s{stride}, lc{n}, lw{n}, cpl{n}, + * pf{n}, rps{n}, tb{n}, flags [bf16]). */ +rocke_status_t rocke_dw_wgrad_lds_kernel_name(const rocke_dw_wgrad_lds_spec_t* spec, + char* out, + size_t out_cap); +bool rocke_dw_wgrad_lds_is_valid_spec(const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch, + char* reason, + size_t reason_cap); +rocke_kernel_def_t* rocke_build_dw_wgrad_lds(rocke_ir_builder_t* b, + const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch); +rocke_kernel_def_t* rocke_build_dw_wgrad_lds_new(rocke_ir_builder_t* b, + const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch); +rocke_status_t rocke_dw_wgrad_lds_lower_to_llvm(const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* ROCKE_INSTANCE_CONV_DIRECT_DW_WGRAD_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_gather.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_gather.h new file mode 100644 index 000000000000..0e1b1ab92134 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_gather.h @@ -0,0 +1,100 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_direct_wgrad_gather.h -- C++ engine port of the dense + * small-C direct wgrad in library/kernels/common/conv_direct_wgrad_gather.py. + * + * One GEMM per filter row r: M = k, N = (s, c) (KW * C columns, padded to 16), + * reduced over the output columns of a 32-wide strip by mfma_f32_16x16x32. The + * block streams output rows, DMAs S X rows and one dY row per step into LDS + * rings, transpose-reads dY, and gathers X (or transpose-reads it when + * S * C keeps every lane 8-byte aligned). It stores one partial slot per block, + * ws[part][(r, s, c)][k], folded into dW by the DwWgradFoldSpec kernel + * (rocke/instance_conv_direct_dw_wgrad.h) with taps = KH * KW * C and + * channels = K. + * + * Python (conv_direct_wgrad_gather.py) C (this header) + * ------------------------------------ ----------------------------------- + * @dataclass(frozen) WgradGatherSpec rocke_direct_wgrad_gather_spec_t + * host-side geometry rocke_direct_wgrad_gather_{threads, + * lds_bytes,n_partitions,fold_taps, + * workspace_bytes,launch_geometry} + * .kernel_name() rocke_direct_wgrad_gather_kernel_name + * is_valid_wgrad_gather_spec(spec, arch) rocke_direct_wgrad_gather_is_valid_spec + * build_wgrad_gather(spec, arch) rocke_build_direct_wgrad_gather(_new) + */ +#ifndef ROCKE_INSTANCE_CONV_DIRECT_WGRAD_GATHER_H +#define ROCKE_INSTANCE_CONV_DIRECT_WGRAD_GATHER_H + +#include +#include + +#include "rocke/instance_conv_direct_grouped.h" /* rocke_direct_conv_problem_t */ +#include "rocke/ir.h" +#include "rocke/lower_llvm.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* ===================================================================== * + * WgradGatherSpec + * + * @dataclass(frozen=True) + * class WgradGatherSpec: + * problem: DirectConvProblem # groups must be 1 + * name: str = "direct_wgrad_gather" + * wave_k16: int = 2 + * waves_k: int = 1 + * prefetch: int = 2 + * rows_per_seg: int = 16 + * target_blocks: int = 1024 + * ===================================================================== */ +typedef struct rocke_direct_wgrad_gather_spec +{ + rocke_direct_conv_problem_t problem; + const char* name; /* default "direct_wgrad_gather" */ + int wave_k16; /* default 2 */ + int waves_k; /* default 1 */ + int prefetch; /* default 2 */ + int rows_per_seg; /* default 16 */ + int target_blocks; /* default 1024 */ +} rocke_direct_wgrad_gather_spec_t; + +rocke_direct_wgrad_gather_spec_t rocke_direct_wgrad_gather_spec_default(void); +int rocke_direct_wgrad_gather_threads(const rocke_direct_wgrad_gather_spec_t* spec); +long rocke_direct_wgrad_gather_lds_bytes(const rocke_direct_wgrad_gather_spec_t* spec); +int rocke_direct_wgrad_gather_n_partitions(const rocke_direct_wgrad_gather_spec_t* spec); +/* fold_taps() = KH * KW * C; fold_channels() = K. */ +int rocke_direct_wgrad_gather_fold_taps(const rocke_direct_wgrad_gather_spec_t* spec); +long rocke_direct_wgrad_gather_workspace_bytes(const rocke_direct_wgrad_gather_spec_t* spec); +void rocke_direct_wgrad_gather_launch_geometry(const rocke_direct_wgrad_gather_spec_t* spec, + int grid[3], + int block[3]); +/* kernel_name_join(name, f{KH}x{KW}p{PAD}s{stride}c{cpg}, wk{wk16}x{waves_k}, + * pf{n}, rps{n}, tb{n}, flags [bf16]). */ +rocke_status_t rocke_direct_wgrad_gather_kernel_name(const rocke_direct_wgrad_gather_spec_t* spec, + char* out, + size_t out_cap); +/* is_valid_wgrad_gather_spec(spec, arch) -> (ok, reason). arch NULL => "gfx950". */ +bool rocke_direct_wgrad_gather_is_valid_spec(const rocke_direct_wgrad_gather_spec_t* spec, + const char* arch, + char* reason, + size_t reason_cap); +rocke_kernel_def_t* rocke_build_direct_wgrad_gather(rocke_ir_builder_t* b, + const rocke_direct_wgrad_gather_spec_t* spec, + const char* arch); +rocke_kernel_def_t* rocke_build_direct_wgrad_gather_new( + rocke_ir_builder_t* b, const rocke_direct_wgrad_gather_spec_t* spec, const char* arch); +rocke_status_t rocke_direct_wgrad_gather_lower_to_llvm(const rocke_direct_wgrad_gather_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* ROCKE_INSTANCE_CONV_DIRECT_WGRAD_GATHER_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_internal.hpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_internal.hpp new file mode 100644 index 000000000000..7347c3d80986 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_internal.hpp @@ -0,0 +1,443 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_direct_wgrad_internal.hpp -- PRIVATE helpers shared by the + * C++ ports of the direct wgrad kernels (conv_direct_wgrad_tiled.cpp, + * conv_direct_dw_wgrad.cpp, conv_direct_wgrad_gather.cpp). Emits no IR of its + * own and declares no public API. + * + * BYTE-IDENTITY CONTRACT. The Python builders emit one op per builder call, in + * Python's left-to-right argument order; C++ leaves argument evaluation order + * unspecified. The `Em` helpers below therefore take at most ONE argument that + * emits: `e.addi(x, 3)` is Python `b.add(x, ci(3))` (the constant, then the + * add). Every call site in the ports keeps to the same rule -- a call may nest + * one emitting argument, never two. + */ +#ifndef ROCKE_INSTANCE_CONV_DIRECT_WGRAD_INTERNAL_HPP +#define ROCKE_INSTANCE_CONV_DIRECT_WGRAD_INTERNAL_HPP + +#include +#include +#include + +#include +#include + +#include "rocke/error_boundary.hpp" /* ckc::guard_builder */ +#include "rocke/instance_conv_direct_grouped.h" /* rocke_direct_conv_problem_t */ +#include "rocke/ir.h" +#include "rocke/ir_internal.h" /* rocke_i_set_err_msg */ +#include "rocke/lower_llvm.h" + +namespace rocke_dwg +{ + +using V = rocke_value_t*; + +constexpr int kWave = 64; + +inline int cdiv(int a, int b) +{ + return (a + b - 1) / b; +} + +/* Python `a % m` for m > 0 (never negative). */ +inline int pymod(int a, int m) +{ + int r = a % m; + return r < 0 ? r + m : r; +} + +/* Python `a // m` for m > 0 (floors). */ +inline int pyfloordiv(int a, int m) +{ + int q = a / m; + return (a % m != 0 && (a < 0) != (m < 0)) ? q - 1 : q; +} + +/* A NULL dtype is not defaulted here: the spec-default constructors set + * "fp16", and the validators reject NULL like Python rejects dtype=None. */ +inline bool problem_is_bf16(const rocke_direct_conv_problem_t* p) +{ + return p->dtype != NULL && strcmp(p->dtype, "bf16") == 0; +} + +inline bool problem_dtype_ok(const rocke_direct_conv_problem_t* p) +{ + const char* d = p->dtype; + return d != NULL && (strcmp(d, "fp16") == 0 || strcmp(d, "bf16") == 0); +} + +/* DirectConvProblem.Ho / .Wo. */ +inline int problem_ho(const rocke_direct_conv_problem_t* p) +{ + return pyfloordiv(p->H + 2 * p->PAD - (p->KH - 1) * p->dil_h - 1, p->stride) + 1; +} + +inline int problem_wo(const rocke_direct_conv_problem_t* p) +{ + return pyfloordiv(p->W + 2 * p->PAD - (p->KW - 1) * p->dil_w - 1, p->stride) + 1; +} + +inline void set_reason(char* out, size_t cap, const char* msg) +{ + if(out != NULL && cap > 0) + { + snprintf(out, cap, "%s", msg); + } +} + +/* conv_direct_grouped.dilation_reason. */ +inline bool dilation_reason(const rocke_direct_conv_problem_t* p, char* why, size_t cap) +{ + if(p->dil_h != 1 || p->dil_w != 1) + { + ROCKE_ERR_SNPRINTF(why, + cap, + "this direct conv kernel has no dilation (got %dx%d); only the " + "non-grouped forward kernel computes a dilated convolution", + p->dil_h, + p->dil_w); + return true; + } + return false; +} + +/* `raise ValueError(f"{prefix}: {reason}")` on the builder. */ +inline void fail_invalid(rocke_ir_builder_t* b, const char* prefix, const char* reason) +{ + if(b->status == ROCKE_OK) + { + std::string msg = std::string(prefix) + ": " + reason; + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, msg.c_str()); + } +} + +/* Emission helpers (see the contract at the top of this header). */ +struct Em +{ + rocke_ir_builder_t* b; + + V ci(int64_t v) const + { + return rocke_b_const_i32(b, v); + } + + V add(V x, V y) const + { + return rocke_b_add(b, x, y); + } + V sub(V x, V y) const + { + return rocke_b_sub(b, x, y); + } + V mul(V x, V y) const + { + return rocke_b_mul(b, x, y); + } + V div(V x, V y) const + { + return rocke_b_div(b, x, y); + } + V mod(V x, V y) const + { + return rocke_b_mod(b, x, y); + } + V smin(V x, V y) const + { + return rocke_b_smin(b, x, y); + } + V smax(V x, V y) const + { + return rocke_b_smax(b, x, y); + } + V land(V x, V y) const + { + return rocke_b_land(b, x, y); + } + V lor(V x, V y) const + { + return rocke_b_lor(b, x, y); + } + V lnot(V x) const + { + return rocke_b_lnot(b, x); + } + V bxor(V x, V y) const + { + return rocke_b_xor(b, x, y); + } + V lt(V x, V y) const + { + return rocke_b_cmp_lt(b, x, y); + } + V ge(V x, V y) const + { + return rocke_b_cmp_ge(b, x, y); + } + V eq(V x, V y) const + { + return rocke_b_cmp_eq(b, x, y); + } + V select(V c, V x, V y) const + { + return rocke_b_select(b, c, x, y); + } + V zext64(V x) const + { + return rocke_b_zext(b, x, rocke_i64()); + } + + /* Python b.op(x, ci(v)): the constant first, then the op. */ + V addi(V x, int64_t v) const + { + V k = ci(v); + return add(x, k); + } + V subi(V x, int64_t v) const + { + V k = ci(v); + return sub(x, k); + } + V muli(V x, int64_t v) const + { + V k = ci(v); + return mul(x, k); + } + V divi(V x, int64_t v) const + { + V k = ci(v); + return div(x, k); + } + V modi(V x, int64_t v) const + { + V k = ci(v); + return mod(x, k); + } + V smini(V x, int64_t v) const + { + V k = ci(v); + return smin(x, k); + } + V lti(V x, int64_t v) const + { + V k = ci(v); + return lt(x, k); + } + V gei(V x, int64_t v) const + { + V k = ci(v); + return ge(x, k); + } + V eqi(V x, int64_t v) const + { + V k = ci(v); + return eq(x, k); + } + V lshri(V x, int64_t v) const + { + V k = ci(v); + return rocke_b_lshr(b, x, k); + } + V shli(V x, int64_t v) const + { + V k = ci(v); + return rocke_b_shl(b, x, k); + } + + /* Python b.div(b.add(x, ci(d - 1)), ci(d)): ceil-divide by a constant. */ + V cdivi(V x, int64_t d) const + { + V t = addi(x, d - 1); + return divi(t, d); + } + + /* (2-D smem) Python b.smem_load_vN(smem, c0, idx, dtype=..., n=...). */ + V smem_load(V smem, V c0, V idx, const rocke_type_t* dtype, int n) const + { + V ix[2] = {c0, idx}; + return rocke_b_smem_load_vN(b, smem, ix, 2, dtype, n); + } + V smem_load_f32(V smem, V c0, V idx, int n) const + { + V ix[2] = {c0, idx}; + return rocke_b_smem_load_vN_f32(b, smem, ix, 2, n); + } + void smem_store(V smem, V c0, V idx, V value, int n) const + { + V ix[2] = {c0, idx}; + rocke_b_smem_store_vN(b, smem, ix, 2, value, n); + } + void smem_store_f32(V smem, V c0, V idx, V value, int n) const + { + V ix[2] = {c0, idx}; + rocke_b_smem_store_vN_f32(b, smem, ix, 2, value, n); + } + V ds_read_tr(V smem, V c0, V idx, const rocke_type_t* dtype) const + { + V ix[2] = {c0, idx}; + return rocke_b_ds_read_tr16_b64(b, smem, ix, 2, dtype); + } + V smem_alloc_1d(const rocke_type_t* elem, int n, const char* hint) const + { + int shape[2] = {1, n}; + return rocke_b_smem_alloc(b, elem, shape, 2, hint); + } + + /* Python swz_val(row, bits): XOR term 2 * sum(((row >> bits[j]) & 1) << j), + * or NULL for no swizzle. */ + V swz_val(V row, const std::vector& bits, V c1) const + { + V v = NULL; + for(size_t j = 0; j < bits.size(); ++j) + { + V t = lshri(row, bits[j]); + t = land(t, c1); + t = shli(t, (int64_t)j + 1); + v = (v == NULL) ? t : lor(v, t); + } + return v; + } + + /* _mfma(b, dtype, "16x16x32", a, bb, acc). */ + V mfma16x16x32(bool bf16, V a, V bb, V acc) const + { + return bf16 ? rocke_b_mfma_f32_16x16x32_bf16(b, a, bb, acc) + : rocke_b_mfma_f32_16x16x32_f16(b, a, bb, acc); + } + + /* `with b.scf_if(cond): body()`. */ + template + void if_(V cond, Fn&& body) const + { + rocke_if_t f = rocke_b_scf_if(b, cond); + if(f.then_region == NULL) + { + return; + } + rocke_b_region_enter(b, f.then_region); + body(); + rocke_b_region_leave(b); + } + + /* `loop = b.scf_for_iter(lo, hi, step, iter_args, iv_name=..., + * elide_trailing_barrier=False)`; `with loop as (iv, carried): body`. + * body(iv, carried) returns the yielded values; returns loop.results. */ + template + std::vector for_iter(V lo, + V hi, + V step, + const std::vector& names, + const std::vector& inits, + const char* iv_name, + Fn&& body) const + { + std::vector args(names.size()); + for(size_t i = 0; i < names.size(); ++i) + { + args[i].name = names[i].c_str(); + args[i].init = inits[i]; + } + rocke_for_t f = rocke_b_scf_for_iter( + b, lo, hi, step, args.data(), (int)args.size(), iv_name, false, false); + if(f.op == NULL || f.body == NULL) + { + return std::vector(); + } + std::vector carried(f.iter_vars, f.iter_vars + f.num_iter_vars); + rocke_b_region_enter(b, f.body); + std::vector out = body(f.iv, carried); + rocke_b_scf_yield(b, out.data(), (int)out.size()); + rocke_b_region_leave(b); + return std::vector(f.op->results, f.op->results + f.op->num_results); + } +}; + +/* The `_new` entry: zero the builder, initialise it with the spec's kernel + * name, then build -- the Python IRBuilder(spec.kernel_name()) + build pair. + * name_fn(out, cap) -> rocke_status_t; build_fn() -> rocke_kernel_def_t*. */ +template +rocke_kernel_def_t* build_new( + rocke_ir_builder_t* b, bool have_spec, const char* who, NameFn&& name_fn, BuildFn&& build_fn) +{ + if(b != NULL) + { + memset(b, 0, sizeof(*b)); + } + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + char name[512]; + if(b == NULL) + { + return NULL; + } + if(!have_spec) + { + std::string msg = std::string(who) + ": null spec"; + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, msg.c_str()); + return NULL; + } + if(name_fn(name, sizeof(name)) != ROCKE_OK) + { + std::string msg = std::string(who) + ": kernel name does not fit"; + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, msg.c_str()); + return NULL; + } + if(rocke_ir_builder_init(b, name) != ROCKE_OK) + { + return NULL; + } + return build_fn(); + }); +} + +/* The `lower_to_llvm` convenience entry: new_fn(&b) builds into a fresh + * builder; the kernel is lowered at `flavor` for `arch`. */ +template +rocke_status_t lower_to_llvm(bool have_spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap, + const char* who, + NewFn&& new_fn) +{ + rocke_ir_builder_t b = {}; + if(out_ll != NULL) + { + *out_ll = NULL; + } + if(!have_spec || out_ll == NULL) + { + set_reason(err, err_cap, "lower_to_llvm: null spec/out"); + return ROCKE_ERR_VALUE; + } + rocke_kernel_def_t* kernel = new_fn(&b); + if(kernel == NULL) + { + const char* m = rocke_ir_builder_error(&b); + rocke_status_t st = rocke_ir_builder_status(&b); + set_reason(err, err_cap, (m != NULL && m[0] != '\0') ? m : who); + rocke_ir_builder_free(&b); + return (st == ROCKE_OK) ? ROCKE_ERR_VALUE : st; + } + rocke_status_t st = rocke_lower_kernel_to_llvm_ex(kernel, flavor, arch, out_ll, err, err_cap); + rocke_ir_builder_free(&b); + return st; +} + +/* Python [(f"{prefix}{i}", v) for i, v in enumerate(values)] names. */ +inline std::vector iter_names(const char* prefix, size_t n) +{ + std::vector out; + out.reserve(n); + for(size_t i = 0; i < n; ++i) + { + out.push_back(std::string(prefix) + std::to_string(i)); + } + return out; +} + +} // namespace rocke_dwg + +#endif /* ROCKE_INSTANCE_CONV_DIRECT_WGRAD_INTERNAL_HPP */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_tiled.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_tiled.h new file mode 100644 index 000000000000..8e605f1a5f3a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_wgrad_tiled.h @@ -0,0 +1,145 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_direct_wgrad_tiled.h -- C++ engine port of the tiled + * direct wgrad builder in library/kernels/common/conv_direct_wgrad_tiled.py. + * + * A stride-1 NHWC weight gradient: a wave owns wave_k16*16 x wave_c16*16 + * channels for all KH*KW taps, the block streams input rows of a 32-column + * strip through LDS ring buffers filled by buffer_load ... lds, with an XOR + * chunk swizzle that keeps the ds_read_b64_tr_b16 fragment reads bank-conflict + * free, and atomically adds its fp32 accumulators into dW once at the end. + * + * Python (conv_direct_wgrad_tiled.py) C (this header) + * ----------------------------------- ------------------------------------ + * @dataclass(frozen) DirectWgradTiledSpec rocke_direct_wgrad_tiled_spec_t + * derived @property geometry rocke_direct_wgrad_tiled_{threads_per_block, + * unroll,acc_vgprs,lds_bytes,n_channel_tiles, + * n_segments,n_splits,launch_geometry} + * .kernel_name() rocke_direct_wgrad_tiled_kernel_name + * is_valid_wgrad_tiled_spec(spec, arch) rocke_direct_wgrad_tiled_is_valid_spec + * pick_swizzle(chunks, shifts, ...) rocke_direct_wgrad_pick_swizzle + * build_direct_wgrad_tiled(spec, arch) rocke_build_direct_wgrad_tiled(_new) + * (+ convenience: build -> lower .ll) rocke_direct_wgrad_tiled_lower_to_llvm + * + * AOT: the kernel declares the direct-conv kernarg block of + * rocke_conv_direct_arg_names("wgrad"); N, H, W, Ho, Wo, groups, the channel + * totals and the strides are runtime, the filter, PAD, dtype and tile geometry + * are baked. dW is fp32 and the caller zero-initialises it. + */ +#ifndef ROCKE_INSTANCE_CONV_DIRECT_WGRAD_TILED_H +#define ROCKE_INSTANCE_CONV_DIRECT_WGRAD_TILED_H + +#include +#include + +#include "rocke/instance_conv_direct_grouped.h" /* rocke_direct_conv_problem_t */ +#include "rocke/ir.h" +#include "rocke/lower_llvm.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* ===================================================================== * + * DirectWgradTiledSpec + * + * @dataclass(frozen=True) + * class DirectWgradTiledSpec: + * problem: DirectConvProblem + * name: str = "direct_wgrad_tiled" + * wave_k16: int = 2 + * wave_c16: int = 2 + * waves_k: int = 4 + * waves_c: int = 2 + * rows_per_seg: int = 30 + * prefetch: int = 2 + * target_blocks: int = 256 + * row_unroll: int = 0 + * unfold_n: int = 1 + * q_groups: int = 1 + * pingpong: bool = False + * ===================================================================== */ +typedef struct rocke_direct_wgrad_tiled_spec +{ + rocke_direct_conv_problem_t problem; + const char* name; /* default "direct_wgrad_tiled" */ + int wave_k16; /* default 2 */ + int wave_c16; /* default 2 */ + int waves_k; /* default 4 */ + int waves_c; /* default 2 */ + int rows_per_seg; /* default 30 */ + int prefetch; /* default 2 */ + int target_blocks; /* default 256 */ + int row_unroll; /* default 0 (= lcm(prefetch + 1, KH)) */ + int unfold_n; /* default 1 */ + int q_groups; /* default 1 */ + bool pingpong; /* default false */ +} rocke_direct_wgrad_tiled_spec_t; + +/* Spec with the Python dataclass defaults; problem == problem_default(). */ +rocke_direct_wgrad_tiled_spec_t rocke_direct_wgrad_tiled_spec_default(void); + +/* Derived geometry (Python @property / host-side methods). */ +int rocke_direct_wgrad_tiled_threads_per_block(const rocke_direct_wgrad_tiled_spec_t* spec); +int rocke_direct_wgrad_tiled_unroll(const rocke_direct_wgrad_tiled_spec_t* spec); +int rocke_direct_wgrad_tiled_acc_vgprs(const rocke_direct_wgrad_tiled_spec_t* spec); +/* lds_bytes(): the padded dY / X rings plus the q-group reduction buffers. */ +long rocke_direct_wgrad_tiled_lds_bytes(const rocke_direct_wgrad_tiled_spec_t* spec); +int rocke_direct_wgrad_tiled_n_channel_tiles(const rocke_direct_wgrad_tiled_spec_t* spec); +int rocke_direct_wgrad_tiled_n_segments(const rocke_direct_wgrad_tiled_spec_t* spec); +int rocke_direct_wgrad_tiled_n_splits(const rocke_direct_wgrad_tiled_spec_t* spec); +/* launch_geometry() -> (grid, block). */ +void rocke_direct_wgrad_tiled_launch_geometry(const rocke_direct_wgrad_tiled_spec_t* spec, + int grid[3], + int block[3]); + +/* kernel_name(): kernel_name_join(name, f{KH}x{KW}p{PAD}, wt{wk}x{wc}, + * w{waves_k}x{waves_c}, rps{n}, pf{n}, tb{n}, flags [bf16, u{n}, un{n}, + * qg{n}, pp]). */ +rocke_status_t rocke_direct_wgrad_tiled_kernel_name(const rocke_direct_wgrad_tiled_spec_t* spec, + char* out, + size_t out_cap); + +/* is_valid_wgrad_tiled_spec(spec, arch) -> (ok, reason). arch NULL => "gfx950". */ +bool rocke_direct_wgrad_tiled_is_valid_spec(const rocke_direct_wgrad_tiled_spec_t* spec, + const char* arch, + char* reason, + size_t reason_cap); + +/* pick_swizzle(chunks, shifts, wu, seg, col_origins): the bit selection + * `bits` of the chunk XOR 2 * sum(((row >> bits[j]) & 1) << j) that makes + * every fragment read conflict free (or none). Writes up to 6 bits into + * bits_out and returns how many. Shared with conv_direct_wgrad_gather. */ +int rocke_direct_wgrad_pick_swizzle(int chunks, + const int* shifts, + int n_shifts, + int wu, + int seg, + const int* col_origins, + int n_origins, + int bits_out[6]); + +/* build_direct_wgrad_tiled(spec, arch) into an already-initialised builder. */ +rocke_kernel_def_t* rocke_build_direct_wgrad_tiled(rocke_ir_builder_t* b, + const rocke_direct_wgrad_tiled_spec_t* spec, + const char* arch); + +/* Initialise `b` with the spec's kernel name, then build. */ +rocke_kernel_def_t* rocke_build_direct_wgrad_tiled_new(rocke_ir_builder_t* b, + const rocke_direct_wgrad_tiled_spec_t* spec, + const char* arch); + +/* Convenience: build + lower to LLVM IR text (*out_ll malloc'd). */ +rocke_status_t rocke_direct_wgrad_tiled_lower_to_llvm(const rocke_direct_wgrad_tiled_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* ROCKE_INSTANCE_CONV_DIRECT_WGRAD_TILED_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_dw_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_dw_wgrad.cpp new file mode 100644 index 000000000000..25b993389530 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_dw_wgrad.cpp @@ -0,0 +1,1192 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * conv_direct_dw_wgrad.cpp -- C++ engine port of + * library/kernels/common/conv_direct_dw_wgrad.py: the partial-slot fold + * (build_dw_wgrad_fold), the small-C depthwise wgrad with lanes along output + * columns (build_dw_wgrad_cols) and the depthwise wgrad with LDS row rings + * (build_dw_wgrad_lds). + * + * Every IR op (including each arith.constant) is emitted in the order the + * Python builder emits it; see the contract in + * rocke/instance_conv_direct_wgrad_internal.hpp. + */ +#include "rocke/instance_conv_direct_dw_wgrad.h" + +#include +#include +#include + +#include +#include +#include + +#include "rocke/arch_target.h" /* rocke_set_unknown_arch_reason */ +#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx */ +#include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/instance_conv_direct_grouped_internal.h" /* rocke_dconv_emit_params */ +#include "rocke/instance_conv_direct_wgrad_internal.hpp" + +using namespace rocke_dwg; + +namespace +{ + +constexpr int kFoldSlices = 16; +constexpr int kFoldCols = 64; +/* Partition loads each slice keeps in flight. */ +constexpr int kFoldUnroll = 4; + +const char* bf16_flag[1] = {"bf16"}; + +} // namespace + +/* ===================================================================== * + * Fold: dW[c, tap] = sum_part ws[part][tap][c] + * ===================================================================== */ + +extern "C" rocke_dw_wgrad_fold_spec_t rocke_dw_wgrad_fold_spec_default(void) +{ + rocke_dw_wgrad_fold_spec_t s; + memset(&s, 0, sizeof(s)); + s.name = "dw_wgrad_fold"; + return s; +} + +extern "C" rocke_status_t + rocke_dw_wgrad_fold_kernel_name(const rocke_dw_wgrad_fold_spec_t* s, char* out, size_t out_cap) +{ + if(s == NULL || out == NULL || out_cap == 0) + { + return ROCKE_ERR_VALUE; + } + char t_buf[24]; + snprintf(t_buf, sizeof t_buf, "t%d", s->taps); + const char* parts[1] = {t_buf}; + return rocke_kernel_name_join(s->name, parts, 1, NULL, NULL, 0, out, out_cap, NULL); +} + +extern "C" void rocke_dw_wgrad_fold_launch_geometry(const rocke_dw_wgrad_fold_spec_t* s, + int channels, + int grid[3], + int block[3]) +{ + grid[0] = cdiv(channels * s->taps, kFoldCols); + grid[1] = 1; + grid[2] = 1; + block[0] = kFoldSlices * kWave; + block[1] = 1; + block[2] = 1; +} + +namespace +{ + +/* Each block owns 64 outputs; 16 waves split the partition axis, meet in LDS. */ +rocke_kernel_def_t* build_fold(rocke_ir_builder_t* b, const rocke_dw_wgrad_fold_spec_t* spec) +{ + Em e{b}; + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", kFoldSlices * kWave); + rocke_param_opts_t ro; + memset(&ro, 0, sizeof(ro)); + ro.noalias = ro.noalias_set = true; + ro.readonly = ro.readonly_set = true; + ro.align = 16; + ro.align_set = true; + rocke_param_opts_t wo; + memset(&wo, 0, sizeof(wo)); + wo.noalias = wo.noalias_set = true; + wo.writeonly = wo.writeonly_set = true; + wo.align = 4; + wo.align_set = true; + V ws = rocke_b_param(b, "A", rocke_ptr_type(b, rocke_f32(), "global"), &ro); + V dw = rocke_b_param(b, "D", rocke_ptr_type(b, rocke_f32(), "global"), &wo); + rocke_b_param(b, "A_bytes", rocke_i32(), NULL); + rocke_b_param(b, "D_bytes", rocke_i32(), NULL); + V parts = rocke_b_param(b, "parts", rocke_i32(), NULL); + V chans = rocke_b_param(b, "channels", rocke_i32(), NULL); + if(!rocke_ir_builder_ok(b)) + { + return NULL; + } + + V tid = rocke_b_thread_id_x(b); + V lane = e.modi(tid, kWave); + V slc = e.divi(tid, kWave); + V n_out = e.muli(chans, spec->taps); + V bx = rocke_b_block_id_x(b); + V o = e.add(e.muli(bx, kFoldCols), lane); + V o_ok = e.lt(o, n_out); + V o_safe = e.select(o_ok, o, e.ci(0)); + V zero = rocke_b_const_f32(b, 0.0); + const int U = kFoldUnroll; + /* Partition p = slc + SLICES * (U * it + j). */ + V n_it = e.cdivi(parts, kFoldSlices * U); + V lo = e.ci(0); + V step = e.ci(1); + std::vector names; + for(int j = 0; j < U; ++j) + { + names.push_back("s" + std::to_string(j)); + } + std::vector res = e.for_iter( + lo, n_it, step, names, std::vector(U, zero), "it", [&](V it, std::vector acc) { + for(int j = 0; j < U; ++j) + { + V pidx = e.addi(e.muli(it, U), j); + pidx = e.add(slc, e.muli(pidx, kFoldSlices)); + V ok = e.lt(pidx, parts); + V sel = e.select(ok, pidx, e.ci(0)); + V idx = e.add(e.mul(sel, n_out), o_safe); + V v = rocke_b_global_load_f32(b, ws, idx, 4); + acc[j] = rocke_b_fadd(b, acc[j], e.select(ok, v, zero)); + } + return acc; + }); + if((int)res.size() != U) + { + return NULL; + } + V tot = res[0]; + for(int j = 1; j < U; ++j) + { + tot = rocke_b_fadd(b, tot, res[j]); + } + V red = e.smem_alloc_1d(rocke_f32(), kFoldSlices * kWave, "fold"); + e.smem_store_f32(red, e.ci(0), tid, tot, 1); + rocke_b_sync(b); + V cond = e.land(e.eqi(slc, 0), o_ok); + e.if_(cond, [&]() { + V t = NULL; + for(int q = 0; q < kFoldSlices; ++q) + { + V c0 = e.ci(0); + V v = e.smem_load_f32(red, c0, e.addi(lane, q * kWave), 1); + v = rocke_b_vec_extract(b, v, 0); + t = (t == NULL) ? v : rocke_b_fadd(b, t, v); + } + V tap = e.div(o, chans); + V c = e.mod(o, chans); + V off = e.add(e.muli(c, spec->taps), tap); + rocke_b_global_store(b, dw, off, t, 4); + }); + return rocke_ir_builder_ok(b) ? b->kernel : NULL; +} + +} // namespace + +extern "C" rocke_kernel_def_t* rocke_build_dw_wgrad_fold(rocke_ir_builder_t* b, + const rocke_dw_wgrad_fold_spec_t* spec, + const char* arch) +{ + (void)arch; + if(b == NULL || spec == NULL) + { + return NULL; + } + return build_fold(b, spec); +} + +extern "C" rocke_kernel_def_t* rocke_build_dw_wgrad_fold_new(rocke_ir_builder_t* b, + const rocke_dw_wgrad_fold_spec_t* spec, + const char* arch) +{ + return build_new( + b, + spec != NULL, + "build_dw_wgrad_fold", + [&](char* out, size_t cap) { return rocke_dw_wgrad_fold_kernel_name(spec, out, cap); }, + [&]() { return rocke_build_dw_wgrad_fold(b, spec, arch); }); +} + +extern "C" rocke_status_t rocke_dw_wgrad_fold_lower_to_llvm(const rocke_dw_wgrad_fold_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + if(arch == NULL) + { + arch = "gfx950"; + } + return lower_to_llvm( + spec != NULL, + arch, + flavor, + out_ll, + err, + err_cap, + "build_dw_wgrad_fold failed", + [&](rocke_ir_builder_t* b) { return rocke_build_dw_wgrad_fold_new(b, spec, arch); }); +} + +/* ===================================================================== * + * Small-C depthwise wgrad: lanes along output columns + * + * A lane owns filter row r (block x) and a strip of `strip` output columns, + * keeps KW x C accumulators, and per output row loads the dY strip and the + * X row segment it touches as whole 16-byte vectors at a build-time phase. + * Block y is a partition over (image, output row, column group) units. + * ===================================================================== */ + +extern "C" rocke_dw_wgrad_cols_spec_t rocke_dw_wgrad_cols_spec_default(void) +{ + rocke_dw_wgrad_cols_spec_t s; + memset(&s, 0, sizeof(s)); + s.problem = rocke_direct_conv_problem_default(); + s.name = "dw_wgrad_cols"; + s.strip = 8; + s.threads = 64; + s.target_blocks = 2048; + return s; +} + +static int cols_n_units(const rocke_dw_wgrad_cols_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + return cdiv(p->N * problem_ho(p) * (problem_wo(p) / s->strip), s->threads); +} + +extern "C" int rocke_dw_wgrad_cols_n_partitions(const rocke_dw_wgrad_cols_spec_t* s) +{ + int a = cols_n_units(s); + int c = cdiv(s->target_blocks, s->problem.KH); + return a < c ? a : c; +} + +extern "C" long rocke_dw_wgrad_cols_workspace_bytes(const rocke_dw_wgrad_cols_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + return (long)rocke_dw_wgrad_cols_n_partitions(s) * p->KH * p->KW * p->groups * 4; +} + +extern "C" void rocke_dw_wgrad_cols_launch_geometry(const rocke_dw_wgrad_cols_spec_t* s, + int grid[3], + int block[3]) +{ + grid[0] = s->problem.KH; + grid[1] = rocke_dw_wgrad_cols_n_partitions(s); + grid[2] = 1; + block[0] = s->threads; + block[1] = 1; + block[2] = 1; +} + +extern "C" long rocke_dw_wgrad_cols_lds_bytes(const rocke_dw_wgrad_cols_spec_t* s) +{ + /* The block reduction buffer: KW x C fp32 accumulators per thread. */ + return (long)s->threads * s->problem.KW * s->problem.groups * 4; +} + +extern "C" rocke_status_t + rocke_dw_wgrad_cols_kernel_name(const rocke_dw_wgrad_cols_spec_t* s, char* out, size_t out_cap) +{ + if(s == NULL || out == NULL || out_cap == 0) + { + return ROCKE_ERR_VALUE; + } + const rocke_direct_conv_problem_t* p = &s->problem; + char f_buf[64], st_buf[24], t_buf[24], tb_buf[24]; + snprintf(f_buf, sizeof f_buf, "f%dx%dp%ds%dc%d", p->KH, p->KW, p->PAD, p->stride, p->groups); + snprintf(st_buf, sizeof st_buf, "st%d", s->strip); + snprintf(t_buf, sizeof t_buf, "t%d", s->threads); + snprintf(tb_buf, sizeof tb_buf, "tb%d", s->target_blocks); + const char* parts[4] = {f_buf, st_buf, t_buf, tb_buf}; + int on[1] = {problem_is_bf16(p) ? 1 : 0}; + return rocke_kernel_name_join(s->name, parts, 4, bf16_flag, on, 1, out, out_cap, NULL); +} + +extern "C" bool rocke_dw_wgrad_cols_is_valid_spec(const rocke_dw_wgrad_cols_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ + (void)arch; + char msg[ROCKE_ERR_MSG_CAP]; + if(s == NULL) + { + set_reason(reason, reason_cap, "null spec"); + return false; + } + const rocke_direct_conv_problem_t* p = &s->problem; + if(dilation_reason(p, msg, sizeof msg)) + { + set_reason(reason, reason_cap, msg); + return false; + } + const char* why = NULL; + int C = p->groups; + if(p->cpg != 1 || p->kpg != 1) + { + why = "depthwise only"; + } + else if(!problem_dtype_ok(p)) + { + why = "fp16/bf16 only"; + } + else if(p->groups > 8) + { + why = "small-C kernel: C <= 8"; + } + else if(p->KW * p->groups > 96) + { + why = "too many accumulators"; + } + else if(s->threads < kWave || s->threads % kWave || s->threads > 1024) + { + why = "threads must be a multiple of 64"; + } + /* Whole 16-byte vectors at a build-time phase: the strip and the rows must + * keep every vector start on an 8-element boundary. */ + else if((s->strip * C) % 8 || (s->strip * p->stride * C) % 8) + { + why = "strip * C (and strip * stride * C) must be multiples of 8"; + } + else if((p->W * C) % 8 || (problem_wo(p) * C) % 8 || s->strip <= 0 || problem_wo(p) % s->strip) + { + why = "W*C and Wo*C must be multiples of 8 and Wo of the strip"; + } + else if(s->target_blocks < 1) + { + why = "target_blocks >= 1"; + } + else if(rocke_dw_wgrad_cols_lds_bytes(s) > 160 * 1024) + { + ROCKE_ERR_SNPRINTF( + msg, sizeof msg, "LDS %ld B > 160 KiB", rocke_dw_wgrad_cols_lds_bytes(s)); + why = msg; + } + if(why != NULL) + { + set_reason(reason, reason_cap, why); + return false; + } + set_reason(reason, reason_cap, "ok"); + return true; +} + +namespace +{ + +rocke_kernel_def_t* + build_cols(rocke_ir_builder_t* b, const rocke_dw_wgrad_cols_spec_t* spec, const char* arch) +{ + char reason[ROCKE_ERR_MSG_CAP]; + if(!rocke_dw_wgrad_cols_is_valid_spec(spec, arch, reason, sizeof reason)) + { + fail_invalid(b, "invalid DwWgradColsSpec", reason); + return NULL; + } + Em e{b}; + const rocke_direct_conv_problem_t& p = spec->problem; + const int KH = p.KH, KW = p.KW, PAD = p.PAD, S = p.stride, C = p.groups; + const int WS = spec->strip; + const int T = spec->threads; + const bool bf16 = problem_is_bf16(&p); + const rocke_type_t* io_type = bf16 ? rocke_bf16() : rocke_f16(); + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", T); + rocke_dconv_params_t P; + rocke_dconv_emit_params(b, &P, "wgrad", io_type); + if(!rocke_ir_builder_ok(b)) + { + return NULL; + } + + V c0 = e.ci(0); + V c1 = e.ci(1); + V c2 = e.ci(2); + V oob = e.ci(((int64_t)1 << 31) - 1); + V tid = rocke_b_thread_id_x(b); + V r_tap = rocke_b_block_id_x(b); /* filter row */ + V part = rocke_b_block_id_y(b); + V strips_per_row = e.divi(P.p_Wo, WS); + V n_strips = e.mul(e.mul(P.p_N, P.p_Ho), strips_per_row); + V n_units = e.cdivi(n_strips, T); + V n_part = e.smin(n_units, e.ci(cdiv(spec->target_blocks, KH))); + V u_lo = e.div(e.mul(part, n_units), n_part); + V u_hi = e.div(e.mul(e.add(part, c1), n_units), n_part); + + V a_rsrc = rocke_b_buffer_rsrc(b, P.A, P.A_bytes); /* dY */ + V b_rsrc = rocke_b_buffer_rsrc(b, P.Bp, P.B_bytes); /* X */ + const int X_COLS = (WS - 1) * S + KW; + /* Window element e sits at vector element e + PH: the window starts PAD + * columns left of an 8-aligned strip. */ + const int PH = pymod(-PAD * C, 8); + const int n_dy_vec = WS * C / 8; + const int n_x_vec = cdiv(PH + X_COLS * C, 8); + + auto load8 = [&](V rsrc, V elem_off, V ok) { + V voff = e.select(ok, e.mul(elem_off, c2), oob); + return bf16 ? rocke_b_buffer_load_vN_bf16(b, rsrc, voff, c0, 4) + : rocke_b_buffer_load_vN_f16(b, rsrc, voff, c0, 4); + }; + auto elem = [&](const std::vector& vecs, int el) { + return rocke_b_cast_to_f32(b, rocke_b_vec_extract(b, vecs[el / 8], el % 8)); + }; + + V zero = rocke_b_const_f32(b, 0.0); + const int n_acc = KW * C; + std::vector accs = e.for_iter( + u_lo, + u_hi, + c1, + iter_names("acc", n_acc), + std::vector(n_acc, zero), + "unit", + [&](V unit, std::vector acc) { + V strip = e.add(e.muli(unit, T), tid); + V live = e.lt(strip, n_strips); + V cs = e.mod(strip, strips_per_row); + V rest = e.div(strip, strips_per_row); + V ho = e.mod(rest, P.p_Ho); + V n_img = e.div(rest, P.p_Ho); + V wo0 = e.muli(cs, WS); + V hi = e.muli(ho, S); + hi = e.add(hi, e.subi(r_tap, PAD)); + V hi_ok = e.ge(hi, c0); + hi_ok = e.land(hi_ok, e.lt(hi, P.p_Hi)); + hi_ok = e.land(live, hi_ok); + V wi0 = e.subi(e.muli(wo0, S), PAD); + V dy_base = e.mul(n_img, P.p_A_stride_n); + dy_base = e.add(dy_base, e.mul(ho, P.p_A_stride_hi)); + dy_base = e.add(dy_base, e.mul(wo0, P.p_A_stride_wi)); + V x_base = e.mul(n_img, P.p_B_stride_n); + x_base = e.add(x_base, e.mul(hi, P.p_B_stride_hi)); + x_base = e.add(x_base, e.mul(wi0, P.p_B_stride_wi)); + std::vector dy_vecs, x_vecs; + for(int i = 0; i < n_dy_vec; ++i) + { + dy_vecs.push_back(load8(a_rsrc, e.addi(dy_base, 8 * i), live)); + } + for(int i = 0; i < n_x_vec; ++i) + { + x_vecs.push_back(load8(b_rsrc, e.addi(x_base, 8 * i - PH), hi_ok)); + } + std::vector> dy(WS); + for(int j = 0; j < WS; ++j) + { + for(int c = 0; c < C; ++c) + { + dy[j].push_back(elem(dy_vecs, j * C + c)); + } + } + std::vector> xs(X_COLS); + for(int col = 0; col < X_COLS; ++col) + { + V wi = e.addi(wi0, col); + V okc = e.ge(wi, c0); + okc = e.land(okc, e.lt(wi, P.p_Wi)); + for(int c = 0; c < C; ++c) + { + xs[col].push_back(e.select(okc, elem(x_vecs, PH + col * C + c), zero)); + } + } + for(int j = 0; j < WS; ++j) + { + for(int s = 0; s < KW; ++s) + { + for(int c = 0; c < C; ++c) + { + int i = s * C + c; + acc[i] = rocke_b_fma(b, dy[j][c], xs[j * S + s][c], acc[i]); + } + } + } + return acc; + }); + if((int)accs.size() != n_acc) + { + return NULL; + } + + /* Block reduction: every lane holds this block's r row; sum over lanes. */ + V red = e.smem_alloc_1d(rocke_f32(), T * n_acc, "red"); + for(int i = 0; i < n_acc; ++i) + { + V ki = e.ci(i); + V idx = e.add(e.muli(ki, T), tid); + e.smem_store_f32(red, c0, idx, accs[i], 1); + } + rocke_b_sync(b); + /* Thread tid sums outputs tid, tid + T, ... (KW * C may exceed T). */ + for(int k = 0; k < cdiv(n_acc, T); ++k) + { + V o = k == 0 ? tid : e.addi(tid, k * T); + e.if_(e.lti(o, n_acc), [&]() { + V tot = NULL; + for(int q = 0; q < T; ++q) + { + V idx = e.addi(e.muli(o, T), q); + V v = rocke_b_vec_extract(b, e.smem_load_f32(red, c0, idx, 1), 0); + tot = (tot == NULL) ? v : rocke_b_fadd(b, tot, v); + } + /* ws[part][tap][c], tap = r*KW + s; output o is (s, c) = divmod(o, C). */ + V s_i = e.divi(o, C); + V c_i = e.modi(o, C); + V tap = e.add(e.muli(r_tap, KW), s_i); + V off = e.add(e.muli(part, KH * KW), tap); + off = e.add(e.muli(off, C), c_i); + rocke_b_global_store(b, P.D, off, tot, 4); + }); + } + return rocke_ir_builder_ok(b) ? b->kernel : NULL; +} + +} // namespace + +extern "C" rocke_kernel_def_t* rocke_build_dw_wgrad_cols(rocke_ir_builder_t* b, + const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch) +{ + if(b == NULL || spec == NULL) + { + return NULL; + } + return build_cols(b, spec, arch ? arch : "gfx950"); +} + +extern "C" rocke_kernel_def_t* rocke_build_dw_wgrad_cols_new(rocke_ir_builder_t* b, + const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch) +{ + return build_new( + b, + spec != NULL, + "build_dw_wgrad_cols", + [&](char* out, size_t cap) { return rocke_dw_wgrad_cols_kernel_name(spec, out, cap); }, + [&]() { return rocke_build_dw_wgrad_cols(b, spec, arch); }); +} + +extern "C" rocke_status_t rocke_dw_wgrad_cols_lower_to_llvm(const rocke_dw_wgrad_cols_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + if(arch == NULL) + { + arch = "gfx950"; + } + return lower_to_llvm( + spec != NULL, + arch, + flavor, + out_ll, + err, + err_cap, + "build_dw_wgrad_cols failed", + [&](rocke_ir_builder_t* b) { return rocke_build_dw_wgrad_cols_new(b, spec, arch); }); +} + +/* ===================================================================== * + * Depthwise wgrad, LDS row rings + * + * The block owns a channel tile and a column tile and streams output rows: + * per step it DMAs S new X rows and one dY row into LDS rings (prefetch steps + * ahead), and every lane reads its window out of LDS. Ring slots fold to + * constants by unrolling U = prefetch + ceil(KH / S) steps: the X ring has + * S*U row slots and the dY ring U. + * ===================================================================== */ + +extern "C" rocke_dw_wgrad_lds_spec_t rocke_dw_wgrad_lds_spec_default(void) +{ + rocke_dw_wgrad_lds_spec_t s; + memset(&s, 0, sizeof(s)); + s.problem = rocke_direct_conv_problem_default(); + s.name = "dw_wgrad_lds"; + s.lanes_c = 8; + s.lanes_w = 32; + s.cols_per_lane = 1; + s.prefetch = 2; + s.rows_per_seg = 32; + s.target_blocks = 1024; + return s; +} + +namespace +{ + +int lds_block_c(const rocke_dw_wgrad_lds_spec_t* s) +{ + return 8 * s->lanes_c; +} + +int lds_block_wo(const rocke_dw_wgrad_lds_spec_t* s) +{ + return s->lanes_w * s->cols_per_lane; +} + +int lds_unroll(const rocke_dw_wgrad_lds_spec_t* s) +{ + return s->prefetch + cdiv(s->problem.KH, s->problem.stride); +} + +int lds_x_cols(const rocke_dw_wgrad_lds_spec_t* s) +{ + return (lds_block_wo(s) - 1) * s->problem.stride + s->problem.KW; +} + +int lds_rps(const rocke_dw_wgrad_lds_spec_t* s) +{ + return cdiv(s->rows_per_seg, lds_unroll(s)) * lds_unroll(s); +} + +int lds_n_channel_tiles(const rocke_dw_wgrad_lds_spec_t* s) +{ + return cdiv(s->problem.groups, lds_block_c(s)); +} + +/* 16-byte DMA instructions per wave for one LDS row of `cols` pixels. */ +int lds_dma_instrs(const rocke_dw_wgrad_lds_spec_t* s, int cols) +{ + int n_waves = s->lanes_c * s->lanes_w / kWave; + return cdiv(cdiv(cols * (lds_block_c(s) / 8), kWave), n_waves); +} + +/* Elements of one X / dY ring slot: whole DMA instructions of every wave. */ +int lds_x_row_elems(const rocke_dw_wgrad_lds_spec_t* s) +{ + return s->lanes_c * s->lanes_w * lds_dma_instrs(s, lds_x_cols(s)) * 8; +} + +int lds_dy_row_elems(const rocke_dw_wgrad_lds_spec_t* s) +{ + return s->lanes_c * s->lanes_w * lds_dma_instrs(s, lds_block_wo(s)) * 8; +} + +} // namespace + +extern "C" int rocke_dw_wgrad_lds_threads(const rocke_dw_wgrad_lds_spec_t* s) +{ + return s->lanes_c * s->lanes_w; +} + +extern "C" long rocke_dw_wgrad_lds_lds_bytes(const rocke_dw_wgrad_lds_spec_t* s) +{ + /* The X and dY rings as the builder allocates them, plus the reduction. */ + long U = lds_unroll(s); + long x = (long)s->problem.stride * U * lds_x_row_elems(s) * 2; + long dy = U * lds_dy_row_elems(s) * 2; + long red = (long)s->lanes_w * lds_block_c(s) * 4; + return x + dy + red; +} + +extern "C" int rocke_dw_wgrad_lds_n_partitions(const rocke_dw_wgrad_lds_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + int n_units = p->N * cdiv(problem_wo(p), lds_block_wo(s)) * cdiv(problem_ho(p), lds_rps(s)); + int c = cdiv(s->target_blocks, lds_n_channel_tiles(s)); + return n_units < c ? n_units : c; +} + +extern "C" long rocke_dw_wgrad_lds_workspace_bytes(const rocke_dw_wgrad_lds_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + return (long)rocke_dw_wgrad_lds_n_partitions(s) * p->KH * p->KW * p->groups * 4; +} + +extern "C" void rocke_dw_wgrad_lds_launch_geometry(const rocke_dw_wgrad_lds_spec_t* s, + int grid[3], + int block[3]) +{ + grid[0] = lds_n_channel_tiles(s); + grid[1] = rocke_dw_wgrad_lds_n_partitions(s); + grid[2] = 1; + block[0] = rocke_dw_wgrad_lds_threads(s); + block[1] = 1; + block[2] = 1; +} + +extern "C" rocke_status_t + rocke_dw_wgrad_lds_kernel_name(const rocke_dw_wgrad_lds_spec_t* s, char* out, size_t out_cap) +{ + if(s == NULL || out == NULL || out_cap == 0) + { + return ROCKE_ERR_VALUE; + } + const rocke_direct_conv_problem_t* p = &s->problem; + char f_buf[64], lc_buf[24], lw_buf[24], cpl_buf[24], pf_buf[24], rps_buf[24], tb_buf[24]; + snprintf(f_buf, sizeof f_buf, "f%dx%dp%ds%d", p->KH, p->KW, p->PAD, p->stride); + snprintf(lc_buf, sizeof lc_buf, "lc%d", s->lanes_c); + snprintf(lw_buf, sizeof lw_buf, "lw%d", s->lanes_w); + snprintf(cpl_buf, sizeof cpl_buf, "cpl%d", s->cols_per_lane); + snprintf(pf_buf, sizeof pf_buf, "pf%d", s->prefetch); + snprintf(rps_buf, sizeof rps_buf, "rps%d", s->rows_per_seg); + snprintf(tb_buf, sizeof tb_buf, "tb%d", s->target_blocks); + const char* parts[7] = {f_buf, lc_buf, lw_buf, cpl_buf, pf_buf, rps_buf, tb_buf}; + int on[1] = {problem_is_bf16(p) ? 1 : 0}; + return rocke_kernel_name_join(s->name, parts, 7, bf16_flag, on, 1, out, out_cap, NULL); +} + +extern "C" bool rocke_dw_wgrad_lds_is_valid_spec(const rocke_dw_wgrad_lds_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ + char msg[ROCKE_ERR_MSG_CAP]; + if(s == NULL) + { + set_reason(reason, reason_cap, "null spec"); + return false; + } + if(arch == NULL) + { + arch = "gfx950"; + } + const rocke_archtarget_t* target = rocke_archtarget_from_gfx(arch); + if(target == NULL) + { + rocke_set_unknown_arch_reason(reason, reason_cap, arch); + return false; + } + const rocke_direct_conv_problem_t* p = &s->problem; + if(dilation_reason(p, msg, sizeof msg)) + { + set_reason(reason, reason_cap, msg); + return false; + } + int threads = rocke_dw_wgrad_lds_threads(s); + const char* why = NULL; + if(target->wave_size != 64 || rocke_archtarget_async_lds_max_dwords(target) < 4) + { + ROCKE_ERR_SNPRINTF( + msg, sizeof msg, "needs wave64 + 16-byte buffer_load_lds (absent on %s)", arch); + why = msg; + } + else if(p->cpg != 1 || p->kpg != 1) + { + why = "depthwise only"; + } + else if(!problem_dtype_ok(p)) + { + why = "fp16/bf16 only"; + } + else if(p->groups % 8) + { + why = "C must be a multiple of 8"; + } + else if(threads < kWave || threads % kWave || threads > 1024) + { + why = "threads must be a multiple of 64, at most 1024"; + } + else if(p->KH * p->KW * 8 > 128) + { + why = "too many accumulators"; + } + else if(s->prefetch < 1) + { + why = "prefetch >= 1"; + } + else if(rocke_dw_wgrad_lds_lds_bytes(s) > 160 * 1024) + { + ROCKE_ERR_SNPRINTF(msg, sizeof msg, "LDS %ld B > 160 KiB", rocke_dw_wgrad_lds_lds_bytes(s)); + why = msg; + } + else if(s->rows_per_seg < 1) + { + why = "rows_per_seg >= 1"; + } + else if(s->target_blocks < 1) + { + why = "target_blocks >= 1"; + } + if(why != NULL) + { + set_reason(reason, reason_cap, why); + return false; + } + set_reason(reason, reason_cap, "ok"); + return true; +} + +namespace +{ + +struct LdsLane +{ + V col, ch, ok; +}; + +struct LdsOff +{ + V off, ok; +}; + +rocke_kernel_def_t* + build_lds(rocke_ir_builder_t* b, const rocke_dw_wgrad_lds_spec_t* spec, const char* arch) +{ + char reason[ROCKE_ERR_MSG_CAP]; + if(!rocke_dw_wgrad_lds_is_valid_spec(spec, arch, reason, sizeof reason)) + { + fail_invalid(b, "invalid DwWgradLdsSpec", reason); + return NULL; + } + Em e{b}; + const rocke_direct_conv_problem_t& p = spec->problem; + const int KH = p.KH, KW = p.KW, PAD = p.PAD, S = p.stride; + const int TAPS = KH * KW; + const int LC = spec->lanes_c, LW = spec->lanes_w, WL = spec->cols_per_lane; + const int BC = lds_block_c(spec), BWO = lds_block_wo(spec); + const int CH = BC / 8; /* 16-byte chunks per pixel */ + const int XC = lds_x_cols(spec); + const int T = rocke_dw_wgrad_lds_threads(spec); + const int D = spec->prefetch; + const int U = lds_unroll(spec); + const int NX = S * U, NY = U; + const int RPS = lds_rps(spec); + const int PRE + = KH - S > 0 ? KH - S : 0; /* X rows a segment needs before its first step's own */ + const bool bf16 = problem_is_bf16(&p); + const rocke_type_t* io_type = bf16 ? rocke_bf16() : rocke_f16(); + + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", T); + rocke_attr_set_bool(b, &b->kernel->attrs, "lds_alias_scopes", true); + rocke_dconv_params_t P; + rocke_dconv_emit_params(b, &P, "wgrad", io_type); + if(!rocke_ir_builder_ok(b)) + { + return NULL; + } + + V c0 = e.ci(0); + V c1 = e.ci(1); + V c2 = e.ci(2); + V oob = e.ci(((int64_t)1 << 31) - 1); + V p_C = P.p_groups; + V tid = rocke_b_thread_id_x(b); + V lane = e.modi(tid, kWave); + V wave = rocke_b_readfirstlane(b, e.divi(tid, kWave)); + V lc = e.modi(tid, LC); + V lw = e.divi(tid, LC); + V ctile = rocke_b_block_id_x(b); + V part = rocke_b_block_id_y(b); + V c_blk0 = e.muli(ctile, BC); + V n_ct = e.cdivi(P.p_Wo, BWO); + V n_hb = e.cdivi(P.p_Ho, RPS); + V n_units = e.mul(e.mul(P.p_N, n_ct), n_hb); + V n_tiles = e.cdivi(p_C, BC); + V tb = e.ci(spec->target_blocks); + V np = e.add(tb, e.sub(n_tiles, c1)); + np = e.div(np, n_tiles); + V n_part = e.smin(n_units, np); + + V a_rsrc = rocke_b_buffer_rsrc(b, P.A, P.A_bytes); /* dY */ + V b_rsrc = rocke_b_buffer_rsrc(b, P.Bp, P.B_bytes); /* X */ + + /* ---- LDS rings: one allocation per slot so the waitcnt pass can tell them apart */ + const int IX = lds_dma_instrs(spec, XC); /* DMA instrs per wave per X row */ + const int IY = lds_dma_instrs(spec, BWO); + const int X_ELEMS = lds_x_row_elems(spec); + const int Y_ELEMS = lds_dy_row_elems(spec); + std::vector x_bufs, y_bufs, x_addr, y_addr; + for(int i = 0; i < NX; ++i) + { + std::string h = "xr" + std::to_string(i); + x_bufs.push_back(e.smem_alloc_1d(io_type, X_ELEMS, h.c_str())); + } + for(int i = 0; i < NY; ++i) + { + std::string h = "yr" + std::to_string(i); + y_bufs.push_back(e.smem_alloc_1d(io_type, Y_ELEMS, h.c_str())); + } + for(V s : x_bufs) + { + x_addr.push_back(rocke_b_smem_addr_of(b, s)); + } + for(V s : y_bufs) + { + y_addr.push_back(rocke_b_smem_addr_of(b, s)); + } + + /* DMA lane geometry: slot L = (wave * I + i) * 64 + lane -> (col, chunk). */ + auto dma_lanes = [&](int n_instr, int cols) { + std::vector out; + for(int i = 0; i < n_instr; ++i) + { + V slot = e.muli(wave, n_instr); + slot = e.addi(slot, i); + slot = e.muli(slot, kWave); + slot = e.add(slot, lane); + LdsLane l; + l.col = e.divi(slot, CH); + l.ch = e.muli(e.modi(slot, CH), 8); + l.ok = e.lti(l.col, cols); + out.push_back(l); + } + return out; + }; + std::vector x_lanes = dma_lanes(IX, XC); + std::vector y_lanes = dma_lanes(IY, BWO); + + auto dma = [&](V rsrc, V addr, int n_instr, int i, V voff) { + V o = e.muli(wave, n_instr); + o = e.addi(o, i); + o = e.muli(o, kWave * 16); + V lds = rocke_b_smem_ptr_add(b, addr, e.zext64(o)); + rocke_b_async_buffer_load_lds(b, rsrc, lds, voff, c0, 4, 0); + }; + + /* Fragment read element indices (lane: chunk lc of a column). */ + auto col_index = [&](V col) { + V a = e.muli(col, BC); + return e.add(a, e.muli(lc, 8)); + }; + auto cvt8 = [&](V v) { + std::vector out; + for(int el = 0; el < 8; ++el) + { + out.push_back(rocke_b_cast_to_f32(b, rocke_b_vec_extract(b, v, el))); + } + return out; + }; + + V zero = rocke_b_const_f32(b, 0.0); + const int n_acc = TAPS * 8; + const int PER_STEP = S * IX + IY; + + std::vector accs = e.for_iter( + part, + n_units, + n_part, + iter_names("acc", n_acc), + std::vector(n_acc, zero), + "unit", + [&](V unit, std::vector carried) { + V hb = e.mod(unit, n_hb); + V rest = e.div(unit, n_hb); + V ct = e.mod(rest, n_ct); + V n_img = e.div(rest, n_ct); + V h0 = e.muli(hb, RPS); + V q0 = e.muli(ct, BWO); + V xi0 = e.subi(e.muli(q0, S), PAD); /* first X column of the tile */ + V x_row0 = e.subi(e.muli(h0, S), PAD); /* X row of relative index 0 */ + V k_rps = e.ci(RPS); + V n_steps = e.smin(k_rps, e.sub(P.p_Ho, h0)); + + /* Per-lane DMA offsets of this unit (rows added per issue). */ + std::vector x_offs, y_offs; + for(const LdsLane& l : x_lanes) + { + V wi = e.add(xi0, l.col); + V a = e.ge(wi, c0); + a = e.land(a, e.lt(wi, P.p_Wi)); + V okc = e.land(l.ok, a); + V cc = e.add(c_blk0, l.ch); + okc = e.land(okc, e.lt(cc, p_C)); + V off = e.mul(n_img, P.p_B_stride_n); + off = e.add(off, e.mul(wi, P.p_B_stride_wi)); + off = e.add(off, cc); + x_offs.push_back({off, okc}); + } + for(const LdsLane& l : y_lanes) + { + V wo = e.add(q0, l.col); + V cc = e.add(c_blk0, l.ch); + V a = e.lt(wo, P.p_Wo); + a = e.land(a, e.lt(cc, p_C)); + V oky = e.land(l.ok, a); + V off = e.mul(n_img, P.p_A_stride_n); + off = e.add(off, e.mul(wo, P.p_A_stride_wi)); + off = e.add(off, cc); + y_offs.push_back({off, oky}); + } + + auto issue_x = [&](V xr, int slot) { + V hi = e.add(x_row0, xr); + /* Rows of later steps than the segment's own stay masked. */ + V a = e.ge(hi, c0); + a = e.land(a, e.lt(hi, P.p_Hi)); + V lim = e.addi(e.muli(n_steps, S), PRE); + V ok_r = e.land(a, e.lt(xr, lim)); + for(int i = 0; i < (int)x_offs.size(); ++i) + { + V okk = e.land(x_offs[i].ok, ok_r); + V row = e.mul(hi, P.p_B_stride_hi); + V voff = e.mul(e.add(x_offs[i].off, row), c2); + voff = e.select(okk, voff, oob); + dma(b_rsrc, x_addr[slot], IX, i, voff); + } + }; + auto issue_y = [&](V t, int slot) { + V ho = e.add(h0, t); + V ok_r = e.lt(t, n_steps); + for(int i = 0; i < (int)y_offs.size(); ++i) + { + V okk = e.land(y_offs[i].ok, ok_r); + V row = e.mul(ho, P.p_A_stride_hi); + V voff = e.mul(e.add(y_offs[i].off, row), c2); + voff = e.select(okk, voff, oob); + dma(a_rsrc, y_addr[slot], IY, i, voff); + } + }; + /* DMA the S X rows that step t adds and dY row t (tm = t mod U). */ + auto issue_step = [&](V t, int tm) { + for(int j = 0; j < S; ++j) + { + V xr = e.addi(e.muli(t, S), PRE + j); + issue_x(xr, (S * tm + PRE + j) % NX); + } + issue_y(t, tm % NY); + }; + + /* Prologue: the PRE leading X rows, then steps 0..D-1. */ + for(int j = 0; j < PRE; ++j) + { + issue_x(e.ci(j), j % NX); + } + for(int t = 0; t < D; ++t) + { + issue_step(e.ci(t), t); + } + + auto compute = [&](int tm, std::vector& la) { + std::vector> dyv; + for(int jw = 0; jw < WL; ++jw) + { + V col = e.addi(e.muli(lw, WL), jw); + dyv.push_back( + cvt8(e.smem_load(y_bufs[tm % NY], c0, col_index(col), io_type, 8))); + } + for(int r = 0; r < KH; ++r) + { + int slot = (S * tm + r) % NX; + std::map> win; + for(int jw = 0; jw < WL; ++jw) + { + for(int s = 0; s < KW; ++s) + { + int xc = jw * S + s; + if(win.find(xc) == win.end()) + { + V col = e.addi(e.muli(lw, WL * S), xc); + win[xc] = cvt8( + e.smem_load(x_bufs[slot], c0, col_index(col), io_type, 8)); + } + const std::vector& xv = win[xc]; + for(int el = 0; el < 8; ++el) + { + int i = (r * KW + s) * 8 + el; + la[i] = rocke_b_fma(b, dyv[jw][el], xv[el], la[i]); + } + } + } + } + }; + + V trips = e.cdivi(n_steps, U); + trips = e.smax(trips, c1); + std::vector accs_out + = e.for_iter(c0, + trips, + c1, + iter_names("a", carried.size()), + carried, + "rb", + [&](V rb, std::vector la) { + for(int u = 0; u < U; ++u) + { + V t = e.addi(e.muli(rb, U), u); + rocke_b_s_waitcnt(b, (D - 1) * PER_STEP, 0, -1); + rocke_b_s_barrier_bare(b); + issue_step(e.addi(t, D), (u + D) % U); + compute(u, la); + } + return la; + }); + rocke_b_s_waitcnt(b, 0, 0, -1); + rocke_b_s_barrier_bare(b); + return accs_out; + }); + if((int)accs.size() != n_acc) + { + return NULL; + } + + /* ---- reduce over the LW column lanes, store ws[part][tap][c] ---- */ + V red = e.smem_alloc_1d(rocke_f32(), LW * BC, "red"); + for(int t = 0; t < TAPS; ++t) + { + V base = e.muli(lw, BC); + base = e.add(base, e.muli(lc, 8)); + for(int e0 = 0; e0 < 8; e0 += 4) + { + V comps[4]; + for(int el = 0; el < 4; ++el) + { + comps[el] = accs[t * 8 + e0 + el]; + } + V vec = rocke_b_vec_pack(b, comps, 4, rocke_f32()); + e.smem_store_f32(red, c0, e.addi(base, e0), vec, 4); + } + rocke_b_sync(b); + for(int k = 0; k < cdiv(BC, T); ++k) + { + V o = e.addi(tid, k * T); + V o_ok = e.lti(o, BC); + V tot = NULL; + for(int q = 0; q < LW; ++q) + { + V kq = e.ci(q * BC); + V idx = e.select(o_ok, e.add(kq, o), c0); + V v = rocke_b_vec_extract(b, e.smem_load_f32(red, c0, idx, 1), 0); + tot = (tot == NULL) ? v : rocke_b_fadd(b, tot, v); + } + V c_glb = e.add(c_blk0, o); + V okw = e.land(o_ok, e.lt(c_glb, p_C)); + V off = e.addi(e.muli(part, TAPS), t); + off = e.add(e.mul(off, p_C), c_glb); + e.if_(okw, [&]() { rocke_b_global_store(b, P.D, off, tot, 4); }); + } + rocke_b_sync(b); + } + return rocke_ir_builder_ok(b) ? b->kernel : NULL; +} + +} // namespace + +extern "C" rocke_kernel_def_t* rocke_build_dw_wgrad_lds(rocke_ir_builder_t* b, + const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch) +{ + if(b == NULL || spec == NULL) + { + return NULL; + } + return build_lds(b, spec, arch ? arch : "gfx950"); +} + +extern "C" rocke_kernel_def_t* rocke_build_dw_wgrad_lds_new(rocke_ir_builder_t* b, + const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch) +{ + return build_new( + b, + spec != NULL, + "build_dw_wgrad_lds", + [&](char* out, size_t cap) { return rocke_dw_wgrad_lds_kernel_name(spec, out, cap); }, + [&]() { return rocke_build_dw_wgrad_lds(b, spec, arch); }); +} + +extern "C" rocke_status_t rocke_dw_wgrad_lds_lower_to_llvm(const rocke_dw_wgrad_lds_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + if(arch == NULL) + { + arch = "gfx950"; + } + return lower_to_llvm( + spec != NULL, + arch, + flavor, + out_ll, + err, + err_cap, + "build_dw_wgrad_lds failed", + [&](rocke_ir_builder_t* b) { return rocke_build_dw_wgrad_lds_new(b, spec, arch); }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_wgrad_gather.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_wgrad_gather.cpp new file mode 100644 index 000000000000..ae12ce5a0e18 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_wgrad_gather.cpp @@ -0,0 +1,758 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * conv_direct_wgrad_gather.cpp -- C++ engine port of build_wgrad_gather + * (library/kernels/common/conv_direct_wgrad_gather.py): the dense small-C + * direct wgrad, one GEMM per filter row with N = (s, c). + * + * Every IR op (including each arith.constant) is emitted in the order the + * Python builder emits it; see the contract in + * rocke/instance_conv_direct_wgrad_internal.hpp. + */ +#include "rocke/instance_conv_direct_wgrad_gather.h" + +#include +#include +#include + +#include +#include + +#include "rocke/arch_target.h" /* rocke_set_unknown_arch_reason */ +#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx */ +#include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/instance_conv_direct_grouped_internal.h" /* rocke_dconv_emit_params */ +#include "rocke/instance_conv_direct_wgrad_internal.hpp" +#include "rocke/instance_conv_direct_wgrad_tiled.h" /* rocke_direct_wgrad_pick_swizzle */ + +using namespace rocke_dwg; + +namespace +{ + +constexpr int kAtom = 16; +constexpr int kWo = 32; +constexpr int kVec = 8; + +int block_k(const rocke_direct_wgrad_gather_spec_t* s) +{ + return kAtom * s->wave_k16 * s->waves_k; +} + +int n_atoms(const rocke_direct_wgrad_gather_spec_t* s) +{ + return cdiv(s->problem.KW * s->problem.cpg, kAtom); +} + +int unroll(const rocke_direct_wgrad_gather_spec_t* s) +{ + return s->prefetch + cdiv(s->problem.KH, s->problem.stride); +} + +int rps(const rocke_direct_wgrad_gather_spec_t* s) +{ + return cdiv(s->rows_per_seg, unroll(s)) * unroll(s); +} + +int x_phase(const rocke_direct_wgrad_gather_spec_t* s) +{ + return pymod(-s->problem.PAD * s->problem.cpg, kVec); +} + +/* Elements of one staged X row: phase + the strip's columns. */ +int x_elems(const rocke_direct_wgrad_gather_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + return x_phase(s) + ((kWo - 1) * p->stride + p->KW) * p->cpg; +} + +int n_k_tiles(const rocke_direct_wgrad_gather_spec_t* s) +{ + return cdiv(s->problem.kpg, block_k(s)); +} + +/* WgradGatherSpec.dma_instrs(): 16-byte DMA instructions per wave and row. */ +int dma_instrs_x(const rocke_direct_wgrad_gather_spec_t* s) +{ + return cdiv(cdiv(cdiv(x_elems(s), kVec), kWave), s->waves_k); +} + +int dma_instrs_dy(const rocke_direct_wgrad_gather_spec_t* s) +{ + return cdiv(cdiv(kWo * block_k(s) / kVec, kWave), s->waves_k); +} + +/* Elements of one X ring slot: whole DMA instructions of every wave, +16 + * because a transpose read of the last column atom runs past it. */ +int x_slot_elems(const rocke_direct_wgrad_gather_spec_t* s) +{ + return s->waves_k * kWave * dma_instrs_x(s) * kVec + 16; +} + +int dy_slot_elems(const rocke_direct_wgrad_gather_spec_t* s) +{ + return s->waves_k * kWave * dma_instrs_dy(s) * kVec; +} + +} // namespace + +extern "C" rocke_direct_wgrad_gather_spec_t rocke_direct_wgrad_gather_spec_default(void) +{ + rocke_direct_wgrad_gather_spec_t s; + memset(&s, 0, sizeof(s)); + s.problem = rocke_direct_conv_problem_default(); + s.name = "direct_wgrad_gather"; + s.wave_k16 = 2; + s.waves_k = 1; + s.prefetch = 2; + s.rows_per_seg = 16; + s.target_blocks = 1024; + return s; +} + +extern "C" int rocke_direct_wgrad_gather_threads(const rocke_direct_wgrad_gather_spec_t* s) +{ + return s->waves_k * kWave; +} + +extern "C" long rocke_direct_wgrad_gather_lds_bytes(const rocke_direct_wgrad_gather_spec_t* s) +{ + /* The X and dY rings as the builder allocates them. */ + long U = unroll(s); + return ((long)s->problem.stride * U * x_slot_elems(s) + U * dy_slot_elems(s)) * 2; +} + +extern "C" int rocke_direct_wgrad_gather_n_partitions(const rocke_direct_wgrad_gather_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + int n_units = p->N * cdiv(problem_wo(p), kWo) * cdiv(problem_ho(p), rps(s)); + int c = cdiv(s->target_blocks, n_k_tiles(s)); + return n_units < c ? n_units : c; +} + +extern "C" int rocke_direct_wgrad_gather_fold_taps(const rocke_direct_wgrad_gather_spec_t* s) +{ + return s->problem.KH * s->problem.KW * s->problem.cpg; +} + +extern "C" long rocke_direct_wgrad_gather_workspace_bytes(const rocke_direct_wgrad_gather_spec_t* s) +{ + return (long)rocke_direct_wgrad_gather_n_partitions(s) * rocke_direct_wgrad_gather_fold_taps(s) + * s->problem.kpg * 4; +} + +extern "C" void rocke_direct_wgrad_gather_launch_geometry(const rocke_direct_wgrad_gather_spec_t* s, + int grid[3], + int block[3]) +{ + grid[0] = n_k_tiles(s); + grid[1] = rocke_direct_wgrad_gather_n_partitions(s); + grid[2] = 1; + block[0] = rocke_direct_wgrad_gather_threads(s); + block[1] = 1; + block[2] = 1; +} + +extern "C" rocke_status_t rocke_direct_wgrad_gather_kernel_name( + const rocke_direct_wgrad_gather_spec_t* s, char* out, size_t out_cap) +{ + if(s == NULL || out == NULL || out_cap == 0) + { + return ROCKE_ERR_VALUE; + } + const rocke_direct_conv_problem_t* p = &s->problem; + char f_buf[64], wk_buf[32], pf_buf[24], rps_buf[24], tb_buf[24]; + snprintf(f_buf, sizeof f_buf, "f%dx%dp%ds%dc%d", p->KH, p->KW, p->PAD, p->stride, p->cpg); + snprintf(wk_buf, sizeof wk_buf, "wk%dx%d", s->wave_k16, s->waves_k); + snprintf(pf_buf, sizeof pf_buf, "pf%d", s->prefetch); + snprintf(rps_buf, sizeof rps_buf, "rps%d", s->rows_per_seg); + snprintf(tb_buf, sizeof tb_buf, "tb%d", s->target_blocks); + const char* parts[5] = {f_buf, wk_buf, pf_buf, rps_buf, tb_buf}; + const char* flag_names[1] = {"bf16"}; + int on[1] = {problem_is_bf16(p) ? 1 : 0}; + return rocke_kernel_name_join(s->name, parts, 5, flag_names, on, 1, out, out_cap, NULL); +} + +extern "C" bool rocke_direct_wgrad_gather_is_valid_spec(const rocke_direct_wgrad_gather_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ + char msg[ROCKE_ERR_MSG_CAP]; + if(s == NULL) + { + set_reason(reason, reason_cap, "null spec"); + return false; + } + if(arch == NULL) + { + arch = "gfx950"; + } + const rocke_archtarget_t* target = rocke_archtarget_from_gfx(arch); + if(target == NULL) + { + rocke_set_unknown_arch_reason(reason, reason_cap, arch); + return false; + } + const rocke_direct_conv_problem_t* p = &s->problem; + if(dilation_reason(p, msg, sizeof msg)) + { + set_reason(reason, reason_cap, msg); + return false; + } + const char* why = NULL; + int acc = p->KH * s->wave_k16 * n_atoms(s) * 4; + if(p->groups != 1) + { + why = "dense layers only"; + } + else if(!problem_dtype_ok(p)) + { + why = "fp16/bf16 only"; + } + else if(!target->memory.has_ds_read_tr) + { + ROCKE_ERR_SNPRINTF(msg, sizeof msg, "needs ds_read_b64_tr_b16 (absent on %s)", arch); + why = msg; + } + else if(target->wave_size != 64 || rocke_archtarget_async_lds_max_dwords(target) < 4) + { + ROCKE_ERR_SNPRINTF( + msg, sizeof msg, "needs wave64 + 16-byte buffer_load_lds (absent on %s)", arch); + why = msg; + } + else if(p->kpg % kVec) + { + why = "K must be a multiple of 8"; + } + else if(p->KW * p->cpg > 96) + { + why = "KW * C must be at most 96"; + } + /* Whole 16-byte chunks from an aligned start at a build-time phase. */ + else if((p->W * p->cpg) % kVec || (kWo * p->stride * p->cpg) % kVec) + { + why = "W * C and 32 * stride * C must be multiples of 8"; + } + else if(acc > 192) + { + ROCKE_ERR_SNPRINTF(msg, sizeof msg, "%d accumulator VGPRs (> 192)", acc); + why = msg; + } + else if(s->wave_k16 < 1 || s->waves_k < 1 || rocke_direct_wgrad_gather_threads(s) > 1024) + { + why = "wave_k16 >= 1 and waves_k in 1..16"; + } + else if(s->prefetch < 1) + { + why = "prefetch >= 1"; + } + else if(rocke_direct_wgrad_gather_lds_bytes(s) > 160 * 1024) + { + ROCKE_ERR_SNPRINTF( + msg, sizeof msg, "LDS %ld B > 160 KiB", rocke_direct_wgrad_gather_lds_bytes(s)); + why = msg; + } + else if(s->rows_per_seg < 1) + { + why = "rows_per_seg >= 1"; + } + else if(s->target_blocks < 1) + { + why = "target_blocks >= 1"; + } + if(why != NULL) + { + set_reason(reason, reason_cap, why); + return false; + } + set_reason(reason, reason_cap, "ok"); + return true; +} + +namespace +{ + +struct Chunk +{ + V a, b, ok; /* X: (element, -, ok); dY: (row, k element, ok) */ +}; + +struct Off +{ + V off, ok; +}; + +rocke_kernel_def_t* + build(rocke_ir_builder_t* b, const rocke_direct_wgrad_gather_spec_t* spec, const char* arch) +{ + char reason[ROCKE_ERR_MSG_CAP]; + if(!rocke_direct_wgrad_gather_is_valid_spec(spec, arch, reason, sizeof reason)) + { + std::string prefix = std::string("invalid WgradGatherSpec for ") + arch; + fail_invalid(b, prefix.c_str(), reason); + return NULL; + } + Em e{b}; + const rocke_direct_conv_problem_t& p = spec->problem; + const int KH = p.KH, KW = p.KW, PAD = p.PAD, S = p.stride, C = p.cpg; + const bool bf16 = problem_is_bf16(&p); + const rocke_type_t* io_type = bf16 ? rocke_bf16() : rocke_f16(); + const int WK16 = spec->wave_k16; + const int NA = n_atoms(spec); + const int BK = block_k(spec); + const int CHK = BK / kVec; + const int PH = x_phase(spec); + const int XE = x_elems(spec); + const int XCH = cdiv(XE, kVec); /* 16-byte chunks of a staged X row */ + const int D = spec->prefetch; + const int U = unroll(spec); + const int NX = S * U, NY = U; + const int RPS = rps(spec); + const int PRE = KH - S > 0 ? KH - S : 0; + const int TAPS = KH * KW * C; + const int THREADS = rocke_direct_wgrad_gather_threads(spec); + + const int IX = dma_instrs_x(spec); + const int IY = dma_instrs_dy(spec); + const int X_ELEMS = x_slot_elems(spec); + /* When a pixel row of the strip (S * C elements) and the phase keep every + * lane's 8-byte read aligned, the X operand is a transpose read too. */ + const bool X_TR = (S * C) % 4 == 0 && PH % 4 == 0; + const int Y_ELEMS = dy_slot_elems(spec); + const int PER_STEP = S * IX + IY; + + std::vector k_origins; + for(int o = 0; o < BK; o += kAtom) + { + k_origins.push_back(o); + } + std::vector swz_k; + { + int zero_shift = 0; + int bits[6]; + int n = rocke_direct_wgrad_pick_swizzle( + CHK, &zero_shift, 1, kWo, kWo, k_origins.data(), (int)k_origins.size(), bits); + swz_k.assign(bits, bits + n); + } + + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", THREADS); + rocke_attr_set_bool(b, &b->kernel->attrs, "lds_alias_scopes", true); + rocke_dconv_params_t P; + rocke_dconv_emit_params(b, &P, "wgrad", io_type); + if(!rocke_ir_builder_ok(b)) + { + return NULL; + } + + V c0 = e.ci(0); + V c1 = e.ci(1); + V c2 = e.ci(2); + V oob = e.ci(((int64_t)1 << 31) - 1); + V p_K = P.p_total_k; + V tid = rocke_b_thread_id_x(b); + V lane = e.modi(tid, kWave); + V wave = rocke_b_readfirstlane(b, e.divi(tid, kWave)); + V k_tile = rocke_b_block_id_x(b); + V part = rocke_b_block_id_y(b); + V k_blk0 = e.muli(k_tile, BK); + V n_q = e.cdivi(P.p_Wo, kWo); + V n_hb = e.cdivi(P.p_Ho, RPS); + V n_units = e.mul(e.mul(P.p_N, n_q), n_hb); + V n_kt = e.cdivi(p_K, BK); + V tb = e.ci(spec->target_blocks); + V np = e.add(tb, e.sub(n_kt, c1)); + np = e.div(np, n_kt); + V n_part = e.smin(n_units, np); + + V a_rsrc = rocke_b_buffer_rsrc(b, P.A, P.A_bytes); /* dY */ + V b_rsrc = rocke_b_buffer_rsrc(b, P.Bp, P.B_bytes); /* X */ + + std::vector x_bufs, y_bufs, x_addr, y_addr; + for(int i = 0; i < NX; ++i) + { + std::string h = "xg" + std::to_string(i); + x_bufs.push_back(e.smem_alloc_1d(io_type, X_ELEMS, h.c_str())); + } + for(int i = 0; i < NY; ++i) + { + std::string h = "yg" + std::to_string(i); + y_bufs.push_back(e.smem_alloc_1d(io_type, Y_ELEMS, h.c_str())); + } + for(V s : x_bufs) + { + x_addr.push_back(rocke_b_smem_addr_of(b, s)); + } + for(V s : y_bufs) + { + y_addr.push_back(rocke_b_smem_addr_of(b, s)); + } + + auto dma = [&](V rsrc, V addr, int n_instr, int i, V voff) { + V o = e.muli(wave, n_instr); + o = e.addi(o, i); + o = e.muli(o, kWave * 16); + V lds = rocke_b_smem_ptr_add(b, addr, e.zext64(o)); + rocke_b_async_buffer_load_lds(b, rsrc, lds, voff, c0, 4, 0); + }; + auto dma_slot = [&](int n_instr, int i) { + V slot = e.muli(wave, n_instr); + slot = e.addi(slot, i); + slot = e.muli(slot, kWave); + return e.add(slot, lane); + }; + + /* X: chunk L of a row is elements [8 L, 8 L + 8) from the aligned start. */ + std::vector x_chunks; + for(int i = 0; i < IX; ++i) + { + V slot = dma_slot(IX, i); + Chunk ch; + ch.a = e.muli(slot, kVec); + ch.b = NULL; + ch.ok = e.lti(slot, XCH); + x_chunks.push_back(ch); + } + /* dY: slot -> (pixel row, swizzled k chunk). */ + std::vector y_chunks; + for(int i = 0; i < IY; ++i) + { + V slot = dma_slot(IY, i); + V row = e.divi(slot, CHK); + V phys = e.modi(slot, CHK); + V sw = e.swz_val(row, swz_k, c1); + V log = (sw == NULL) ? phys : e.bxor(phys, sw); + Chunk ch; + ch.a = row; + ch.b = e.muli(log, kVec); + ch.ok = e.lti(row, kWo); + y_chunks.push_back(ch); + } + + /* ---- fragment addresses ---- */ + V lane_hi = e.divi(lane, 16); + V lane_q = e.modi(e.divi(lane, 4), 4); + V lane_c4 = e.modi(lane, 4); + V lane_n = e.modi(lane, 16); + V wave_k0 = e.muli(wave, WK16 * kAtom); + + auto dy_index = [&](int rd, V col0) { + V row = e.add(e.muli(lane_hi, 8), lane_q); + row = e.addi(row, rd * 4); + V col = e.add(col0, e.muli(lane_c4, 4)); + V chunk = e.divi(col, kVec); + V sw = e.swz_val(row, swz_k, c1); + V phys = (sw == NULL) ? chunk : e.bxor(chunk, sw); + V idx = e.muli(row, CHK); + idx = e.add(idx, phys); + idx = e.muli(idx, kVec); + return e.add(idx, e.modi(col, kVec)); + }; + std::vector> dy_idx(WK16); + for(int mi = 0; mi < WK16; ++mi) + { + for(int rd = 0; rd < 2; ++rd) + { + V col0 = e.addi(wave_k0, mi * kAtom); + dy_idx[mi].push_back(dy_index(rd, col0)); + } + } + /* X gather base per atom: element PH + (S * px0 + s) * C + c, px0 = 8 * (lane / 16). */ + std::vector x_base; + for(int nj = 0; nj < NA; ++nj) + { + V n = e.addi(lane_n, nj * kAtom); + V s_i = e.divi(n, C); + s_i = e.smini(s_i, KW - 1); /* n >= KW*C: any in-range column */ + V c_i = e.modi(n, C); + V kph = e.ci(PH); + V t = e.add(e.muli(lane_hi, 8 * S), s_i); + t = e.add(e.muli(t, C), c_i); + x_base.push_back(e.add(kph, t)); + } + std::vector> x_tr_idx(NA); + for(int nj = 0; nj < NA; ++nj) + { + for(int rd = 0; rd < 2; ++rd) + { + V k0 = e.ci(PH + nj * kAtom); + V t = e.add(e.muli(lane_hi, 8), lane_q); + t = e.addi(t, rd * 4); + t = e.muli(t, S * C); + t = e.add(t, e.muli(lane_c4, 4)); + x_tr_idx[nj].push_back(e.add(k0, t)); + } + } + + auto read_dy = [&](int buf) { + std::vector out; + for(int mi = 0; mi < WK16; ++mi) + { + V lo = e.ds_read_tr(y_bufs[buf], c0, dy_idx[mi][0], io_type); + V hi = e.ds_read_tr(y_bufs[buf], c0, dy_idx[mi][1], io_type); + out.push_back(rocke_b_vec_concat(b, lo, hi)); + } + return out; + }; + auto read_x = [&](int buf, int nj) { + if(X_TR) + { + V lo = e.ds_read_tr(x_bufs[buf], c0, x_tr_idx[nj][0], io_type); + V hi = e.ds_read_tr(x_bufs[buf], c0, x_tr_idx[nj][1], io_type); + return rocke_b_vec_concat(b, lo, hi); + } + V vals[8]; + for(int j = 0; j < 8; ++j) + { + V idx = e.addi(x_base[nj], j * S * C); + vals[j] = rocke_b_vec_extract(b, e.smem_load(x_bufs[buf], c0, idx, io_type, 1), 0); + } + return rocke_b_vec_pack(b, vals, 8, io_type); + }; + + V zero_acc = rocke_b_zero_vec_f32(b, 4); + V zero_elem = rocke_b_vec_extract(b, rocke_b_zero_vec(b, io_type, 2), 0); + const int n_acc = KH * WK16 * NA; + auto acc_index = [&](int r, int mi, int nj) { return (r * WK16 + mi) * NA + nj; }; + + std::vector accs = e.for_iter( + part, + n_units, + n_part, + iter_names("acc", n_acc), + std::vector(n_acc, zero_acc), + "unit", + [&](V unit, std::vector carried) { + V hb = e.mod(unit, n_hb); + V rest = e.div(unit, n_hb); + V q = e.mod(rest, n_q); + V n_img = e.div(rest, n_q); + V h0 = e.muli(hb, RPS); + V q0 = e.muli(q, kWo); + V col0 = e.subi(e.muli(q0, S), PAD); /* first X column of the strip */ + V x_row0 = e.subi(e.muli(h0, S), PAD); + V k_rps = e.ci(RPS); + V n_steps = e.smin(k_rps, e.sub(P.p_Ho, h0)); + /* Strip X start, 8-aligned: element col0 * C sits at PH. */ + V xs0 = e.mul(n_img, P.p_B_stride_n); + V x_start = e.add(xs0, e.subi(e.muli(col0, C), PH)); + std::vector y_offs; + for(const Chunk& ch : y_chunks) + { + V wo = e.add(q0, ch.a); + V kk = e.add(k_blk0, ch.b); + V a = e.lt(wo, P.p_Wo); + a = e.land(a, e.lt(kk, p_K)); + V ok = e.land(ch.ok, a); + V off = e.mul(n_img, P.p_A_stride_n); + off = e.add(off, e.mul(wo, P.p_A_stride_wi)); + y_offs.push_back({e.add(off, kk), ok}); + } + /* Padding columns staged from neighbouring rows: [lo, hi) element ranges + * of the row to zero, for the edge strips. */ + V pad_lo = e.muli(e.sub(c0, col0), C); + pad_lo = e.smax(pad_lo, c0); /* cols < 0 */ + V cols_in = e.sub(P.p_Wi, col0); /* cols >= W start here (relative) */ + V kph = e.ci(PH); + V pad_hi0 = e.add(kph, e.muli(cols_in, C)); + V edge = e.lt(col0, c0); + edge = e.lor(edge, e.lti(pad_hi0, XE)); + + auto issue_x = [&](V xr, int slot) { + V hi = e.add(x_row0, xr); + V a = e.ge(hi, c0); + a = e.land(a, e.lt(hi, P.p_Hi)); + V lim = e.addi(e.muli(n_steps, S), PRE); + V ok_r = e.land(a, e.lt(xr, lim)); + V row_off = e.add(x_start, e.mul(hi, P.p_B_stride_hi)); + for(int i = 0; i < (int)x_chunks.size(); ++i) + { + V okk = e.land(x_chunks[i].ok, ok_r); + V voff = e.mul(e.add(row_off, x_chunks[i].a), c2); + voff = e.select(okk, voff, oob); + dma(b_rsrc, x_addr[slot], IX, i, voff); + } + }; + auto issue_y = [&](V t, int slot) { + V ho = e.add(h0, t); + V ok_r = e.lt(t, n_steps); + for(int i = 0; i < (int)y_offs.size(); ++i) + { + V okk = e.land(y_offs[i].ok, ok_r); + V row = e.mul(ho, P.p_A_stride_hi); + V voff = e.mul(e.add(y_offs[i].off, row), c2); + voff = e.select(okk, voff, oob); + dma(a_rsrc, y_addr[slot], IY, i, voff); + } + }; + auto issue_step = [&](V t, int tm) { + for(int j = 0; j < S; ++j) + { + V xr = e.addi(e.muli(t, S), PRE + j); + issue_x(xr, (S * tm + PRE + j) % NX); + } + issue_y(t, tm % NY); + }; + /* Zero the padding columns of freshly landed X rows (edge strips). */ + auto fix_padding = [&](const std::vector& slots) { + e.if_(edge, [&]() { + for(int slot : slots) + { + for(int k = 0; k < cdiv(XE, THREADS); ++k) + { + V el = e.addi(tid, k * THREADS); + V rel = e.subi(el, PH); + V in_row = e.lti(el, XE); + V lo = e.lt(rel, pad_lo); + V hi = e.ge(el, pad_hi0); + V is_pad = e.land(in_row, e.lor(lo, hi)); + e.if_(is_pad, + [&]() { e.smem_store(x_bufs[slot], c0, el, zero_elem, 1); }); + } + } + rocke_b_s_waitcnt(b, -1, 0, -1); + rocke_b_s_barrier_bare(b); + }); + }; + + for(int j = 0; j < PRE; ++j) + { + issue_x(e.ci(j), j % NX); + } + for(int t = 0; t < D; ++t) + { + issue_step(e.ci(t), t); + } + + V trips = e.cdivi(n_steps, U); + trips = e.smax(trips, c1); + std::vector accs_out + = e.for_iter(c0, + trips, + c1, + iter_names("a", carried.size()), + carried, + "rb", + [&](V rb, std::vector la) { + for(int tm = 0; tm < U; ++tm) + { + V t = e.addi(e.muli(rb, U), tm); + rocke_b_s_waitcnt(b, (D - 1) * PER_STEP, 0, -1); + rocke_b_s_barrier_bare(b); + std::vector new_slots; + if(tm == 0) + { + /* Step 0 of a segment also consumes the PRE leading rows. */ + for(int j = 0; j < PRE; ++j) + { + new_slots.push_back(j % NX); + } + } + for(int j = 0; j < S; ++j) + { + new_slots.push_back((S * tm + PRE + j) % NX); + } + fix_padding(new_slots); + issue_step(e.addi(t, D), (tm + D) % U); + std::vector dyf = read_dy(tm % NY); + for(int r = 0; r < KH; ++r) + { + int slot = (S * tm + r) % NX; + for(int nj = 0; nj < NA; ++nj) + { + V xf = read_x(slot, nj); + for(int mi = 0; mi < WK16; ++mi) + { + int i = acc_index(r, mi, nj); + la[i] = e.mfma16x16x32(bf16, dyf[mi], xf, la[i]); + } + } + } + } + return la; + }); + rocke_b_s_waitcnt(b, 0, 0, -1); + rocke_b_s_barrier_bare(b); + return accs_out; + }); + if((int)accs.size() != n_acc) + { + return NULL; + } + + /* ---- epilogue: partial slot ws[part][(r, s, c)][k] ---- + * acc lane layout: k = (lane / 16) * 4 + slot of atom mi, n = lane % 16 of atom nj. */ + V k_lane0 = e.add(k_blk0, wave_k0); + k_lane0 = e.add(k_lane0, e.muli(lane_hi, 4)); + for(int r = 0; r < KH; ++r) + { + for(int nj = 0; nj < NA; ++nj) + { + V n = e.addi(lane_n, nj * kAtom); + V n_ok = e.lti(n, KW * C); + V kt = e.ci(r * KW * C); + V tap = e.add(kt, n); /* (r, s, c) flattened */ + V row_base = e.add(e.muli(part, TAPS), tap); + row_base = e.mul(row_base, p_K); + for(int mi = 0; mi < WK16; ++mi) + { + for(int slot = 0; slot < 4; ++slot) + { + V kk = e.addi(k_lane0, mi * kAtom + slot); + V ok = e.land(n_ok, e.lt(kk, p_K)); + V val = rocke_b_vec_extract(b, accs[acc_index(r, mi, nj)], slot); + e.if_(ok, [&]() { rocke_b_global_store(b, P.D, e.add(row_base, kk), val, 4); }); + } + } + } + } + return rocke_ir_builder_ok(b) ? b->kernel : NULL; +} + +} // namespace + +extern "C" rocke_kernel_def_t* rocke_build_direct_wgrad_gather( + rocke_ir_builder_t* b, const rocke_direct_wgrad_gather_spec_t* spec, const char* arch) +{ + if(b == NULL || spec == NULL) + { + return NULL; + } + return build(b, spec, arch ? arch : "gfx950"); +} + +extern "C" rocke_kernel_def_t* rocke_build_direct_wgrad_gather_new( + rocke_ir_builder_t* b, const rocke_direct_wgrad_gather_spec_t* spec, const char* arch) +{ + return build_new( + b, + spec != NULL, + "build_wgrad_gather", + [&](char* out, size_t cap) { + return rocke_direct_wgrad_gather_kernel_name(spec, out, cap); + }, + [&]() { return rocke_build_direct_wgrad_gather(b, spec, arch); }); +} + +extern "C" rocke_status_t + rocke_direct_wgrad_gather_lower_to_llvm(const rocke_direct_wgrad_gather_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + if(arch == NULL) + { + arch = "gfx950"; + } + return lower_to_llvm( + spec != NULL, + arch, + flavor, + out_ll, + err, + err_cap, + "build_wgrad_gather failed", + [&](rocke_ir_builder_t* b) { return rocke_build_direct_wgrad_gather_new(b, spec, arch); }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_wgrad_tiled.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_wgrad_tiled.cpp new file mode 100644 index 000000000000..136217df1251 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_wgrad_tiled.cpp @@ -0,0 +1,1241 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * conv_direct_wgrad_tiled.cpp -- C++ engine port of build_direct_wgrad_tiled + * (library/kernels/common/conv_direct_wgrad_tiled.py): the stride-1 direct + * wgrad with multi-atom wave tiles and an LDS-DMA row pipeline. + * + * Every IR op (including each arith.constant) is emitted in the order the + * Python builder emits it; see the contract in + * rocke/instance_conv_direct_wgrad_internal.hpp. The functions below follow + * the Python closures one to one (seg_lane_offsets, issue_row, frag_index, + * step, segment). + */ +#include "rocke/instance_conv_direct_wgrad_tiled.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "rocke/arch_target.h" /* rocke_set_unknown_arch_reason */ +#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx */ +#include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/instance_conv_direct_grouped_internal.h" /* rocke_dconv_emit_params */ +#include "rocke/instance_conv_direct_wgrad_internal.hpp" + +using namespace rocke_dwg; + +namespace +{ + +constexpr int kAtom = 16; +constexpr int kWo = 32; +constexpr int kVec = 8; +/* Bank model used to pick the swizzle: 64 four-byte banks, a 32-lane half of a + * ds_read_b64 serviced per pass. */ +constexpr int kBanks = 64; +constexpr int kHalf = 32; + +int wave_k(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return kAtom * s->wave_k16; +} + +int wave_c(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return kAtom * s->wave_c16; +} + +int n_waves(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return s->waves_k * s->waves_c; +} + +int strip_w(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return kWo / s->unfold_n; +} + +/* DirectWgradTiledSpec.dma_instrs(): 1 KiB DMA instructions per wave and row. */ +int dma_instrs_dy(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return cdiv(cdiv(kWo * wave_k(s) * s->waves_k / kVec, kWave), n_waves(s)); +} + +int dma_instrs_x(const rocke_direct_wgrad_tiled_spec_t* s) +{ + int seg = strip_w(s) + s->problem.KW - 1; + return cdiv(cdiv(s->unfold_n * seg * wave_c(s) * s->waves_c / kVec, kWave), n_waves(s)); +} + +/* ---- pick_swizzle ------------------------------------------------------ */ + +/* _frag_row: LDS row a ds_read_b64_tr_b16 lane reads for reduction block rd. */ +int frag_row(int lane, int rd, int shift, int wu, int seg) +{ + int kk0 = (lane / 16) * 8; + return (kk0 / wu) * seg + kk0 % wu + rd * 4 + (lane / 4) % 4 + shift; +} + +int swz_of(const std::vector& bits, int row) +{ + int v = 0; + for(size_t j = 0; j < bits.size(); ++j) + { + v |= ((row >> bits[j]) & 1) << j; + } + return 2 * v; +} + +/* _swizzle_ok: no two lanes of a 32-lane half hit one bank, for every read. */ +bool swizzle_ok(const std::vector& bits, + int chunks, + const int* shifts, + int n_shifts, + int wu, + int seg, + const int* origins, + int n_origins) +{ + for(int si = 0; si < n_shifts; ++si) + { + for(int oi = 0; oi < n_origins; ++oi) + { + for(int rd = 0; rd < 2; ++rd) + { + for(int half = 0; half < 2; ++half) + { + std::set banks; + for(int lane = half * kHalf; lane < (half + 1) * kHalf; ++lane) + { + int row = frag_row(lane, rd, shifts[si], wu, seg); + int col = origins[oi] + (lane % 4) * 4; + int phys = (col / kVec) ^ swz_of(bits, row); + int byte = (row * chunks + phys) * 16 + (col % kVec) * 2; + for(int w = 0; w < 2; ++w) + { + banks.insert((byte / 4 + w) % kBanks); + } + } + if((int)banks.size() != 2 * kHalf) + { + return false; + } + } + } + } + } + return true; +} + +/* itertools.permutations(range(6), nbits) in its (lexicographic) order; + * stops at the first selection `ok` accepts. */ +template +bool first_permutation(int nbits, std::vector& cur, std::vector& used, Ok&& ok) +{ + if((int)cur.size() == nbits) + { + return ok(cur); + } + for(int v = 0; v < 6; ++v) + { + if(used[v]) + { + continue; + } + used[v] = true; + cur.push_back(v); + if(first_permutation(nbits, cur, used, ok)) + { + return true; + } + cur.pop_back(); + used[v] = false; + } + return false; +} + +std::vector pick_swizzle( + int chunks, const std::vector& shifts, int wu, int seg, const std::vector& origins) +{ + int out[6]; + int n = rocke_direct_wgrad_pick_swizzle(chunks, + shifts.data(), + (int)shifts.size(), + wu, + seg, + origins.data(), + (int)origins.size(), + out); + return std::vector(out, out + n); +} + +} // namespace + +extern "C" int rocke_direct_wgrad_pick_swizzle(int chunks, + const int* shifts, + int n_shifts, + int wu, + int seg, + const int* col_origins, + int n_origins, + int bits_out[6]) +{ + auto ok = [&](const std::vector& bits) { + return swizzle_ok(bits, chunks, shifts, n_shifts, wu, seg, col_origins, n_origins); + }; + std::vector none; + if(ok(none)) + { + return 0; + } + int n_pairs = chunks / 2; + int width = 0; + if(n_pairs > 1) + { + /* int(math.log2(n_pairs)) */ + while((2 << width) <= n_pairs) + { + ++width; + } + } + for(int nbits = 1; nbits <= width; ++nbits) + { + std::vector cur; + std::vector used(6, false); + if(first_permutation(nbits, cur, used, ok)) + { + for(int j = 0; j < nbits; ++j) + { + bits_out[j] = cur[j]; + } + return nbits; + } + } + return 0; +} + +/* ===================================================================== * + * Spec defaults / derived geometry / kernel_name / is_valid + * ===================================================================== */ + +extern "C" rocke_direct_wgrad_tiled_spec_t rocke_direct_wgrad_tiled_spec_default(void) +{ + rocke_direct_wgrad_tiled_spec_t s; + memset(&s, 0, sizeof(s)); + s.problem = rocke_direct_conv_problem_default(); + s.name = "direct_wgrad_tiled"; + s.wave_k16 = 2; + s.wave_c16 = 2; + s.waves_k = 4; + s.waves_c = 2; + s.rows_per_seg = 30; + s.prefetch = 2; + s.target_blocks = 256; + s.row_unroll = 0; + s.unfold_n = 1; + s.q_groups = 1; + s.pingpong = false; + return s; +} + +extern "C" int rocke_direct_wgrad_tiled_threads_per_block(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return n_waves(s) * s->q_groups * kWave; +} + +extern "C" int rocke_direct_wgrad_tiled_unroll(const rocke_direct_wgrad_tiled_spec_t* s) +{ + if(s->row_unroll) + { + return s->row_unroll; + } + return std::lcm(s->prefetch + 1, s->problem.KH); +} + +extern "C" int rocke_direct_wgrad_tiled_acc_vgprs(const rocke_direct_wgrad_tiled_spec_t* s) +{ + return s->problem.KH * s->problem.KW * s->wave_k16 * s->wave_c16 * 4; +} + +extern "C" long rocke_direct_wgrad_tiled_lds_bytes(const rocke_direct_wgrad_tiled_spec_t* s) +{ + /* The padded dY / X ring buffers, plus the q-group reduction buffers. */ + long slot + = (long)s->q_groups * n_waves(s) * (dma_instrs_dy(s) + dma_instrs_x(s)) * kWave * kVec * 2; + long red = 0; + if(s->q_groups > 1) + { + red = 2L * n_waves(s) * s->wave_k16 * s->wave_c16 * kWave * 4 * 4; + } + return (s->prefetch + 1) * slot + red; +} + +extern "C" int rocke_direct_wgrad_tiled_n_channel_tiles(const rocke_direct_wgrad_tiled_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + int block_k = wave_k(s) * s->waves_k; + int block_c = wave_c(s) * s->waves_c; + return p->groups * cdiv(p->kpg, block_k) * cdiv(p->cpg, block_c); +} + +extern "C" int rocke_direct_wgrad_tiled_n_segments(const rocke_direct_wgrad_tiled_spec_t* s) +{ + const rocke_direct_conv_problem_t* p = &s->problem; + return cdiv(p->N, s->unfold_n) * cdiv(problem_wo(p), strip_w(s)) * cdiv(p->H, s->rows_per_seg); +} + +extern "C" int rocke_direct_wgrad_tiled_n_splits(const rocke_direct_wgrad_tiled_spec_t* s) +{ + int a = cdiv(rocke_direct_wgrad_tiled_n_segments(s), s->q_groups); + int c = cdiv(s->target_blocks, rocke_direct_wgrad_tiled_n_channel_tiles(s)); + return a < c ? a : c; +} + +extern "C" void rocke_direct_wgrad_tiled_launch_geometry(const rocke_direct_wgrad_tiled_spec_t* s, + int grid[3], + int block[3]) +{ + grid[0] = rocke_direct_wgrad_tiled_n_channel_tiles(s) * rocke_direct_wgrad_tiled_n_splits(s); + grid[1] = 1; + grid[2] = 1; + block[0] = rocke_direct_wgrad_tiled_threads_per_block(s); + block[1] = 1; + block[2] = 1; +} + +extern "C" rocke_status_t rocke_direct_wgrad_tiled_kernel_name( + const rocke_direct_wgrad_tiled_spec_t* s, char* out, size_t out_cap) +{ + if(s == NULL || out == NULL || out_cap == 0) + { + return ROCKE_ERR_VALUE; + } + const rocke_direct_conv_problem_t* p = &s->problem; + char f_buf[48], wt_buf[32], w_buf[32], rps_buf[24], pf_buf[24], tb_buf[24]; + char u_buf[24], un_buf[24], qg_buf[24]; + snprintf(f_buf, sizeof f_buf, "f%dx%dp%d", p->KH, p->KW, p->PAD); + snprintf(wt_buf, sizeof wt_buf, "wt%dx%d", wave_k(s), wave_c(s)); + snprintf(w_buf, sizeof w_buf, "w%dx%d", s->waves_k, s->waves_c); + snprintf(rps_buf, sizeof rps_buf, "rps%d", s->rows_per_seg); + snprintf(pf_buf, sizeof pf_buf, "pf%d", s->prefetch); + snprintf(tb_buf, sizeof tb_buf, "tb%d", s->target_blocks); + snprintf(u_buf, sizeof u_buf, "u%d", s->row_unroll); + snprintf(un_buf, sizeof un_buf, "un%d", s->unfold_n); + snprintf(qg_buf, sizeof qg_buf, "qg%d", s->q_groups); + const char* parts[6] = {f_buf, wt_buf, w_buf, rps_buf, pf_buf, tb_buf}; + const char* flag_names[5] = {"bf16", u_buf, un_buf, qg_buf, "pp"}; + int flag_on[5] = {problem_is_bf16(p) ? 1 : 0, + s->row_unroll ? 1 : 0, + s->unfold_n > 1 ? 1 : 0, + s->q_groups > 1 ? 1 : 0, + s->pingpong ? 1 : 0}; + return rocke_kernel_name_join(s->name, parts, 6, flag_names, flag_on, 5, out, out_cap, NULL); +} + +extern "C" bool rocke_direct_wgrad_tiled_is_valid_spec(const rocke_direct_wgrad_tiled_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ + char msg[ROCKE_ERR_MSG_CAP]; +#define TILED_REJECT(...) \ + do \ + { \ + ROCKE_ERR_SNPRINTF(msg, sizeof(msg), __VA_ARGS__); \ + set_reason(reason, reason_cap, msg); \ + return false; \ + } while(0) + + if(s == NULL) + { + set_reason(reason, reason_cap, "null spec"); + return false; + } + if(arch == NULL) + { + arch = "gfx950"; + } + const rocke_archtarget_t* target = rocke_archtarget_from_gfx(arch); + if(target == NULL) + { + rocke_set_unknown_arch_reason(reason, reason_cap, arch); + return false; + } + const rocke_direct_conv_problem_t* p = &s->problem; + if(dilation_reason(p, msg, sizeof msg)) + { + set_reason(reason, reason_cap, msg); + return false; + } + if(!problem_dtype_ok(p)) + { + if(p->dtype == NULL) + { + TILED_REJECT("unsupported dtype None"); + } + TILED_REJECT("unsupported dtype '%s'", p->dtype); + } + if(p->stride != 1) + { + TILED_REJECT("stride-1 algorithm (got stride=%d)", p->stride); + } + if(target->wave_size != 64 || !target->memory.has_ds_read_tr) + { + TILED_REJECT("needs wave64 + ds_read_b64_tr_b16 (absent on %s)", arch); + } + if(rocke_archtarget_async_lds_max_dwords(target) < 4) + { + TILED_REJECT("needs 16-byte buffer_load_lds (absent on %s)", arch); + } + const char* ab = problem_is_bf16(p) ? "bf16" : "f16"; + if(!rocke_mma_catalog_has_shape( + rocke_archtarget_mma(target), "mma", ab, ab, "fp32", 16, 16, 32)) + { + TILED_REJECT("missing mfma_f32_16x16x32_%s on %s", ab, arch); + } + auto in124 = [](int v) { return v == 1 || v == 2 || v == 4; }; + if(!in124(s->wave_k16) || !in124(s->wave_c16)) + { + TILED_REJECT("wave_k16 / wave_c16 must be 1, 2 or 4"); + } + if(s->waves_k < 1 || (s->waves_k & (s->waves_k - 1))) + { + TILED_REJECT("waves_k must be a power of two"); + } + if(s->waves_c < 1 || (s->waves_c & (s->waves_c - 1))) + { + TILED_REJECT("waves_c must be a power of two"); + } + if(!in124(s->unfold_n)) + { + TILED_REJECT("unfold_n must be 1, 2 or 4"); + } + if(s->q_groups != 1 && s->q_groups != 2) + { + TILED_REJECT("q_groups must be 1 or 2"); + } + if(s->pingpong && s->q_groups != 2) + { + TILED_REJECT("pingpong needs q_groups == 2"); + } + if(rocke_direct_wgrad_tiled_threads_per_block(s) + > rocke_archtarget_max_threads_per_block(target)) + { + TILED_REJECT("too many threads per block"); + } + if(p->cpg % kVec || p->kpg % kVec) + { + TILED_REJECT("cpg and kpg must be multiples of 8 (16-byte chunks)"); + } + if(p->kpg < wave_k(s) || p->cpg < wave_c(s)) + { + TILED_REJECT("wave tile wider than the group's channels"); + } + if(rocke_direct_wgrad_tiled_acc_vgprs(s) > 192) + { + TILED_REJECT("accumulators need %d VGPRs (> 192)", rocke_direct_wgrad_tiled_acc_vgprs(s)); + } + if(s->row_unroll < 0) + { + TILED_REJECT("row_unroll must be >= 0"); + } + if(s->target_blocks < 1) + { + TILED_REJECT("target_blocks must be >= 1"); + } + /* (Python divides by prefetch + 1 here; a negative prefetch, which raises + * there, falls through to the prefetch check.) */ + if(s->prefetch >= 0 && s->row_unroll % (s->prefetch + 1)) + { + TILED_REJECT("row_unroll must be a multiple of prefetch + 1"); + } + if(s->prefetch < 1) + { + TILED_REJECT("prefetch must be >= 1"); + } + { + int u = rocke_direct_wgrad_tiled_unroll(s); + if(s->rows_per_seg <= 0 || u <= 0 || s->rows_per_seg % u) + { + TILED_REJECT("rows_per_seg must be a positive multiple of %d", u); + } + } + if(p->KH > 8 || p->KW > 8) + { + TILED_REJECT("filter extent above 8"); + } + if(rocke_direct_wgrad_tiled_lds_bytes(s) > 160 * 1024) + { + TILED_REJECT("LDS %ld B > 160 KiB", rocke_direct_wgrad_tiled_lds_bytes(s)); + } +#undef TILED_REJECT + set_reason(reason, reason_cap, "ok"); + return true; +} + +/* ===================================================================== * + * Builder + * ===================================================================== */ +namespace +{ + +struct DmaLane +{ + V img, col, ch, okr; +}; + +struct LaneOff +{ + V off, ok; +}; + +struct Ctx +{ + Em e; + const rocke_direct_wgrad_tiled_spec_t* spec; + rocke_direct_conv_problem_t p; + const rocke_type_t* io_type; + bool bf16; + + int KH, KW, PAD, NW, WK16, WC16, BK, BC, CHK, CHC, UN, WU, SEG, STRIP, D, NB, HPB, U, T_PRO; + int IK, IC, N_INSTR, QG, GDY_ELEMS, GX_ELEMS; + std::vector swz_k, swz_c; + + rocke_dconv_params_t params; + V c0, c1, c_two, oob; + V c_cpg, c_kpg; + V lane, qg, wave; + V k_blk0, c_glb0, c_blk0, group; + V a_rsrc, b_rsrc; + std::vector dy_bufs, x_bufs, dy_addr, x_addr; + V qg_dy_bytes = NULL, qg_x_bytes = NULL; + std::vector dy_lanes, x_lanes; + std::vector> dy_idx; /* [mi][rd] */ + std::vector>> x_idx; /* [s][ni][rd] */ + V zero_frag; + int n_acc; +}; + +int acc_index(const Ctx& c, int r, int s, int mi, int ni) +{ + return ((r * c.KW + s) * c.WK16 + mi) * c.WC16 + ni; +} + +std::vector + dma_lanes(Ctx& c, int n_instr, int chunks, int rows, int per_img, const std::vector& bits) +{ + const Em& e = c.e; + std::vector out; + for(int i = 0; i < n_instr; ++i) + { + V slot = e.muli(c.wave, n_instr); + slot = e.addi(slot, i); + slot = e.muli(slot, kWave); + slot = e.add(slot, c.lane); + V row = e.divi(slot, chunks); + V phys = e.modi(slot, chunks); + V sw = e.swz_val(row, bits, c.c1); + V log = (sw == NULL) ? phys : e.bxor(phys, sw); + DmaLane d; + d.img = c.UN > 1 ? e.divi(row, per_img) : c.c0; + d.col = c.UN > 1 ? e.modi(row, per_img) : row; + d.ch = e.muli(log, kVec); + d.okr = e.lti(row, rows); + out.push_back(d); + } + return out; +} + +/* seg_lane_offsets(n0, q0) */ +void seg_lane_offsets(Ctx& c, V n0, V q0, std::vector& dy, std::vector& xs) +{ + const Em& e = c.e; + const rocke_dconv_params_t& P = c.params; + for(const DmaLane& d : c.dy_lanes) + { + V n_i = e.add(n0, d.img); + V wo = e.add(q0, d.col); + V a = e.lt(wo, P.p_Wo); + V bq = e.lt(n_i, P.p_N); + V ok = e.land(d.okr, e.land(a, bq)); + V m1 = e.mul(n_i, P.p_A_stride_n); + V m2 = e.mul(wo, P.p_A_stride_wi); + V s = e.add(m1, m2); + V t = e.add(c.k_blk0, d.ch); + dy.push_back({e.add(s, t), ok}); + } + for(const DmaLane& d : c.x_lanes) + { + V n_i = e.add(n0, d.img); + V wi = e.add(q0, e.subi(d.col, c.PAD)); + V a = e.ge(wi, c.c0); + V bq = e.lt(wi, P.p_Wi); + V ab = e.land(a, bq); + V cn = e.lt(n_i, P.p_N); + V ok = e.land(d.okr, e.land(ab, cn)); + V m1 = e.mul(n_i, P.p_B_stride_n); + V m2 = e.mul(wi, P.p_B_stride_wi); + V s = e.add(m1, m2); + V t = e.add(c.c_glb0, d.ch); + xs.push_back({e.add(s, t), ok}); + } +} + +struct Seg +{ + V h0, h_end, t_end; + std::vector dy, xs; +}; + +/* issue_row(buf, t, h0, h_end, lane_offs, t_end) */ +void issue_row(Ctx& c, const Seg& sg, int buf, V t) +{ + const Em& e = c.e; + const rocke_dconv_params_t& P = c.params; + V hi = e.add(sg.h0, e.subi(t, c.T_PRO)); + V ho = e.addi(hi, c.PAD); + V live = e.lt(t, sg.t_end); + V a = e.ge(ho, c.c0); + V bq = e.lt(ho, P.p_Ho); + V dy_ok = e.land(live, e.land(a, bq)); + V x0 = e.gei(t, c.T_PRO); + V x1 = e.ge(hi, c.c0); + V x2 = e.lt(hi, sg.h_end); + V x12 = e.land(x1, x2); + V x_ok = e.land(live, e.land(x0, x12)); + V dy_row = e.mul(ho, P.p_A_stride_hi); + V x_row = e.mul(hi, P.p_B_stride_hi); + for(int i = 0; i < (int)sg.dy.size(); ++i) + { + V okk = e.land(sg.dy[i].ok, dy_ok); + V voff = e.mul(e.add(sg.dy[i].off, dy_row), c.c_two); + voff = e.select(okk, voff, c.oob); + V base + = c.QG == 1 ? c.dy_addr[buf] : rocke_b_smem_ptr_add(e.b, c.dy_addr[buf], c.qg_dy_bytes); + V o = e.muli(c.wave, c.IK); + o = e.addi(o, i); + o = e.muli(o, kWave * 16); + o = e.zext64(o); + V lds = rocke_b_smem_ptr_add(e.b, base, o); + rocke_b_async_buffer_load_lds(e.b, c.a_rsrc, lds, voff, c.c0, 4, 0); + } + for(int i = 0; i < (int)sg.xs.size(); ++i) + { + V okk = e.land(sg.xs[i].ok, x_ok); + V voff = e.mul(e.add(sg.xs[i].off, x_row), c.c_two); + voff = e.select(okk, voff, c.oob); + V base = c.QG == 1 ? c.x_addr[buf] : rocke_b_smem_ptr_add(e.b, c.x_addr[buf], c.qg_x_bytes); + V o = e.muli(c.wave, c.IC); + o = e.addi(o, i); + o = e.muli(o, kWave * 16); + o = e.zext64(o); + V lds = rocke_b_smem_ptr_add(e.b, base, o); + rocke_b_async_buffer_load_lds(e.b, c.b_rsrc, lds, voff, c.c0, 4, 0); + } +} + +V read_frag(Ctx& c, V smem, const std::vector& idx) +{ + V lo = c.e.ds_read_tr(smem, c.c0, idx[0], c.io_type); + V hi = c.e.ds_read_tr(smem, c.c0, idx[1], c.io_type); + return rocke_b_vec_concat(c.e.b, lo, hi); +} + +typedef std::vector> Ring; /* [KH][WK16] */ + +/* step(buf, t, ring, accs, compute, issue) */ +void step(Ctx& c, const Seg& sg, int buf, V t, Ring& ring, std::vector& accs, bool compute) +{ + const Em& e = c.e; + rocke_b_s_waitcnt(e.b, (c.D - 1) * c.N_INSTR, 0, -1); + rocke_b_s_barrier_bare(e.b); + issue_row(c, sg, (buf + c.D) % c.NB, e.addi(t, c.D)); + std::vector newest; + for(int mi = 0; mi < c.WK16; ++mi) + { + newest.push_back(read_frag(c, c.dy_bufs[buf], c.dy_idx[mi])); + } + ring.erase(ring.begin()); + ring.push_back(newest); + if(!compute) + { + return; + } + std::vector> xf(c.KW); + for(int s = 0; s < c.KW; ++s) + { + for(int ni = 0; ni < c.WC16; ++ni) + { + xf[s].push_back(read_frag(c, c.x_bufs[buf], c.x_idx[s][ni])); + } + } + if(c.spec->pingpong) + { + rocke_b_s_waitcnt(e.b, -1, 0, -1); + rocke_b_s_barrier_bare(e.b); + rocke_b_s_setprio(e.b, 1); + } + for(int r = 0; r < c.KH; ++r) + { + const std::vector& dyf = ring[c.KH - 1 - r]; + for(int s = 0; s < c.KW; ++s) + { + for(int mi = 0; mi < c.WK16; ++mi) + { + for(int ni = 0; ni < c.WC16; ++ni) + { + int i = acc_index(c, r, s, mi, ni); + accs[i] = e.mfma16x16x32(c.bf16, dyf[mi], xf[s][ni], accs[i]); + } + } + } + } + if(c.spec->pingpong) + { + rocke_b_s_setprio(e.b, 0); + } +} + +/* segment(accs, n0, q, h0, h_end, valid=None) */ +std::vector segment(Ctx& c, std::vector accs, V n0, V q, V h0, V h_end, V valid) +{ + const Em& e = c.e; + Seg sg; + sg.h0 = h0; + sg.h_end = h_end; + seg_lane_offsets(c, n0, e.muli(q, c.WU), sg.dy, sg.xs); + sg.t_end = e.ci(c.T_PRO + c.HPB); + if(valid != NULL) + { + sg.t_end = e.select(valid, sg.t_end, c.c0); + } + for(int t = 0; t < c.D; ++t) + { + issue_row(c, sg, t % c.NB, e.ci(t)); + } + Ring ring(c.KH, std::vector(c.WK16, c.zero_frag)); + for(int t = 0; t < c.T_PRO; ++t) + { + step(c, sg, t % c.NB, e.ci(t), ring, accs, false); + } + std::vector names = iter_names("a", accs.size()); + std::vector inits = accs; + for(int j = 1; j < c.KH; ++j) + { + for(int mi = 0; mi < c.WK16; ++mi) + { + names.push_back("r" + std::to_string(j) + "_" + std::to_string(mi)); + inits.push_back(ring[j][mi]); + } + } + V hi_bound = e.ci(c.HPB / c.U); + std::vector res + = e.for_iter(c.c0, hi_bound, c.c1, names, inits, "rb", [&](V rb, std::vector rc) { + std::vector accs_i(rc.begin(), rc.begin() + c.n_acc); + Ring ring_i; + ring_i.push_back(std::vector(c.WK16, c.zero_frag)); + for(int j = 1; j < c.KH; ++j) + { + ring_i.push_back(std::vector(rc.begin() + c.n_acc + (j - 1) * c.WK16, + rc.begin() + c.n_acc + j * c.WK16)); + } + V t_base = e.addi(e.muli(rb, c.U), c.T_PRO); + for(int u = 0; u < c.U; ++u) + { + step(c, sg, (c.T_PRO + u) % c.NB, e.addi(t_base, u), ring_i, accs_i, true); + } + std::vector out = accs_i; + for(int j = 1; j < c.KH; ++j) + { + out.insert(out.end(), ring_i[j].begin(), ring_i[j].end()); + } + return out; + }); + std::vector accs_out(res.begin(), res.begin() + std::min((size_t)c.n_acc, res.size())); + rocke_b_s_waitcnt(e.b, 0, 0, -1); + rocke_b_s_barrier_bare(e.b); + return accs_out; +} + +/* frag_index(shift, rd, col0, chunks, bits, seg, g_elems) */ +struct FragLanes +{ + V lane_q, lane_c4, img_l, jl; +}; + +V frag_index(Ctx& c, + const FragLanes& fl, + int shift, + int rd, + V col0, + int chunks, + const std::vector& bits, + int seg, + int g_elems) +{ + const Em& e = c.e; + V row = e.muli(fl.img_l, seg); + row = e.add(row, e.add(fl.jl, fl.lane_q)); + row = e.addi(row, shift + rd * 4); + V col = e.add(col0, e.muli(fl.lane_c4, 4)); + V chunk = e.divi(col, kVec); + V sw = e.swz_val(row, bits, c.c1); + V phys = (sw == NULL) ? chunk : e.bxor(chunk, sw); + V idx = e.muli(row, chunks); + idx = e.add(idx, phys); + idx = e.muli(idx, kVec); + idx = e.add(idx, e.modi(col, kVec)); + if(c.QG > 1) + { + idx = e.add(idx, e.muli(c.qg, g_elems)); + } + return idx; +} + +rocke_kernel_def_t* + build(rocke_ir_builder_t* b, const rocke_direct_wgrad_tiled_spec_t* spec, const char* arch) +{ + char reason[ROCKE_ERR_MSG_CAP]; + if(!rocke_direct_wgrad_tiled_is_valid_spec(spec, arch, reason, sizeof reason)) + { + std::string prefix = std::string("invalid DirectWgradTiledSpec for ") + arch; + fail_invalid(b, prefix.c_str(), reason); + return NULL; + } + Ctx c; + c.e.b = b; + const Em& e = c.e; + c.spec = spec; + c.p = spec->problem; + const rocke_direct_conv_problem_t& p = c.p; + c.bf16 = problem_is_bf16(&p); + c.io_type = c.bf16 ? rocke_bf16() : rocke_f16(); + c.KH = p.KH; + c.KW = p.KW; + c.PAD = p.PAD; + c.NW = n_waves(spec); + c.WK16 = spec->wave_k16; + c.WC16 = spec->wave_c16; + c.BK = wave_k(spec) * spec->waves_k; + c.BC = wave_c(spec) * spec->waves_c; + c.CHK = c.BK / kVec; + c.CHC = c.BC / kVec; + c.UN = spec->unfold_n; + c.WU = strip_w(spec); + c.SEG = c.WU + c.KW - 1; + c.STRIP = c.UN * c.SEG; + c.D = spec->prefetch; + c.NB = spec->prefetch + 1; + c.HPB = spec->rows_per_seg; + c.U = rocke_direct_wgrad_tiled_unroll(spec); + c.T_PRO = c.KH - 1; + c.IK = dma_instrs_dy(spec); + c.IC = dma_instrs_x(spec); + c.N_INSTR = c.IK + c.IC; + c.QG = spec->q_groups; + c.GDY_ELEMS = c.NW * c.IK * kWave * kVec; + c.GX_ELEMS = c.NW * c.IC * kWave * kVec; + int DY_ELEMS = c.QG * c.GDY_ELEMS; + int X_ELEMS = c.QG * c.GX_ELEMS; + c.n_acc = c.KH * c.KW * c.WK16 * c.WC16; + + std::vector k_origins, c_origins; + for(int wk = 0; wk < spec->waves_k; ++wk) + { + for(int mi = 0; mi < c.WK16; ++mi) + { + k_origins.push_back(wk * wave_k(spec) + mi * kAtom); + } + } + for(int wc = 0; wc < spec->waves_c; ++wc) + { + for(int ni = 0; ni < c.WC16; ++ni) + { + c_origins.push_back(wc * wave_c(spec) + ni * kAtom); + } + } + std::vector shifts_c; + for(int s = 0; s < c.KW; ++s) + { + shifts_c.push_back(s); + } + c.swz_k = pick_swizzle(c.CHK, std::vector{0}, c.WU, c.WU, k_origins); + c.swz_c = pick_swizzle(c.CHC, shifts_c, c.WU, c.SEG, c_origins); + + rocke_attr_set_int(b, + &b->kernel->attrs, + "max_workgroup_size", + rocke_direct_wgrad_tiled_threads_per_block(spec)); + /* Separate scopes per row buffer, or the waitcnt pass drains the DMA ring. */ + rocke_attr_set_bool(b, &b->kernel->attrs, "lds_alias_scopes", true); + /* Two waves per SIMD: the register budget the 32x32 tile is sized for. */ + { + const int64_t wpe[2] = {2, 8}; + rocke_attr_set_int_list(b, &b->kernel->attrs, "waves_per_eu", wpe, 2); + } + rocke_dconv_emit_params(b, &c.params, "wgrad", c.io_type); + if(!rocke_ir_builder_ok(b)) + { + return NULL; + } + const rocke_dconv_params_t& P = c.params; + + c.c0 = e.ci(0); + c.c1 = e.ci(1); + c.c_two = e.ci(2); + c.oob = e.ci(((int64_t)1 << 31) - 1); + + V p_N = P.p_N; + V p_H = P.p_Hi; + V p_Wo = P.p_Wo; + V p_groups = P.p_groups; + c.c_cpg = e.div(P.p_total_c, p_groups); + c.c_kpg = e.div(P.p_total_k, p_groups); + + V tid = rocke_b_thread_id_x(b); + c.lane = e.modi(tid, kWave); + V wave_all = rocke_b_readfirstlane(b, e.divi(tid, kWave)); + if(c.QG > 1) + { + c.qg = e.divi(wave_all, c.NW); + c.wave = e.modi(wave_all, c.NW); + } + else + { + c.qg = c.c0; + c.wave = wave_all; + } + + /* ---- grid decode ---- */ + V n_k_tiles = e.cdivi(c.c_kpg, c.BK); + V n_c_tiles = e.cdivi(c.c_cpg, c.BC); + V n_ch_tiles = e.mul(e.mul(p_groups, n_k_tiles), n_c_tiles); + V n_q = e.cdivi(p_Wo, c.WU); + V n_ng = e.cdivi(p_N, c.UN); + V n_hb = e.cdivi(p_H, c.HPB); + V total_segs = e.mul(e.mul(n_ng, n_q), n_hb); + /* Same formula as DirectWgradTiledSpec.n_splits, which sizes the grid. */ + V sa = e.cdivi(total_segs, spec->q_groups); + V tb = e.ci(spec->target_blocks); + V sb = e.add(tb, e.sub(n_ch_tiles, c.c1)); + sb = e.div(sb, n_ch_tiles); + V n_splits = e.smin(sa, sb); + + /* Renumber block ids so each XCD gets a contiguous run. */ + V lin = rocke_b_block_id_x(b); + V n_blocks = e.mul(n_ch_tiles, n_splits); + V xcd = e.modi(lin, 8); + V pos = e.divi(lin, 8); + V per = e.divi(n_blocks, 8); + V rem = e.modi(n_blocks, 8); + V xp = e.mul(xcd, per); + V xr = e.smin(xcd, rem); + lin = e.add(e.add(xp, xr), pos); + V ch_tile = e.mod(lin, n_ch_tiles); + V split = e.div(lin, n_ch_tiles); + V c_tile = e.mod(ch_tile, n_c_tiles); + V gk = e.div(ch_tile, n_c_tiles); + V k_tile = e.mod(gk, n_k_tiles); + c.group = e.div(gk, n_k_tiles); + { + V g = e.mul(c.group, c.c_kpg); + V k = e.muli(k_tile, c.BK); + c.k_blk0 = e.add(g, k); + } + c.c_blk0 = e.muli(c_tile, c.BC); + c.c_glb0 = e.add(e.mul(c.group, c.c_cpg), c.c_blk0); + + V wave_k_id = e.modi(c.wave, spec->waves_k); + V wave_c_id = e.divi(c.wave, spec->waves_k); + + c.a_rsrc = rocke_b_buffer_rsrc(b, P.A, P.A_bytes); + c.b_rsrc = rocke_b_buffer_rsrc(b, P.Bp, P.B_bytes); + + /* ---- LDS: one allocation per ring buffer (one alias scope each) ---- */ + for(int i = 0; i < c.NB; ++i) + { + std::string h = "dy" + std::to_string(i); + c.dy_bufs.push_back(e.smem_alloc_1d(c.io_type, DY_ELEMS, h.c_str())); + } + for(int i = 0; i < c.NB; ++i) + { + std::string h = "xs" + std::to_string(i); + c.x_bufs.push_back(e.smem_alloc_1d(c.io_type, X_ELEMS, h.c_str())); + } + for(V s : c.dy_bufs) + { + c.dy_addr.push_back(rocke_b_smem_addr_of(b, s)); + } + for(V s : c.x_bufs) + { + c.x_addr.push_back(rocke_b_smem_addr_of(b, s)); + } + if(c.QG > 1) + { + c.qg_dy_bytes = e.zext64(e.muli(c.qg, c.GDY_ELEMS * 2)); + c.qg_x_bytes = e.zext64(e.muli(c.qg, c.GX_ELEMS * 2)); + } + + /* ---- DMA lane geometry ---- */ + c.dy_lanes = dma_lanes(c, c.IK, c.CHK, kWo, c.WU, c.swz_k); + c.x_lanes = dma_lanes(c, c.IC, c.CHC, c.STRIP, c.SEG, c.swz_c); + + /* ---- fragment read addresses ---- */ + V lane_hi = e.divi(c.lane, 16); + FragLanes fl; + fl.lane_q = e.modi(e.divi(c.lane, 4), 4); + fl.lane_c4 = e.modi(c.lane, 4); + V kk0 = e.muli(lane_hi, 8); + if(c.UN > 1) + { + fl.img_l = e.divi(kk0, c.WU); + fl.jl = e.modi(kk0, c.WU); + } + else + { + fl.img_l = c.c0; + fl.jl = kk0; + } + V k_wave0 = e.muli(wave_k_id, wave_k(spec)); + V c_wave0 = e.muli(wave_c_id, wave_c(spec)); + c.dy_idx.resize(c.WK16); + for(int mi = 0; mi < c.WK16; ++mi) + { + for(int rd = 0; rd < 2; ++rd) + { + V col0 = e.addi(k_wave0, mi * kAtom); + c.dy_idx[mi].push_back( + frag_index(c, fl, 0, rd, col0, c.CHK, c.swz_k, c.WU, c.GDY_ELEMS)); + } + } + c.x_idx.assign(c.KW, std::vector>(c.WC16)); + for(int s = 0; s < c.KW; ++s) + { + for(int ni = 0; ni < c.WC16; ++ni) + { + for(int rd = 0; rd < 2; ++rd) + { + V col0 = e.addi(c_wave0, ni * kAtom); + c.x_idx[s][ni].push_back( + frag_index(c, fl, s, rd, col0, c.CHC, c.swz_c, c.SEG, c.GX_ELEMS)); + } + } + } + + V zero_acc = rocke_b_zero_vec_f32(b, 4); + c.zero_frag = rocke_b_zero_vec(b, c.io_type, kVec); + + /* ---- persistent segment loop ---- */ + if(spec->pingpong) + { + /* Group 1 starts one barrier late. */ + e.if_(e.eq(c.qg, c.c1), [&]() { rocke_b_s_barrier_bare(b); }); + } + std::vector acc_names = iter_names("acc", c.n_acc); + std::vector zero_accs(c.n_acc, zero_acc); + std::vector accs; + if(c.QG == 1) + { + accs = e.for_iter(split, + total_segs, + n_splits, + acc_names, + zero_accs, + "seg", + [&](V seg, std::vector carried) { + V hb = e.mod(seg, n_hb); + V rest = e.div(seg, n_hb); + V h0 = e.muli(hb, c.HPB); + V n0 = e.muli(e.div(rest, n_q), c.UN); + V q = e.mod(rest, n_q); + V h_end = e.smin(e.addi(h0, c.HPB), p_H); + return segment(c, carried, n0, q, h0, h_end, NULL); + }); + } + else + { + /* Group g of block `split` takes virtual split split * QG + g. */ + V n_vsplits = e.muli(n_splits, c.QG); + V vsplit0 = e.muli(split, c.QG); + V s1 = e.sub(total_segs, vsplit0); + V s2 = e.sub(n_vsplits, c.c1); + V n_iter = e.div(e.add(s1, s2), n_vsplits); + accs = e.for_iter( + c.c0, n_iter, c.c1, acc_names, zero_accs, "segi", [&](V segi, std::vector carried) { + V a = e.add(vsplit0, c.qg); + V seg = e.add(a, e.mul(segi, n_vsplits)); + V hb = e.mod(seg, n_hb); + V rest = e.div(seg, n_hb); + V h0 = e.muli(hb, c.HPB); + V n0 = e.muli(e.div(rest, n_q), c.UN); + V q = e.mod(rest, n_q); + V h_end = e.smin(e.addi(h0, c.HPB), p_H); + V valid = e.lt(seg, total_segs); + return segment(c, carried, n0, q, h0, h_end, valid); + }); + } + if((int)accs.size() != c.n_acc || !rocke_ir_builder_ok(b)) + { + return NULL; + } + if(spec->pingpong) + { + e.if_(e.eq(c.qg, c.c0), [&]() { rocke_b_s_barrier_bare(b); }); + } + + /* ---- epilogue: atomic add into dW[total_k, KH, KW, cpg] ---- */ + V D_ptr = P.D; + V rs_stride = c.c_cpg; + V k_stride = e.muli(c.c_cpg, c.KH * c.KW); + V k_lane0 = e.add(c.k_blk0, k_wave0); + k_lane0 = e.add(k_lane0, e.muli(lane_hi, 4)); + V k_in_grp0 = e.sub(k_lane0, e.mul(c.group, c.c_kpg)); + V c_lane0 = e.add(c.c_blk0, c_wave0); + c_lane0 = e.add(c_lane0, e.modi(c.lane, 16)); + if(c.QG > 1) + { + /* Tap t is summed and written by group t % 2; the other group hands its + * accumulators over through LDS, double-buffered on the tap parity. */ + int n_atoms = c.WK16 * c.WC16; + V red[2]; + for(int i = 0; i < 2; ++i) + { + int shape[2] = {1, c.NW * n_atoms * kWave * 4}; + std::string h = "red" + std::to_string(i); + red[i] = rocke_b_smem_alloc_f32(b, shape, 2, h.c_str()); + } + auto red_index = [&](int atom) { + V v = e.muli(c.wave, n_atoms); + v = e.addi(v, atom); + v = e.muli(v, kWave); + v = e.add(v, c.lane); + return e.muli(v, 4); + }; + for(int r = 0; r < c.KH; ++r) + { + for(int s = 0; s < c.KW; ++s) + { + int tap = r * c.KW + s; + V buf = red[tap % 2]; + V is_recv = e.eqi(c.qg, tap % 2); + e.if_(e.lnot(is_recv), [&]() { + for(int mi = 0; mi < c.WK16; ++mi) + { + for(int ni = 0; ni < c.WC16; ++ni) + { + V idx = red_index(mi * c.WC16 + ni); + e.smem_store_f32(buf, c.c0, idx, accs[acc_index(c, r, s, mi, ni)], 4); + } + } + }); + rocke_b_s_waitcnt(b, -1, 0, -1); + rocke_b_s_barrier_bare(b); + e.if_(is_recv, [&]() { + for(int ni = 0; ni < c.WC16; ++ni) + { + V c_in = e.addi(c_lane0, ni * kAtom); + V c_ok = e.lt(c_in, c.c_cpg); + for(int mi = 0; mi < c.WK16; ++mi) + { + V ld = e.smem_load_f32(buf, c.c0, red_index(mi * c.WC16 + ni), 4); + V tot = rocke_b_fadd(b, accs[acc_index(c, r, s, mi, ni)], ld); + for(int slot = 0; slot < 4; ++slot) + { + V k_in = e.addi(k_in_grp0, mi * kAtom + slot); + V k_glb = e.addi(k_lane0, mi * kAtom + slot); + V ok = e.land(c_ok, e.lt(k_in, c.c_kpg)); + V off = e.add(e.mul(k_glb, k_stride), c_in); + V tk = e.ci(tap); + off = e.add(off, e.mul(tk, rs_stride)); + e.if_(ok, [&]() { + V val = rocke_b_vec_extract(b, tot, slot); + rocke_b_global_atomic_add(b, D_ptr, off, val, NULL); + }); + } + } + } + }); + } + } + return rocke_ir_builder_ok(b) ? b->kernel : NULL; + } + for(int ni = 0; ni < c.WC16; ++ni) + { + V c_in = e.addi(c_lane0, ni * kAtom); + V c_ok = e.lt(c_in, c.c_cpg); + for(int mi = 0; mi < c.WK16; ++mi) + { + for(int slot = 0; slot < 4; ++slot) + { + V k_in = e.addi(k_in_grp0, mi * kAtom + slot); + V k_glb = e.addi(k_lane0, mi * kAtom + slot); + V ok = e.land(c_ok, e.lt(k_in, c.c_kpg)); + V base = e.add(e.mul(k_glb, k_stride), c_in); + for(int r = 0; r < c.KH; ++r) + { + for(int s = 0; s < c.KW; ++s) + { + V val = rocke_b_vec_extract(b, accs[acc_index(c, r, s, mi, ni)], slot); + V tk = e.ci(r * c.KW + s); + V off = e.add(base, e.mul(tk, rs_stride)); + e.if_(ok, [&]() { rocke_b_global_atomic_add(b, D_ptr, off, val, NULL); }); + } + } + } + } + } + return rocke_ir_builder_ok(b) ? b->kernel : NULL; +} + +} // namespace + +extern "C" rocke_kernel_def_t* rocke_build_direct_wgrad_tiled( + rocke_ir_builder_t* b, const rocke_direct_wgrad_tiled_spec_t* spec, const char* arch) +{ + if(b == NULL || spec == NULL) + { + return NULL; + } + return build(b, spec, arch ? arch : "gfx950"); +} + +extern "C" rocke_kernel_def_t* rocke_build_direct_wgrad_tiled_new( + rocke_ir_builder_t* b, const rocke_direct_wgrad_tiled_spec_t* spec, const char* arch) +{ + return build_new( + b, + spec != NULL, + "build_direct_wgrad_tiled", + [&](char* out, size_t cap) { return rocke_direct_wgrad_tiled_kernel_name(spec, out, cap); }, + [&]() { return rocke_build_direct_wgrad_tiled(b, spec, arch); }); +} + +extern "C" rocke_status_t + rocke_direct_wgrad_tiled_lower_to_llvm(const rocke_direct_wgrad_tiled_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + if(arch == NULL) + { + arch = "gfx950"; + } + return lower_to_llvm( + spec != NULL, + arch, + flavor, + out_ll, + err, + err_cap, + "build_direct_wgrad_tiled failed", + [&](rocke_ir_builder_t* b) { return rocke_build_direct_wgrad_tiled_new(b, spec, arch); }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md index f94c080a5e50..feebce4ae0f8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md @@ -88,6 +88,9 @@ and exact gfx targets. | `build_direct_conv_16c` | [L/common/conv_direct_grouped.py](../../../library/kernels/common/conv_direct_grouped.py) | grouped direct convolution | `mfma_f32_16x16x16_f16`, or `mfma_f32_16x16x32_f16` with K folding | owning validator | | `build_direct_conv_4c` | [L/common/conv_direct_grouped.py](../../../library/kernels/common/conv_direct_grouped.py) | small-channel grouped convolution | direct `mfma_f32_4x4x4_f16` intrinsic outside `MmaCatalog` | validated exact gfx target | | `build_direct_conv_nongrouped` | [L/common/conv_direct_nongrouped.py](../../../library/kernels/common/conv_direct_nongrouped.py) | non-grouped (`groups == 1`) direct convolution, LDS halo reuse | square `mfma_f32_{32x32x16,32x32x8,16x16x32,16x16x16}` | owning validator | +| `build_direct_wgrad_tiled` | [L/common/conv_direct_wgrad_tiled.py](../../../library/kernels/common/conv_direct_wgrad_tiled.py) | direct wgrad, LDS-DMA row pipeline, fp32 atomics | `mfma_f32_16x16x32_{f16,bf16}` | owning validator | +| `build_dw_wgrad_lds`, `build_dw_wgrad_cols`, `build_dw_wgrad_fold` | [L/common/conv_direct_dw_wgrad.py](../../../library/kernels/common/conv_direct_dw_wgrad.py) | depthwise wgrad into partial slots, and their fold | VALU FMA | owning validator | +| `build_wgrad_gather` | [L/common/conv_direct_wgrad_gather.py](../../../library/kernels/common/conv_direct_wgrad_gather.py) | dense small-C wgrad into partial slots | `mfma_f32_16x16x32_{f16,bf16}` | owning validator | | `build_deep_fused_conv_pool` | [L/common/deep_fused_conv_pool.py](../../../library/kernels/common/deep_fused_conv_pool.py), [L/gfx1151/deep_fused_conv_pool.py](../../../library/kernels/gfx1151/deep_fused_conv_pool.py), [L/gfx1201/deep_fused_conv_pool.py](../../../library/kernels/gfx1201/deep_fused_conv_pool.py), [L/gfx950/deep_fused_conv_pool.py](../../../library/kernels/gfx950/deep_fused_conv_pool.py) | convolution plus pooling | target-specific matrix MMA plus VALU/window reduction | exact implementation selected by target | | `build_moe_gate_up_silu_gemm`, `build_moe_interleaved_gate_up_silu_gemm`, `build_moe_down_reduce_gemm`, `build_moe_down_silu_reduce_gemm` | [P/common/moe_gemm_fused.py](../../python/rocke/instances/common/moe_gemm_fused.py) | expert GEMM plus activation or reduction | catalog-selected matrix MMA plus fused VALU/reduction | owning validator | | `build_moe_fused_mega_gemm`, `build_moe_fused_mega_gemm_fp8`, `build_moe_fused_mega_wmma` | [P/common/moe_fused_mega.py](../../python/rocke/instances/common/moe_fused_mega.py), [P/common/moe_fused_mega_fp8.py](../../python/rocke/instances/common/moe_fused_mega_fp8.py), [P/gfx1250/fused_moe_mega_wmma.py](../../python/rocke/instances/gfx1250/fused_moe_mega_wmma.py) | persistent expert GEMM pipeline | target-selected matrix MMA plus fused routing/activation | owning validator or exact-target module | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index be117731ee0f..15e8df4b1c3f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md @@ -5,6 +5,9 @@ This page covers: - `library/kernels/common/conv_implicit_gemm.py` - `library/kernels/common/conv_direct_grouped.py` - `library/kernels/common/conv_direct_nongrouped.py` +- `library/kernels/common/conv_direct_wgrad_tiled.py`, + `library/kernels/common/conv_direct_dw_wgrad.py`, + `library/kernels/common/conv_direct_wgrad_gather.py` (direct wgrad) - `library/kernels/common/img2col.py` - `instances/common/pooling.py` @@ -768,6 +771,54 @@ C++ engine twin: `rocke_build_direct_conv_nongrouped` in `conv_direct_nongrouped/*` in `tests/instances/rocke_ir_parity_harness.py`, binding `rocke_engine.conv_direct_nongrouped_*` / `rocke.core.backend.lower_conv_direct_nongrouped`. +## Direct Wgrad Kernels + +Sources: `library/kernels/common/conv_direct_wgrad_tiled.py`, +`conv_direct_dw_wgrad.py`, `conv_direct_wgrad_gather.py`. Weight-gradient +kernels of the direct-conv AOT sweep (`--directions wgrad`), gfx950; all take +the `conv_direct_arg_names("wgrad")` kernargs (dY on A, X on B) and accumulate +in fp32. + +| Spec | Serves | Output | +|---|---|---| +| `DirectWgradTiledSpec` | dense and grouped, stride 1, `cpg`/`kpg` multiples of 8 | `global_atomic_add` into a zeroed fp32 `dW[K, KH, KW, cpg]` | +| `DwWgradLdsSpec` | depthwise, `C % 8 == 0`, 3x3, stride 1/2 | partial slots `ws[part][tap][c]` | +| `DwWgradColsSpec` | depthwise, `C <= 8` (stem layers), any filter | partial slots `ws[part][tap][c]` | +| `WgradGatherSpec` | dense, small `C` (`KW * C <= 96`), any stride | partial slots `ws[part][(r, s, c)][k]` | +| `DwWgradFoldSpec` | the fold after a partial-slot pass | `dW` written whole (no zeroing) | + +- **Tiled.** A wave owns `wave_k16*16 x wave_c16*16` channels for all taps; + the block DMAs input rows of a 32-column strip into LDS ring buffers + (`buffer_load ... lds`, `prefetch` rows in flight) with an XOR chunk swizzle + that `pick_swizzle` picks at build time so the `ds_read_b64_tr_b16` fragment + reads are bank-conflict free. `unfold_n` packs several images into one strip; + `q_groups=2` (+`pingpong`) runs two wave groups half a row apart. +- **Partial slots + fold.** Same-address fp32 atomics serialize, and a + depthwise `dW` has only a few thousand addresses, so the depthwise and gather + kernels store one partial slot per partition with plain stores and + `DwWgradFoldSpec` sums the partitions. In the sweep they are two-step plans: + the pass writes the scratch `ws`, the fold reads it (kernargs `A, D, A_bytes, + D_bytes, parts, channels`). +- All four row-streaming kernels set the kernel attr `lds_alias_scopes`: every + smem allocation gets its own alias scope, so the waitcnt pass does not drain + the in-flight LDS DMA before each `ds_read`. +- The validators gate on `lds_bytes()`, which counts what the builder + allocates: ring slots rounded up to whole DMA instructions of every wave + (plus the gather's 16-element transpose tail) and the reduction buffers. + The LDS-DMA kernels (tiled, `DwWgradLdsSpec`, gather) also require a wave64 + target with 16-byte `buffer_load ... lds` (`async_lds_max_dwords >= 4`). + +C++ engine twins: `rocke_build_direct_wgrad_tiled` (`cpp/instances/common/conv_direct_wgrad_tiled.cpp`, +also `rocke_direct_wgrad_pick_swizzle`), `rocke_build_dw_wgrad_{fold,cols,lds}` +(`conv_direct_dw_wgrad.cpp`), `rocke_build_direct_wgrad_gather` +(`conv_direct_wgrad_gather.cpp`); byte-identity gated by +`library/tests/parity/conv_direct_grouped_emit.{c,py}` (configs 95+), Python IR +golden `conv_direct_wgrad/*` in `tests/instances/rocke_ir_parity_harness.py`. +On-GPU numeric coverage: `TestDirectWgradAotKernelsCorrectness` in +`library/tests/test_direct_conv_correctness.py` (built on the cache probe, +launched on other shapes); the two-step plans: +`library/tests/test_direct_kernel_cache.py`. + ## Img2Col Source: `library/kernels/common/img2col.py`. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md index f596adbd0916..a79e0aa975ad 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md @@ -43,6 +43,9 @@ implementation. | `conv_implicit_gemm.py` | `ConvProblem`, `ImplicitGemmConvSpec` | `instances/convolution.md` | | `conv_direct_grouped.py` | `DirectConvProblem`, `DirectConv16cSpec`, `DirectConv4cSpec` | `instances/convolution.md` | | `conv_direct_nongrouped.py` | `DirectNongroupedConvSpec` (`groups == 1`) | `instances/convolution.md` | +| `conv_direct_wgrad_tiled.py` | `DirectWgradTiledSpec` (direct wgrad, atomics) | `instances/convolution.md` | +| `conv_direct_dw_wgrad.py` | `DwWgradLdsSpec`, `DwWgradColsSpec`, `DwWgradFoldSpec` (depthwise wgrad) | `instances/convolution.md` | +| `conv_direct_wgrad_gather.py` | `WgradGatherSpec` (dense small-C wgrad) | `instances/convolution.md` | | `img2col.py` | `Img2ColSpec` | `instances/convolution.md` | | `pooling.py` | `PoolingProblem`, `Pooling2DSpec`, `PoolOp` | `instances/convolution.md` | @@ -225,6 +228,7 @@ From `helpers/README.md`: | conv_implicit_gemm | yes (buffer) | full (unmerge+embed+pad) | - | - | - | yes | - | `AsyncTileLoader`, `CoalescedTileLoader`, `CShuffleEpilogue`, `MfmaAtom`, `WarpGrid`, `LdsLayout`, `SchedulePolicy` | | conv_direct_grouped | - | input/output/weight + H/W pad | - | - | - | yes | - | `MfmaAtom` (4x4x4 / 16x16x{16,32}) | | conv_direct_nongrouped | - | - | - | - | - | yes | - | `chiplet_aware_super_tile_dynamic` (`helpers.grid`), `emit_direct_params` (AOT kernargs) | +| conv_direct_wgrad (tiled / dw / gather) | - | - | - | - | - | yes | - | `emit_direct_params` (AOT kernargs), kernel attr `lds_alias_scopes` | | img2col | - | reuses A descriptor | - | yes | - | yes | - | - | | pooling | - | full (input + pad) | - | yes | - | yes | - | - | | attention_unified | - | Q + output + paged-KV | - | - | - | yes | - | `OnlineSoftmaxState`, `PagedKvDescriptor` | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index 68487349eae9..18fcfb9441ce 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -1790,6 +1790,8 @@ def __init__( ) self._smem_globals: List[Tuple[str, SmemType]] = [] self._smem_storage_name: Dict[str, str] = {} # IR value name -> @global name + # Byte size of each allocation placed in the pool. + self._smem_sizes: Dict[str, int] = {} # smem pool: one unified addrspace(3) buffer; per-allocation byte offsets. self._smem_offsets: Dict[str, int] = {} self._smem_pool_size: int = 0 @@ -1801,7 +1803,7 @@ def __init__( # waits vmcnt(0) first; with them it waits only for DMA into the # read's own allocation. self._lds_aa = bool(kernel.attrs.get("lds_alias_scopes", False)) - self._smem_sizes: Dict[str, int] = {} + # i64 LDS address value name -> the allocation it points into. self._lds_prov: Dict[str, str] = {} self._lds_aa_used = False # Cache of per-allocation base pointers already emitted, keyed by @@ -5135,7 +5137,7 @@ def _op_tile_ds_read_tr16_b64(self, op: Op) -> None: raw = self._fresh("tr.raw") self._current().emit( f" {raw} = call <4 x i16> @llvm.amdgcn.ds.read.tr16.b64(" - f"ptr addrspace(3) {base})" + f"ptr addrspace(3) {base}){self._lds_aa_suffix(gname)}" ) elem_ty = _llvm_type(op.result.type.elem) # type: ignore[attr-defined] self._current().emit( @@ -6322,6 +6324,7 @@ def _op_tile_async_buffer_load_lds(self, op: Op) -> None: f"i32 {self._operand(soffset)}, " f"i32 0, " f"i32 {aux})" + f"{self._lds_aa_suffix(self._lds_prov.get(lds_ptr.name))}" ) def _op_tile_buffer_load_lds_async(self, op: Op) -> None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index 8563444346e1..32acf749a114 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -99,6 +99,8 @@ and [packed FP6 input contract](../examples/gfx1250/gemm/FP6.md). | `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `default` epilogue, `split_k=1` | | `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag ops | | `conv_direct_nongrouped` | ✅ | ✅ | ❌ | `groups=1`; gfx942 uses 32x32x8 / 16x16x16 atoms, gfx950 also 32x32x16 / 16x16x32; MFMA-only; the only direct kernel with dilation. Built from `library/kernels`; not yet wired into dispatch | +| `direct_wgrad_tiled` | ❌ | ✅ | ❌ | dense and grouped wgrad, stride 1; needs `ds_read_b64_tr_b16` and `mfma_f32_16x16x32`. Built from `library/kernels`; AOT sweep only, not yet wired into dispatch | +| `dw_wgrad_lds` / `dw_wgrad_cols` / `direct_wgrad_gather` | ❌ | ✅ | ❌ | depthwise (`C % 8 == 0` / `C <= 8`) and dense small-C wgrad, partial slots + `dw_wgrad_fold`; validated on gfx950 only (the LDS kernels DMA 16 bytes per lane). AOT sweep only, not yet wired into dispatch | | `direct_depthwise_tiled` | ✅ | ✅ | ❌ | `cpg=kpg=1`, filters up to 31x31, any padding; scalar FMA, wave64. Built from `library/kernels`; not yet wired into dispatch | --- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index ad125f138984..cd2dfb08b99a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -1162,6 +1162,70 @@ "sha256": "c80479a6bbdb5c5e81decdc1ac5b7bfc773b017e02fd518b9cd4fa0b1d75d50a", "status": "ok" }, + "conv_direct_wgrad/gfx950/dw_cols_c3_7x7_s2": { + "arch": "gfx950", + "bytes": 80331, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_dw_cols_c3_7x7_s2_f7x7p3s2c3_st8_t64_tb2048_bf16", + "sha256": "eb9ba342a46fd966a8fef30e04c461c89715a971f243eb49b6edb5eae7bc6559", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/dw_lds_c64": { + "arch": "gfx950", + "bytes": 323590, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_dw_lds_c64_f3x3p1s1_lc8_lw32_cpl1_pf2_rps32_tb1024_bf16", + "sha256": "3ef756d5ce85d6a9761e59b04681a1a01229ae78a5a9b53ef8e77cb838895611", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/fold_t9": { + "arch": "gfx950", + "bytes": 11967, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_fold_t9_t9", + "sha256": "141a91ac18bd835273d8dc029fa5768e36723419221bcd8f0fcdbc38714beae7", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/gather_c24_tr": { + "arch": "gfx950", + "bytes": 229377, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_gather_c24_tr_f3x3p1s2c24_wk1x6_pf2_rps8_tb256_bf16", + "sha256": "b96c105663423b58b8f3254252625ead91bac203e50c0ab65de2874f2dff490c", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/gather_c3_s2": { + "arch": "gfx950", + "bytes": 143496, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_gather_c3_s2_f3x3p1s2c3_wk2x1_pf2_rps16_tb1024_bf16", + "sha256": "b6d37178d5e31f5c397aaed864dd61d12c9931872fe9b69f4034d7c442f23d61", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_bf16_c128": { + "arch": "gfx950", + "bytes": 193857, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_bf16_c128_f3x3p1_wt32x32_w4x2_rps30_pf2_tb256_bf16", + "sha256": "46c81d2ba55e1130ba0bb3aa34856e0a5859208e49b6e1e06ad0e86024ff6640", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_fp16_unfold4": { + "arch": "gfx950", + "bytes": 188337, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_fp16_unfold4_f3x3p1_wt32x32_w4x2_rps15_pf2_tb256_u3_un4", + "sha256": "b12a6615a8b98ede61d291b8506bdcca374b13037a71d6e89ee43c665bb79828", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_qg2_pingpong": { + "arch": "gfx950", + "bytes": 344288, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_qg2_pingpong_f3x3p1_wt16x32_w2x2_rps12_pf3_tb256_bf16_qg2_pp", + "sha256": "9dacd91ff5ea1ba3156c9c97eebedbb5b0a764ddd616e7326bd9185f9cfdaaab", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53578, @@ -1902,7 +1966,7 @@ "gfx1250": 4, "gfx90a": 5, "gfx942": 42, - "gfx950": 151 + "gfx950": 159 }, "expected_failures": 4, "families": { @@ -1916,6 +1980,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, + "conv_direct_wgrad": 8, "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -1927,7 +1992,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 222 + "ok_cases": 230 } }, "llvm22": { @@ -3092,6 +3157,70 @@ "sha256": "65aecb5e9192cae49254347700f6aefdd11f47c805312d9fd9b4604a67f7755f", "status": "ok" }, + "conv_direct_wgrad/gfx950/dw_cols_c3_7x7_s2": { + "arch": "gfx950", + "bytes": 80427, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_dw_cols_c3_7x7_s2_f7x7p3s2c3_st8_t64_tb2048_bf16", + "sha256": "ac7974519cd6b2db4fb7e4da8fb96a1ac421f36f0db0632773ececa8472dd557", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/dw_lds_c64": { + "arch": "gfx950", + "bytes": 323686, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_dw_lds_c64_f3x3p1s1_lc8_lw32_cpl1_pf2_rps32_tb1024_bf16", + "sha256": "f420503d056420ab4ae36118f2fb032ffd725a6faa5985f4ac92b5bc9f46d385", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/fold_t9": { + "arch": "gfx950", + "bytes": 11974, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_fold_t9_t9", + "sha256": "6eb6f5f3f266d284fe6c83c94beb4883f34cfbc99f0cb32858cf5507d31177a4", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/gather_c24_tr": { + "arch": "gfx950", + "bytes": 229477, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_gather_c24_tr_f3x3p1s2c24_wk1x6_pf2_rps8_tb256_bf16", + "sha256": "2ca1296df7a196c58816527b26f1a19ab64b49c1232bdc6d6b882c7f9ef3165d", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/gather_c3_s2": { + "arch": "gfx950", + "bytes": 143592, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_gather_c3_s2_f3x3p1s2c3_wk2x1_pf2_rps16_tb1024_bf16", + "sha256": "0649e52910a0cae5a02cd8c9257f033d712aec4d9889067e9ab1f2c2dab08583", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_bf16_c128": { + "arch": "gfx950", + "bytes": 193953, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_bf16_c128_f3x3p1_wt32x32_w4x2_rps30_pf2_tb256_bf16", + "sha256": "d40aacacc77c8890cf9a90591b5900a5d7490617f9bec1b0b05391c62d877581", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_fp16_unfold4": { + "arch": "gfx950", + "bytes": 188433, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_fp16_unfold4_f3x3p1_wt32x32_w4x2_rps15_pf2_tb256_u3_un4", + "sha256": "ed19d39776f25ba1979a26173b65c7fb793d268600a60536f2dc10d792f26bd2", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_qg2_pingpong": { + "arch": "gfx950", + "bytes": 344388, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_qg2_pingpong_f3x3p1_wt16x32_w2x2_rps12_pf3_tb256_bf16_qg2_pp", + "sha256": "1fdeb533d3a975eeaa7431efb89457bcce602ad7e865ac20677b5aa35eed6bb1", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53717, @@ -3832,7 +3961,7 @@ "gfx1250": 4, "gfx90a": 5, "gfx942": 42, - "gfx950": 151 + "gfx950": 159 }, "expected_failures": 4, "families": { @@ -3846,6 +3975,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, + "conv_direct_wgrad": 8, "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -3857,7 +3987,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 222 + "ok_cases": 230 } }, "llvm23": { @@ -5022,6 +5152,70 @@ "sha256": "21f519428620240701ef490ea0c66989d993fa27a28f234b3b86069477530a8e", "status": "ok" }, + "conv_direct_wgrad/gfx950/dw_cols_c3_7x7_s2": { + "arch": "gfx950", + "bytes": 80491, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_dw_cols_c3_7x7_s2_f7x7p3s2c3_st8_t64_tb2048_bf16", + "sha256": "0b9f3cabec37a475f60014a5b151c6fca590c06b6045d68be2d6036afb00384a", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/dw_lds_c64": { + "arch": "gfx950", + "bytes": 323750, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_dw_lds_c64_f3x3p1s1_lc8_lw32_cpl1_pf2_rps32_tb1024_bf16", + "sha256": "805e6b9db3d4aca2c6a1656159dcb50a35a35ffe09bbc32bb1d04474047355f7", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/fold_t9": { + "arch": "gfx950", + "bytes": 12038, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_fold_t9_t9", + "sha256": "c3e1983fafb62392d0ba4aefb86eb7e4304fb38551554f10f1b64fde2dadd651", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/gather_c24_tr": { + "arch": "gfx950", + "bytes": 229541, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_gather_c24_tr_f3x3p1s2c24_wk1x6_pf2_rps8_tb256_bf16", + "sha256": "1008db7b8586064083b043c19b159020fd43a6a0a1eb6bb2d950c654d0b00578", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/gather_c3_s2": { + "arch": "gfx950", + "bytes": 143656, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_gather_c3_s2_f3x3p1s2c3_wk2x1_pf2_rps16_tb1024_bf16", + "sha256": "44512065b1f099343444bc836b155b65df48f5cf1f7a25c529afe8faaf64ee39", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_bf16_c128": { + "arch": "gfx950", + "bytes": 194017, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_bf16_c128_f3x3p1_wt32x32_w4x2_rps30_pf2_tb256_bf16", + "sha256": "57717396251cb0791a8c8edf28dc312fd36688394fa3ecc596c00a9d5aa84cb2", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_fp16_unfold4": { + "arch": "gfx950", + "bytes": 188497, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_fp16_unfold4_f3x3p1_wt32x32_w4x2_rps15_pf2_tb256_u3_un4", + "sha256": "3543328889743364b097d2d9d47a5e0dc6cbd1f7a5967821786a74f967528a1c", + "status": "ok" + }, + "conv_direct_wgrad/gfx950/tiled_qg2_pingpong": { + "arch": "gfx950", + "bytes": 344452, + "family": "conv_direct_wgrad", + "kernel_name": "irhash_direct_wgrad_tiled_qg2_pingpong_f3x3p1_wt16x32_w2x2_rps12_pf3_tb256_bf16_qg2_pp", + "sha256": "f926c3ff4d613a6e38cbddb59be8aad725317d4d13c0e4e4789422cac7a6e950", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53781, @@ -5777,7 +5971,7 @@ "gfx1250": 8, "gfx90a": 5, "gfx942": 42, - "gfx950": 151 + "gfx950": 159 }, "expected_failures": 0, "families": { @@ -5791,6 +5985,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, + "conv_direct_wgrad": 8, "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -5802,7 +5997,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 226 + "ok_cases": 234 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 488d394ff62b..c6c5ac94bbfc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -1207,6 +1207,48 @@ def _build(): return _build +def build_direct_wgrad(name, arch, kind, problem, **spec_kw): + """A direct wgrad kernel: ``kind`` is tiled / dw_lds / dw_cols / gather / + fold; ``problem`` overrides the DirectConvProblem fields (unused by the + fold). Configs mirror parity/conv_direct_grouped_emit.* (indices 95+), + where the C++ engine is gated byte-identical by check_byte_identity.py.""" + + def _build(): + from kernels.common import conv_direct_dw_wgrad as dw + from kernels.common import conv_direct_wgrad_gather as wg + from kernels.common import conv_direct_wgrad_tiled as wt + from kernels.common.conv_direct_grouped import DirectConvProblem + + if kind == "fold": + return dw.build_dw_wgrad_fold( + dw.DwWgradFoldSpec(name=name, **spec_kw), arch=arch + ) + kw = dict( + N=2, + H=16, + W=32, + groups=1, + cpg=64, + kpg=64, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + ) + kw.update(problem) + p = DirectConvProblem(**kw) + spec_cls, build = { + "tiled": (wt.DirectWgradTiledSpec, wt.build_direct_wgrad_tiled), + "dw_lds": (dw.DwWgradLdsSpec, dw.build_dw_wgrad_lds), + "dw_cols": (dw.DwWgradColsSpec, dw.build_dw_wgrad_cols), + "gather": (wg.WgradGatherSpec, wg.build_wgrad_gather), + }[kind] + return build(spec_cls(problem=p, name=name, **spec_kw), arch=arch) + + return _build + + def build_direct_depthwise_tiled( name, arch, N, H, W, groups, KH, PAD, stride=1, dtype="bf16", **spec_kw ): @@ -4074,6 +4116,62 @@ def add(family, case_id, arch, build): ), ) + # --- conv_direct_wgrad: the direct wgrad kernels on gfx950 --- + # Tiled (LDS-DMA row pipeline, swizzled transpose reads, atomics), the + # depthwise partial-slot kernels and their fold, and the small-C gather. + # The full branch matrix (and the C++ twins) lives in + # library/tests/parity/conv_direct_grouped_emit.* (indices 95+). + for _case_id, _kind, _prob, _over in ( + ("tiled_bf16_c128", "tiled", dict(cpg=128, kpg=128, H=30), {}), + ( + "tiled_qg2_pingpong", + "tiled", + {}, + dict( + q_groups=2, + pingpong=True, + waves_k=2, + wave_k16=1, + rows_per_seg=12, + prefetch=3, + ), + ), + ( + "tiled_fp16_unfold4", + "tiled", + dict(N=4, H=15, W=40, dtype="fp16"), + dict(unfold_n=4, rows_per_seg=15, row_unroll=3), + ), + ("fold_t9", "fold", {}, dict(taps=9)), + ( + "dw_cols_c3_7x7_s2", + "dw_cols", + dict(groups=3, cpg=1, kpg=1, KH=7, KW=7, PAD=3, stride=2, H=20), + {}, + ), + ("dw_lds_c64", "dw_lds", dict(groups=64, cpg=1, kpg=1), {}), + ( + "gather_c3_s2", + "gather", + dict(cpg=3, kpg=24, stride=2, H=32, W=64), + {}, + ), + ( + "gather_c24_tr", + "gather", + dict(cpg=24, kpg=96, stride=2, H=30, W=40), + dict(wave_k16=1, waves_k=6, rows_per_seg=8, target_blocks=256), + ), + ): + add( + "conv_direct_wgrad", + f"conv_direct_wgrad/gfx950/{_case_id}", + "gfx950", + build_direct_wgrad( + f"irhash_direct_wgrad_{_case_id}", "gfx950", _kind, _prob, **_over + ), + ) + # --- conv_direct_depthwise_tiled: output-stationary depthwise direct conv --- # Runtime filter-row loop, sliding input-row window, one weight row in # registers. Mirrors library/tests/parity/conv_direct_grouped_emit.*