From 38f9d4cea671a88798d28398ba08b0c0020f65cf Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Fri, 30 Oct 2026 21:05:31 +0000 Subject: [PATCH 1/5] feat(rocke): shape-generic (AOT) Winograd forward transforms for gfx1250 The Winograd F(m,3) transform kernels (K1 filter, K2a prescale max, K2 input, K4 output) bake the problem extents in, so every convolution shape needs its own binaries and the pipeline cannot be served from an AOT cache. conv_winograd_fwd_rt builds the same stages, workspace layouts and arithmetic with every extent as a kernel argument (divisions by runtime extents through a host-computed magic number), so one binary serves every problem its knobs fit. The K3 batched GEMMs already take their extents at run time. - kernels/common/conv_winograd_fwd_rt.py: per-stage specs (RtFilterSpec, RtPrescaleSpec, RtInputSpec, RtOutputSpec), stage_launch (grid, block and scalar kernargs), the AOT grid (aot_stage_specs) and the default per-problem pick (default_stage_specs). K2 keeps the register path (with a fused prescale over a runtime lane count) and the LDS region. - C++ engine port with byte-identical IR; the shared Winograd math helpers move into rocke_wino (no IR change to the existing families). - conv_winograd_fwd_rt_emit.{py,c} parity pair and a representative IR golden shard. - Dispatch: plan kind unfused2d_rt and the opt-in candidate conv_winograd_fwd_aot (algorithm="winograd_aot") over the selector's 3x3 plans; algorithm="auto" never picks it. - benchmarks/common/benchmark_winograd_aot.py: --compile-all builds the grid into a cache, --run-from-cache assembles each case's pipeline from it, picking each stage's binary by timing and verifying the pipeline. - Tests: CPU (kernels, dispatch) and on-GPU numerics; docs and support matrix. Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_winograd_aot.py | 720 +++++++ .../library/dispatch/conv_winograd_plan.py | 83 +- .../library/dispatch/grouped_convolution.py | 109 +- .../kernels/common/conv_winograd_fwd_rt.py | 1537 ++++++++++++++ .../kernels/common/conv_winograd_select.py | 2 +- .../test_conv_winograd_aot_dispatch.py | 177 ++ .../tests/parity/conv_winograd_fwd_rt_emit.c | 209 ++ .../tests/parity/conv_winograd_fwd_rt_emit.py | 153 ++ .../tests/test_builder_signature_contract.py | 5 + .../tests/test_conv_winograd_fwd_rt_cpu.py | 168 ++ .../tests/test_conv_winograd_fwd_rt_gpu.py | 253 +++ .../rocke/instance_conv_winograd_fwd_rt.h | 102 + .../rocke/instance_conv_winograd_internal.h | 18 + .../instances/common/conv_winograd_common.cpp | 215 ++ .../instances/common/conv_winograd_fwd.cpp | 189 +- .../instances/common/conv_winograd_fwd_rt.cpp | 1772 +++++++++++++++++ .../dsl_docs/instances/convolution.md | 37 +- .../platform/dsl_docs/instances/index.md | 4 +- .../python/rocke/instances/SUPPORT_MATRIX.md | 2 +- .../conv_winograd_fwd_rt.json | 167 ++ .../instances/rocke_ir_parity_harness.py | 39 + 21 files changed, 5755 insertions(+), 206 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_fwd_rt.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py new file mode 100644 index 000000000000..e29746337837 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py @@ -0,0 +1,720 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""AOT Winograd ``F(m x m, 3 x 3)`` forward: compile once, run every shape from the cache. + +``benchmark_winograd_conv.py`` builds the transform kernels for each problem +(the extents are baked in). This driver uses the shape-generic transforms of +:mod:`kernels.common.conv_winograd_fwd_rt` and the K3 GEMMs (the WMMA batched +GEMM and the implicit-GEMM conv on the probe problem, both shape-generic +already), so the whole pipeline is a set of cached binaries: + +* ``--compile-all --cache-dir DIR``: build every binary of the grid below + into ``DIR//wino_rt`` (CPU only, ``--jobs`` workers); +* ``--run-from-cache DIR --miopen-file F``: for each 3x3 stride-1 case and + each tile (all of :data:`conv_winograd_fwd_rt.AOT_TILES`), pick each stage's + fastest fitting binary by timing it alone (K3: every plan, as + ``benchmark_winograd_conv.py`` sweeps them), verify the assembled pipeline + against the reference conv, and time the whole chain (min over + ``--repeats``). One result row per case (``--results``, JSON Lines) with + every tile's pipeline, its time and its binaries in ``top``. + +The prescale and operand dtype follow ``benchmark_winograd_conv.py``'s +``auto`` rules (bf16 operands only for F(2,3) with bf16 I/O, also tried with +fp16 operands; prescale for bf16 I/O with a tile side >= 4). +""" + +from __future__ import annotations + +import argparse +import dataclasses +import json +import os +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Dict, List, Optional, Tuple + +os.environ.setdefault("ROCKE_CPP_QUIET_FALLBACK", "1") + +VECS = (8, 4, 2, 1) +IOS = ("fp16", "bf16") + + +def _tiles(): + from kernels.common.conv_winograd_fwd_rt import AOT_TILES + + return tuple(AOT_TILES) + + +def _operands(tile, io: str) -> Tuple[str, ...]: + """Operand dtypes timed per tile: the ``auto`` one, and fp16 operands as + well for bf16 F(2,3).""" + from kernels.common.conv_winograd_fwd_rt import pipeline_dtypes + + op, _ = pipeline_dtypes(tile, io) + return (op, "fp16") if op == "bf16" else (op,) + + +def _prescale(tile, io: str) -> bool: + from kernels.common.conv_winograd_fwd_rt import pipeline_dtypes + + return pipeline_dtypes(tile, io)[1] + + +# --------------------------------------------------------------------------- +# The binary grid +# --------------------------------------------------------------------------- + + +def transform_specs() -> list: + """The transform binaries of the grid (``conv_winograd_fwd_rt.aot_stage_specs``).""" + from kernels.common.conv_winograd_fwd_rt import aot_stage_specs + + return aot_stage_specs() + + +def _k3i_name(cfg, dtype: str, vab: int, vc: int) -> str: + from benchmarks.common.benchmark_winograd_conv import k3_label + + return f"wino_rt_k3i_{dtype}_{k3_label(cfg)}_va{vab}_vc{vc}" + + +def _k3i_spec(cfg, dtype: str, vab: int, vc: int, arch: str): + """The implicit-GEMM K3 spec with the given vector widths (probe problem).""" + from benchmarks.common.benchmark_winograd_conv import k3_spec + + fake = SimpleNamespace(cpg=vab, kpg=vc) + return k3_spec(cfg, fake, dtype, arch) + + +def _job_list(arch: str) -> list: + from benchmarks.common import winograd_k3_wmma as K3W + from benchmarks.common.benchmark_winograd_conv import K3_CONFIGS_SHORT + + jobs = [("T", s) for s in transform_specs()] + for dt in IOS: + for knobs in K3W.K3W_CONFIGS: + jobs.append(("W", (dict(knobs), dt))) + for cfg in K3_CONFIGS_SHORT: + for vab in VECS: + for vc in VECS: + if cfg[7] == "default" and vc != 1: + continue + jobs.append(("I", (cfg, dt, vab, vc))) + return jobs + + +def _compile_one(item): + """``(name, meta, hsaco, symbol)`` or ``(name, None, error, None)``.""" + kind, payload, arch = item + from rocke import compile_kernel + + try: + if kind == "T": + from kernels.common import conv_winograd_fwd_rt as W + + spec = payload + name = spec.kernel_name() + art = compile_kernel(W.build_stage(spec, arch=arch), arch=arch) + meta = dict(kind="T", stage=spec.stage, cls=type(spec).__name__, + knobs=dataclasses.asdict(spec)) # fmt: skip + elif kind == "W": + from benchmarks.common import winograd_k3_wmma as K3W + from kernels.common.conv_winograd_gemm_wmma import ( + build_winograd_gemm_wmma, + winograd_gemm_reason, + ) + + knobs, dt = payload + spec = K3W.spec_of(knobs, dt) + name = "wino_rt_k3w_" + spec.kernel_name() + if winograd_gemm_reason(spec, arch) is not None: + return name, None, "invalid", None + art = compile_kernel(build_winograd_gemm_wmma(spec, arch=arch), arch=arch) + meta = dict(kind="W", stage="K3", knobs=knobs, dtype=dt) + else: + from kernels.common.conv_implicit_gemm import ( + build_implicit_gemm_conv, + is_valid_spec, + ) + + cfg, dt, vab, vc = payload + name = _k3i_name(cfg, dt, vab, vc) + spec = _k3i_spec(cfg, dt, vab, vc, arch) + if spec is None or not is_valid_spec(spec, arch)[0]: + return name, None, "invalid", None + art = compile_kernel(build_implicit_gemm_conv(spec, arch=arch), arch=arch) + meta = dict(kind="I", stage="K3", cfg=list(cfg), dtype=dt, vab=vab, vc=vc) + except Exception as e: # noqa: BLE001 -- reported, the grid goes on + return f"{kind}:{payload!r}", None, f"{type(e).__name__}: {e}", None + return name, meta, bytes(art.hsaco), art.kernel_name + + +def compile_all(cache_dir: str, arch: str, jobs: int) -> int: + import multiprocessing + from concurrent.futures import ProcessPoolExecutor + + out = Path(cache_dir) / arch / "wino_rt" + out.mkdir(parents=True, exist_ok=True) + work = [(k, p, arch) for k, p in _job_list(arch)] + print(f"compile-all: {len(work)} binaries -> {out}", flush=True) + t0 = time.time() + index: Dict[str, dict] = {} + n_bad = 0 + try: + ctx = multiprocessing.get_context("forkserver") + except ValueError: + ctx = multiprocessing.get_context("spawn") + from benchmarks.common import benchmark_winograd_aot as _self + + with ProcessPoolExecutor(max_workers=max(1, jobs), mp_context=ctx) as pool: + for i, (name, meta, data, sym) in enumerate( + pool.map(_self._compile_one, work, chunksize=4), 1 + ): + if meta is None: + if data != "invalid": + n_bad += 1 + print(f" [fail] {name}: {data}", flush=True) + continue + (out / f"{name}.hsaco").write_bytes(data) + meta["symbol"] = sym + index[name] = meta + if i % 200 == 0: + print(f" {i}/{len(work)} ({time.time() - t0:.0f}s)", flush=True) + (out / "index.json").write_text(json.dumps(index, indent=1, sort_keys=True)) + print( + f"compile-all done: {len(index)} binaries, {n_bad} failed " + f"({time.time() - t0:.0f}s)", + flush=True, + ) + return 1 if n_bad else 0 + + +# --------------------------------------------------------------------------- +# Cache +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Binary: + name: str + meta: dict + hsaco: bytes + kernel_name: str + spec: object = None + + @property + def art(self): + """The artifact ``K3Plan.bind`` launches (``.hsaco``, ``.kernel_name``).""" + return self + + +def load_cache(cache_dir: str, arch: str) -> List[Binary]: + from kernels.common import conv_winograd_fwd_rt as W + + d = Path(cache_dir) / arch / "wino_rt" + index = json.loads((d / "index.json").read_text()) + out = [] + for name, meta in sorted(index.items()): + spec = None + if meta["kind"] == "T": + knobs = dict(meta["knobs"]) + if "lds" in knobs: + knobs["lds"] = tuple(knobs["lds"]) + spec = getattr(W, meta["cls"])(**knobs) + out.append( + Binary(name, meta, (d / f"{name}.hsaco").read_bytes(), meta["symbol"], spec) + ) + return out + + +# --------------------------------------------------------------------------- +# One case +# --------------------------------------------------------------------------- + + +def _chunks(problem, tile, prescale: bool) -> Optional[int]: + """Fewest N-chunks (a divisor of N) whose buffers and K3 operands fit the + 32-bit byte offsets. (The WMMA K3's own constraints are checked per plan: + a problem it cannot run may still run on the implicit-GEMM K3.)""" + from kernels.common import conv_winograd_fwd_rt as W + + for c in range(1, problem.N + 1): + if problem.N % c: + continue + sub = dataclasses.replace(problem, N=problem.N // c) + geo = W.RtGeometry(sub, *tile) + if geo.check(prescale): + continue + g = _gemm_prob(geo) + elems = (g.T * g.G * g.Cg, g.G * g.Kg * g.Cg, g.T * g.G * g.Kg) + if 2 * max(elems) < 2**31: + return c + return None + + +def _gemm_prob(geo): + from kernels.common.conv_winograd_gemm_wmma import WinogradGemmProblem + + p = geo.problem + return WinogradGemmProblem( + T=geo.num_tiles, G=geo.n2 * p.groups, Cg=p.cpg, Kg=p.K // p.groups + ) + + +def _gemm_conv_problem(geo): + from kernels.common._conv_implicit_gemm_common import ConvProblem + + p = geo.problem + return ConvProblem( + N=geo.num_tiles, + Hi=1, + Wi=1, + C=geo.n2 * p.C, + K=geo.n2 * p.K, + Y=1, + X=1, + groups=geo.n2 * p.groups, + ) + + +class Case: + """Device buffers of one case: x / y (all chunks) and w are shared by the tiles.""" + + def __init__(self, rt, problem, io: str, seed: int, verify: bool): + from benchmarks.common.benchmark_winograd_conv import make_inputs, reference + + self.rt = rt + self.problem = problem + self.io = io + self.x_t, self.w_t, self.y_t = make_inputs(problem, io, seed=seed) + self.ref = reference(self.x_t, self.w_t, problem, io) if verify else None + self.xb = self.x_t.numel() * 2 + self.wb = self.w_t.numel() * 2 + self.yb = self.y_t.numel() * 2 + self.x = int(rt.alloc(max(self.xb, 4))) + self.w = int(rt.alloc(max(self.wb, 4))) + self.y = int(rt.alloc(max(self.yb, 4))) + from benchmarks.common.benchmark_winograd_conv import _u8 + + rt.sync() + for ptr, t in ((self.x, self.x_t), (self.w, self.w_t)): + h = t.contiguous().pin_memory() + rt.memcpy_h2d(ptr, _u8(h), h.nbytes) + rt.sync() + del h + + def read_y(self): + import torch + + from benchmarks.common.benchmark_winograd_conv import _u8 + + self.rt.sync() + h = torch.empty_like(self.y_t, device="cpu").pin_memory() + self.rt.memcpy_d2h(_u8(h), self.y, h.nbytes) + self.rt.sync() + return h.clone() + + def free(self): + from rocke.runtime import synchronize_and_release + + synchronize_and_release(0) + for p in (self.x, self.w, self.y): + self.rt.free(p) + + +class TilePipe: + """Workspaces and bound launches of one tile's pipeline on one case.""" + + def __init__(self, case: Case, tile, op: str, prescale: bool, chunks: int): + from kernels.common import conv_winograd_fwd_rt as W + + self.case = case + self.tile = tuple(tile) + self.op = op + self.prescale = prescale + self.chunks = chunks + self.sub = dataclasses.replace(case.problem, N=case.problem.N // chunks) + self.geo = W.RtGeometry(self.sub, *tile) + sz = self.geo.buffer_bytes(prescale) + self.sizes = sz + rt = case.rt + self.ws = {k: int(rt.alloc(max(sz[k], 4))) for k in ("U", "V", "M", "e")} + self.xc = sz["x"] + self.yc = sz["y"] + + def free(self): + from rocke.runtime import synchronize_and_release + + synchronize_and_release(0) + for p in self.ws.values(): + self.case.rt.free(p) + + def poison(self): + rt = self.case.rt + for k, p in self.ws.items(): + rt.memset(p, 0x7F, self.sizes[k]) + rt.memset(self.case.y, 0x7F, self.case.yb) + + # -- launches -- + + def bind_stage(self, binary: Binary, chunk: int = 0): + from kernels.common import conv_winograd_fwd_rt as W + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + spec = binary.spec + grid, block, vals = W.stage_launch(spec, self.geo) + x = self.case.x + chunk * self.xc + y = self.case.y + chunk * self.yc + bufs = dict( + w=(self.case.w, self.sizes["w"]), + u=(self.ws["U"], self.sizes["U"]), + x=(x, self.xc), + v=(self.ws["V"], self.sizes["V"]), + e=(self.ws["e"], self.sizes["e"]), + mm=(self.ws["M"], self.sizes["M"]), + m=(self.ws["M"], self.sizes["M"]), + y=(y, self.yc), + ) + sig = W.stage_signature(spec) + args = dict(vals) + for a in sig: + nm = a["name"] + if nm in args: + continue + if nm.endswith("_bytes"): + args[nm] = bufs[nm[: -len("_bytes")]][1] + else: + args[nm] = bufs[nm][0] + lau = KernelLauncher( + hsaco=binary.hsaco, kernel_name=binary.kernel_name, signature=sig + ) + return lau.bind(args, config=LaunchConfig(grid=grid, block=(block, 1, 1))) + + def bind_k3(self, plan_or_impl): + """A WMMA ``K3Plan`` or an implicit-GEMM ``(binary, spec)``.""" + from benchmarks.common.winograd_k3_wmma import K3Plan + + if isinstance(plan_or_impl, K3Plan): + return plan_or_impl.bind( + _gemm_prob(self.geo), self.ws["V"], self.ws["U"], self.ws["M"] + ) + from kernels.common.conv_abi import conv_args_signature + from kernels.common.conv_args import ConvArgs + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + binary, spec = plan_or_impl + gp = _gemm_conv_problem(self.geo) + lau = KernelLauncher( + hsaco=binary.hsaco, + kernel_name=binary.kernel_name, + signature=conv_args_signature(self.op, is_3d=False), + ) + S = self.sizes + vals = ConvArgs.from_problem( + gp, tile_m=spec.tile_m, tile_n=spec.tile_n + ).to_launch_values( + self.ws["V"], self.ws["U"], self.ws["M"], S["V"], S["U"], S["M"] + ) + grid = (-(-gp.N_gemm // spec.tile_n), -(-gp.M // spec.tile_m), gp.groups) + if min(grid) <= 0 or grid[1] > 65535 or grid[2] > 65535: + raise ValueError(f"K3 grid {grid} not launchable") + return lau.bind( + vals, config=LaunchConfig(grid=grid, block=(spec.launch_block_size, 1, 1)) + ) + + def chain(self, pick: dict, k3): + """The whole forward: K1 once, then per chunk K2a, K2, K3, K4.""" + fns = [self.bind_stage(pick["K1"])] + k3f = self.bind_k3(k3) + for i in range(self.chunks): + if "K2a" in pick: + fns.append(self.bind_stage(pick["K2a"], i)) + fns.append(self.bind_stage(pick["K2"], i)) + fns.append(k3f) + fns.append(self.bind_stage(pick["K4"], i)) + + def run(_f=tuple(fns)): + for f in _f: + f() + + return run + + +def _fits(binary: Binary, tp: TilePipe) -> bool: + s = binary.spec + m = binary.meta + if m["kind"] != "T": + return False + k = m["knobs"] + if k.get("io_dtype") != tp.case.io: + return False + if m["stage"] != "K2a": + if (k["mh"], k["mw"]) != tp.tile or k["operand_dtype"] != tp.op: + return False + if m["stage"] in ("K2", "K4") and k["prescale"] != tp.prescale: + return False + return s.reason(tp.geo) is None + + +def run_case( + case: Case, tile, op: str, cache: List[Binary], args, log +) -> Optional[dict]: + """The best verified pipeline of one tile: ``{label, ms, binaries, stages}``.""" + from benchmarks.common import winograd_k3_wmma as K3W + from benchmarks.common.benchmark_winograd_conv import verify_error + from rocke.runtime import time_launches + + prescale = _prescale(tile, case.io) + chunks = _chunks(case.problem, tile, prescale) + if chunks is None: + log(f" [skip] F{tile}: no batch split fits") + return None + tp = TilePipe(case, tile, op, prescale, chunks) + try: + + def us(fn): + return time_launches(fn, warmup=args.warmup, iters=args.iters) * 1e3 + + ranked: Dict[str, List[Tuple[float, Binary]]] = {} + for st in ("K1", "K2a", "K2", "K4") if prescale else ("K1", "K2", "K4"): + cands = [b for b in cache if b.meta.get("stage") == st and _fits(b, tp)] + timed = [] + for b in cands: + try: + timed.append((us(tp.bind_stage(b)), b)) + except ValueError as e: + log(f" [skip] {b.name}: {e}") + timed.sort(key=lambda r: r[0]) + ranked[st] = timed + # The prescale's input side: a fused K2, or K2a + a split K2. + if prescale: + fused = [r for r in ranked["K2"] if r[1].spec.fused] + split = [r for r in ranked["K2"] if not r[1].spec.fused] + opts = [] + if fused: + opts.append((fused[0][0], {"K2": fused[0][1]})) + if split and ranked["K2a"]: + k2a = ranked["K2a"][0] + opts.append((split[0][0] + k2a[0], {"K2a": k2a[1], "K2": split[0][1]})) + if not opts: + log(f" [skip] F{tile}: no K2 binary fits") + return None + k2t, k2pick = min(opts, key=lambda o: o[0]) + ranked["K2"] = [(k2t, k2pick["K2"])] + if "K2a" not in k2pick: + del ranked["K2a"] + for st in ("K1", "K2", "K4"): + if not ranked[st]: + log(f" [skip] F{tile}: no {st} binary fits") + return None + stages = tuple(st for st in ("K1", "K2a", "K2", "K4") if st in ranked) + pick = {st: ranked[st][0][1] for st in stages} + if "K2a" in stages: + pick["K2a"] = k2pick["K2a"] + + # K3: every WMMA plan (single + main/tail splits) and every implicit config. + prob = _gemm_prob(tp.geo) + gp = _gemm_conv_problem(tp.geo) + built = [ + (b.meta["knobs"], K3W.spec_of(b.meta["knobs"], tp.op), b.art) + for b in cache + if b.meta["kind"] == "W" and b.meta["dtype"] == tp.op + ] + plans = K3W.candidate_plans(prob, built, args.cus, tails=True) + k3c = [] + for pl in plans: + if pl.reason(prob): + continue + try: + k3c.append((us(tp.bind_k3(pl)), "wmma:" + pl.label(), pl, + [s.art.name for s in pl.steps])) # fmt: skip + except ValueError: + continue + vab = _pow2(gp.cpg) + vc_max = _pow2(gp.kpg) + for b in cache: + m = b.meta + if m["kind"] != "I" or m["dtype"] != tp.op or m["vab"] != vab: + continue + cfg = tuple(m["cfg"]) + want_vc = 1 if cfg[7] == "default" else vc_max + if m["vc"] != want_vc: + continue + spec = _k3i_spec(cfg, tp.op, m["vab"], m["vc"], args.arch) + if spec is None or -(-gp.M // spec.tile_m) > 65535: + continue + try: + k3c.append( + ( + us(tp.bind_k3((b, spec))), + "implicit:" + b.name, + (b, spec), + [b.name], + ) + ) + except ValueError: + continue + if not k3c: + log(f" [skip] F{tile}: no K3 candidate") + return None + k3c.sort(key=lambda r: r[0]) + + # Assemble: the fastest K3 that verifies through the whole pipeline. + best = None + for k3us, k3label, k3, k3bins in k3c[: args.k3_finalists]: + run = tp.chain(pick, k3) + err = None + if case.ref is not None: + tp.poison() + run() + err, _ = verify_error(case.read_y(), case.ref) + if not err < 5e-2: + log(f" F{tile} op={op} K3 {k3label[:70]} FAIL err={err:.2e}") + continue + best = (run, k3label, k3bins, err) + break + if best is None: + return None + run, k3label, k3bins, err = best + total = None + for _ in range(args.repeats): + if args.cooldown > 0: + time.sleep(args.cooldown) + t = us(run) + total = t if total is None else min(total, t) + bins = [pick[st].name for st in stages] + list(k3bins) + label = ( + f"winograd-{case.io}-F{tile[0]}x{tile[1]}-op={op}-ps={int(prescale)}" + f"-k1={pick['K1'].name}-k2={pick['K2'].name}-k3={k3label}-k4={pick['K4'].name}" + ) + log( + f" F{tile[0]}x{tile[1]} op={op} ps={int(prescale)} chunks={chunks}: " + f"{total:9.2f} us err={err if err is not None else float('nan'):.2e} " + + " ".join(f"{st}={ranked[st][0][0]:.1f}" for st in stages) + + f" K3={k3c[0][0]:.1f} [{k3label[:60]}]" + ) + return dict(label=label, ms=total / 1e3, binaries=bins, err=err) + finally: + tp.free() + + +def _pow2(n: int, cap: int = 8) -> int: + v = cap + while v > 1 and n % v: + v //= 2 + return v + + +def run_from_cache(args) -> int: + import torch # noqa: F401 -- claim the HIP context first + + from benchmarks.common.benchmark_implicit_gemm_conv import parse_miopen_cmd + from benchmarks.common.early_stop import case_key + from rocke.runtime.hip_module import Runtime + + cache = load_cache(args.run_from_cache, args.arch) + print(f"cache: {len(cache)} binaries", flush=True) + cmds = [ + ln.strip() + for ln in open(args.miopen_file) + if ln.strip() and not ln.startswith("#") + ] + out = open(args.results, "a" if args.resume else "w") if args.results else None + done = set() + if args.resume and args.results and Path(args.results).exists(): + for ln in open(args.results): + try: + done.add(json.loads(ln)["miopen_cmd"]) + except (ValueError, KeyError): + pass + rt = Runtime() + rc = 0 + for ci, cmd in enumerate(cmds, 1): + if cmd in done: + continue + problem, dtype, forw = parse_miopen_cmd(cmd) + if ( + forw != 1 + or (problem.Y, problem.X) != (3, 3) + or (problem.sH, problem.sW) != (1, 1) + ): + print(f"[skip] not a 3x3 stride-1 forward: {cmd}", flush=True) + continue + print(f"\n== [{ci}/{len(cmds)}] {problem.short()} {dtype}", flush=True) + case = Case(rt, problem, dtype, args.seed, args.verify) + res = [] + try: + for tile in _tiles(): + for op in _operands(tile, dtype): + try: + r = run_case( + case, tile, op, cache, args, lambda m: print(m, flush=True) + ) + except (ValueError, RuntimeError) as e: + print(f" [skip] F{tile} op={op}: {e}", flush=True) + r = None + if r: + res.append(r) + finally: + case.free() + if not res: + print(" no Winograd pipeline passed", flush=True) + rc = 1 + continue + res.sort(key=lambda r: r["ms"]) + best = res[0] + print(f"Best: {best['ms'] * 1e3:.2f} us — {best['label']}", flush=True) + if out: + row = dict( + key=case_key(problem, dtype, "fwd"), + miopen_cmd=cmd, + engine="winograd", + shape=problem.short(), + dtype=dtype, + direction="fwd", + ms=round(best["ms"], 6), + kernel=best["label"], + timing="wino_aot", + top=[[r["label"], round(r["ms"], 6), r["binaries"]] for r in res], + ) + out.write(json.dumps(row, sort_keys=True) + "\n") + out.flush() + if out: + out.close() + return rc + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--arch", default="gfx1250") + ap.add_argument("--compile-all", action="store_true") + ap.add_argument("--cache-dir", default="./kernel_cache") + ap.add_argument("--jobs", type=int, default=os.cpu_count() or 8) + ap.add_argument("--run-from-cache", default=None, metavar="DIR") + ap.add_argument("--miopen-file", default=None) + ap.add_argument("--results", default=None) + ap.add_argument("--resume", action="store_true") + ap.add_argument("--verify", action="store_true") + ap.add_argument("--warmup", type=int, default=3) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument("--repeats", type=int, default=2) + ap.add_argument("--cooldown", type=float, default=0.0) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--k3-finalists", type=int, default=4) + ap.add_argument("--cus", type=int, default=0, help="0: query the device") + args = ap.parse_args(argv) + if args.compile_all: + return compile_all(args.cache_dir, args.arch, args.jobs) + if not args.run_from_cache or not args.miopen_file: + ap.error("give --compile-all, or --run-from-cache with --miopen-file") + if not args.cus: + from benchmarks.common.winograd_k3_wmma import device_cus + + args.cus = device_cus() + return run_from_cache(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py index 412d7a15bc30..11275ce0e793 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py @@ -13,13 +13,17 @@ input transform, K3 batched GEMM (:mod:`kernels.common.conv_winograd_gemm_wmma`) and K4 output transform; a batch split runs K2 / K3 / K4 once per chunk of ``N / chunks`` images on byte offsets of x and y, sharing U; +* ``unfused2d_rt`` (3x3, :mod:`kernels.common.conv_winograd_fwd_rt`): the same + pipeline (without the merged launch) on shape-generic transform binaries, + every problem extent a kernarg, so each binary is compiled once ahead of + time and named by its kernel name; * ``unfused_gen`` (1-D, 7x1, 5x5, segmented long filters, :mod:`kernels.common.conv_winograd_gen`): K1, K2, K3 with ``taps`` segments, K4; * ``fused1d`` (:mod:`kernels.common.conv_winograd_fused_1d_wmma`): the generalized K1 writes U, then one fused kernel goes from x to y. -The steps are :class:`ConvLaunchStep` records with no AOT identity: each +The steps are :class:`ConvLaunchStep` records with no conv-kernel identity: each carries the kernel to build (``kernel.build(arch)``) and its kernargs by name, resolved by :meth:`ConvLaunchPlan.values` against the buffers ``x`` / ``w`` / ``y`` and the plan's workspaces (``U``, ``V``, ``M``, ``e``). Transform knobs, @@ -309,6 +313,76 @@ def stage_step(name, chunk): return tuple(steps), workspaces +def _rt_kernargs(signature, vals: dict, sizes: dict, offsets: dict) -> tuple: + """Kernargs of a shape-generic transform step: its scalar values from + ``stage_launch``, ``*_bytes`` the chunk's buffer sizes, pointers at the + chunk's byte offsets.""" + out = [] + for a in signature: + nm = a["name"] + if nm in vals: + out.append(("i32", int(vals[nm]))) + elif nm.endswith("_bytes"): + out.append(("i32", int(sizes[_BUF_OF[nm[: -len("_bytes")]]]))) + else: + buf = _BUF_OF[nm] + out.append(("buf", buf, int(offsets.get(buf, 0)))) + return tuple(out) + + +def _unfused2d_rt(problem: ConvProblem, dtype: str, plan: WinogradPlan): + """The 3x3 pipeline on the shape-generic (AOT) transform binaries of + :mod:`kernels.common.conv_winograd_fwd_rt`: each step is named by its + binary's kernel name (its AOT identity) and its problem extents are + kernargs. The binaries are :func:`default_stage_specs`; K3 is the plan's + WMMA GEMM, as in :func:`_unfused2d`.""" + from kernels.common import conv_winograd_fwd_rt as R + from kernels.common.conv_winograd_gemm_wmma import ( + WinogradGemmProblem, + WinogradGemmWmmaSpec, + ) + + mh, mw = plan.tile + chunks = max(1, int(plan.chunks)) + if problem.N % chunks: + raise ValueError(f"batch split {chunks} does not divide N={problem.N}") + sub = dataclasses.replace(problem, N=problem.N // chunks) + geo = R.RtGeometry(sub, mh, mw) + operand, prescale = R.pipeline_dtypes((mh, mw), dtype) + why = geo.check(prescale) + if why: + raise ValueError(why) + specs = R.default_stage_specs(sub, (mh, mw), dtype) + sizes = dict(geo.buffer_bytes(prescale)) + x_chunk, y_chunk = sizes["x"], sizes["y"] + + def stage_step(stage, chunk): + spec = specs[stage] + grid, block, vals = R.stage_launch(spec, geo) + sig = R.stage_signature(spec) + offs = {"x": chunk * x_chunk, "y": chunk * y_chunk} + return _step( + spec.kernel_name(), + functools.partial(_build_with_arch, R.build_stage, spec), + grid, + block, + sig, + _rt_kernargs(sig, vals, sizes, offs), + ) + + steps = [stage_step("K1", 0)] + k3 = WinogradGemmWmmaSpec(dtype=operand, **_k3_knobs(plan.k3)) + prob = WinogradGemmProblem.from_conv(sub, (mh, mw)) + for i in range(chunks): + for k in ("K2a", "K2"): + if k in specs: + steps.append(stage_step(k, i)) + steps.append(_k3_step(k3, prob)) + steps.append(stage_step("K4", i)) + workspaces = {k: int(max(sizes[k], 4)) for k in ("U", "V", "M", "e")} + return tuple(steps), workspaces + + def _unfused_gen(problem: ConvProblem, dtype: str, plan: WinogradPlan): from kernels.common import conv_winograd_gen as WG from kernels.common.conv_winograd_gemm_wmma import WinogradGemmWmmaSpec @@ -407,7 +481,12 @@ def _fused1d(problem: ConvProblem, dtype: str, plan: WinogradPlan): return tuple(steps), {"U": int(max(gsizes["U"], 4))} -_KINDS = {"unfused2d": _unfused2d, "unfused_gen": _unfused_gen, "fused1d": _fused1d} +_KINDS = { + "unfused2d": _unfused2d, + "unfused2d_rt": _unfused2d_rt, + "unfused_gen": _unfused_gen, + "fused1d": _fused1d, +} @functools.lru_cache(maxsize=256) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py index bfc9a3837160..2c9aeb5ebbdc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py @@ -23,6 +23,10 @@ built per problem (no AOT binary), its eligibility and plan come from :mod:`kernels.common.conv_winograd_select`, and its :class:`ConvLaunchPlan` runs the transform, GEMM and output kernels in order over named workspaces. +Its AOT twin, ``conv_winograd_fwd_aot`` (``algorithm="winograd_aot"``, also +opt-in), runs the same 3x3 plans on the shape-generic transform binaries of +:mod:`kernels.common.conv_winograd_fwd_rt` (plan kind ``unfused2d_rt``): +every kernel of its plans is compiled once and named by its kernel name. Every record there spells the :class:`KernelIdentity` of one AOT kernel (format: @@ -1195,7 +1199,9 @@ def _winograd_model(arch: str) -> WinogradCostModel: return WinogradCostModel(num_cus=_ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS)) -def _winograd_spec(req: ConvGroupedRequest, plan) -> ConvWinogradSpec: +def _winograd_spec( + req: ConvGroupedRequest, plan, name: str = _WINOGRAD_NAME +) -> ConvWinogradSpec: """The spec of ``plan`` on ``req``; raises ``ValueError`` when a kernel of the plan cannot run the problem or does not lower at the request's LLVM flavor (the K3 GEMM lowers at llvm23 only).""" @@ -1204,9 +1210,9 @@ def _winograd_spec(req: ConvGroupedRequest, plan) -> ConvWinogradSpec: steps, _ = winograd_steps(problem, dtype, plan) why = winograd_flavor_reason(problem, dtype, plan, req.llvm_flavor) if why is not None: - raise ValueError(f"{_WINOGRAD_NAME}:{plan.label}: {why}") + raise ValueError(f"{name}:{plan.label}: {why}") return ConvWinogradSpec( - name=f"{_WINOGRAD_NAME}:{plan.label}", + name=f"{name}:{plan.label}", plan=plan, problem=problem, dtype=dtype, @@ -1266,12 +1272,81 @@ def _winograd_selection( return None, str(e) +# The same 3x3 plans on the shape-generic transform binaries +# (:mod:`kernels.common.conv_winograd_fwd_rt`): every kernel of the pipeline is +# compiled once, ahead of time, and named by its kernel name. +_WINOGRAD_AOT_NAME = "conv_winograd_fwd_aot" + + +def _as_aot(plan): + """The ``unfused2d_rt`` twin of a 3x3 ``unfused2d`` plan, else ``None`` + (the 1-D and generalized plans have no shape-generic kernels).""" + if plan.kind != "unfused2d": + return None + return dataclasses.replace(plan, kind="unfused2d_rt", merge="") + + +@functools.lru_cache(maxsize=256) +def _winograd_aot_sweep(req: ConvGroupedRequest) -> tuple[ConvWinogradSpec, ...]: + """Every safe 3x3 Winograd plan of ``req`` on the AOT binaries, in model + cost order.""" + if _winograd_problem_reason(req) is not None: + return () + dtype = req.dtype.lower() + out = [] + for plan in winograd_plans( + _conv_problem(req), dtype, model=_winograd_model(req.arch) + ): + aot = _as_aot(plan) + if aot is None: + continue + try: + out.append(_winograd_spec(req, aot, _WINOGRAD_AOT_NAME)) + except ValueError: + continue + return tuple(out) + + +@functools.lru_cache(maxsize=256) +def _winograd_aot_selection( + req: ConvGroupedRequest, +) -> tuple[ConvWinogradSpec | None, str]: + """The selector's plan for ``req`` on the AOT binaries (the cheapest safe + 3x3 plan when the selector picks a 1-D / generalized one, or when a + binary of its plan cannot run the problem), or ``(None, why)``.""" + why = _winograd_problem_reason(req) + if why is not None: + return None, why + sel = select_winograd( + _conv_problem(req), req.dtype.lower(), model=_winograd_model(req.arch) + ) + if sel.kind == "none": + return None, sel.reason + aot = _as_aot(sel) + if aot is not None: + try: + return _winograd_spec(req, aot, _WINOGRAD_AOT_NAME), "ok" + except ValueError as e: + why = str(e) + else: + why = f"the selected {sel.label} plan has no AOT kernels" + rest = _winograd_aot_sweep(req) + if rest: + return rest[0], "ok" + return None, why + + def _winograd_last_step(spec: ConvWinogradSpec): steps, _ = winograd_steps(spec.problem, spec.dtype, spec.plan) return steps[-1] -def _make_winograd_candidate() -> KernelCandidate: +def _make_winograd_candidate( + name: str = _WINOGRAD_NAME, + algorithm: str = "winograd", + selection=_winograd_selection, + sweep=_winograd_sweep, +) -> KernelCandidate: def support(req: OperatorRequest) -> tuple[bool, str]: errors = _request_errors(req) if errors: @@ -1280,27 +1355,27 @@ def support(req: OperatorRequest) -> tuple[bool, str]: ok, why = selector_matches(req, candidate) if not ok: return False, why - spec, why = _winograd_selection(req) + spec, why = selection(req) return spec is not None, why def select(req: OperatorRequest) -> ConvWinogradSpec: ok, why = candidate.admits(req) if not ok: - raise ValueError(f"{_WINOGRAD_NAME} does not support request: {why}") + raise ValueError(f"{name} does not support request: {why}") assert isinstance(req, ConvGroupedRequest) - spec, why = _winograd_selection(req) + spec, why = selection(req) if spec is None: - raise ValueError(f"{_WINOGRAD_NAME} does not support request: {why}") + raise ValueError(f"{name} does not support request: {why}") return spec def build(spec: ConvWinogradSpec, arch: str): return _winograd_last_step(spec).kernel.build(arch) candidate = KernelCandidate( - name=_WINOGRAD_NAME, + name=name, family=_FAMILIES["fwd"], - algorithm="winograd", - spec_id=_WINOGRAD_NAME, + algorithm=algorithm, + spec_id=name, abi_version=CONV_WINOGRAD_ABI_VERSION, # After every listed instance; opt-in, so never picked under "auto". priority=1_000_000, @@ -1314,9 +1389,7 @@ def build(spec: ConvWinogradSpec, arch: str): signature=lambda spec: _winograd_last_step(spec).signature, grid=lambda spec, req: _winograd_last_step(spec).grid, block=lambda spec: _winograd_last_step(spec).block, - sweep_space=lambda req: ( - _winograd_sweep(req) if candidate.admits(req)[0] else () - ), + sweep_space=lambda req: (sweep(req) if candidate.admits(req)[0] else ()), build=build, opt_in=True, ) @@ -1326,6 +1399,14 @@ def build(spec: ConvWinogradSpec, arch: str): _WINOGRAD_CANDIDATE = _make_winograd_candidate() _REGISTRIES["fwd"].register(_WINOGRAD_CANDIDATE) _EXECUTION_REGISTRIES["fwd"].register(_WINOGRAD_CANDIDATE) +_WINOGRAD_AOT_CANDIDATE = _make_winograd_candidate( + _WINOGRAD_AOT_NAME, + "winograd_aot", + _winograd_aot_selection, + _winograd_aot_sweep, +) +_REGISTRIES["fwd"].register(_WINOGRAD_AOT_CANDIDATE) +_EXECUTION_REGISTRIES["fwd"].register(_WINOGRAD_AOT_CANDIDATE) # --------------------------------------------------------------------------- diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py new file mode 100644 index 000000000000..afe965dc924b --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py @@ -0,0 +1,1537 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Shape-generic Winograd ``F(mh x mw, 3 x 3)`` forward transform kernels (AOT). + +The kernels of :mod:`kernels.common.conv_winograd_fwd` bake the problem +extents in, so every conv shape is its own set of binaries. The kernels here +compute the same K1 / K2a / K2 / K4 stages (same workspace layouts, same +arithmetic, same OOB-select addressing) but take every extent as a kernel +argument, so one binary serves every shape its knobs fit: they can be compiled +ahead of time like the direct and implicit-GEMM conv instances. K3 needs no +counterpart -- the WMMA batched GEMM +(:mod:`kernels.common.conv_winograd_gemm_wmma`) already takes its extents at +run time. + +What stays compile time (the instance's knobs): the tile, the I/O and operand +dtypes, the channel vector width (the channel extent of the stage must be a +multiple of it), the workgroup size, the K2 run length or LDS region, and +whether the prescale is on. Divisions by runtime extents use a multiply-high +by a host-computed magic number (:func:`magic_div`), exact for the 31-bit +indices the kernels use. + +Differences from the shape-specialised kernels: + +* the fused prescale (``RtInputSpec.fused``, register path) reduces the + tile maxima over a runtime number of channel lanes (<= ``block_size``, so + ``C <= block_size * vec``), through xor shuffles and LDS; the split form + (K2a per-pixel maxima, then K2) serves any C; +* K2's register path has no ``wg_tile_rows`` stacking; +* the LDS path uses a fixed number of channel lanes (``lds_lanes``) and loads + ``vec``-channel pieces; +* the prescale limit (which depends on ``C / groups``) is a kernel argument. + +Each stage's spec has a ``kernel_name()`` that is its AOT identity, and +:func:`stage_launch` turns (spec, problem) into the grid, block and kernarg +values of one launch. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, List, Optional, Sequence, Tuple + +from rocke.core.ir import F32, I32, IRBuilder, KernelDef, PtrType, Value +from rocke.helpers.spec import SignatureBuilder, ceil_div, kernel_name_join +from kernels.common._conv_implicit_gemm_common import ConvProblem +from kernels.common.conv_winograd_fwd import ( + _CPOL_NT, + _MAX_BUFFER_BYTES, + _OOB, + _PRESCALE_LOG2_LIMIT, + PRESCALE_BETA, + _abs_bits_vec, + _bits_to_f32, + _hmax_bits, + _IROps, + _ir_dtype, + _lincomb, + _load_raw, + _max_bits, + _pack_vec_f32, + _pow2_ceil, + _pow2_f32, + _raw_to_f32, + _store_packed, + _store_vec_f32, + at_1d, + bt_1d, + filter_matrix_int, +) + +_R = 3 +_FILTER_BLOCK = 128 +_PRESCALE_BLOCK = 256 +_IO = ("fp16", "bf16") +_TILES = (2, 4, 6) +_PRESCALE_MAX_EXP = 100 + + +# --------------------------------------------------------------------------- +# Runtime division +# --------------------------------------------------------------------------- + + +def magic_div(d: int) -> Tuple[int, int]: + """``(magic, shift)`` with ``n // d == umulhi(n, magic) >> shift`` for + ``0 <= n < 2**31``; ``d == 1`` is ``(0, 0)`` (the kernel selects ``n``). + + ``l = ceil(log2 d)``, ``magic = ceil(2**(31 + l) / d)`` (< 2**32 for + ``d >= 2``): the rounding error ``magic * d - 2**(31 + l) < d <= 2**l`` + times ``n < 2**31`` stays below ``2**(31 + l)``, so the floor is exact. + """ + if d < 1: + raise ValueError(f"divisor must be >= 1 (got {d})") + if d == 1: + return 0, 0 + lg = (d - 1).bit_length() + magic = -(-(1 << (31 + lg)) // d) + assert magic < (1 << 32) + return magic, lg - 1 + + +def _i32(v: int) -> int: + """An unsigned 32-bit kernarg as the signed value the launcher packs.""" + v &= 0xFFFFFFFF + return v - (1 << 32) if v >= (1 << 31) else v + + +class _Div: + """A runtime divisor: kernargs ````, ``_mg``, ``_sh``.""" + + def __init__(self, b: IRBuilder, name: str): + self.b = b + self.d = b.param(name, I32) + self.mg = b.param(f"{name}_mg", I32) + self.sh = b.param(f"{name}_sh", I32) + + def div(self, n: Value) -> Value: + b = self.b + q = b.lshr(b.umul_hi_i32(n, self.mg), self.sh) + return b.select(b.cmp_eq(self.mg, b.const_i32(0)), n, q) + + def divmod(self, n: Value) -> Tuple[Value, Value]: + q = self.div(n) + return q, self.b.sub(n, self.b.mul(q, self.d)) + + +def _div_values(name: str, d: int) -> dict: + mg, sh = magic_div(d) + return {name: d, f"{name}_mg": _i32(mg), f"{name}_sh": sh} + + +def _div_sig(sb: SignatureBuilder, name: str) -> SignatureBuilder: + return ( + sb.scalar(name, "i32").scalar(f"{name}_mg", "i32").scalar(f"{name}_sh", "i32") + ) + + +def _oob(b: IRBuilder, ok: Value, off: Value) -> Value: + return b.select(ok, off, b.const_i32(_OOB)) + + +def _c(b: IRBuilder, v: int) -> Value: + return b.const_i32(int(v)) + + +def _mul_c(b: IRBuilder, v: Value, k: int) -> Value: + """``v * k`` for a small constant ``k`` (0 and 1 folded).""" + if k == 0: + return b.const_i32(0) + if k == 1: + return v + return b.mul(v, b.const_i32(k)) + + +# --------------------------------------------------------------------------- +# Geometry (host side) +# --------------------------------------------------------------------------- + + +def _split_tile(m) -> Tuple[int, int]: + if isinstance(m, (tuple, list)): + return int(m[0]), int(m[1]) + return int(m), int(m) + + +@dataclass(frozen=True) +class RtGeometry: + """Host-side extents of one problem (one batch chunk) for tile ``(mh, mw)``.""" + + problem: ConvProblem + mh: int + mw: int + + @property + def nh(self) -> int: + return self.mh + _R - 1 + + @property + def nw(self) -> int: + return self.mw + _R - 1 + + @property + def n2(self) -> int: + return self.nh * self.nw + + @property + def tiles_h(self) -> int: + p = self.problem + return ceil_div(ceil_div(p.Ho, p.dH), self.mh) + + @property + def tiles_w(self) -> int: + p = self.problem + return ceil_div(ceil_div(p.Wo, p.dW), self.mw) + + @property + def num_tiles(self) -> int: + p = self.problem + return p.N * p.dH * p.dW * self.tiles_h * self.tiles_w + + @property + def prescale_scale(self) -> float: + """``beta / 2**L`` of :func:`conv_winograd_fwd._prescale_exp`.""" + beta = math.sqrt(PRESCALE_BETA[self.mh] * PRESCALE_BETA[self.mw]) + lim = _PRESCALE_LOG2_LIMIT - math.ceil(math.log2(self.problem.cpg) / 2) + return beta / float(2**lim) + + def buffer_bytes(self, prescale: bool) -> dict: + p = self.problem + t = self.num_tiles + return { + "x": p.N * p.Hi * p.Wi * p.C * 2, + "w": p.K * 9 * p.cpg * 2, + "y": p.N * p.Ho * p.Wo * p.K * 2, + "U": self.n2 * p.K * p.cpg * 2, + "V": t * self.n2 * p.C * 2, + "M": t * self.n2 * p.K * 2, + "e": (t * 4 + p.N * p.Hi * p.Wi * 4) if prescale else 4, + } + + def check(self, prescale: bool) -> Optional[str]: + p = self.problem + if p.is_3d or (p.Y, p.X) != (3, 3) or (p.sH, p.sW) != (1, 1): + return "Winograd F(m,3) needs a 2-D 3x3 stride-1 conv" + if p.Ho <= 0 or p.Wo <= 0: + return "empty output" + for nm, nbytes in self.buffer_bytes(prescale).items(): + if nbytes > _MAX_BUFFER_BYTES: + return f"{nm} needs {nbytes} bytes (32-bit buffer offsets)" + return None + + +def _ps_bits(geo: RtGeometry) -> int: + import struct + + return struct.unpack(" Value: + """``clamp(ceil(log2(amax * ps_scale)), 0, 100)``; ``ps_scale = beta / 2**L``.""" + bits = b.bitcast(b.fmul(amax, ps_scale), I32) + expo = b.sub(b.land(b.lshr(bits, _c(b, 23)), _c(b, 0xFF)), _c(b, 127)) + frac = b.land(bits, _c(b, 0x7FFFFF)) + e_val = b.add(expo, b.select(b.cmp_ne(frac, _c(b, 0)), _c(b, 1), _c(b, 0))) + return b.smin(b.smax(e_val, _c(b, 0)), _c(b, _PRESCALE_MAX_EXP)) + + +# --------------------------------------------------------------------------- +# Stage specs +# --------------------------------------------------------------------------- + + +def _check_common(mh, mw, io, op, vec): + if mh not in _TILES or mw not in _TILES: + raise ValueError(f"tile must be in {_TILES} (got {mh}x{mw})") + if io not in _IO or op not in _IO: + raise ValueError("dtypes must be fp16 / bf16") + if op == "bf16" and ((mh, mw) != (2, 2) or io != "bf16"): + raise ValueError("bf16 operands only for F(2,3) with bf16 I/O") + if vec not in (1, 2, 4, 8): + raise ValueError(f"vec must be 1/2/4/8 (got {vec})") + + +def _tile_tag(mh, mw) -> str: + return f"f{mh}" if mh == mw else f"f{mh}x{mw}" + + +@dataclass(frozen=True) +class RtFilterSpec: + """K1: ``w -> U``; ``vec`` divides ``C / groups``.""" + + mh: int = 4 + mw: int = 4 + io_dtype: str = "fp16" + operand_dtype: str = "fp16" + vec: int = 8 + filter_rows: int = 1 + + def __post_init__(self): + _check_common(self.mh, self.mw, self.io_dtype, self.operand_dtype, self.vec) + if (self.mh + 2) % self.filter_rows: + raise ValueError("filter_rows must divide n") + + stage = "K1" + + def kernel_name(self) -> str: + return kernel_name_join( + "wino_rt_filter", + _tile_tag(self.mh, self.mw), + self.io_dtype, + f"op{self.operand_dtype}", + f"v{self.vec}", + f"fr{self.filter_rows}", + ) + + def reason(self, geo: RtGeometry) -> Optional[str]: + if (geo.mh, geo.mw) != (self.mh, self.mw): + return "tile mismatch" + if geo.problem.cpg % self.vec: + return f"C/groups={geo.problem.cpg} not a multiple of vec {self.vec}" + return None + + +@dataclass(frozen=True) +class RtPrescaleSpec: + """K2a: per-pixel ``max |x|`` over C; ``vec`` divides C.""" + + io_dtype: str = "fp16" + vec: int = 8 + + def __post_init__(self): + if self.io_dtype not in _IO or self.vec not in (1, 2, 4, 8): + raise ValueError("bad K2a spec") + + stage = "K2a" + + def kernel_name(self) -> str: + return kernel_name_join("wino_rt_pmax", self.io_dtype, f"v{self.vec}") + + def reason(self, geo: RtGeometry) -> Optional[str]: + if geo.problem.C % self.vec: + return f"C={geo.problem.C} not a multiple of vec {self.vec}" + return None + + +@dataclass(frozen=True) +class RtInputSpec: + """K2: ``x -> V`` (and ``e``); ``vec`` divides C. + + ``lds = (bh, bw)``: stage ``bh x bw`` tiles x ``lds_lanes * vec`` channels + through LDS, ``tiles_per_thread`` of them per thread; ``(0, 0)``: the + register path, a run of ``tiles_per_thread`` tiles along W per thread. + """ + + mh: int = 4 + mw: int = 4 + io_dtype: str = "fp16" + operand_dtype: str = "fp16" + vec: int = 8 + block_size: int = 256 + prescale: bool = False + tiles_per_thread: int = 1 + lds: Tuple[int, int] = (0, 0) + lds_lanes: int = 16 + nt_stores: bool = True + # Fused prescale: the workgroup holds all channels of its runs and reduces + # the tile maxima itself (no K2a); register path only. + fused: bool = False + + def __post_init__(self): + _check_common(self.mh, self.mw, self.io_dtype, self.operand_dtype, self.vec) + if self.block_size < 32 or self.block_size & (self.block_size - 1): + raise ValueError("block_size must be a power of two >= 32") + if self.fused and (not self.prescale or self.input_lds): + raise ValueError("fused needs the prescale and the register path") + if self.tiles_per_thread < 1: + raise ValueError("tiles_per_thread must be >= 1") + if self.input_lds: + bh, bw = self.lds + if (bh * bw) % self.tiles_per_thread: + raise ValueError("tiles_per_thread must divide bh * bw") + if self.lds_block > 1024: + raise ValueError("LDS region too large for one workgroup") + + stage = "K2" + + @property + def input_lds(self) -> bool: + return self.lds[0] > 0 + + @property + def nh(self) -> int: + return self.mh + 2 + + @property + def nw(self) -> int: + return self.mw + 2 + + @property + def lds_block(self) -> int: + bh, bw = self.lds + return bh * bw * self.lds_lanes // self.tiles_per_thread + + def kernel_name(self) -> str: + tags = [ + _tile_tag(self.mh, self.mw), + self.io_dtype, + f"op{self.operand_dtype}", + f"v{self.vec}", + ] + if self.input_lds: + tags += [ + "lds{}x{}".format(*self.lds), + f"ll{self.lds_lanes}", + f"t{self.tiles_per_thread}", + ] + else: + tags += [f"b{self.block_size}", f"r{self.tiles_per_thread}"] + if self.prescale: + tags.append("psf" if self.fused else "ps") + if not self.nt_stores: + tags.append("rt") + return kernel_name_join("wino_rt_input", *tags) + + def reason(self, geo: RtGeometry) -> Optional[str]: + if (geo.mh, geo.mw) != (self.mh, self.mw): + return "tile mismatch" + if geo.problem.C % self.vec: + return f"C={geo.problem.C} not a multiple of vec {self.vec}" + if self.fused and _pow2_ceil(geo.problem.C // self.vec) > self.block_size: + return f"fused prescale needs C/vec <= {self.block_size}" + return None + + +@dataclass(frozen=True) +class RtOutputSpec: + """K4: ``M -> y`` (``x 2**e``); ``vec`` divides K.""" + + mh: int = 4 + mw: int = 4 + io_dtype: str = "fp16" + operand_dtype: str = "fp16" + vec: int = 8 + block_size: int = 256 + prescale: bool = False + nt_stores: bool = True + + def __post_init__(self): + _check_common(self.mh, self.mw, self.io_dtype, self.operand_dtype, self.vec) + if self.block_size < 32 or self.block_size & (self.block_size - 1): + raise ValueError("block_size must be a power of two >= 32") + + stage = "K4" + + def kernel_name(self) -> str: + tags = [ + _tile_tag(self.mh, self.mw), + self.io_dtype, + f"op{self.operand_dtype}", + f"v{self.vec}", + f"b{self.block_size}", + ] + if self.prescale: + tags.append("ps") + if not self.nt_stores: + tags.append("rt") + return kernel_name_join("wino_rt_output", *tags) + + def reason(self, geo: RtGeometry) -> Optional[str]: + if (geo.mh, geo.mw) != (self.mh, self.mw): + return "tile mismatch" + if geo.problem.K % self.vec: + return f"K={geo.problem.K} not a multiple of vec {self.vec}" + return None + + +# --------------------------------------------------------------------------- +# K1: filter transform +# --------------------------------------------------------------------------- + + +def filter_signature(spec: RtFilterSpec) -> list: + sb = ( + SignatureBuilder() + .ptr("w", spec.io_dtype) + .ptr("u", spec.operand_dtype) + .scalar("w_bytes", "i32") + .scalar("u_bytes", "i32") + .scalar("K", "i32") + .scalar("cg", "i32") + .scalar("items", "i32") + ) + return _div_sig(sb, "cgv").build() + + +def build_filter( + spec: RtFilterSpec, *, arch: str = "gfx1250" +) -> KernelDef: # noqa: ARG001 + """K1 of :func:`conv_winograd_fwd.build_winograd_filter_transform`, extents at run time.""" + vec = spec.vec + n, nw = spec.mh + 2, spec.mw + 2 + gint, den = filter_matrix_int(spec.mh) + gint_w, den_w = filter_matrix_int(spec.mw) + scale = 1.0 / float(den * den_w) + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = _FILTER_BLOCK + w = b.param( + "w", PtrType(_ir_dtype(spec.io_dtype), "global"), noalias=True, readonly=True + ) + u = b.param( + "u", + PtrType(_ir_dtype(spec.operand_dtype), "global"), + noalias=True, + writeonly=True, + ) + w_bytes = b.param("w_bytes", I32) + u_bytes = b.param("u_bytes", I32) + K = b.param("K", I32) + cg = b.param("cg", I32) + items = b.param("items", I32) + cgv = _Div(b, "cgv") + w_rsrc = b.buffer_rsrc(w, w_bytes) + u_rsrc = b.buffer_rsrc(u, u_bytes) + + gid = b.add(b.mul(b.block_id_x(), _c(b, _FILTER_BLOCK)), b.thread_id_x()) + row_y = b.block_id_y() + ok = b.cmp_lt(gid, items) + k, cm = cgv.divmod(gid) + c = _mul_c(b, cm, vec) + cg2 = b.mul(cg, _c(b, 2)) + wbase = _oob(b, ok, b.mul(b.add(b.mul(k, b.mul(cg, _c(b, 9))), c), _c(b, 2))) + taps = [] + off = wbase + for ky in range(3): + row = [] + for kx in range(3): + if ky or kx: + off = b.add(off, cg2) + row.append( + _raw_to_f32(b, _load_raw(b, w_rsrc, off, spec.io_dtype, vec), vec) + ) + taps.append(row) + sc = b.const_f32(scale) + nr = spec.filter_rows + kcg2 = b.mul(K, cg2) # bytes of one U[bin] + for rr in range(nr): + if nr == n: + row_i = _c(b, rr) + coef = [b.const_f32(float(gint[rr][ky])) for ky in range(3)] + else: + row_i = b.add(_mul_c(b, row_y, nr), _c(b, rr)) + coef = [] + for ky in range(3): + cval = b.const_f32(float(gint[0][ky])) + for i in range(1, n): + cval = b.select( + b.cmp_eq(row_i, _c(b, i)), b.const_f32(float(gint[i][ky])), cval + ) + coef.append(cval) + tmp = [ + [ + b.fma( + coef[2], + taps[2][kx][ch], + b.fma(coef[1], taps[1][kx][ch], b.fmul(coef[0], taps[0][kx][ch])), + ) + for ch in range(vec) + ] + for kx in range(3) + ] + # U[(row_i * nw + j)][k][c] + ubase = _oob( + b, + ok, + b.add( + b.mul(_mul_c(b, row_i, nw), kcg2), + b.mul(b.add(b.mul(k, cg), c), _c(b, 2)), + ), + ) + uoff = ubase + for j in range(nw): + if j: + uoff = b.add(uoff, kcg2) + vals = [ + b.fmul( + _lincomb(b, [(gint_w[j][kx], tmp[kx][ch]) for kx in range(3)]), sc + ) + for ch in range(vec) + ] + _store_vec_f32(b, u_rsrc, uoff, vals, spec.operand_dtype) + return b.kernel + + +def filter_launch(spec: RtFilterSpec, geo: RtGeometry): + p = geo.problem + cgv = p.cpg // spec.vec + items = p.K * cgv + vals = dict(K=p.K, cg=p.cpg, items=items, **_div_values("cgv", cgv)) + grid = (ceil_div(items, _FILTER_BLOCK), (spec.mh + 2) // spec.filter_rows, 1) + return grid, _FILTER_BLOCK, vals + + +# --------------------------------------------------------------------------- +# K2a: per-pixel max |x| over C +# --------------------------------------------------------------------------- + + +def prescale_signature(spec: RtPrescaleSpec) -> list: + return ( + SignatureBuilder() + .ptr("x", spec.io_dtype) + .ptr("e", "i32") + .scalar("x_bytes", "i32") + .scalar("e_bytes", "i32") + .scalar("C", "i32") + .scalar("nvec", "i32") + .scalar("tpc_lg", "i32") + .scalar("loads", "i32") + .scalar("npix", "i32") + .scalar("T", "i32") + .build() + ) + + +def build_prescale( + spec: RtPrescaleSpec, *, arch: str = "gfx1250" +) -> KernelDef: # noqa: ARG001 + """K2a of :func:`conv_winograd_fwd.build_winograd_prescale`, extents at run time. + + ``tpc = 2**tpc_lg <= 32`` lanes per pixel, each reduces ``loads`` + channel vectors (a runtime loop), then xor shuffles inside the lane group. + """ + vec = spec.vec + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = _PRESCALE_BLOCK + x = b.param( + "x", PtrType(_ir_dtype(spec.io_dtype), "global"), noalias=True, readonly=True + ) + e_ptr = b.param("e", PtrType(I32, "global"), noalias=True, writeonly=True) + x_bytes = b.param("x_bytes", I32) + e_bytes = b.param("e_bytes", I32) + C = b.param("C", I32) + nvec = b.param("nvec", I32) + tpc_lg = b.param("tpc_lg", I32) + loads = b.param("loads", I32) + npix = b.param("npix", I32) + T = b.param("T", I32) + x_rsrc = b.buffer_rsrc(x, x_bytes) + e_rsrc = b.buffer_rsrc(e_ptr, e_bytes) + + tid = b.thread_id_x() + tpc = b.shl(_c(b, 1), tpc_lg) + lc = b.land(tid, b.sub(tpc, _c(b, 1))) + lp = b.lshr(tid, tpc_lg) + ppw = b.lshr(_c(b, _PRESCALE_BLOCK), tpc_lg) + pix = b.add(b.mul(b.block_id_x(), ppw), lp) + pok = b.cmp_lt(pix, npix) + base = b.mul(b.add(b.mul(pix, C), _mul_c(b, lc, vec)), _c(b, 2)) + step_b = b.mul(tpc, _c(b, vec * 2)) + zero = _c(b, 0) + loop = b.scf_for_iter(zero, loads, _c(b, 1), [("pm_acc", zero)], iv_name="pm_q") + with loop as (q, (acc,)): + idx = b.add(lc, b.mul(q, tpc)) + ok = b.land(pok, b.cmp_lt(idx, nvec)) + off = _oob(b, ok, b.add(base, b.mul(q, step_b))) + a_ = _hmax_bits( + b, _abs_bits_vec(b, _load_raw(b, x_rsrc, off, spec.io_dtype, vec), vec), vec + ) + b.scf_yield(b.smax(acc, a_)) + mx = loop.results[0] + s = 1 + while s < 32: + sh = b.warp_shuffle_xor(mx, s) + mx = b.select(b.cmp_lt(_c(b, s), tpc), b.smax(mx, sh), mx) + s *= 2 + ok_s = b.land(pok, b.cmp_eq(lc, _c(b, 0))) + b.buffer_store_f32( + e_rsrc, + _oob(b, ok_s, b.mul(b.add(pix, T), _c(b, 4))), + _c(b, 0), + b.bitcast(mx, F32), + ) + return b.kernel + + +def prescale_launch(spec: RtPrescaleSpec, geo: RtGeometry): + p = geo.problem + nv = p.C // spec.vec + tpc = min(_pow2_ceil(nv), 32) + npix = p.N * p.Hi * p.Wi + vals = dict( + C=p.C, + nvec=nv, + tpc_lg=tpc.bit_length() - 1, + loads=ceil_div(nv, tpc), + npix=npix, + T=geo.num_tiles, + ) + return (ceil_div(npix, _PRESCALE_BLOCK // tpc), 1, 1), _PRESCALE_BLOCK, vals + + +# --------------------------------------------------------------------------- +# K2 / K4 thread mapping (register path) +# --------------------------------------------------------------------------- + + +@dataclass +class _Coords: + ok: Value + cv: Value + c_ok: Value + n: Value + a: Value + bsub: Value + ti: Value + tj0: Value + t0: Value + tiles_w: Value + tpc: Value + + +_RUN_SCALARS = ("nvec", "tpc_lg", "rpw", "num_runs", "tiles_w") +_RUN_DIVS = ("cx", "rpr", "th", "ddw", "ddh") + + +def _run_sig(sb: SignatureBuilder) -> SignatureBuilder: + for nm in _RUN_SCALARS: + sb = sb.scalar(nm, "i32") + for nm in _RUN_DIVS: + sb = _div_sig(sb, nm) + return sb + + +def _run_coords(b: IRBuilder, tiles: int) -> _Coords: + """Thread -> (run of ``tiles`` tiles along W, channel vector), all runtime.""" + nvec = b.param("nvec", I32) + tpc_lg = b.param("tpc_lg", I32) + rpw = b.param("rpw", I32) + num_runs = b.param("num_runs", I32) + tiles_w = b.param("tiles_w", I32) + cx = _Div(b, "cx") + rpr = _Div(b, "rpr") + th = _Div(b, "th") + ddw = _Div(b, "ddw") + ddh = _Div(b, "ddh") + tid = b.thread_id_x() + bid = b.block_id_x() + tpc = b.shl(_c(b, 1), tpc_lg) + lc = b.land(tid, b.sub(tpc, _c(b, 1))) + lr = b.lshr(tid, tpc_lg) + rb, cxi = cx.divmod(bid) + cv = b.add(b.mul(cxi, tpc), lc) + c_ok = b.cmp_lt(cv, nvec) + rho = b.add(b.mul(rb, rpw), lr) + ok = b.cmp_lt(rho, num_runs) + row, jr = rpr.divmod(rho) + tj0 = _mul_c(b, jr, tiles) + t0 = b.add(b.mul(row, tiles_w), tj0) + r, ti = th.divmod(row) + r, bsub = ddw.divmod(r) + r, a = ddh.divmod(r) + return _Coords(ok, cv, c_ok, r, a, bsub, ti, tj0, t0, tiles_w, tpc) + + +def _run_values( + geo: RtGeometry, nvec: int, tiles: int, block: int +) -> Tuple[dict, tuple]: + p = geo.problem + tpc = min(_pow2_ceil(nvec), block) + cx = ceil_div(nvec, tpc) + rpw = block // tpc + rpr = ceil_div(geo.tiles_w, tiles) + rows = geo.num_tiles // geo.tiles_w + num_runs = rows * rpr + vals = dict( + nvec=nvec, + tpc_lg=tpc.bit_length() - 1, + rpw=rpw, + num_runs=num_runs, + tiles_w=geo.tiles_w, + **_div_values("cx", cx), + **_div_values("rpr", rpr), + **_div_values("th", geo.tiles_h), + **_div_values("ddw", p.dW), + **_div_values("ddh", p.dH), + ) + return vals, (ceil_div(num_runs, rpw) * cx, 1, 1) + + +def _group_max_rt( + b: IRBuilder, vals: List[Value], tpc: Value, block: int +) -> List[Value]: + """Max of each i32 in ``vals`` over the ``tpc`` (runtime power of two, + ``<= block``) threads of a run: xor shuffles inside a 32-lane row, then + across the row maxima through LDS (``conv_winograd_fwd._group_max_bits`` + with a runtime lane count).""" + out = [] + for v in vals: + s = 1 + while s < 32: + sh = b.warp_shuffle_xor(v, s) + v = b.select(b.cmp_lt(_c(b, s), tpc), b.smax(v, sh), v) + s *= 2 + out.append(v) + rows = block // 32 + if rows <= 1: + return out + nv = len(out) + red = b.smem_alloc_f32([rows * nv], name_hint="wino_rt_amax") + tid = b.thread_id_x() + row_id = b.lshr(tid, _c(b, 5)) + with b.scf_if(b.cmp_eq(b.land(tid, _c(b, 31)), _c(b, 0))): + for i, v in enumerate(out): + idx = b.add(_mul_c(b, row_id, nv), _c(b, i)) + b.smem_store_vN_f32(red, [idx], b.vec_pack([b.bitcast(v, F32)], F32), 1) + b.sync() + # The rows of one run: ``per = tpc / 32`` rows from ``first`` (tpc > 32); + # with tpc <= 32 a run lives in one row and the shuffles were enough. + per = b.lshr(tpc, _c(b, 5)) + big = b.cmp_lt(_c(b, 32), tpc) + first = b.mul(b.select(big, b.div(row_id, b.select(big, per, _c(b, 1))), row_id), + b.select(big, per, _c(b, 1))) # fmt: skip + res = [] + for i in range(nv): + acc = out[i] + for r in range(rows): + idx = b.add(_mul_c(b, b.add(first, _c(b, r)), nv), _c(b, i)) + ok = b.land(big, b.cmp_lt(_c(b, r), per)) + idx = b.select(ok, idx, _c(b, 0)) + x = b.bitcast(b.vec_extract(b.smem_load_vN_f32(red, idx, n=1), 0), I32) + acc = b.select(ok, b.smax(acc, x), acc) + res.append(acc) + return res + + +# --------------------------------------------------------------------------- +# K2: input transform +# --------------------------------------------------------------------------- + +_INPUT_GEOM = ("C", "Hi", "Wi", "pH", "pW", "dH", "dW", "dHm", "dWm", "T", "ps_bits") + + +def input_signature(spec: RtInputSpec) -> list: + sb = ( + SignatureBuilder() + .ptr("x", spec.io_dtype) + .ptr("v", spec.operand_dtype) + .ptr("e", "i32") + .scalar("x_bytes", "i32") + .scalar("v_bytes", "i32") + .scalar("e_bytes", "i32") + ) + for nm in _INPUT_GEOM: + sb = sb.scalar(nm, "i32") + if spec.input_lds: + sb = sb.scalar("nvec", "i32").scalar("tiles_h", "i32").scalar("tiles_w", "i32") + for nm in ("cx", "nbw", "nbh", "ddw", "ddh"): + sb = _div_sig(sb, nm) + else: + sb = _run_sig(sb) + return sb.build() + + +def _input_tiles_transform(b, spec, raw, R, tile_ok, scales, vbase0, v_rsrc, cb2): + """``V = B^T d B`` of ``R`` tiles adjacent along W; ``cb2 = C * 2`` (runtime).""" + n, m = spec.nh, spec.mh + nw, mw = spec.nw, spec.mw + vec = spec.vec + ops = _IROps(b) + vstride = _mul_c(b, cb2, n * nw) + colt = {} + pending = [] + b.sched_barrier(0) + + def col_transform(vv): + if vv not in colt: + data = [_raw_to_f32(b, raw[uu][vv], vec) for uu in range(n)] + colt[vv] = [ + bt_1d(ops, m, [data[uu][ch] for uu in range(n)]) for ch in range(vec) + ] + return colt[vv] + + for r in range(R): + cols = [col_transform(mw * r + vv) for vv in range(nw)] + base = vbase0 if r == 0 else b.add(vbase0, _mul_c(b, vstride, r)) + off = _oob(b, tile_ok[r], base) + first = True + for i in range(n): + outs = [ + bt_1d(ops, mw, [cols[vv][ch][i] for vv in range(nw)]) + for ch in range(vec) + ] + for j in range(nw): + if not first: + off = b.add(off, cb2) + first = False + vals = [outs[ch][j] for ch in range(vec)] + if scales[r] is not None: + vals = [b.fmul(x_, scales[r]) for x_ in vals] + pending.append((off, _pack_vec_f32(b, vals, spec.operand_dtype))) + b.sched_barrier(0) + cpol = _CPOL_NT if spec.nt_stores else 0 + for off, packed in pending: + _store_packed(b, v_rsrc, off, packed, spec.operand_dtype, cpol) + + +def build_input( + spec: RtInputSpec, *, arch: str = "gfx1250" +) -> KernelDef: # noqa: ARG001 + """K2 of :func:`conv_winograd_fwd.build_winograd_input_transform`, extents at run time.""" + if spec.input_lds: + return _build_input_lds(spec) + n, m = spec.nh, spec.mh + nw, mw = spec.nw, spec.mw + vec = spec.vec + R = spec.tiles_per_thread + ncols = mw * R + (nw - mw) + ps = spec.prescale + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = spec.block_size + x = b.param( + "x", PtrType(_ir_dtype(spec.io_dtype), "global"), noalias=True, readonly=True + ) + v = b.param( + "v", + PtrType(_ir_dtype(spec.operand_dtype), "global"), + noalias=True, + writeonly=True, + ) + e_ptr = b.param("e", PtrType(I32, "global"), noalias=True) + x_bytes = b.param("x_bytes", I32) + v_bytes = b.param("v_bytes", I32) + e_bytes = b.param("e_bytes", I32) + g = {nm: b.param(nm, I32) for nm in _INPUT_GEOM} + x_rsrc = b.buffer_rsrc(x, x_bytes) + v_rsrc = b.buffer_rsrc(v, v_bytes) + e_rsrc = b.buffer_rsrc(e_ptr, e_bytes) if ps else None + rc = _run_coords(b, R) + + C, Hi, Wi = g["C"], g["Hi"], g["Wi"] + cb2 = b.mul(C, _c(b, 2)) + c = _mul_c(b, rc.cv, vec) + h0 = b.add(b.sub(rc.a, g["pH"]), b.mul(rc.ti, g["dHm"])) + w0 = b.add(b.sub(rc.bsub, g["pW"]), b.mul(rc.tj0, g["dWm"])) + img_row0 = b.mul(rc.n, Hi) + row_ok, row_pix = [], [] + hi = h0 + for uu in range(n): + if uu: + hi = b.add(hi, g["dH"]) + row_ok.append(b.land(rc.ok, b.cmp_ult(hi, Hi))) + row_pix.append(b.mul(b.add(img_row0, hi), Wi)) + col_ok, col_w = [], [] + wi = w0 + for vv in range(ncols): + if vv: + wi = b.add(wi, g["dW"]) + col_ok.append(b.cmp_ult(wi, Wi)) + col_w.append(wi) + c2 = _mul_c(b, c, 2) + xrow = [ + _oob(b, b.land(row_ok[uu], rc.c_ok), b.add(b.mul(row_pix[uu], cb2), c2)) + for uu in range(n) + ] + xcol = [_oob(b, col_ok[vv], b.mul(col_w[vv], cb2)) for vv in range(ncols)] + raw = [ + [ + _load_raw(b, x_rsrc, b.add(xrow[uu], xcol[vv]), spec.io_dtype, vec) + for vv in range(ncols) + ] + for uu in range(n) + ] + base_ok = b.land(rc.ok, rc.c_ok) + tile_ok = [] + for r in range(R): + if r == 0: + tile_ok.append(base_ok) + else: + tile_ok.append( + b.land(base_ok, b.cmp_lt(b.add(rc.tj0, _c(b, r)), rc.tiles_w)) + ) + + scales: List[Optional[Value]] = [None] * R + if ps and spec.fused: + colmax = [] + for vv in range(ncols): + acc = None + for uu in range(n): + a_ = _abs_bits_vec(b, raw[uu][vv], vec) + acc = a_ if acc is None else _max_bits(b, acc, a_, vec) + colmax.append(acc) + tmax = [] + for r in range(R): + acc = colmax[mw * r] + for vv in range(1, nw): + acc = _max_bits(b, acc, colmax[mw * r + vv], vec) + tmax.append(_hmax_bits(b, acc, vec)) + tmax = _group_max_rt(b, tmax, rc.tpc, spec.block_size) + ps_scale = b.bitcast(g["ps_bits"], F32) + lane0 = b.cmp_eq(rc.cv, _c(b, 0)) + for r in range(R): + e_r = _prescale_exp_rt(b, _bits_to_f32(b, tmax[r], spec.io_dtype), ps_scale) + t_r = b.add(rc.t0, _c(b, r)) if r else rc.t0 + b.buffer_store_f32( + e_rsrc, + _oob(b, b.land(lane0, tile_ok[r]), b.mul(t_r, _c(b, 4))), + _c(b, 0), + b.bitcast(e_r, F32), + ) + scales[r] = _pow2_f32(b, b.sub(_c(b, 0), e_r)) + elif ps: + T = g["T"] + prow = [ + _oob(b, row_ok[uu], b.mul(b.add(row_pix[uu], T), _c(b, 4))) + for uu in range(n) + ] + pcol = [_oob(b, col_ok[vv], b.mul(col_w[vv], _c(b, 4))) for vv in range(ncols)] + colmax = [] + for vv in range(ncols): + acc = None + for uu in range(n): + pm = b.buffer_load(e_rsrc, b.add(prow[uu], pcol[vv]), _c(b, 0), I32) + acc = pm if acc is None else b.smax(acc, pm) + colmax.append(acc) + ps_scale = b.bitcast(g["ps_bits"], F32) + lane0 = b.cmp_eq(rc.cv, _c(b, 0)) + for r in range(R): + acc = colmax[mw * r] + for vv in range(1, nw): + acc = b.smax(acc, colmax[mw * r + vv]) + e_r = _prescale_exp_rt(b, _bits_to_f32(b, acc, spec.io_dtype), ps_scale) + t_r = b.add(rc.t0, _c(b, r)) if r else rc.t0 + b.buffer_store_f32( + e_rsrc, + _oob(b, b.land(lane0, tile_ok[r]), b.mul(t_r, _c(b, 4))), + _c(b, 0), + b.bitcast(e_r, F32), + ) + scales[r] = _pow2_f32(b, b.sub(_c(b, 0), e_r)) + + vbase0 = b.add(b.mul(rc.t0, _mul_c(b, cb2, spec.nh * spec.nw)), c2) + _input_tiles_transform(b, spec, raw, R, tile_ok, scales, vbase0, v_rsrc, cb2) + return b.kernel + + +def input_launch(spec: RtInputSpec, geo: RtGeometry): + p = geo.problem + vals = dict( + C=p.C, + Hi=p.Hi, + Wi=p.Wi, + pH=p.pH, + pW=p.pW, + dH=p.dH, + dW=p.dW, + dHm=p.dH * spec.mh, + dWm=p.dW * spec.mw, + T=geo.num_tiles, + ps_bits=_ps_bits(geo), + ) + nv = p.C // spec.vec + if spec.input_lds: + bh, bw = spec.lds + tpc = spec.lds_lanes + cx = ceil_div(nv, tpc) + nbh = ceil_div(geo.tiles_h, bh) + nbw = ceil_div(geo.tiles_w, bw) + vals.update( + nvec=nv, + tiles_h=geo.tiles_h, + tiles_w=geo.tiles_w, + **_div_values("cx", cx), + **_div_values("nbw", nbw), + **_div_values("nbh", nbh), + **_div_values("ddw", p.dW), + **_div_values("ddh", p.dH), + ) + grid = (p.N * p.dH * p.dW * nbh * nbw * cx, 1, 1) + return grid, spec.lds_block, vals + rv, grid = _run_values(geo, nv, spec.tiles_per_thread, spec.block_size) + vals.update(rv) + return grid, spec.block_size, vals + + +def _build_input_lds(spec: RtInputSpec) -> KernelDef: + """K2 through LDS (``conv_winograd_fwd._build_input_lds``), extents at run time. + + A workgroup stages ``bh x bw`` tiles x ``lds_lanes * vec`` channels; every + channel lane past C is masked (a small-C problem leaves lanes idle). + """ + n, m = spec.nh, spec.mh + nw, mw = spec.nw, spec.mw + vec = spec.vec + bh, bw = spec.lds + tpc = spec.lds_lanes + cc = tpc * vec + lv = vec + rh = m * bh + n - m + rw = mw * bw + nw - mw + block = spec.lds_block + tpt = spec.tiles_per_thread + ppx = cc // lv + npx = rh * rw + loads = ceil_div(npx * ppx, block) + ps = spec.prescale + dt = _ir_dtype(spec.io_dtype) + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = block + x = b.param("x", PtrType(dt, "global"), noalias=True, readonly=True) + v = b.param( + "v", + PtrType(_ir_dtype(spec.operand_dtype), "global"), + noalias=True, + writeonly=True, + ) + e_ptr = b.param("e", PtrType(I32, "global"), noalias=True) + x_bytes = b.param("x_bytes", I32) + v_bytes = b.param("v_bytes", I32) + e_bytes = b.param("e_bytes", I32) + g = {nm: b.param(nm, I32) for nm in _INPUT_GEOM} + nvec = b.param("nvec", I32) + tiles_h = b.param("tiles_h", I32) + tiles_w = b.param("tiles_w", I32) + cx = _Div(b, "cx") + nbw = _Div(b, "nbw") + nbh = _Div(b, "nbh") + ddw = _Div(b, "ddw") + ddh = _Div(b, "ddh") + x_rsrc = b.buffer_rsrc(x, x_bytes) + v_rsrc = b.buffer_rsrc(v, v_bytes) + e_rsrc = b.buffer_rsrc(e_ptr, e_bytes) if ps else None + C, Hi, Wi = g["C"], g["Hi"], g["Wi"] + cb2 = b.mul(C, _c(b, 2)) + + tid = b.thread_id_x() + bid = b.block_id_x() + r, cxi = cx.divmod(bid) + r, bwi = nbw.divmod(r) + sub, bhi = nbh.divmod(r) + r, bsub = ddw.divmod(sub) + img, a = ddh.divmod(r) + ti0 = _mul_c(b, bhi, bh) + tj0 = _mul_c(b, bwi, bw) + h0 = b.add(b.sub(a, g["pH"]), b.mul(ti0, g["dHm"])) + w0 = b.add(b.sub(bsub, g["pW"]), b.mul(tj0, g["dWm"])) + img_row0 = b.mul(img, Hi) + c0 = _mul_c(b, cxi, cc) + + lds = b.smem_alloc(dt, [loads * block * lv], name_hint="wino_xin") + for k in range(loads): + q = b.add(tid, _c(b, k * block)) if k else tid + px = b.lshr(q, _c(b, ppx.bit_length() - 1)) if ppx > 1 else q + pc = b.land(q, _c(b, ppx - 1)) if ppx > 1 else _c(b, 0) + rr = b.div(px, _c(b, rw)) + cc_ = b.sub(px, b.mul(rr, _c(b, rw))) + hi = b.add(h0, b.mul(rr, g["dH"])) + wi = b.add(w0, b.mul(cc_, g["dW"])) + ch = b.add(c0, _mul_c(b, pc, lv)) + ok = b.land(b.cmp_ult(hi, Hi), b.cmp_ult(wi, Wi)) + if (k + 1) * block > npx * ppx: + ok = b.land(ok, b.cmp_lt(px, _c(b, npx))) + ok = b.land(ok, b.cmp_lt(ch, C)) + pix = b.add(b.mul(b.add(img_row0, hi), Wi), wi) + off = _oob(b, ok, b.add(b.mul(pix, cb2), _mul_c(b, ch, 2))) + val = _load_raw(b, x_rsrc, off, spec.io_dtype, lv) + b.smem_store_vN(lds, [_mul_c(b, q, lv)], val, lv) + pm_lds = None + if ps: + T = g["T"] + pm_lds = b.smem_alloc( + I32, [ceil_div(npx, block) * block], name_hint="wino_pmax" + ) + for k in range(ceil_div(npx, block)): + px = b.add(tid, _c(b, k * block)) if k else tid + rr = b.div(px, _c(b, rw)) + cc_ = b.sub(px, b.mul(rr, _c(b, rw))) + hi = b.add(h0, b.mul(rr, g["dH"])) + wi = b.add(w0, b.mul(cc_, g["dW"])) + ok = b.land(b.cmp_ult(hi, Hi), b.cmp_ult(wi, Wi)) + if (k + 1) * block > npx: + ok = b.land(ok, b.cmp_lt(px, _c(b, npx))) + pix = b.add(b.mul(b.add(img_row0, hi), Wi), wi) + off = _oob(b, ok, b.mul(b.add(pix, T), _c(b, 4))) + b.smem_store_vN(pm_lds, [px], b.buffer_load(e_rsrc, off, _c(b, 0), I32), 1) + b.sync() + + cl = b.land(tid, _c(b, tpc - 1)) if tpc > 1 else _c(b, 0) + tb0 = b.lshr(tid, _c(b, tpc.bit_length() - 1)) if tpc > 1 else tid + cv = b.add(_mul_c(b, cxi, tpc), cl) + c_ok = b.cmp_lt(cv, nvec) + ps_scale = b.bitcast(g["ps_bits"], F32) if ps else None + for rep in range(tpt): + tb = b.add(tb0, _c(b, rep * (block // tpc))) if rep else tb0 + bi = b.div(tb, _c(b, bw)) if bw > 1 else tb + bj = b.sub(tb, _mul_c(b, bi, bw)) if bw > 1 else _c(b, 0) + ti = b.add(ti0, bi) + tj = b.add(tj0, bj) + tile_ok = b.land( + b.land(b.cmp_lt(ti, tiles_h), b.cmp_lt(tj, tiles_w)), + c_ok, + ) + t = b.add(b.mul(b.add(b.mul(sub, tiles_h), ti), tiles_w), tj) + pbase = b.add(_mul_c(b, bi, m * rw), _mul_c(b, bj, mw)) + raw = [] + for uu in range(n): + row = [] + for vv in range(nw): + idx = b.add( + _mul_c(b, b.add(pbase, _c(b, uu * rw + vv)), cc), _mul_c(b, cl, vec) + ) + ld = b.smem_load_vN(lds, idx, dtype=dt, n=vec) + row.append(b.vec_extract(ld, 0) if vec == 1 else ld) + raw.append(row) + scales: List[Optional[Value]] = [None] + if ps: + acc = None + for uu in range(n): + for vv in range(nw): + pm = b.vec_extract( + b.smem_load_vN( + pm_lds, b.add(pbase, _c(b, uu * rw + vv)), dtype=I32, n=1 + ), + 0, + ) + acc = pm if acc is None else b.smax(acc, pm) + e_t = _prescale_exp_rt(b, _bits_to_f32(b, acc, spec.io_dtype), ps_scale) + ok_e = b.land(tile_ok, b.cmp_eq(cv, _c(b, 0))) + b.buffer_store_f32( + e_rsrc, _oob(b, ok_e, b.mul(t, _c(b, 4))), _c(b, 0), b.bitcast(e_t, F32) + ) + scales = [_pow2_f32(b, b.sub(_c(b, 0), e_t))] + vbase0 = b.add( + b.mul(t, _mul_c(b, cb2, spec.nh * spec.nw)), _mul_c(b, cv, vec * 2) + ) + _input_tiles_transform(b, spec, raw, 1, [tile_ok], scales, vbase0, v_rsrc, cb2) + return b.kernel + + +# --------------------------------------------------------------------------- +# K4: output transform +# --------------------------------------------------------------------------- + +_OUTPUT_GEOM = ("K", "Ho", "Wo", "dH", "dW", "dHm", "dWm") + + +def output_signature(spec: RtOutputSpec) -> list: + sb = ( + SignatureBuilder() + .ptr("mm", spec.operand_dtype) + .ptr("e", "i32") + .ptr("y", spec.io_dtype) + .scalar("m_bytes", "i32") + .scalar("y_bytes", "i32") + ) + for nm in _OUTPUT_GEOM: + sb = sb.scalar(nm, "i32") + return _run_sig(sb).build() + + +def build_output( + spec: RtOutputSpec, *, arch: str = "gfx1250" +) -> KernelDef: # noqa: ARG001 + """K4 of :func:`conv_winograd_fwd.build_winograd_output_transform`, extents at run time.""" + n, m = spec.mh + 2, spec.mh + nw, mw = spec.mw + 2, spec.mw + vec = spec.vec + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = spec.block_size + mm = b.param( + "mm", + PtrType(_ir_dtype(spec.operand_dtype), "global"), + noalias=True, + readonly=True, + ) + e_ptr = b.param("e", PtrType(I32, "global"), noalias=True, readonly=True) + y = b.param( + "y", PtrType(_ir_dtype(spec.io_dtype), "global"), noalias=True, writeonly=True + ) + m_bytes = b.param("m_bytes", I32) + y_bytes = b.param("y_bytes", I32) + g = {nm: b.param(nm, I32) for nm in _OUTPUT_GEOM} + m_rsrc = b.buffer_rsrc(mm, m_bytes) + y_rsrc = b.buffer_rsrc(y, y_bytes) + rc = _run_coords(b, 1) + + K, Ho, Wo = g["K"], g["Ho"], g["Wo"] + kb2 = b.mul(K, _c(b, 2)) + t = rc.t0 + ok = b.land(rc.ok, rc.c_ok) + k = _mul_c(b, rc.cv, vec) + k2 = _mul_c(b, k, 2) + mbase = _oob(b, ok, b.add(b.mul(t, _mul_c(b, kb2, n * nw)), k2)) + raw = [] + off = mbase + for uu in range(n): + row = [] + for vv in range(nw): + if uu or vv: + off = b.add(off, kb2) + row.append(_load_raw(b, m_rsrc, off, spec.operand_dtype, vec)) + raw.append(row) + scale = None + if spec.prescale: + t_safe = b.select(ok, t, _c(b, 0)) + scale = _pow2_f32(b, b.global_load_i32(e_ptr, t_safe)) + + ho0 = b.add(rc.a, b.mul(rc.ti, g["dHm"])) + wo0 = b.add(rc.bsub, b.mul(rc.tj0, g["dWm"])) + img_row0 = b.mul(rc.n, Ho) + yrow = [] + ho = ho0 + for i in range(m): + if i: + ho = b.add(ho, g["dH"]) + okr = b.land(ok, b.cmp_lt(ho, Ho)) + yrow.append(_oob(b, okr, b.add(b.mul(b.mul(b.add(img_row0, ho), Wo), kb2), k2))) + ycol = [] + wo = wo0 + for j in range(mw): + if j: + wo = b.add(wo, g["dW"]) + ycol.append(_oob(b, b.cmp_lt(wo, Wo), b.mul(wo, kb2))) + + ops = _IROps(b) + b.sched_barrier(0) + colt = [] + for vv in range(nw): + data = [_raw_to_f32(b, raw[uu][vv], vec) for uu in range(n)] + colt.append( + [at_1d(ops, m, [data[uu][ch] for uu in range(n)]) for ch in range(vec)] + ) + pending = [] + for i in range(m): + outs = [ + at_1d(ops, mw, [colt[vv][ch][i] for vv in range(nw)]) for ch in range(vec) + ] + for j in range(mw): + vals = [outs[ch][j] for ch in range(vec)] + if scale is not None: + vals = [b.fmul(x_, scale) for x_ in vals] + pending.append( + (b.add(yrow[i], ycol[j]), _pack_vec_f32(b, vals, spec.io_dtype)) + ) + b.sched_barrier(0) + cpol = _CPOL_NT if spec.nt_stores else 0 + for off_, packed in pending: + _store_packed(b, y_rsrc, off_, packed, spec.io_dtype, cpol) + return b.kernel + + +def output_launch(spec: RtOutputSpec, geo: RtGeometry): + p = geo.problem + vals = dict( + K=p.K, + Ho=p.Ho, + Wo=p.Wo, + dH=p.dH, + dW=p.dW, + dHm=p.dH * spec.mh, + dWm=p.dW * spec.mw, + ) + rv, grid = _run_values(geo, p.K // spec.vec, 1, spec.block_size) + vals.update(rv) + return grid, spec.block_size, vals + + +# --------------------------------------------------------------------------- +# Stage table +# --------------------------------------------------------------------------- + +_STAGES = { + RtFilterSpec: (build_filter, filter_signature, filter_launch), + RtPrescaleSpec: (build_prescale, prescale_signature, prescale_launch), + RtInputSpec: (build_input, input_signature, input_launch), + RtOutputSpec: (build_output, output_signature, output_launch), +} + + +def build_stage(spec, *, arch: str = "gfx1250") -> KernelDef: + return _STAGES[type(spec)][0](spec, arch=arch) + + +def stage_signature(spec) -> list: + return _STAGES[type(spec)][1](spec) + + +def stage_launch( + spec, geo: RtGeometry +) -> Tuple[Tuple[int, int, int], int, Dict[str, int]]: + """``(grid, block, scalar kernarg values)`` of one launch; the buffer + pointers and ``_bytes`` sizes are the caller's.""" + why = spec.reason(geo) + if why: + raise ValueError(f"{spec.kernel_name()}: {why}") + grid, block, vals = _STAGES[type(spec)][2](spec, geo) + if min(grid) <= 0: + raise ValueError("empty grid") + if grid[0] >= 2**31: + raise ValueError(f"grid {grid} too large") + return grid, block, vals + + +# --------------------------------------------------------------------------- +# The AOT grid and the default pick +# --------------------------------------------------------------------------- + +#: Tiles of the AOT grid: every tile the 3x3 selector considers +#: (``conv_winograd_fwd.AUTO_TILES``). +AOT_TILES: Tuple[Tuple[int, int], ...] = ( + (4, 4), + (6, 4), + (2, 2), + (6, 2), + (4, 2), + (2, 4), + (2, 6), + (4, 6), +) +AOT_VECS: Tuple[int, ...] = (8, 4, 2, 1) +#: K2's LDS region for tiles with fewer than 36 bins (the JIT pipeline's +#: tuned LDS knobs) and the fill below which the register path is used. +_LDS_REGION = (4, 8) +_LDS_TPT = 4 +_LDS_LANES = 16 +_LDS_MIN_FILL = 0.75 + + +def pipeline_dtypes(tile, io_dtype: str) -> Tuple[str, bool]: + """``(operand dtype, prescale)`` of a pipeline, as the JIT pipeline's + ``auto`` rules: bf16 operands only for F(2,3) with bf16 I/O; the prescale + for bf16 I/O with a tile side of 4 or more.""" + mh, mw = _split_tile(tile) + operand = "bf16" if ((mh, mw) == (2, 2) and io_dtype == "bf16") else "fp16" + return operand, io_dtype == "bf16" and max(mh, mw) >= 4 + + +def _small_tile(mh: int, mw: int) -> bool: + return (mh + 2) * (mw + 2) < 36 + + +def aot_stage_specs() -> list: + """Every stage spec of the AOT grid: per I/O dtype, K2a at each vector + width; per tile and operand dtype (bf16 F(2,3) also with fp16 operands), + K1, K2 and K4 at each vector width -- K2 on the register path, with the + fused prescale at blocks 256 and 1024 when the tile takes the prescale, + and for tiles under 36 bins also with 2-tile runs and through the LDS + region.""" + out = [] + for io in _IO: + for v in AOT_VECS: + out.append(RtPrescaleSpec(io, v)) + for mh, mw in AOT_TILES: + _, ps = pipeline_dtypes((mh, mw), io) + ops = ( + ("bf16", "fp16") if ((mh, mw) == (2, 2) and io == "bf16") else ("fp16",) + ) + fr = 2 if mh >= 4 else mh + 2 + for op in ops: + for v in AOT_VECS: + out.append(RtFilterSpec(mh, mw, io, op, v, fr)) + out.append(RtOutputSpec(mh, mw, io, op, v, 256, prescale=ps)) + out.append(RtInputSpec(mh, mw, io, op, v, 256, prescale=ps)) + if ps: + for blk in (256, 1024): + out.append( + RtInputSpec( + mh, mw, io, op, v, blk, prescale=True, fused=True + ) + ) + if _small_tile(mh, mw): + out.append( + RtInputSpec( + mh, mw, io, op, v, 256, prescale=ps, tiles_per_thread=2 + ) + ) + out.append( + RtInputSpec( + mh, mw, io, op, v, prescale=ps, lds=_LDS_REGION, + tiles_per_thread=_LDS_TPT, lds_lanes=_LDS_LANES, + ) + ) # fmt: skip + return out + + +def _pow2_div(n: int, cap: int) -> int: + v = cap + while v > 1 and n % v: + v //= 2 + return v + + +def default_stage_specs(problem: ConvProblem, tile, io_dtype: str) -> dict: + """The AOT binaries a pipeline uses by default, ``{"K1", ("K2a",) "K2", + "K4"}`` -> spec: the JIT pipeline's tuned knobs (8-wide channel vectors + for F(2,3)-sized tiles, 4-wide for larger ones, narrowed to divide the + channel extent; the fused prescale while a workgroup can hold all + channels, else the split one; the LDS region for small tiles whose tile + grid fills it). A pure function of the problem; every spec is in + :func:`aot_stage_specs`.""" + mh, mw = _split_tile(tile) + operand, ps = pipeline_dtypes((mh, mw), io_dtype) + p = problem + cap = 4 if max(mh, mw) >= 4 else 8 + vc = _pow2_div(p.C, cap) + out = { + "K1": RtFilterSpec( + mh, mw, io_dtype, operand, _pow2_div(p.cpg, cap), 2 if mh >= 4 else mh + 2 + ), + "K4": RtOutputSpec( + mh, mw, io_dtype, operand, _pow2_div(p.K, cap), 256, prescale=ps + ), + } + if ps: + lanes = _pow2_ceil(p.C // vc) + if lanes <= 1024: + blk = 256 if lanes <= 256 else 1024 + out["K2"] = RtInputSpec( + mh, mw, io_dtype, operand, vc, blk, prescale=True, fused=True + ) + else: + out["K2a"] = RtPrescaleSpec(io_dtype, _pow2_div(p.C, 8)) + out["K2"] = RtInputSpec(mh, mw, io_dtype, operand, vc, 256, prescale=True) + return out + geo = RtGeometry(p, mh, mw) + bh, bw = _LDS_REGION + th, tw = geo.tiles_h, geo.tiles_w + fill = th * tw / (ceil_div(th, bh) * bh * ceil_div(tw, bw) * bw) + if _small_tile(mh, mw) and fill >= _LDS_MIN_FILL: + out["K2"] = RtInputSpec( + mh, mw, io_dtype, operand, vc, lds=_LDS_REGION, + tiles_per_thread=_LDS_TPT, lds_lanes=_LDS_LANES, + ) # fmt: skip + else: + out["K2"] = RtInputSpec(mh, mw, io_dtype, operand, vc, 256) + return out + + +__all__: Sequence[str] = ( + "RtFilterSpec", + "RtGeometry", + "RtInputSpec", + "RtOutputSpec", + "RtPrescaleSpec", + "aot_stage_specs", + "build_stage", + "default_stage_specs", + "pipeline_dtypes", + "magic_div", + "stage_launch", + "stage_signature", +) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_select.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_select.py index 1d53a1863dce..ef76aca6fce4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_select.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_select.py @@ -205,7 +205,7 @@ def label(self) -> str: generalized plans).""" if self.kind == "none": return "none" - if self.kind == "unfused2d": + if self.kind in ("unfused2d", "unfused2d_rt"): return f"{self.kind}:F({self.tile[0]}x{self.tile[1]},3x3)" return f"{self.kind}:{self.variant}" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py new file mode 100644 index 000000000000..b571fc8e9134 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py @@ -0,0 +1,177 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""CPU tests of the opt-in AOT Winograd forward route of the grouped-convolution +dispatcher (``conv_winograd_fwd_aot``, plan kind ``unfused2d_rt``). + +* ``algorithm="auto"`` never picks it; ``algorithm="winograd_aot"`` selects a + 3x3 plan on the shape-generic transform binaries; +* every step of its launch plans is an AOT binary: the transforms are members + of :func:`conv_winograd_fwd_rt.aot_stage_specs` and are named by their + kernel names, every kernarg is spelled, every kernel builds; +* the batch split re-launches the same binaries on byte offsets of x and y; +* requests without a 3x3 plan are refused with a reason. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +import dispatch.grouped_convolution as gc +from dispatch.conv_winograd_plan import ConvWinogradSpec, winograd_launch_plan +from kernels.common import conv_winograd_fwd_rt as R +from kernels.common.conv_winograd_gemm_wmma import WinogradGemmProblem +from kernels.common.conv_winograd_select import WinogradPlan + +_ARCH = "gfx1250" +_PTRS = { + "x": 1 << 20, + "w": 2 << 20, + "y": 3 << 20, + "U": 4 << 20, + "V": 5 << 20, + "M": 6 << 20, + "e": 7 << 20, +} +_GRID = {s.kernel_name() for s in R.aot_stage_specs()} + + +def _req(**kw): + fields = dict(N=2, C=64, K=64, Hi=16, Wi=16, Y=3, X=3, pad_h=1, pad_w=1) + fields.update(kw) + fields.setdefault("arch", _ARCH) + fields.setdefault("llvm_flavor", "llvm23") # the K3 GEMM lowers at llvm23 + return gc.ConvGroupedRequest(**fields) + + +def _pin(req): + return dataclasses.replace(req, algorithm="winograd_aot") + + +def _check_plan(spec: ConvWinogradSpec, *, build: bool = True): + plan = gc.launch_plan_for(None, spec) + assert plan.steps + for step in plan.steps: + assert step.identity is None + assert len(step.kernargs) == len(step.signature) + vals = plan.values(step, _PTRS, {}) + assert list(vals) == [a["name"] for a in step.signature] + if not step.kernel.name.startswith("K3"): + assert step.kernel.name in _GRID + assert min(step.grid) >= 1 and step.block[0] >= 32 + if build: + k = step.kernel.build(_ARCH) + assert [p.name for p in k.params] == [a["name"] for a in step.signature] + return plan + + +def _aot_spec(req, tile): + for s in gc.conv_grouped_sweep_space(req): + if s.plan.kind == "unfused2d_rt" and s.plan.tile == tile: + return s + raise AssertionError(f"no AOT {tile} plan") + + +def test_auto_never_picks_the_aot_route(): + with pytest.raises(ValueError): + gc.dispatch_conv_grouped(_req()) + + +def test_pinned_request_selects_an_aot_plan(): + res = gc.dispatch_conv_grouped(_pin(_req())) + assert isinstance(res.spec, ConvWinogradSpec) + assert res.candidate.name == "conv_winograd_fwd_aot" and res.candidate.opt_in + assert res.candidate.algorithm == "winograd_aot" + assert res.spec.plan.kind == "unfused2d_rt" and not res.spec.plan.merge + assert res.spec.name.startswith("conv_winograd_fwd_aot:unfused2d_rt:F(") + last = gc.launch_plan_for(_pin(_req()), res.spec).steps[-1] + assert res.grid == last.grid and res.block == last.block + assert last.kernel.name.startswith("wino_rt_output_") + + +def test_sweep_lists_both_routes(): + specs = gc.conv_grouped_sweep_space(_req()) + kinds = {s.plan.kind for s in specs} + assert {"unfused2d", "unfused2d_rt"} <= kinds + aot = [s for s in specs if s.plan.kind == "unfused2d_rt"] + jit = [s for s in specs if s.plan.kind == "unfused2d"] + assert {s.plan.tile for s in aot} == {s.plan.tile for s in jit} + + +def test_fp16_plan_steps(): + req = _req() + spec = _aot_spec(req, (4, 4)) + plan = _check_plan(spec) + names = [st.kernel.name for st in plan.steps] + assert names[0].startswith("wino_rt_filter_f4_fp16") + assert names[1].startswith("wino_rt_input_f4_fp16") and names[2] == "K3" + geo = R.RtGeometry(spec.problem, 4, 4) + sizes = geo.buffer_bytes(False) + assert plan.workspaces == { + "U": sizes["U"], + "V": sizes["V"], + "M": sizes["M"], + "e": 4, + } + vals = plan.values(plan.steps[2], _PTRS, {}) + prob = WinogradGemmProblem.from_conv(spec.problem, (4, 4)) + assert (vals["T"], vals["G"], vals["Cg"], vals["Kg"]) == ( + prob.T, + prob.G, + prob.Cg, + prob.Kg, + ) + # The input transform reads the problem from its kernargs. + k2 = plan.values(plan.steps[1], _PTRS, {}) + assert (k2["C"], k2["Hi"], k2["Wi"], k2["pH"], k2["T"]) == ( + 64, + 16, + 16, + 1, + geo.num_tiles, + ) + + +def test_small_tile_uses_the_lds_region(): + spec = _aot_spec(_req(Hi=32, Wi=32), (2, 2)) + plan = _check_plan(spec) + assert "_lds4x8_" in plan.steps[1].kernel.name + + +def test_bf16_prescale_plans(): + req = dataclasses.replace(_req(K=96), dtype="bf16") + fused = _check_plan(_aot_spec(req, (4, 4))) + assert fused.steps[1].kernel.name.endswith("_psf") + assert fused.workspaces["e"] > 4 + assert ("buf", "e", 0) in fused.steps[-1].kernargs + wide = dataclasses.replace(_req(C=16384, K=64, Hi=8, Wi=8), dtype="bf16") + split = _check_plan(_aot_spec(wide, (4, 4)), build=False) + assert split.steps[1].kernel.name.startswith("wino_rt_pmax_bf16") + assert split.steps[2].kernel.name.endswith("_ps") + + +def test_batch_split_relaunches_the_same_binaries(): + p = _req(N=4, Hi=8, Wi=8) + plan = WinogradPlan( + kind="unfused2d_rt", tile=(2, 2), k3="128x128x128_w2x2_f4x4_s2_tdm", chunks=2 + ) + split = winograd_launch_plan(gc._conv_problem(p), "fp16", plan, _ARCH) + names = [st.kernel.name for st in split.steps] + assert len(names) == 7 and names[2] == names[5] == "K3" + assert names[1] == names[4] and names[3] == names[6] + sub = dataclasses.replace(gc._conv_problem(p), N=2) + sz = R.RtGeometry(sub, 2, 2).buffer_bytes(False) + v = split.values(split.steps[4], _PTRS, {}) + assert v["x"] == _PTRS["x"] + sz["x"] and v["x_bytes"] == sz["x"] + v = split.values(split.steps[6], _PTRS, {}) + assert v["y"] == _PTRS["y"] + sz["y"] and v["y_bytes"] == sz["y"] + + +def test_no_3x3_plan_is_refused(): + req = _pin(_req(Y=1, X=11, Hi=1, Wi=200, pad_h=0, pad_w=5)) + with pytest.raises(ValueError, match="has no AOT kernels"): + gc.dispatch_conv_grouped(req) + strided = _pin(_req(stride_h=2, stride_w=2)) + with pytest.raises(ValueError): + gc.dispatch_conv_grouped(strided) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c new file mode 100644 index 000000000000..e109721e53c9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c @@ -0,0 +1,209 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * tests/parity/conv_winograd_fwd_rt_emit.c -- C-side emitter for the + * shape-generic Winograd forward transform parity harness. Selects one of N + * sampled configs by argv[1], builds the rocke_winograd_rt_spec_t identically + * to the Python emitter conv_winograd_fwd_rt_emit.py (same table, same + * order), builds the stage (rocke_build_winograd_rt_new) and lowers it for the + * config's arch, printing the .ll (or the serialized IR / verifier output) to + * stdout. + * + * Config index table: see the Python emitter (must stay in sync). + */ +#include +#include +#include + +#include "rocke/instance_conv_winograd_fwd_rt.h" +#include "rocke/ir.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_llvm.h" +#include "rocke/verify.h" + +typedef struct cfg +{ + const char* stage; + const char* arch; + int mh, mw; + const char* io; + const char* op; + int vec, fr, block, ps, tpt, bh, bw, lanes, nt, fused; +} cfg_t; + +/* clang-format off */ +static const cfg_t CONFIGS[] = { + {"K1", "gfx1250", 4, 4, "fp16", "fp16", 8, 2, 256, 0, 1, 0, 0, 16, 1, 0}, /* 0: F4 fp16 K1 (2 filter rows) */ + {"K1", "gfx1250", 2, 2, "bf16", "bf16", 8, 4, 256, 0, 1, 0, 0, 16, 1, 0}, /* 1: F2 bf16 operands K1 (all rows) */ + {"K1", "gfx1250", 6, 4, "fp16", "fp16", 4, 2, 256, 0, 1, 0, 0, 16, 1, 0}, /* 2: F(6x4) K1 vec 4 */ + {"K1", "gfx1250", 2, 6, "fp16", "fp16", 1, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 3: F(2x6) K1 vec 1, 1 filter row */ + {"K1", "gfx1250", 4, 2, "bf16", "fp16", 2, 3, 256, 0, 1, 0, 0, 16, 1, 0}, /* 4: F(4x2) bf16 I/O K1 vec 2 */ + {"K2a", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 5: K2a fp16 vec 8 */ + {"K2a", "gfx1250", 4, 4, "bf16", "fp16", 1, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 6: K2a bf16 vec 1 */ + {"K2a", "gfx1250", 4, 4, "bf16", "fp16", 4, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 7: K2a bf16 vec 4 */ + {"K2", "gfx1250", 4, 4, "fp16", "fp16", 4, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 8: F4 K2 register, 1 tile */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 2, 0, 0, 16, 1, 0}, /* 9: F2 K2 register, 2 tiles per thread */ + {"K2", "gfx1250", 6, 4, "fp16", "fp16", 4, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 10: F(6x4) K2 register vec 4 */ + {"K2", "gfx1250", 2, 4, "fp16", "fp16", 2, 1, 256, 0, 3, 0, 0, 16, 1, 0}, /* 11: F(2x4) K2 register 3 tiles per thread (partial runs) */ + {"K2", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 64, 0, 1, 0, 0, 16, 0, 0}, /* 12: F4 K2 regular stores, block 64 */ + {"K2", "gfx1250", 4, 4, "bf16", "fp16", 4, 1, 256, 1, 1, 0, 0, 16, 1, 0}, /* 13: F4 bf16 K2 split prescale */ + {"K2", "gfx1250", 4, 6, "bf16", "fp16", 2, 1, 256, 1, 2, 0, 0, 16, 1, 0}, /* 14: F(4x6) bf16 K2 split prescale 2 tiles */ + {"K2", "gfx1250", 4, 4, "bf16", "fp16", 4, 1, 256, 1, 1, 0, 0, 16, 1, 1}, /* 15: F4 bf16 K2 fused prescale block 256 (LDS group max) */ + {"K2", "gfx1250", 6, 4, "bf16", "fp16", 2, 1, 1024, 1, 1, 0, 0, 16, 1, 1}, /* 16: F(6x4) bf16 K2 fused prescale block 1024 */ + {"K2", "gfx1250", 4, 4, "bf16", "fp16", 1, 1, 32, 1, 1, 0, 0, 16, 1, 1}, /* 17: F4 bf16 K2 fused prescale block 32 (shuffles only) */ + {"K2", "gfx1250", 2, 2, "bf16", "bf16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 18: F2 bf16 operands K2 register */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 4, 4, 8, 16, 1, 0}, /* 19: F2 K2 LDS 4x8, 4 tiles per thread, 16 lanes */ + {"K2", "gfx1250", 2, 2, "bf16", "bf16", 1, 1, 256, 0, 4, 4, 8, 16, 1, 0}, /* 20: F2 bf16 K2 LDS 4x8 vec 1 */ + {"K2", "gfx1250", 4, 4, "bf16", "fp16", 4, 1, 256, 1, 1, 2, 4, 8, 1, 0}, /* 21: F4 bf16 K2 LDS 2x4 + split prescale, 8 lanes */ + {"K2", "gfx1250", 4, 2, "fp16", "fp16", 2, 1, 256, 0, 1, 4, 1, 1, 1, 0}, /* 22: F(4x2) K2 LDS 4x1 (one tile column), 1 lane */ + {"K4", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 23: F4 K4 */ + {"K4", "gfx1250", 2, 2, "fp16", "fp16", 1, 1, 64, 0, 1, 0, 0, 16, 1, 0}, /* 24: F2 K4 vec 1 block 64 */ + {"K4", "gfx1250", 6, 4, "bf16", "fp16", 4, 1, 256, 1, 1, 0, 0, 16, 1, 0}, /* 25: F(6x4) bf16 K4 prescale */ + {"K4", "gfx1250", 2, 6, "fp16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 0, 0}, /* 26: F(2x6) K4 regular stores */ + {"K4", "gfx1250", 6, 6, "fp16", "fp16", 2, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 27: F6 K4 */ + {"K2", "gfx950", 4, 4, "fp16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 28: gfx950: F4 K2 register */ + {"K2", "gfx950", 4, 4, "bf16", "fp16", 8, 1, 256, 1, 1, 0, 0, 16, 1, 1}, /* 29: gfx950: F4 bf16 K2 fused prescale */ + {"K2", "gfx950", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 4, 4, 8, 16, 1, 0}, /* 30: gfx950: F2 K2 LDS */ + {"K1", "gfx950", 4, 4, "fp16", "fp16", 8, 2, 256, 0, 1, 0, 0, 16, 1, 0}, /* 31: gfx950: F4 K1 */ + {"K4", "gfx950", 4, 4, "bf16", "fp16", 8, 1, 256, 1, 1, 0, 0, 16, 1, 0}, /* 32: gfx950: F4 bf16 K4 prescale */ + {"K2a", "gfx950", 4, 4, "bf16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 33: gfx950: K2a */ + {"K2", "gfx1250", 3, 4, "fp16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 34: rejected: tile 3 */ + {"K1", "gfx1250", 4, 4, "bf16", "bf16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 35: rejected: bf16 operands with F4 */ + {"K4", "gfx1250", 4, 4, "fp16", "fp16", 3, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 36: rejected: vec 3 */ + {"K2", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 256, 0, 1, 0, 0, 16, 1, 1}, /* 37: rejected: fused without prescale */ + {"K2", "gfx1250", 2, 2, "bf16", "fp16", 8, 1, 256, 1, 4, 4, 8, 16, 1, 1}, /* 38: rejected: fused on the LDS path */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 4, 2, 3, 16, 1, 0}, /* 39: rejected: tiles_per_thread does not divide the LDS block */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 1, 8, 8, 32, 1, 0}, /* 40: rejected: LDS region over 1024 threads */ + {"K1", "gfx1250", 4, 4, "fp16", "fp16", 8, 4, 256, 0, 1, 0, 0, 16, 1, 0}, /* 41: rejected: filter_rows does not divide n */ + {"K4", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 96, 0, 1, 0, 0, 16, 1, 0}, /* 42: rejected: block not a power of two */ + {"K2a", "gfx1250", 4, 4, "fp16", "fp16", 16, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 43: rejected: K2a vec 16 */ +}; +/* clang-format on */ + +#define NUM_CONFIGS ((int)(sizeof(CONFIGS) / sizeof(CONFIGS[0]))) + +static int stage_of(const char* name, rocke_winograd_rt_stage_t* out) +{ + if(strcmp(name, "K1") == 0) + *out = ROCKE_WINOGRAD_RT_K1; + else if(strcmp(name, "K2a") == 0) + *out = ROCKE_WINOGRAD_RT_K2A; + else if(strcmp(name, "K2") == 0) + *out = ROCKE_WINOGRAD_RT_K2; + else if(strcmp(name, "K4") == 0) + *out = ROCKE_WINOGRAD_RT_K4; + else + return 0; + return 1; +} + +int main(int argc, char** argv) +{ + if(argc < 2) + { + fprintf(stderr, "usage: %s [ll|ir|verify]\n", argv[0]); + return 2; + } + int idx = atoi(argv[1]); + const char* mode = (argc > 2) ? argv[2] : "ll"; + if(idx < 0 || idx >= NUM_CONFIGS) + { + fprintf(stderr, "unknown config index %d\n", idx); + return 2; + } + const cfg_t* c = &CONFIGS[idx]; + rocke_winograd_rt_stage_t stage; + if(!stage_of(c->stage, &stage)) + { + fprintf(stderr, "bad stage %s\n", c->stage); + return 2; + } + rocke_winograd_rt_spec_t s = rocke_winograd_rt_spec_default(stage); + s.io_dtype = c->io; + s.vec = c->vec; + if(stage != ROCKE_WINOGRAD_RT_K2A) + { + s.mh = c->mh; + s.mw = c->mw; + s.operand_dtype = c->op; + } + if(stage == ROCKE_WINOGRAD_RT_K1) + s.filter_rows = c->fr; + if(stage == ROCKE_WINOGRAD_RT_K2 || stage == ROCKE_WINOGRAD_RT_K4) + { + s.block_size = c->block; + s.prescale = c->ps != 0; + s.nt_stores = c->nt != 0; + } + if(stage == ROCKE_WINOGRAD_RT_K2) + { + s.tiles_per_thread = c->tpt; + s.lds[0] = c->bh; + s.lds[1] = c->bw; + s.lds_lanes = c->lanes; + s.fused = c->fused != 0; + } + + rocke_ir_builder_t b; + rocke_kernel_def_t* kernel = rocke_build_winograd_rt_new(&b, &s, c->arch); + if(kernel == NULL) + { + const char* m = rocke_ir_builder_error(&b); + fprintf(stderr, "build failed: %s\n", m ? m : "(no message)"); + rocke_ir_builder_free(&b); + return 1; + } + + int ret = 0; + if(strcmp(mode, "ll") == 0) + { + char* llvm_text = NULL; + rocke_status_t st + = rocke_lower_kernel_to_llvm(kernel, ROCKE_LLVM_FLAVOR_AUTO, c->arch, &llvm_text); + if(st != ROCKE_OK || !llvm_text) + { + fprintf(stderr, "lower failed: status=%d\n", (int)st); + rocke_ir_builder_free(&b); + return 1; + } + fputs(llvm_text, stdout); + free(llvm_text); + } + else if(strcmp(mode, "ir") == 0) + { + char* ir_text = NULL; + rocke_status_t st = rocke_ir_serialize(kernel, &ir_text); + if(st != ROCKE_OK || !ir_text) + { + fprintf(stderr, "serialize failed: status=%d\n", (int)st); + rocke_ir_builder_free(&b); + return 1; + } + fputs(ir_text, stdout); + free(ir_text); + } + else if(strcmp(mode, "verify") == 0) + { + rocke_diag_t* diags = NULL; + size_t ndiag = 0; + rocke_verify(kernel, &diags, &ndiag); + for(size_t i = 0; i < ndiag; i++) + { + char* t = rocke_diag_to_string(&diags[i]); + if(t) + { + fputs(t, stdout); + fputc('\n', stdout); + free(t); + } + } + rocke_diags_free(diags, ndiag); + } + else + { + fprintf(stderr, "unknown mode '%s' (expected ll|ir|verify)\n", mode); + ret = 2; + } + rocke_ir_builder_free(&b); + return ret; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py new file mode 100644 index 000000000000..ef59c469773a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +# +# tests/parity/conv_winograd_fwd_rt_emit.py -- Python reference emitter for the +# shape-generic Winograd F(mh x mw, 3 x 3) forward transform kernels (K1 +# filter, K2a per-pixel max, K2 input, K4 output) of conv_winograd_fwd_rt. +# Selects one of N sampled configs by argv[1], builds the stage's spec and +# kernel (build_stage) and prints the LLVM IR, to be byte-compared with the C +# emitter conv_winograd_fwd_rt_emit.c (same table, same order). +# +# A config row: the stage ("K1" | "K2a" | "K2" | "K4"), the arch, and the +# spec knobs (a stage reads the ones its spec class has). Index map: +# 0 F4 fp16 K1 (2 filter rows) [gfx1250] +# 1 F2 bf16 operands K1 (all rows) [gfx1250] +# 2 F(6x4) K1 vec 4 [gfx1250] +# 3 F(2x6) K1 vec 1, 1 filter row [gfx1250] +# 4 F(4x2) bf16 I/O K1 vec 2 [gfx1250] +# 5 K2a fp16 vec 8 [gfx1250] +# 6 K2a bf16 vec 1 [gfx1250] +# 7 K2a bf16 vec 4 [gfx1250] +# 8 F4 K2 register, 1 tile [gfx1250] +# 9 F2 K2 register, 2 tiles per thread [gfx1250] +# 10 F(6x4) K2 register vec 4 [gfx1250] +# 11 F(2x4) K2 register 3 tiles per thread (partial runs) [gfx1250] +# 12 F4 K2 regular stores, block 64 [gfx1250] +# 13 F4 bf16 K2 split prescale [gfx1250] +# 14 F(4x6) bf16 K2 split prescale 2 tiles [gfx1250] +# 15 F4 bf16 K2 fused prescale block 256 (LDS group max) [gfx1250] +# 16 F(6x4) bf16 K2 fused prescale block 1024 [gfx1250] +# 17 F4 bf16 K2 fused prescale block 32 (shuffles only) [gfx1250] +# 18 F2 bf16 operands K2 register [gfx1250] +# 19 F2 K2 LDS 4x8, 4 tiles per thread, 16 lanes [gfx1250] +# 20 F2 bf16 K2 LDS 4x8 vec 1 [gfx1250] +# 21 F4 bf16 K2 LDS 2x4 + split prescale, 8 lanes [gfx1250] +# 22 F(4x2) K2 LDS 4x1 (one tile column), 1 lane [gfx1250] +# 23 F4 K4 [gfx1250] +# 24 F2 K4 vec 1 block 64 [gfx1250] +# 25 F(6x4) bf16 K4 prescale [gfx1250] +# 26 F(2x6) K4 regular stores [gfx1250] +# 27 F6 K4 [gfx1250] +# 28 gfx950: F4 K2 register [gfx950] +# 29 gfx950: F4 bf16 K2 fused prescale [gfx950] +# 30 gfx950: F2 K2 LDS [gfx950] +# 31 gfx950: F4 K1 [gfx950] +# 32 gfx950: F4 bf16 K4 prescale [gfx950] +# 33 gfx950: K2a [gfx950] +# 34 rejected: tile 3 [gfx1250] +# 35 rejected: bf16 operands with F4 [gfx1250] +# 36 rejected: vec 3 [gfx1250] +# 37 rejected: fused without prescale [gfx1250] +# 38 rejected: fused on the LDS path [gfx1250] +# 39 rejected: tiles_per_thread does not divide the LDS block [gfx1250] +# 40 rejected: LDS region over 1024 threads [gfx1250] +# 41 rejected: filter_rows does not divide n [gfx1250] +# 42 rejected: block not a power of two [gfx1250] +# 43 rejected: K2a vec 16 [gfx1250] +from _emit_common import run_emit +from kernels.common.conv_winograd_fwd_rt import ( + RtFilterSpec, + RtInputSpec, + RtOutputSpec, + RtPrescaleSpec, + build_stage, +) + +# fmt: off +CONFIGS = [ + ('K1', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 2, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 0: F4 fp16 K1 (2 filter rows) + ('K1', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'bf16', 'op': 'bf16', 'vec': 8, 'fr': 4, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 1: F2 bf16 operands K1 (all rows) + ('K1', 'gfx1250', {'mh': 6, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 4, 'fr': 2, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 2: F(6x4) K1 vec 4 + ('K1', 'gfx1250', {'mh': 2, 'mw': 6, 'io': 'fp16', 'op': 'fp16', 'vec': 1, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 3: F(2x6) K1 vec 1, 1 filter row + ('K1', 'gfx1250', {'mh': 4, 'mw': 2, 'io': 'bf16', 'op': 'fp16', 'vec': 2, 'fr': 3, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 4: F(4x2) bf16 I/O K1 vec 2 + ('K2a', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 5: K2a fp16 vec 8 + ('K2a', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 1, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 6: K2a bf16 vec 1 + ('K2a', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 7: K2a bf16 vec 4 + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 8: F4 K2 register, 1 tile + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 2, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 9: F2 K2 register, 2 tiles per thread + ('K2', 'gfx1250', {'mh': 6, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 10: F(6x4) K2 register vec 4 + ('K2', 'gfx1250', {'mh': 2, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 2, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 3, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 11: F(2x4) K2 register 3 tiles per thread (partial runs) + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 64, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 0, 'fused': 0}), # 12: F4 K2 regular stores, block 64 + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 13: F4 bf16 K2 split prescale + ('K2', 'gfx1250', {'mh': 4, 'mw': 6, 'io': 'bf16', 'op': 'fp16', 'vec': 2, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 2, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 14: F(4x6) bf16 K2 split prescale 2 tiles + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 1}), # 15: F4 bf16 K2 fused prescale block 256 (LDS group max) + ('K2', 'gfx1250', {'mh': 6, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 2, 'fr': 1, 'block': 1024, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 1}), # 16: F(6x4) bf16 K2 fused prescale block 1024 + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 1, 'fr': 1, 'block': 32, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 1}), # 17: F4 bf16 K2 fused prescale block 32 (shuffles only) + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'bf16', 'op': 'bf16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 18: F2 bf16 operands K2 register + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 4, 'bh': 4, 'bw': 8, 'lanes': 16, 'nt': 1, 'fused': 0}), # 19: F2 K2 LDS 4x8, 4 tiles per thread, 16 lanes + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'bf16', 'op': 'bf16', 'vec': 1, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 4, 'bh': 4, 'bw': 8, 'lanes': 16, 'nt': 1, 'fused': 0}), # 20: F2 bf16 K2 LDS 4x8 vec 1 + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 1, 'bh': 2, 'bw': 4, 'lanes': 8, 'nt': 1, 'fused': 0}), # 21: F4 bf16 K2 LDS 2x4 + split prescale, 8 lanes + ('K2', 'gfx1250', {'mh': 4, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 2, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 4, 'bw': 1, 'lanes': 1, 'nt': 1, 'fused': 0}), # 22: F(4x2) K2 LDS 4x1 (one tile column), 1 lane + ('K4', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 23: F4 K4 + ('K4', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 1, 'fr': 1, 'block': 64, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 24: F2 K4 vec 1 block 64 + ('K4', 'gfx1250', {'mh': 6, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 4, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 25: F(6x4) bf16 K4 prescale + ('K4', 'gfx1250', {'mh': 2, 'mw': 6, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 0, 'fused': 0}), # 26: F(2x6) K4 regular stores + ('K4', 'gfx1250', {'mh': 6, 'mw': 6, 'io': 'fp16', 'op': 'fp16', 'vec': 2, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 27: F6 K4 + ('K2', 'gfx950', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 28: gfx950: F4 K2 register + ('K2', 'gfx950', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 1}), # 29: gfx950: F4 bf16 K2 fused prescale + ('K2', 'gfx950', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 4, 'bh': 4, 'bw': 8, 'lanes': 16, 'nt': 1, 'fused': 0}), # 30: gfx950: F2 K2 LDS + ('K1', 'gfx950', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 2, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 31: gfx950: F4 K1 + ('K4', 'gfx950', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 32: gfx950: F4 bf16 K4 prescale + ('K2a', 'gfx950', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 33: gfx950: K2a + ('K2', 'gfx1250', {'mh': 3, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 34: rejected: tile 3 + ('K1', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'bf16', 'op': 'bf16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 35: rejected: bf16 operands with F4 + ('K4', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 3, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 36: rejected: vec 3 + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 1}), # 37: rejected: fused without prescale + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'bf16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 1, 'tpt': 4, 'bh': 4, 'bw': 8, 'lanes': 16, 'nt': 1, 'fused': 1}), # 38: rejected: fused on the LDS path + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 4, 'bh': 2, 'bw': 3, 'lanes': 16, 'nt': 1, 'fused': 0}), # 39: rejected: tiles_per_thread does not divide the LDS block + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 8, 'bw': 8, 'lanes': 32, 'nt': 1, 'fused': 0}), # 40: rejected: LDS region over 1024 threads + ('K1', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 4, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 41: rejected: filter_rows does not divide n + ('K4', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 96, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 42: rejected: block not a power of two + ('K2a', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 16, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 43: rejected: K2a vec 16 +] +# fmt: on + + +def _spec(idx: int): + if not 0 <= idx < len(CONFIGS): + raise SystemExit(f"unknown config index {idx}") + stage, arch, d = CONFIGS[idx] + if stage == "K1": + spec = RtFilterSpec( + mh=d["mh"], mw=d["mw"], io_dtype=d["io"], operand_dtype=d["op"], + vec=d["vec"], filter_rows=d["fr"], + ) # fmt: skip + elif stage == "K2a": + spec = RtPrescaleSpec(io_dtype=d["io"], vec=d["vec"]) + elif stage == "K2": + spec = RtInputSpec( + mh=d["mh"], mw=d["mw"], io_dtype=d["io"], operand_dtype=d["op"], + vec=d["vec"], block_size=d["block"], prescale=bool(d["ps"]), + tiles_per_thread=d["tpt"], lds=(d["bh"], d["bw"]), + lds_lanes=d["lanes"], nt_stores=bool(d["nt"]), fused=bool(d["fused"]), + ) # fmt: skip + else: + spec = RtOutputSpec( + mh=d["mh"], mw=d["mw"], io_dtype=d["io"], operand_dtype=d["op"], + vec=d["vec"], block_size=d["block"], prescale=bool(d["ps"]), + nt_stores=bool(d["nt"]), + ) # fmt: skip + return spec, arch + + +def main() -> int: + return run_emit( + _spec, + build_stage, + usage="usage: conv_winograd_fwd_rt_emit.py [ll|ir|verify]\n", + ) + + +if __name__ == "__main__": + raise SystemExit(main()) 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 b8539da5f34f..8bf3941e13b1 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 @@ -364,6 +364,11 @@ "kernels.common.conv_winograd_fwd.WinogradFwdSpec": ("problem",), "kernels.common.conv_winograd_gemm_wmma.WinogradGemmReduceSpec": (), "kernels.common.conv_winograd_gemm_wmma.WinogradGemmWmmaSpec": (), + # Winograd forward: shape-generic (AOT) transforms, every field defaulted + "kernels.common.conv_winograd_fwd_rt.RtFilterSpec": (), + "kernels.common.conv_winograd_fwd_rt.RtInputSpec": (), + "kernels.common.conv_winograd_fwd_rt.RtOutputSpec": (), + "kernels.common.conv_winograd_fwd_rt.RtPrescaleSpec": (), # Winograd forward: merged transform launch "kernels.common.conv_winograd_fwd.WinogradMergedSpec": ("fwd",), # Generalized Winograd forward diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py new file mode 100644 index 000000000000..4d4df5ea40fa --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py @@ -0,0 +1,168 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""CPU tests of the shape-generic Winograd transform kernels +(:mod:`kernels.common.conv_winograd_fwd_rt`). + +* ``magic_div``: the multiply-high division the kernels use is exact for the + 31-bit indices they divide; +* spec validation refuses what the kernels cannot build; +* the AOT grid: unique kernel names, every spec builds, and the default pick of + any problem is a member of it; +* ``stage_launch`` spells every scalar kernarg of the stage's signature, and + the kernel's parameters are that signature; +* refusals of problems a binary cannot run. +""" + +from __future__ import annotations + +import random + +import pytest + +from kernels.common import conv_winograd_fwd_rt as R +from kernels.common._conv_implicit_gemm_common import ConvProblem + + +def _q(n: int, d: int) -> int: + magic, shift = R.magic_div(d) + if magic == 0: + return n + return ((n * magic) >> 32) >> shift + + +@pytest.mark.parametrize("d", [1, 2, 3, 5, 7, 9, 31, 32, 33, 96, 255, 641, 1000, 65537]) +def test_magic_div_exact(d): + rng = random.Random(d) + ns = [0, 1, d - 1, d, d + 1, 2**31 - 1, 2**31 - d, 2**30] + ns += [rng.randrange(0, 2**31) for _ in range(2000)] + for n in ns: + assert _q(n, d) == n // d, (n, d) + + +def test_magic_div_every_small_divisor(): + rng = random.Random(0) + for d in range(1, 4097): + magic, _ = R.magic_div(d) + assert 0 <= magic < 2**32 + for n in (rng.randrange(0, 2**31) for _ in range(16)): + assert _q(n, d) == n // d + + +def test_magic_div_rejects_zero(): + with pytest.raises(ValueError): + R.magic_div(0) + + +@pytest.mark.parametrize( + "make", + [ + lambda: R.RtFilterSpec(mh=3), + lambda: R.RtFilterSpec(io_dtype="bf16", operand_dtype="bf16"), + lambda: R.RtFilterSpec(filter_rows=4), + lambda: R.RtOutputSpec(vec=3), + lambda: R.RtOutputSpec(block_size=96), + lambda: R.RtInputSpec(fused=True), + lambda: R.RtInputSpec(mh=2, mw=2, lds=(2, 3), tiles_per_thread=4), + lambda: R.RtInputSpec(mh=2, mw=2, lds=(8, 8), lds_lanes=32), + lambda: R.RtPrescaleSpec(vec=16), + ], +) +def test_invalid_specs_are_refused(make): + with pytest.raises(ValueError): + make() + + +def test_aot_grid_names_are_unique(): + specs = R.aot_stage_specs() + names = [s.kernel_name() for s in specs] + assert len(names) == len(set(names)) + stages = {s.stage for s in specs} + assert stages == {"K1", "K2a", "K2", "K4"} + + +def test_aot_grid_builds(): + """Every binary of the grid builds, with its signature as parameters.""" + for spec in R.aot_stage_specs(): + k = R.build_stage(spec, arch="gfx1250") + assert [p.name for p in k.params] == [ + a["name"] for a in R.stage_signature(spec) + ] + assert k.name == spec.kernel_name() + + +_PROBLEMS = [ + ConvProblem(N=2, Hi=20, Wi=24, C=64, K=96, Y=3, X=3, pH=1, pW=1), + ConvProblem(N=1, Hi=13, Wi=11, C=36, K=20, Y=3, X=3, pH=0, pW=2), + ConvProblem(N=1, Hi=40, Wi=40, C=128, K=64, Y=3, X=3, pH=2, pW=2, dH=2, dW=2), + ConvProblem(N=2, Hi=16, Wi=16, C=128, K=128, Y=3, X=3, pH=1, pW=1, groups=4), + ConvProblem(N=1, Hi=8, Wi=8, C=16384, K=64, Y=3, X=3, pH=1, pW=1), + ConvProblem(N=1, Hi=5, Wi=7, C=3, K=5, Y=3, X=3, pH=1, pW=1), +] + + +@pytest.mark.parametrize("problem", _PROBLEMS) +@pytest.mark.parametrize("io", ["fp16", "bf16"]) +def test_default_pick_is_in_the_grid(problem, io): + grid = {s.kernel_name() for s in R.aot_stage_specs()} + for tile in R.AOT_TILES: + specs = R.default_stage_specs(problem, tile, io) + _, prescale = R.pipeline_dtypes(tile, io) + assert ("K2a" in specs) == (prescale and not specs["K2"].fused) + geo = R.RtGeometry(problem, *tile) + for spec in specs.values(): + assert spec.kernel_name() in grid + assert spec.reason(geo) is None + + +def test_fused_prescale_falls_back_to_split_for_wide_c(): + p = _PROBLEMS[4] # C / vec = 4096 lanes: more than one workgroup + specs = R.default_stage_specs(p, (4, 4), "bf16") + assert "K2a" in specs and not specs["K2"].fused + narrow = R.default_stage_specs(_PROBLEMS[0], (4, 4), "bf16") + assert "K2a" not in narrow and narrow["K2"].fused + + +@pytest.mark.parametrize("problem", _PROBLEMS[:4]) +def test_stage_launch_spells_every_scalar(problem): + for tile in ((4, 4), (2, 2), (6, 4)): + geo = R.RtGeometry(problem, *tile) + for io in ("fp16", "bf16"): + for spec in R.default_stage_specs(problem, tile, io).values(): + grid, block, vals = R.stage_launch(spec, geo) + assert min(grid) >= 1 and 32 <= block <= 1024 + sig = [a["name"] for a in R.stage_signature(spec)] + pointers = {"w", "u", "x", "v", "e", "mm", "y"} + scalars = [ + n for n in sig if n not in pointers and not n.endswith("_bytes") + ] + assert sorted(scalars) == sorted(vals) + for name, val in vals.items(): + assert -(2**31) <= val < 2**31, name + if name.endswith("_mg"): + d = vals[name[: -len("_mg")]] + magic, shift = R.magic_div(d) + assert val % 2**32 == magic and vals[name[:-3] + "_sh"] == shift + + +def test_stage_launch_refuses_unfit_problems(): + p = ConvProblem(N=1, Hi=8, Wi=8, C=36, K=20, Y=3, X=3, pH=1, pW=1) + geo = R.RtGeometry(p, 4, 4) + with pytest.raises(ValueError, match="multiple of vec"): + R.stage_launch(R.RtInputSpec(vec=8), geo) + with pytest.raises(ValueError, match="multiple of vec"): + R.stage_launch(R.RtOutputSpec(vec=8), geo) + with pytest.raises(ValueError, match="tile mismatch"): + R.stage_launch(R.RtFilterSpec(mh=2, mw=2, vec=4), geo) + wide = ConvProblem(N=1, Hi=8, Wi=8, C=4096, K=64, Y=3, X=3, pH=1, pW=1) + with pytest.raises(ValueError, match="fused prescale"): + R.stage_launch( + R.RtInputSpec(io_dtype="bf16", vec=1, prescale=True, fused=True), + R.RtGeometry(wide, 4, 4), + ) + + +def test_geometry_refuses_non_3x3_and_huge_buffers(): + p5 = ConvProblem(N=1, Hi=8, Wi=8, C=8, K=8, Y=5, X=5, pH=2, pW=2) + assert R.RtGeometry(p5, 4, 4).check(False) + big = ConvProblem(N=64, Hi=512, Wi=512, C=256, K=256, Y=3, X=3, pH=1, pW=1) + assert "32-bit" in R.RtGeometry(big, 4, 4).check(False) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py new file mode 100644 index 000000000000..ef6b13a82d38 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py @@ -0,0 +1,253 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""On-GPU numeric test of the AOT Winograd forward route: the shape-generic +transform binaries of :mod:`kernels.common.conv_winograd_fwd_rt` through +``dispatch_conv_grouped`` (pinned ``algorithm="winograd_aot"``) -> +``launch_plan_for`` -> compile every step -> launch the steps in order -> +compare y with the conv reference (max|out - ref| / max(max|ref|, 1) < 5e-2, +the conv benchmarks' metric). + +The binaries are compiled once per kernel name and reused across the cases -- +the point of the shape-generic kernels. Besides the default pick, the cases +force the other binaries of the AOT grid (register runs of 2 tiles, the LDS +region, the split prescale, 1-wide channel vectors) on the same pipeline. + +Workspaces and y are poisoned (bytes 0x7F) first, so an element no kernel +writes shows up. Requires a gfx1250 GPU and torch. +""" + +from __future__ import annotations + +import dataclasses +import importlib.util + +import pytest + +_HAS_TORCH = importlib.util.find_spec("torch") is not None +if _HAS_TORCH: + import torch + + torch.cuda.is_available() # let torch claim the HIP context first + +from rocke.runtime.hip_module import get_device_arch # noqa: E402 + +GPU_ARCH = get_device_arch(0) + +pytestmark = pytest.mark.skipif( + not _HAS_TORCH or GPU_ARCH != "gfx1250", + reason="needs torch and a gfx1250 GPU", +) + +_TOL = 5e-2 +_ARTIFACTS: dict = {} + + +def _u8(t): + import ctypes + + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + +def _artifact(step): + """One compile per kernel name: a shape-generic binary serves every case.""" + from rocke import compile_kernel + + name = step.kernel.name + if name.startswith("K3"): + name = f"K3:{step.signature!r}:{step.block}" # the GEMM config's binary + if name not in _ARTIFACTS: + _ARTIFACTS[name] = compile_kernel(step.kernel.build(GPU_ARCH), arch=GPU_ARCH) + return _ARTIFACTS[name] + + +def _run(req, spec): + """Run ``spec`` (an AOT Winograd plan selected for ``req``); ``err``.""" + from benchmarks.common.benchmark_winograd_conv import ( + make_inputs, + reference, + verify_error, + ) + from dispatch.grouped_convolution import launch_plan_for + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + problem, dtype = spec.problem, spec.dtype + plan = launch_plan_for(req, spec) + launchers = [] + for step in plan.steps: + art = _artifact(step) + launchers.append( + KernelLauncher( + hsaco=art.hsaco, kernel_name=art.kernel_name, signature=step.signature + ) + ) + x_t, w_t, y_t = make_inputs(problem, dtype, seed=7) + ref = reference(x_t, w_t, problem, dtype) + sizes = {"x": x_t.nbytes, "w": w_t.nbytes, "y": y_t.nbytes, **plan.workspaces} + rt = Runtime() + ptrs = {k: int(rt.alloc(max(int(v), 4))) for k, v in sizes.items()} + try: + rt.sync() + for key, t in (("x", x_t), ("w", w_t)): + h = t.contiguous().pin_memory() + rt.memcpy_h2d(ptrs[key], _u8(h), h.nbytes) + rt.sync() + del h + for key in (*plan.workspaces, "y"): + rt.memset(ptrs[key], 0x7F, sizes[key]) + rt.sync() + for step, lau in zip(plan.steps, launchers): + lau( + plan.values(step, ptrs, sizes), + config=LaunchConfig(grid=step.grid, block=step.block, fence=True), + ) + rt.sync() + out = torch.empty_like(y_t, device="cpu").pin_memory() + rt.memcpy_d2h(_u8(out), ptrs["y"], out.nbytes) + rt.sync() + finally: + rt.sync() + for p in ptrs.values(): + rt.free(p) + err, _ = verify_error(out, ref) + return err + + +def _request(dtype="fp16", **kw): + from dispatch.grouped_convolution import ConvGroupedRequest + from kernels.conv.kernel_identity import current_llvm_flavor + + return ConvGroupedRequest( + arch=GPU_ARCH, + algorithm="winograd_aot", + llvm_flavor=current_llvm_flavor(), + dtype=dtype, + **kw, + ) + + +def _plans(req, tiles=None): + from dispatch.grouped_convolution import conv_grouped_sweep_space + + out = [s for s in conv_grouped_sweep_space(req) if s.plan.kind == "unfused2d_rt"] + if tiles is not None: + out = [s for s in out if s.plan.tile in tiles] + assert out, "no AOT plan" + return out + + +_CASES = [ + ("fp16", dict(N=2, C=64, K=96, Hi=20, Wi=24, Y=3, X=3, pad_h=1, pad_w=1)), + ("bf16", dict(N=2, C=64, K=96, Hi=20, Wi=24, Y=3, X=3, pad_h=1, pad_w=1)), + ( + "fp16", + dict( + N=1, + C=128, + K=64, + Hi=40, + Wi=40, + Y=3, + X=3, + pad_h=2, + pad_w=2, + dilation_h=2, + dilation_w=2, + ), + ), # fmt: skip + ("bf16", dict(N=2, C=128, K=128, Hi=16, Wi=16, Y=3, X=3, pad_h=1, pad_w=1, G=4)), + ("fp16", dict(N=1, C=64, K=64, Hi=17, Wi=23, Y=3, X=3, pad_h=0, pad_w=2)), +] + + +@pytest.mark.parametrize("dtype,geom", _CASES) +def test_aot_route_every_tile(dtype, geom): + req = _request(dtype, **geom) + for spec in _plans(req): + err = _run(req, spec) + assert err < _TOL, f"{spec.name}: err {err:.3e}" + + +def test_aot_selection_fp16_and_bf16(): + from dispatch.grouped_convolution import dispatch_conv_grouped + + for dtype in ("fp16", "bf16"): + req = _request(dtype, N=4, C=96, K=64, Hi=28, Wi=28, Y=3, X=3, pad_h=1, pad_w=1) + res = dispatch_conv_grouped(req) + assert res.spec.plan.kind == "unfused2d_rt" + assert _run(req, res.spec) < _TOL + + +def test_batch_split(): + from dispatch.conv_winograd_plan import winograd_steps + + req = _request(N=4, C=64, K=64, Hi=12, Wi=12, Y=3, X=3, pad_h=1, pad_w=1) + spec = _plans(req, tiles={(2, 2)})[0] + split = dataclasses.replace(spec, plan=dataclasses.replace(spec.plan, chunks=2)) + winograd_steps(split.problem, split.dtype, split.plan) # builds the plan + assert _run(req, split) < _TOL + + +_VARIANTS = { + "register_r2": lambda s, p: dataclasses.replace( + s, tiles_per_thread=2, lds=(0, 0), block_size=256 + ), + "lds": lambda s, p: dataclasses.replace( + s, tiles_per_thread=4, lds=(4, 8), lds_lanes=16 + ), + "vec1": lambda s, p: dataclasses.replace(s, vec=1), +} + + +@pytest.mark.parametrize("variant", sorted(_VARIANTS)) +@pytest.mark.parametrize("tile", [(2, 2), (4, 2)]) +def test_forced_input_variants(monkeypatch, variant, tile): + """The same pipeline with the K2 binary swapped for another grid member.""" + from dispatch.conv_winograd_plan import winograd_steps + from kernels.common import conv_winograd_fwd_rt as R + + pick = R.default_stage_specs + + def forced(problem, t, io): + specs = dict(pick(problem, t, io)) + specs["K2"] = _VARIANTS[variant](specs["K2"], problem) + if variant == "vec1": + specs["K1"] = dataclasses.replace(specs["K1"], vec=1) + specs["K4"] = dataclasses.replace(specs["K4"], vec=1) + return specs + + monkeypatch.setattr(R, "default_stage_specs", forced) + winograd_steps.cache_clear() + try: + req = _request(N=2, C=64, K=64, Hi=26, Wi=30, Y=3, X=3, pad_h=1, pad_w=1) + spec = _plans(req, tiles={tile})[0] + assert _run(req, spec) < _TOL + finally: + winograd_steps.cache_clear() + + +def test_split_prescale(monkeypatch): + """bf16 with the per-pixel max pass (K2a) instead of the fused prescale.""" + from dispatch.conv_winograd_plan import winograd_steps + from kernels.common import conv_winograd_fwd_rt as R + + pick = R.default_stage_specs + + def split(problem, t, io): + specs = dict(pick(problem, t, io)) + k2 = specs["K2"] + if k2.prescale: + specs["K2"] = dataclasses.replace(k2, fused=False, block_size=256) + specs["K2a"] = R.RtPrescaleSpec(io, 8) + return specs + + monkeypatch.setattr(R, "default_stage_specs", split) + winograd_steps.cache_clear() + try: + req = _request( + "bf16", N=2, C=64, K=96, Hi=20, Wi=24, Y=3, X=3, pad_h=1, pad_w=1 + ) + for spec in _plans(req, tiles={(4, 4), (6, 4)}): + assert _run(req, spec) < _TOL, spec.name + finally: + winograd_steps.cache_clear() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_fwd_rt.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_fwd_rt.h new file mode 100644 index 000000000000..08cffa794a69 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_fwd_rt.h @@ -0,0 +1,102 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_winograd_fwd_rt.h -- C++ engine port of the + * shape-generic Winograd F(mh x mw, 3 x 3) forward transform kernels in + * library/kernels/common/conv_winograd_fwd_rt.py. + * + * Python (conv_winograd_fwd_rt.py) C (this header) + * ------------------------------------ ------------------------------------- + * RtFilterSpec / RtPrescaleSpec / rocke_winograd_rt_spec_t (+ .stage) + * RtInputSpec / RtOutputSpec + * __post_init__ rocke_winograd_rt_is_valid_spec + * .kernel_name() rocke_winograd_rt_kernel_name + * build_stage(spec) rocke_build_winograd_rt(_new) + * magic_div(d) rocke_winograd_rt_magic_div + * + * Every problem extent is a kernel argument (divisions by runtime extents use + * a host-computed magic number), so one binary serves every problem its knobs + * fit. The launch geometry and kernarg values are host-side Python + * (stage_launch); the AOT cache stores the binaries built here. + */ +#ifndef ROCKE_INSTANCE_CONV_WINOGRAD_FWD_RT_H +#define ROCKE_INSTANCE_CONV_WINOGRAD_FWD_RT_H + +#include +#include +#include + +#include "rocke/ir.h" +#include "rocke/lower_llvm.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* The stage a spec builds (the Python spec class). */ +typedef enum rocke_winograd_rt_stage +{ + ROCKE_WINOGRAD_RT_K1 = 0, /* RtFilterSpec: filter transform */ + ROCKE_WINOGRAD_RT_K2A = 1, /* RtPrescaleSpec: per-pixel max |x| */ + ROCKE_WINOGRAD_RT_K2 = 2, /* RtInputSpec: input transform */ + ROCKE_WINOGRAD_RT_K4 = 3 /* RtOutputSpec: output transform */ +} rocke_winograd_rt_stage_t; + +/* The union of the four spec classes' fields (defaults in the comments); a + * stage reads only the fields its Python class has. */ +typedef struct rocke_winograd_rt_spec +{ + rocke_winograd_rt_stage_t stage; + int mh; /* 4 (K1, K2, K4) */ + int mw; /* 4 (K1, K2, K4) */ + const char* io_dtype; /* "fp16" */ + const char* operand_dtype; /* "fp16" (K1, K2, K4) */ + int vec; /* 8 */ + int filter_rows; /* 1 (K1) */ + int block_size; /* 256 (K2, K4) */ + bool prescale; /* false (K2, K4) */ + int tiles_per_thread; /* 1 (K2) */ + int lds[2]; /* (0, 0) (K2) */ + int lds_lanes; /* 16 (K2) */ + bool nt_stores; /* true (K2, K4) */ + bool fused; /* false (K2) */ +} rocke_winograd_rt_spec_t; + +/* Spec with the dataclass defaults for `stage`. */ +rocke_winograd_rt_spec_t rocke_winograd_rt_spec_default(rocke_winograd_rt_stage_t stage); + +/* __post_init__ of the stage's spec class: (ok, reason). */ +bool rocke_winograd_rt_is_valid_spec(const rocke_winograd_rt_spec_t* spec, + char* reason, + size_t reason_cap); + +/* spec.kernel_name(): the AOT identity of the binary. */ +rocke_status_t + rocke_winograd_rt_kernel_name(const rocke_winograd_rt_spec_t* spec, char* out, size_t out_cap); + +/* magic_div(d): (magic, shift) with n / d == umulhi(n, magic) >> shift for + * 0 <= n < 2**31; d == 1 gives (0, 0). False for d < 1. */ +bool rocke_winograd_rt_magic_div(int64_t d, uint32_t* magic, int* shift); + +/* build_stage(spec) into an initialised builder / a fresh one named like + * Python. NULL (with the builder's error set) for an invalid spec. */ +rocke_kernel_def_t* rocke_build_winograd_rt(rocke_ir_builder_t* b, + const rocke_winograd_rt_spec_t* spec, + const char* arch); +rocke_kernel_def_t* rocke_build_winograd_rt_new(rocke_ir_builder_t* b, + const rocke_winograd_rt_spec_t* spec, + const char* arch); + +/* Build + lower to LLVM IR. *out_ll is malloc'd (caller frees). */ +rocke_status_t rocke_winograd_rt_lower_to_llvm(const rocke_winograd_rt_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_WINOGRAD_FWD_RT_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_internal.h index f1e0206558d8..bab2959b5361 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_winograd_internal.h @@ -119,6 +119,24 @@ rocke_value_t* lincomb(rocke_ir_builder_t* b, rocke_value_t* lincomb_f(rocke_ir_builder_t* b, const std::vector>& terms); +/* The 1-D Winograd transforms with shared subexpressions (Python bt_1d / + * at_1d over _IROps): B^T d and A^T d for one length-n column, fp32. */ +std::vector + bt_1d(rocke_ir_builder_t* b, int m, const std::vector& d); +std::vector + at_1d(rocke_ir_builder_t* b, int m, const std::vector& d); +/* filter_matrix_int(m): integer numerators of G (row-major n x 3) and their + * common denominator. */ +int64_t filter_matrix_int(int m, std::vector& gi); +/* conv_winograd_fwd's prescale bit helpers: 2**e from an i32 exponent; |raw| + * of 16-bit floats as integers; their (vector) max; the horizontal max as an + * i32; a 16-bit float pattern as fp32. */ +rocke_value_t* pow2_f32(rocke_ir_builder_t* b, rocke_value_t* x); +rocke_value_t* abs_bits_vec(rocke_ir_builder_t* b, rocke_value_t* raw, int n); +rocke_value_t* max_bits(rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y, int n); +rocke_value_t* hmax_bits(rocke_ir_builder_t* b, rocke_value_t* v, int n); +rocke_value_t* bits_to_f32(rocke_ir_builder_t* b, rocke_value_t* bits, const char* dt); + } // namespace rocke_wino #endif /* ROCKE_INSTANCE_CONV_WINOGRAD_INTERNAL_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_common.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_common.cpp index 983a5de64cee..2c1f81412634 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_common.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_common.cpp @@ -10,6 +10,7 @@ */ #include "rocke/instance_conv_winograd_internal.h" +#include #include namespace rocke_wino @@ -928,4 +929,218 @@ rocke_value_t* lincomb_f(rocke_ir_builder_t* b, return acc != nullptr ? acc : rocke_b_const_f32(b, 0.0); } +using V = rocke_value_t*; +using VV = std::vector; + +namespace +{ + +/* _IROps: the scalar fp32 ops of the 1-D transforms. */ +V o_add(rocke_ir_builder_t* b, V x, V y) +{ + return rocke_b_fadd(b, x, y); +} +V o_sub(rocke_ir_builder_t* b, V x, V y) +{ + return rocke_b_fsub(b, x, y); +} +V o_fma(rocke_ir_builder_t* b, double c, V x, V y) +{ + V k = rocke_b_const_f32(b, c); + return rocke_b_fma(b, k, x, y); +} + +/* _matvec_1d over rows x cols of `mat` (row stride `cols`). */ +VV matvec_1d(rocke_ir_builder_t* b, const Rational* mat, int rows, int cols, const VV& d) +{ + VV out; + for(int i = 0; i < rows; ++i) + { + std::vector> unit, rest; + for(int k = 0; k < cols; ++k) + { + const Rational& r = mat[i * cols + k]; + if(r.is_zero()) + continue; + double c = r.value(); + if(fabs(c) == 1.0) + unit.push_back({c, d[k]}); + else + rest.push_back({c, d[k]}); + } + V acc = nullptr; + for(auto& t : unit) + { + if(acc == nullptr) + { + if(t.first > 0) + acc = t.second; + else + { + V z = o_sub(b, t.second, t.second); + acc = o_sub(b, z, t.second); + } + } + else + acc = t.first > 0 ? o_add(b, acc, t.second) : o_sub(b, acc, t.second); + } + for(auto& t : rest) + { + if(acc != nullptr) + acc = o_fma(b, t.first, t.second, acc); + else + { + V z = o_sub(b, t.second, t.second); + acc = o_fma(b, t.first, t.second, z); + } + } + out.push_back(acc != nullptr ? acc : o_sub(b, d[0], d[0])); + } + return out; +} + +} // namespace + +VV bt_1d(rocke_ir_builder_t* b, int m, const VV& d) +{ + if(m == 2) + { + V a = o_sub(b, d[2], d[0]); + V b1 = o_add(b, d[1], d[2]); + V c = o_sub(b, d[2], d[1]); + V dd = o_sub(b, d[3], d[1]); + return {a, b1, c, dd}; + } + if(m == 4) + { + V s12 = o_add(b, d[1], d[2]); + V s34 = o_add(b, d[3], d[4]); + V d12 = o_sub(b, d[1], d[2]); + V d43 = o_sub(b, d[4], d[3]); + V d31 = o_sub(b, d[3], d[1]); + V d42 = o_sub(b, d[4], d[2]); + VV out; + V t = o_fma(b, -5.0, d[2], d[4]); + out.push_back(o_fma(b, 4.0, d[0], t)); + out.push_back(o_fma(b, -4.0, s12, s34)); + out.push_back(o_fma(b, 4.0, d12, d43)); + out.push_back(o_fma(b, 2.0, d31, d42)); + out.push_back(o_fma(b, -2.0, d31, d42)); + t = o_fma(b, -5.0, d[3], d[5]); + out.push_back(o_fma(b, 4.0, d[1], t)); + return out; + } + const ToomCook* tc = toom_cook(m, 3); + return matvec_1d(b, tc->bt, tc->n, tc->n, d); +} + +VV at_1d(rocke_ir_builder_t* b, int m, const VV& d) +{ + if(m == 2) + { + V t = o_add(b, d[0], d[1]); + V a = o_add(b, t, d[2]); + V u = o_sub(b, d[1], d[2]); + V c = o_add(b, u, d[3]); + return {a, c}; + } + if(m == 4) + { + V s12 = o_add(b, d[1], d[2]); + V d12 = o_sub(b, d[1], d[2]); + V s34 = o_add(b, d[3], d[4]); + V d34 = o_sub(b, d[3], d[4]); + VV out; + V t = o_add(b, d[0], s12); + out.push_back(o_add(b, t, s34)); + out.push_back(o_fma(b, 2.0, d34, d12)); + out.push_back(o_fma(b, 4.0, s34, s12)); + t = o_fma(b, 8.0, d34, d12); + out.push_back(o_add(b, t, d[5])); + return out; + } + const ToomCook* tc = toom_cook(m, 3); + return matvec_1d(b, tc->at, tc->m, tc->n, d); +} + +V pow2_f32(rocke_ir_builder_t* b, V x) +{ + V c127 = rocke_b_const_i32(b, 127); + V a = rocke_b_add(b, x, c127); + V c23 = rocke_b_const_i32(b, 23); + V bits = rocke_b_shl(b, a, c23); + return rocke_b_bitcast(b, bits, rocke_f32()); +} + +V abs_bits_vec(rocke_ir_builder_t* b, V raw, int n) +{ + if(n == 1) + { + V bc = rocke_b_bitcast(b, raw, rocke_i16()); + V z = rocke_b_zext(b, bc, rocke_i32()); + V m = rocke_b_const_i32(b, 0x7FFF); + return rocke_b_land(b, z, m); + } + V bits = rocke_b_vec_bitcast(b, raw, rocke_vector_type(b, rocke_i16(), n)); + V c = rocke_b_const_i32(b, 0x7FFF); + V t = rocke_b_trunc(b, c, rocke_i16()); + V mask = rocke_b_vector_splat(b, t, n); + return rocke_b_vector_and(b, bits, mask); +} + +V max_bits(rocke_ir_builder_t* b, V x, V y, int n) +{ + return n == 1 ? rocke_b_smax(b, x, y) : rocke_b_vector_smax(b, x, y); +} + +V hmax_bits(rocke_ir_builder_t* b, V v, int n) +{ + if(n == 1) + return v; + V acc = nullptr; + for(int i = 0; i < n; ++i) + { + V x = rocke_b_vec_extract(b, v, i); + V z = rocke_b_zext(b, x, rocke_i32()); + acc = acc == nullptr ? z : rocke_b_smax(b, acc, z); + } + return acc; +} + +V bits_to_f32(rocke_ir_builder_t* b, V bits, const char* dt) +{ + if(strcmp(dt, "bf16") == 0) + { + V c16 = rocke_b_const_i32(b, 16); + V sh = rocke_b_shl(b, bits, c16); + return rocke_b_bitcast(b, sh, rocke_f32()); + } + V t = rocke_b_trunc(b, bits, rocke_i16()); + V h = rocke_b_bitcast(b, t, rocke_f16()); + return rocke_b_cast_to_f32(b, h); +} + +/* filter_matrix_int(m): integer numerators of G and their denominator. */ +int64_t filter_matrix_int(int m, std::vector& gi) +{ + const ToomCook* tc = toom_cook(m, 3); + int64_t den = 1; + for(int i = 0; i < tc->n * 3; ++i) + { + int64_t d = tc->g[i].den; + int64_t a = den, bb = d; + while(bb) + { + int64_t t = a % bb; + a = bb; + bb = t; + } + den = den * d / a; + } + gi.resize(tc->n * 3); + for(int i = 0; i < tc->n * 3; ++i) + gi[i] = tc->g[i].num * (den / tc->g[i].den); + return den; +} + } // namespace rocke_wino diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd.cpp index 3e6cd130f6b3..3d4059e08aaf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd.cpp @@ -499,130 +499,15 @@ void set_mwgs(E& e, int v) rocke_attr_set_int(e.b, &rocke_ir_builder_kernel(e.b)->attrs, "max_workgroup_size", v); } -/* _IROps */ -V o_add(E& e, V x, V y) -{ - return rocke_b_fadd(e.b, x, y); -} -V o_sub(E& e, V x, V y) -{ - return rocke_b_fsub(e.b, x, y); -} -V o_fma(E& e, double c, V x, V y) -{ - V k = cf(e, c); - return rocke_b_fma(e.b, k, x, y); -} - -/* _matvec_1d over rows x cols of `mat` (row stride `cols`). */ -VV matvec_1d(E& e, const rocke_wino::Rational* mat, int rows, int cols, const VV& d) -{ - VV out; - for(int i = 0; i < rows; ++i) - { - std::vector> unit, rest; - for(int k = 0; k < cols; ++k) - { - const rocke_wino::Rational& r = mat[i * cols + k]; - if(r.is_zero()) - continue; - double c = r.value(); - if(fabs(c) == 1.0) - unit.push_back({c, d[k]}); - else - rest.push_back({c, d[k]}); - } - V acc = nullptr; - for(auto& t : unit) - { - if(acc == nullptr) - { - if(t.first > 0) - acc = t.second; - else - { - V z = o_sub(e, t.second, t.second); - acc = o_sub(e, z, t.second); - } - } - else - acc = t.first > 0 ? o_add(e, acc, t.second) : o_sub(e, acc, t.second); - } - for(auto& t : rest) - { - if(acc != nullptr) - acc = o_fma(e, t.first, t.second, acc); - else - { - V z = o_sub(e, t.second, t.second); - acc = o_fma(e, t.first, t.second, z); - } - } - out.push_back(acc != nullptr ? acc : o_sub(e, d[0], d[0])); - } - return out; -} - +/* _IROps (the math helpers are shared: rocke_wino::bt_1d / at_1d / ...). */ VV bt_1d(E& e, int m, const VV& d) { - if(m == 2) - { - V a = o_sub(e, d[2], d[0]); - V b1 = o_add(e, d[1], d[2]); - V c = o_sub(e, d[2], d[1]); - V dd = o_sub(e, d[3], d[1]); - return {a, b1, c, dd}; - } - if(m == 4) - { - V s12 = o_add(e, d[1], d[2]); - V s34 = o_add(e, d[3], d[4]); - V d12 = o_sub(e, d[1], d[2]); - V d43 = o_sub(e, d[4], d[3]); - V d31 = o_sub(e, d[3], d[1]); - V d42 = o_sub(e, d[4], d[2]); - VV out; - V t = o_fma(e, -5.0, d[2], d[4]); - out.push_back(o_fma(e, 4.0, d[0], t)); - out.push_back(o_fma(e, -4.0, s12, s34)); - out.push_back(o_fma(e, 4.0, d12, d43)); - out.push_back(o_fma(e, 2.0, d31, d42)); - out.push_back(o_fma(e, -2.0, d31, d42)); - t = o_fma(e, -5.0, d[3], d[5]); - out.push_back(o_fma(e, 4.0, d[1], t)); - return out; - } - const rocke_wino::ToomCook* tc = rocke_wino::toom_cook(m, 3); - return matvec_1d(e, tc->bt, tc->n, tc->n, d); + return rocke_wino::bt_1d(e.b, m, d); } VV at_1d(E& e, int m, const VV& d) { - if(m == 2) - { - V t = o_add(e, d[0], d[1]); - V a = o_add(e, t, d[2]); - V u = o_sub(e, d[1], d[2]); - V c = o_add(e, u, d[3]); - return {a, c}; - } - if(m == 4) - { - V s12 = o_add(e, d[1], d[2]); - V d12 = o_sub(e, d[1], d[2]); - V s34 = o_add(e, d[3], d[4]); - V d34 = o_sub(e, d[3], d[4]); - VV out; - V t = o_add(e, d[0], s12); - out.push_back(o_add(e, t, s34)); - out.push_back(o_fma(e, 2.0, d34, d12)); - out.push_back(o_fma(e, 4.0, s34, s12)); - t = o_fma(e, 8.0, d34, d12); - out.push_back(o_add(e, t, d[5])); - return out; - } - const rocke_wino::ToomCook* tc = rocke_wino::toom_cook(m, 3); - return matvec_1d(e, tc->at, tc->m, tc->n, d); + return rocke_wino::at_1d(e.b, m, d); } V oob_sel(E& e, V ok, V off) @@ -636,62 +521,27 @@ V load_raw(E& e, V rsrc, V off, const char* dt, int n) V pow2_f32(E& e, V x) { - rocke_ir_builder_t* b = e.b; - V c127 = ci(e, 127); - V a = rocke_b_add(b, x, c127); - V c23 = ci(e, 23); - V bits = rocke_b_shl(b, a, c23); - return rocke_b_bitcast(b, bits, rocke_f32()); + return rocke_wino::pow2_f32(e.b, x); } V abs_bits_vec(E& e, V raw, int n) { - rocke_ir_builder_t* b = e.b; - if(n == 1) - { - V bc = rocke_b_bitcast(b, raw, rocke_i16()); - V z = rocke_b_zext(b, bc, rocke_i32()); - V m = ci(e, 0x7FFF); - return rocke_b_land(b, z, m); - } - V bits = rocke_b_vec_bitcast(b, raw, rocke_vector_type(b, rocke_i16(), n)); - V c = ci(e, 0x7FFF); - V t = rocke_b_trunc(b, c, rocke_i16()); - V mask = rocke_b_vector_splat(b, t, n); - return rocke_b_vector_and(b, bits, mask); + return rocke_wino::abs_bits_vec(e.b, raw, n); } V max_bits(E& e, V x, V y, int n) { - return n == 1 ? rocke_b_smax(e.b, x, y) : rocke_b_vector_smax(e.b, x, y); + return rocke_wino::max_bits(e.b, x, y, n); } V hmax_bits(E& e, V v, int n) { - if(n == 1) - return v; - V acc = nullptr; - for(int i = 0; i < n; ++i) - { - V x = rocke_b_vec_extract(e.b, v, i); - V z = rocke_b_zext(e.b, x, rocke_i32()); - acc = acc == nullptr ? z : rocke_b_smax(e.b, acc, z); - } - return acc; + return rocke_wino::hmax_bits(e.b, v, n); } V bits_to_f32(E& e, V bits, const char* dt) { - rocke_ir_builder_t* b = e.b; - if(strcmp(dt, "bf16") == 0) - { - V c16 = ci(e, 16); - V sh = rocke_b_shl(b, bits, c16); - return rocke_b_bitcast(b, sh, rocke_f32()); - } - V t = rocke_b_trunc(b, bits, rocke_i16()); - V h = rocke_b_bitcast(b, t, rocke_f16()); - return rocke_b_cast_to_f32(b, h); + return rocke_wino::bits_to_f32(e.b, bits, dt); } V prescale_exp(E& e, V amax) @@ -951,28 +801,7 @@ RunCoords run_coords(E& e, const RunMap& rm) /* ------------------------------------------------------------ K1 */ -/* filter_matrix_int(m): integer numerators of G and their denominator. */ -int64_t filter_matrix_int(int m, std::vector& gi) -{ - const rocke_wino::ToomCook* tc = rocke_wino::toom_cook(m, 3); - int64_t den = 1; - for(int i = 0; i < tc->n * 3; ++i) - { - int64_t d = tc->g[i].den; - int64_t a = den, bb = d; - while(bb) - { - int64_t t = a % bb; - a = bb; - bb = t; - } - den = den * d / a; - } - gi.resize(tc->n * 3); - for(int i = 0; i < tc->n * 3; ++i) - gi[i] = tc->g[i].num * (den / tc->g[i].den); - return den; -} +using rocke_wino::filter_matrix_int; void emit_filter(E& e) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp new file mode 100644 index 000000000000..f77db1cb495e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp @@ -0,0 +1,1772 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * conv_winograd_fwd_rt.cpp -- C++ port of the shape-generic Winograd + * F(mh x mw, 3 x 3) forward transform kernels + * (library/kernels/common/conv_winograd_fwd_rt.py): K1 filter transform, K2a + * per-pixel max, K2 input transform (register path with a fused or split + * prescale, LDS-region path) and K4 output transform, every problem extent a + * kernel argument. See the header and the Python module. + * + * Every emitter runs in PYTHON EVALUATION ORDER (arguments left to right, + * innermost call first): nested builder calls are hoisted into ordered + * locals, since C++ leaves the evaluation order of call arguments + * unspecified. + */ +#include "rocke/instance_conv_winograd_fwd_rt.h" + +#include +#include +#include + +#include +#include + +#include "rocke/error_boundary.hpp" /* ckc::guard_builder */ +#include "rocke/instance_conv_winograd_internal.h" +#include "rocke/ir_internal.h" /* rocke_i_set_err_msg */ +#include "rocke/lower_llvm.h" + +namespace +{ + +using rocke_wino::kCpolNt; +using Spec = rocke_winograd_rt_spec_t; +using V = rocke_value_t*; +using VV = std::vector; + +constexpr int kFilterBlock = 128; +constexpr int kPrescaleBlock = 256; +constexpr int kPrescaleMaxExp = 100; +constexpr int NOALIAS = 1, RO = 2, WO = 4; + +void set_reason(char* reason, size_t cap, const std::string& msg) +{ + if(reason != NULL && cap > 0) + snprintf(reason, cap, "%s", msg.c_str()); +} + +bool is_io(const char* d) +{ + return d != NULL && (strcmp(d, "fp16") == 0 || strcmp(d, "bf16") == 0); +} + +bool is_tile(int m) +{ + return m == 2 || m == 4 || m == 6; +} + +bool is_vec(int v) +{ + return v == 1 || v == 2 || v == 4 || v == 8; +} + +/* _check_common(mh, mw, io, op, vec) */ +bool check_common(const Spec* s, std::string* err) +{ + if(!is_tile(s->mh) || !is_tile(s->mw)) + { + *err = "tile must be in (2, 4, 6) (got " + std::to_string(s->mh) + "x" + + std::to_string(s->mw) + ")"; + return false; + } + if(!is_io(s->io_dtype) || !is_io(s->operand_dtype)) + { + *err = "dtypes must be fp16 / bf16"; + return false; + } + if(strcmp(s->operand_dtype, "bf16") == 0 + && (s->mh != 2 || s->mw != 2 || strcmp(s->io_dtype, "bf16") != 0)) + { + *err = "bf16 operands only for F(2,3) with bf16 I/O"; + return false; + } + if(!is_vec(s->vec)) + { + *err = "vec must be 1/2/4/8 (got " + std::to_string(s->vec) + ")"; + return false; + } + return true; +} + +bool check_block(const Spec* s, std::string* err) +{ + if(s->block_size < 32 || (s->block_size & (s->block_size - 1))) + { + *err = "block_size must be a power of two >= 32"; + return false; + } + return true; +} + +bool input_lds(const Spec* s) +{ + return s->lds[0] > 0; +} + +int lds_block(const Spec* s) +{ + return s->lds[0] * s->lds[1] * s->lds_lanes / s->tiles_per_thread; +} + +bool validate(const Spec* s, std::string* err) +{ + switch(s->stage) + { + case ROCKE_WINOGRAD_RT_K1: + if(!check_common(s, err)) + return false; + if(s->filter_rows < 1 || (s->mh + 2) % s->filter_rows) + { + *err = "filter_rows must divide n"; + return false; + } + return true; + case ROCKE_WINOGRAD_RT_K2A: + if(!is_io(s->io_dtype) || !is_vec(s->vec)) + { + *err = "bad K2a spec"; + return false; + } + return true; + case ROCKE_WINOGRAD_RT_K2: + if(!check_common(s, err) || !check_block(s, err)) + return false; + if(s->fused && (!s->prescale || input_lds(s))) + { + *err = "fused needs the prescale and the register path"; + return false; + } + if(s->tiles_per_thread < 1) + { + *err = "tiles_per_thread must be >= 1"; + return false; + } + if(input_lds(s)) + { + if((s->lds[0] * s->lds[1]) % s->tiles_per_thread) + { + *err = "tiles_per_thread must divide bh * bw"; + return false; + } + if(lds_block(s) > 1024) + { + *err = "LDS region too large for one workgroup"; + return false; + } + } + return true; + case ROCKE_WINOGRAD_RT_K4: + return check_common(s, err) && check_block(s, err); + } + *err = "unknown stage"; + return false; +} + +std::string tile_tag(const Spec* s) +{ + if(s->mh == s->mw) + return "f" + std::to_string(s->mh); + return "f" + std::to_string(s->mh) + "x" + std::to_string(s->mw); +} + +std::string join(const std::vector& parts) +{ + std::string out; + for(const std::string& p : parts) + { + if(p.empty()) + continue; + if(!out.empty()) + out += "_"; + out += p; + } + return out; +} + +std::string kernel_name(const Spec* s) +{ + std::string v = "v" + std::to_string(s->vec); + std::string op = std::string("op") + s->operand_dtype; + switch(s->stage) + { + case ROCKE_WINOGRAD_RT_K1: + return join({"wino_rt_filter", + tile_tag(s), + s->io_dtype, + op, + v, + "fr" + std::to_string(s->filter_rows)}); + case ROCKE_WINOGRAD_RT_K2A: + return join({"wino_rt_pmax", s->io_dtype, v}); + case ROCKE_WINOGRAD_RT_K2: + { + std::vector tags = {"wino_rt_input", tile_tag(s), s->io_dtype, op, v}; + if(input_lds(s)) + { + tags.push_back("lds" + std::to_string(s->lds[0]) + "x" + std::to_string(s->lds[1])); + tags.push_back("ll" + std::to_string(s->lds_lanes)); + tags.push_back("t" + std::to_string(s->tiles_per_thread)); + } + else + { + tags.push_back("b" + std::to_string(s->block_size)); + tags.push_back("r" + std::to_string(s->tiles_per_thread)); + } + if(s->prescale) + tags.push_back(s->fused ? "psf" : "ps"); + if(!s->nt_stores) + tags.push_back("rt"); + return join(tags); + } + case ROCKE_WINOGRAD_RT_K4: + { + std::vector tags = { + "wino_rt_output", tile_tag(s), s->io_dtype, op, v, "b" + std::to_string(s->block_size)}; + if(s->prescale) + tags.push_back("ps"); + if(!s->nt_stores) + tags.push_back("rt"); + return join(tags); + } + } + return ""; +} + +/* ------------------------------------------------------------ IR helpers */ + +struct E +{ + rocke_ir_builder_t* b; + const Spec* s; +}; + +V ci(E& e, int64_t v) +{ + return rocke_b_const_i32(e.b, v); +} + +/* b.param(name, t, noalias=..., readonly=..., writeonly=...) */ +V param(E& e, const char* name, const rocke_type_t* t, int flags) +{ + if(flags == 0) + return rocke_b_param(e.b, name, t, NULL); + rocke_param_opts_t o; + memset(&o, 0, sizeof o); + if(flags & NOALIAS) + { + o.noalias = true; + o.noalias_set = true; + } + if(flags & RO) + { + o.readonly = true; + o.readonly_set = true; + } + if(flags & WO) + { + o.writeonly = true; + o.writeonly_set = true; + } + return rocke_b_param(e.b, name, t, &o); +} + +V p32(E& e, const char* name) +{ + return param(e, name, rocke_i32(), 0); +} + +const rocke_type_t* gptr(E& e, const rocke_type_t* elem) +{ + return rocke_ptr_type(e.b, elem, "global"); +} + +void set_mwgs(E& e, int v) +{ + rocke_attr_set_int(e.b, &rocke_ir_builder_kernel(e.b)->attrs, "max_workgroup_size", v); +} + +V oob(E& e, V ok, V off) +{ + return rocke_wino::oob_sel(e.b, ok, off); +} + +/* _mul_c: v * k for a small constant k (0 and 1 folded). */ +V mul_c(E& e, V v, int64_t k) +{ + if(k == 0) + return ci(e, 0); + if(k == 1) + return v; + V c = ci(e, k); + return rocke_b_mul(e.b, v, c); +} + +/* _Div: a runtime divisor (kernargs , _mg, _sh). */ +struct Div +{ + V d, mg, sh; +}; + +Div div_params(E& e, const char* name) +{ + Div dv; + std::string n = name; + dv.d = p32(e, n.c_str()); + dv.mg = p32(e, (n + "_mg").c_str()); + dv.sh = p32(e, (n + "_sh").c_str()); + return dv; +} + +V div_q(E& e, const Div& dv, V n) +{ + rocke_ir_builder_t* b = e.b; + V hi = rocke_b_umul_hi_i32(b, n, dv.mg); + V q = rocke_b_lshr(b, hi, dv.sh); + V z = ci(e, 0); + V is1 = rocke_b_cmp_eq(b, dv.mg, z); + return rocke_b_select(b, is1, n, q); +} + +/* divmod(n) -> (n / d, n % d) */ +void div_mod(E& e, const Div& dv, V n, V* q, V* r) +{ + *q = div_q(e, dv, n); + V m = rocke_b_mul(e.b, *q, dv.d); + *r = rocke_b_sub(e.b, n, m); +} + +/* _prescale_exp_rt: clamp(ceil(log2(amax * ps_scale)), 0, 100). */ +V prescale_exp_rt(E& e, V amax, V ps_scale) +{ + rocke_ir_builder_t* b = e.b; + V sv = rocke_b_fmul(b, amax, ps_scale); + V bits = rocke_b_bitcast(b, sv, rocke_i32()); + V c23 = ci(e, 23); + V sh = rocke_b_lshr(b, bits, c23); + V cff = ci(e, 0xFF); + V an = rocke_b_land(b, sh, cff); + V c127 = ci(e, 127); + V expo = rocke_b_sub(b, an, c127); + V cm = ci(e, 0x7FFFFF); + V frac = rocke_b_land(b, bits, cm); + V c0 = ci(e, 0); + V ne = rocke_b_cmp_ne(b, frac, c0); + V c1 = ci(e, 1); + V c0b = ci(e, 0); + V sel = rocke_b_select(b, ne, c1, c0b); + V e_val = rocke_b_add(b, expo, sel); + V z = ci(e, 0); + V mx = rocke_b_smax(b, e_val, z); + V cmax = ci(e, kPrescaleMaxExp); + return rocke_b_smin(b, mx, cmax); +} + +/* ------------------------------------------------------------ K1 */ + +void emit_filter(E& e) +{ + rocke_ir_builder_t* b = e.b; + const Spec* s = e.s; + int vec = s->vec; + int n = s->mh + 2, nwv = s->mw + 2; + std::vector gint, gint_w; + int64_t den = rocke_wino::filter_matrix_int(s->mh, gint); + int64_t den_w = rocke_wino::filter_matrix_int(s->mw, gint_w); + double scale = 1.0 / (double)(den * den_w); + + set_mwgs(e, kFilterBlock); + V w = param(e, "w", gptr(e, rocke_wino::ir_dtype(s->io_dtype)), NOALIAS | RO); + V u = param(e, "u", gptr(e, rocke_wino::ir_dtype(s->operand_dtype)), NOALIAS | WO); + V w_bytes = p32(e, "w_bytes"); + V u_bytes = p32(e, "u_bytes"); + V K = p32(e, "K"); + V cg = p32(e, "cg"); + V items = p32(e, "items"); + Div cgv = div_params(e, "cgv"); + V w_rsrc = rocke_b_buffer_rsrc(b, w, w_bytes); + V u_rsrc = rocke_b_buffer_rsrc(b, u, u_bytes); + + V gid; + { + V bx = rocke_b_block_id_x(b); + V k = ci(e, kFilterBlock); + V mu = rocke_b_mul(b, bx, k); + V tx = rocke_b_thread_id_x(b); + gid = rocke_b_add(b, mu, tx); + } + V row_y = rocke_b_block_id_y(b); + V ok = rocke_b_cmp_lt(b, gid, items); + V kk, cm; + div_mod(e, cgv, gid, &kk, &cm); + V c = mul_c(e, cm, vec); + V cg2; + { + V k2 = ci(e, 2); + cg2 = rocke_b_mul(b, cg, k2); + } + V wbase; + { + V k9 = ci(e, 9); + V cg9 = rocke_b_mul(b, cg, k9); + V m1 = rocke_b_mul(b, kk, cg9); + V a1 = rocke_b_add(b, m1, c); + V k2 = ci(e, 2); + V m2 = rocke_b_mul(b, a1, k2); + wbase = oob(e, ok, m2); + } + std::vector> taps(3); + V off = wbase; + for(int ky = 0; ky < 3; ++ky) + { + for(int kx = 0; kx < 3; ++kx) + { + if(ky || kx) + off = rocke_b_add(b, off, cg2); + V raw = rocke_wino::load_raw(b, w_rsrc, off, s->io_dtype, vec); + taps[ky].push_back(rocke_wino::raw_to_f32(b, raw, vec)); + } + } + V sc = rocke_b_const_f32(b, scale); + int nr = s->filter_rows; + V kcg2 = rocke_b_mul(b, K, cg2); + for(int rr = 0; rr < nr; ++rr) + { + V row_i; + VV coef; + if(nr == n) + { + row_i = ci(e, rr); + for(int ky = 0; ky < 3; ++ky) + coef.push_back(rocke_b_const_f32(b, (double)gint[rr * 3 + ky])); + } + else + { + V m1 = mul_c(e, row_y, nr); + V k = ci(e, rr); + row_i = rocke_b_add(b, m1, k); + for(int ky = 0; ky < 3; ++ky) + { + V cval = rocke_b_const_f32(b, (double)gint[ky]); + for(int i = 1; i < n; ++i) + { + V ki = ci(e, i); + V eq = rocke_b_cmp_eq(b, row_i, ki); + V cv2 = rocke_b_const_f32(b, (double)gint[i * 3 + ky]); + cval = rocke_b_select(b, eq, cv2, cval); + } + coef.push_back(cval); + } + } + std::vector tmp(3); + for(int kx = 0; kx < 3; ++kx) + { + for(int ch = 0; ch < vec; ++ch) + { + V f0 = rocke_b_fmul(b, coef[0], taps[0][kx][ch]); + V f1 = rocke_b_fma(b, coef[1], taps[1][kx][ch], f0); + tmp[kx].push_back(rocke_b_fma(b, coef[2], taps[2][kx][ch], f1)); + } + } + V ubase; + { + V rn = mul_c(e, row_i, nwv); + V m1 = rocke_b_mul(b, rn, kcg2); + V m2 = rocke_b_mul(b, kk, cg); + V a2 = rocke_b_add(b, m2, c); + V k2 = ci(e, 2); + V m3 = rocke_b_mul(b, a2, k2); + V a3 = rocke_b_add(b, m1, m3); + ubase = oob(e, ok, a3); + } + V uoff = ubase; + for(int j = 0; j < nwv; ++j) + { + if(j) + uoff = rocke_b_add(b, uoff, kcg2); + VV vals; + for(int ch = 0; ch < vec; ++ch) + { + std::vector> terms; + for(int kx = 0; kx < 3; ++kx) + terms.push_back({gint_w[j * 3 + kx], tmp[kx][ch]}); + V lc = rocke_wino::lincomb(b, terms); + vals.push_back(rocke_b_fmul(b, lc, sc)); + } + rocke_wino::store_vec_f32(b, u_rsrc, uoff, vals, s->operand_dtype); + } + } +} + +/* ------------------------------------------------------------ K2a */ + +void emit_prescale(E& e) +{ + rocke_ir_builder_t* b = e.b; + const Spec* s = e.s; + int vec = s->vec; + set_mwgs(e, kPrescaleBlock); + V x = param(e, "x", gptr(e, rocke_wino::ir_dtype(s->io_dtype)), NOALIAS | RO); + V e_ptr = param(e, "e", gptr(e, rocke_i32()), NOALIAS | WO); + V x_bytes = p32(e, "x_bytes"); + V e_bytes = p32(e, "e_bytes"); + V C = p32(e, "C"); + V nvec = p32(e, "nvec"); + V tpc_lg = p32(e, "tpc_lg"); + V loads = p32(e, "loads"); + V npix = p32(e, "npix"); + V T = p32(e, "T"); + V x_rsrc = rocke_b_buffer_rsrc(b, x, x_bytes); + V e_rsrc = rocke_b_buffer_rsrc(b, e_ptr, e_bytes); + + V tid = rocke_b_thread_id_x(b); + V tpc; + { + V c1 = ci(e, 1); + tpc = rocke_b_shl(b, c1, tpc_lg); + } + V lc; + { + V c1 = ci(e, 1); + V t1 = rocke_b_sub(b, tpc, c1); + lc = rocke_b_land(b, tid, t1); + } + V lp = rocke_b_lshr(b, tid, tpc_lg); + V ppw; + { + V cb = ci(e, kPrescaleBlock); + ppw = rocke_b_lshr(b, cb, tpc_lg); + } + V pix; + { + V bx = rocke_b_block_id_x(b); + V mu = rocke_b_mul(b, bx, ppw); + pix = rocke_b_add(b, mu, lp); + } + V pok = rocke_b_cmp_lt(b, pix, npix); + V base; + { + V m1 = rocke_b_mul(b, pix, C); + V m2 = mul_c(e, lc, vec); + V a1 = rocke_b_add(b, m1, m2); + V k2 = ci(e, 2); + base = rocke_b_mul(b, a1, k2); + } + V step_b; + { + V k = ci(e, vec * 2); + step_b = rocke_b_mul(b, tpc, k); + } + V zero = ci(e, 0); + V one = ci(e, 1); + rocke_iter_arg_t ia[1]; + ia[0].name = "pm_acc"; + ia[0].init = zero; + rocke_for_t loop = rocke_b_scf_for_iter(b, zero, loads, one, ia, 1, "pm_q", false, true); + rocke_b_region_enter(b, loop.body); + { + V q = loop.iv; + V acc = loop.iter_vars[0]; + V mq = rocke_b_mul(b, q, tpc); + V idx = rocke_b_add(b, lc, mq); + V lt = rocke_b_cmp_lt(b, idx, nvec); + V ok = rocke_b_land(b, pok, lt); + V ms = rocke_b_mul(b, q, step_b); + V ad = rocke_b_add(b, base, ms); + V off = oob(e, ok, ad); + V raw = rocke_wino::load_raw(b, x_rsrc, off, s->io_dtype, vec); + V ab = rocke_wino::abs_bits_vec(b, raw, vec); + V a_ = rocke_wino::hmax_bits(b, ab, vec); + V nx = rocke_b_smax(b, acc, a_); + rocke_b_scf_yield(b, &nx, 1); + } + rocke_b_region_leave(b); + V mx = loop.op->results[0]; + for(int sft = 1; sft < 32; sft *= 2) + { + V shv = rocke_b_warp_shuffle_xor(b, mx, sft); + V ks = ci(e, sft); + V lt = rocke_b_cmp_lt(b, ks, tpc); + V m = rocke_b_smax(b, mx, shv); + mx = rocke_b_select(b, lt, m, mx); + } + V ok_s; + { + V c0 = ci(e, 0); + V eq = rocke_b_cmp_eq(b, lc, c0); + ok_s = rocke_b_land(b, pok, eq); + } + V a1 = rocke_b_add(b, pix, T); + V k4 = ci(e, 4); + V m4 = rocke_b_mul(b, a1, k4); + V off = oob(e, ok_s, m4); + V z = ci(e, 0); + V bc = rocke_b_bitcast(b, mx, rocke_f32()); + rocke_b_buffer_store_f32(b, e_rsrc, off, z, bc); +} + +/* ------------------------------------------------------------ run coords */ + +struct Coords +{ + V ok, cv, c_ok, n, a, bsub, ti, tj0, t0, tiles_w, tpc; +}; + +/* _run_coords(b, tiles): the kernargs of the run map and the thread's run. */ +Coords run_coords(E& e, int tiles) +{ + rocke_ir_builder_t* b = e.b; + Coords rc; + V nvec = p32(e, "nvec"); + V tpc_lg = p32(e, "tpc_lg"); + V rpw = p32(e, "rpw"); + V num_runs = p32(e, "num_runs"); + V tiles_w = p32(e, "tiles_w"); + Div cx = div_params(e, "cx"); + Div rpr = div_params(e, "rpr"); + Div th = div_params(e, "th"); + Div ddw = div_params(e, "ddw"); + Div ddh = div_params(e, "ddh"); + V tid = rocke_b_thread_id_x(b); + V bid = rocke_b_block_id_x(b); + V tpc; + { + V c1 = ci(e, 1); + tpc = rocke_b_shl(b, c1, tpc_lg); + } + V lc; + { + V c1 = ci(e, 1); + V t1 = rocke_b_sub(b, tpc, c1); + lc = rocke_b_land(b, tid, t1); + } + V lr = rocke_b_lshr(b, tid, tpc_lg); + V rb, cxi; + div_mod(e, cx, bid, &rb, &cxi); + { + V mu = rocke_b_mul(b, cxi, tpc); + rc.cv = rocke_b_add(b, mu, lc); + } + rc.c_ok = rocke_b_cmp_lt(b, rc.cv, nvec); + V rho; + { + V mu = rocke_b_mul(b, rb, rpw); + rho = rocke_b_add(b, mu, lr); + } + rc.ok = rocke_b_cmp_lt(b, rho, num_runs); + V row, jr; + div_mod(e, rpr, rho, &row, &jr); + rc.tj0 = mul_c(e, jr, tiles); + { + V mu = rocke_b_mul(b, row, tiles_w); + rc.t0 = rocke_b_add(b, mu, rc.tj0); + } + V r; + div_mod(e, th, row, &r, &rc.ti); + div_mod(e, ddw, r, &r, &rc.bsub); + div_mod(e, ddh, r, &r, &rc.a); + rc.n = r; + rc.tiles_w = tiles_w; + rc.tpc = tpc; + return rc; +} + +/* _group_max_rt: max of each i32 over the tpc (runtime) threads of a run. */ +VV group_max_rt(E& e, const VV& vals, V tpc, int block) +{ + rocke_ir_builder_t* b = e.b; + VV out; + for(V v : vals) + { + for(int sft = 1; sft < 32; sft *= 2) + { + V shv = rocke_b_warp_shuffle_xor(b, v, sft); + V ks = ci(e, sft); + V lt = rocke_b_cmp_lt(b, ks, tpc); + V m = rocke_b_smax(b, v, shv); + v = rocke_b_select(b, lt, m, v); + } + out.push_back(v); + } + int rows = block / 32; + if(rows <= 1) + return out; + int nv = (int)out.size(); + int shape[1] = {rows * nv}; + V red = rocke_b_smem_alloc_f32(b, shape, 1, "wino_rt_amax"); + V tid = rocke_b_thread_id_x(b); + V row_id; + { + V c5 = ci(e, 5); + row_id = rocke_b_lshr(b, tid, c5); + } + { + V c31 = ci(e, 31); + V la = rocke_b_land(b, tid, c31); + V c0 = ci(e, 0); + V cond = rocke_b_cmp_eq(b, la, c0); + rocke_if_t g = rocke_b_scf_if(b, cond); + rocke_b_region_enter(b, g.then_region); + for(int i = 0; i < nv; ++i) + { + V m = mul_c(e, row_id, nv); + V k = ci(e, i); + V idx = rocke_b_add(b, m, k); + V bc = rocke_b_bitcast(b, out[i], rocke_f32()); + V pk = rocke_b_vec_pack(b, &bc, 1, rocke_f32()); + rocke_b_smem_store_vN_f32(b, red, &idx, 1, pk, 1); + } + rocke_b_region_leave(b); + } + rocke_b_sync(b); + V per; + { + V c5 = ci(e, 5); + per = rocke_b_lshr(b, tpc, c5); + } + V big; + { + V c32 = ci(e, 32); + big = rocke_b_cmp_lt(b, c32, tpc); + } + V first; + { + V c1 = ci(e, 1); + V sp = rocke_b_select(b, big, per, c1); + V dv = rocke_b_div(b, row_id, sp); + V s1 = rocke_b_select(b, big, dv, row_id); + V c1b = ci(e, 1); + V sp2 = rocke_b_select(b, big, per, c1b); + first = rocke_b_mul(b, s1, sp2); + } + VV res; + for(int i = 0; i < nv; ++i) + { + V acc = out[i]; + for(int r = 0; r < rows; ++r) + { + V kr = ci(e, r); + V fr = rocke_b_add(b, first, kr); + V m = mul_c(e, fr, nv); + V ki = ci(e, i); + V idx = rocke_b_add(b, m, ki); + V kr2 = ci(e, r); + V lt = rocke_b_cmp_lt(b, kr2, per); + V okr = rocke_b_land(b, big, lt); + V z = ci(e, 0); + idx = rocke_b_select(b, okr, idx, z); + V ld = rocke_b_smem_load_vN_f32(b, red, &idx, 1, 1); + V ex = rocke_b_vec_extract(b, ld, 0); + V xv = rocke_b_bitcast(b, ex, rocke_i32()); + V mm = rocke_b_smax(b, acc, xv); + acc = rocke_b_select(b, okr, mm, acc); + } + res.push_back(acc); + } + return res; +} + +/* ------------------------------------------------------------ K2 */ + +const char* const kInputGeom[] + = {"C", "Hi", "Wi", "pH", "pW", "dH", "dW", "dHm", "dWm", "T", "ps_bits"}; +enum +{ + G_C, + G_HI, + G_WI, + G_PH, + G_PW, + G_DH, + G_DW, + G_DHM, + G_DWM, + G_T, + G_PS, + G_N +}; + +/* _input_tiles_transform: V = B^T d B of R tiles adjacent along W. */ +void input_tiles_transform(E& e, + const std::vector& raw, + int R, + const VV& tile_ok, + const VV& scales, + V vbase0, + V v_rsrc, + V cb2) +{ + rocke_ir_builder_t* b = e.b; + const Spec* s = e.s; + int n = s->mh + 2, m = s->mh; + int nwv = s->mw + 2, mw = s->mw; + int vec = s->vec; + V vstride = mul_c(e, cb2, n * nwv); + std::vector> colt(raw.empty() ? 0 : raw[0].size()); + std::vector have(colt.size(), false); + std::vector> pending; + rocke_b_sched_barrier(b, 0); + + auto col_transform = [&](int vv) -> const std::vector& { + if(!have[vv]) + { + std::vector data; + for(int uu = 0; uu < n; ++uu) + data.push_back(rocke_wino::raw_to_f32(b, raw[uu][vv], vec)); + std::vector per_ch; + for(int ch = 0; ch < vec; ++ch) + { + VV col; + for(int uu = 0; uu < n; ++uu) + col.push_back(data[uu][ch]); + per_ch.push_back(rocke_wino::bt_1d(b, m, col)); + } + colt[vv] = per_ch; + have[vv] = true; + } + return colt[vv]; + }; + + for(int r = 0; r < R; ++r) + { + std::vector> cols; + for(int vv = 0; vv < nwv; ++vv) + cols.push_back(col_transform(mw * r + vv)); + V base = vbase0; + if(r != 0) + { + V ms = mul_c(e, vstride, r); + base = rocke_b_add(b, vbase0, ms); + } + V off = oob(e, tile_ok[r], base); + bool first = true; + for(int i = 0; i < n; ++i) + { + std::vector outs; + for(int ch = 0; ch < vec; ++ch) + { + VV row; + for(int vv = 0; vv < nwv; ++vv) + row.push_back(cols[vv][ch][i]); + outs.push_back(rocke_wino::bt_1d(b, mw, row)); + } + for(int j = 0; j < nwv; ++j) + { + if(!first) + off = rocke_b_add(b, off, cb2); + first = false; + VV vals; + for(int ch = 0; ch < vec; ++ch) + vals.push_back(outs[ch][j]); + if(scales[r] != nullptr) + { + for(auto& x_ : vals) + x_ = rocke_b_fmul(b, x_, scales[r]); + } + pending.push_back({off, rocke_wino::pack_vec_f32(b, vals, s->operand_dtype)}); + } + } + } + rocke_b_sched_barrier(b, 0); + int cpol = s->nt_stores ? kCpolNt : 0; + for(auto& pr : pending) + rocke_wino::store_packed(b, v_rsrc, pr.first, pr.second, s->operand_dtype, cpol); +} + +/* K2's leading kernargs: the buffers, their sizes, the geometry. */ +struct InputParams +{ + V x, v, e_ptr, x_bytes, v_bytes, e_bytes; + V g[G_N]; +}; + +InputParams input_params(E& e) +{ + const Spec* s = e.s; + InputParams ip; + ip.x = param(e, "x", gptr(e, rocke_wino::ir_dtype(s->io_dtype)), NOALIAS | RO); + ip.v = param(e, "v", gptr(e, rocke_wino::ir_dtype(s->operand_dtype)), NOALIAS | WO); + ip.e_ptr = param(e, "e", gptr(e, rocke_i32()), NOALIAS); + ip.x_bytes = p32(e, "x_bytes"); + ip.v_bytes = p32(e, "v_bytes"); + ip.e_bytes = p32(e, "e_bytes"); + for(int i = 0; i < G_N; ++i) + ip.g[i] = p32(e, kInputGeom[i]); + return ip; +} + +/* The prescale's tile store: e[t] (lane 0 of a tile) and the 2**-e scale. */ +V store_e(E& e, V e_rsrc, V lane0, V tok, V t, V e_t) +{ + rocke_ir_builder_t* b = e.b; + V la = rocke_b_land(b, lane0, tok); + V k4 = ci(e, 4); + V m4 = rocke_b_mul(b, t, k4); + V off = oob(e, la, m4); + V z = ci(e, 0); + V bc = rocke_b_bitcast(b, e_t, rocke_f32()); + rocke_b_buffer_store_f32(b, e_rsrc, off, z, bc); + V z2 = ci(e, 0); + V ne = rocke_b_sub(b, z2, e_t); + return rocke_wino::pow2_f32(b, ne); +} + +bool emit_input_lds(E& e, std::string* err); + +bool emit_input(E& e, std::string* err) +{ + const Spec* s = e.s; + if(input_lds(s)) + return emit_input_lds(e, err); + rocke_ir_builder_t* b = e.b; + int n = s->mh + 2; + int nwv = s->mw + 2, mw = s->mw; + int vec = s->vec; + int R = s->tiles_per_thread; + int ncols = mw * R + (nwv - mw); + bool ps = s->prescale; + + set_mwgs(e, s->block_size); + InputParams ip = input_params(e); + V* g = ip.g; + V x_rsrc = rocke_b_buffer_rsrc(b, ip.x, ip.x_bytes); + V v_rsrc = rocke_b_buffer_rsrc(b, ip.v, ip.v_bytes); + V e_rsrc = ps ? rocke_b_buffer_rsrc(b, ip.e_ptr, ip.e_bytes) : nullptr; + Coords rc = run_coords(e, R); + + V C = g[G_C], Hi = g[G_HI], Wi = g[G_WI]; + V cb2; + { + V k2 = ci(e, 2); + cb2 = rocke_b_mul(b, C, k2); + } + V c = mul_c(e, rc.cv, vec); + V h0; + { + V s1 = rocke_b_sub(b, rc.a, g[G_PH]); + V m1 = rocke_b_mul(b, rc.ti, g[G_DHM]); + h0 = rocke_b_add(b, s1, m1); + } + V w0; + { + V s1 = rocke_b_sub(b, rc.bsub, g[G_PW]); + V m1 = rocke_b_mul(b, rc.tj0, g[G_DWM]); + w0 = rocke_b_add(b, s1, m1); + } + V img_row0 = rocke_b_mul(b, rc.n, Hi); + VV row_ok, row_pix; + V hi = h0; + for(int uu = 0; uu < n; ++uu) + { + if(uu) + hi = rocke_b_add(b, hi, g[G_DH]); + V lt = rocke_b_cmp_ult(b, hi, Hi); + row_ok.push_back(rocke_b_land(b, rc.ok, lt)); + V a1 = rocke_b_add(b, img_row0, hi); + row_pix.push_back(rocke_b_mul(b, a1, Wi)); + } + VV col_ok, col_w; + V wi = w0; + for(int vv = 0; vv < ncols; ++vv) + { + if(vv) + wi = rocke_b_add(b, wi, g[G_DW]); + col_ok.push_back(rocke_b_cmp_ult(b, wi, Wi)); + col_w.push_back(wi); + } + V c2 = mul_c(e, c, 2); + VV xrow, xcol; + for(int uu = 0; uu < n; ++uu) + { + V la = rocke_b_land(b, row_ok[uu], rc.c_ok); + V m1 = rocke_b_mul(b, row_pix[uu], cb2); + V a1 = rocke_b_add(b, m1, c2); + xrow.push_back(oob(e, la, a1)); + } + for(int vv = 0; vv < ncols; ++vv) + { + V m1 = rocke_b_mul(b, col_w[vv], cb2); + xcol.push_back(oob(e, col_ok[vv], m1)); + } + std::vector raw(n); + for(int uu = 0; uu < n; ++uu) + { + for(int vv = 0; vv < ncols; ++vv) + { + V ad = rocke_b_add(b, xrow[uu], xcol[vv]); + raw[uu].push_back(rocke_wino::load_raw(b, x_rsrc, ad, s->io_dtype, vec)); + } + } + V base_ok = rocke_b_land(b, rc.ok, rc.c_ok); + VV tile_ok; + for(int r = 0; r < R; ++r) + { + if(r == 0) + { + tile_ok.push_back(base_ok); + continue; + } + V kr = ci(e, r); + V a1 = rocke_b_add(b, rc.tj0, kr); + V lt = rocke_b_cmp_lt(b, a1, rc.tiles_w); + tile_ok.push_back(rocke_b_land(b, base_ok, lt)); + } + + VV scales(R, nullptr); + if(ps && s->fused) + { + VV colmax; + for(int vv = 0; vv < ncols; ++vv) + { + V acc = nullptr; + for(int uu = 0; uu < n; ++uu) + { + V a_ = rocke_wino::abs_bits_vec(b, raw[uu][vv], vec); + acc = acc == nullptr ? a_ : rocke_wino::max_bits(b, acc, a_, vec); + } + colmax.push_back(acc); + } + VV tmax; + for(int r = 0; r < R; ++r) + { + V acc = colmax[mw * r]; + for(int vv = 1; vv < nwv; ++vv) + acc = rocke_wino::max_bits(b, acc, colmax[mw * r + vv], vec); + tmax.push_back(rocke_wino::hmax_bits(b, acc, vec)); + } + tmax = group_max_rt(e, tmax, rc.tpc, s->block_size); + V ps_scale = rocke_b_bitcast(b, g[G_PS], rocke_f32()); + V lane0; + { + V c0 = ci(e, 0); + lane0 = rocke_b_cmp_eq(b, rc.cv, c0); + } + for(int r = 0; r < R; ++r) + { + V f = rocke_wino::bits_to_f32(b, tmax[r], s->io_dtype); + V e_r = prescale_exp_rt(e, f, ps_scale); + V t_r = rc.t0; + if(r) + { + V kr = ci(e, r); + t_r = rocke_b_add(b, rc.t0, kr); + } + scales[r] = store_e(e, e_rsrc, lane0, tile_ok[r], t_r, e_r); + } + } + else if(ps) + { + V T = g[G_T]; + VV prow, pcol; + for(int uu = 0; uu < n; ++uu) + { + V a1 = rocke_b_add(b, row_pix[uu], T); + V k4 = ci(e, 4); + V m1 = rocke_b_mul(b, a1, k4); + prow.push_back(oob(e, row_ok[uu], m1)); + } + for(int vv = 0; vv < ncols; ++vv) + { + V k4 = ci(e, 4); + V m1 = rocke_b_mul(b, col_w[vv], k4); + pcol.push_back(oob(e, col_ok[vv], m1)); + } + VV colmax; + for(int vv = 0; vv < ncols; ++vv) + { + V acc = nullptr; + for(int uu = 0; uu < n; ++uu) + { + V ad = rocke_b_add(b, prow[uu], pcol[vv]); + V z = ci(e, 0); + V pm = rocke_b_buffer_load(b, e_rsrc, ad, z, rocke_i32()); + acc = acc == nullptr ? pm : rocke_b_smax(b, acc, pm); + } + colmax.push_back(acc); + } + V ps_scale = rocke_b_bitcast(b, g[G_PS], rocke_f32()); + V lane0; + { + V c0 = ci(e, 0); + lane0 = rocke_b_cmp_eq(b, rc.cv, c0); + } + for(int r = 0; r < R; ++r) + { + V acc = colmax[mw * r]; + for(int vv = 1; vv < nwv; ++vv) + acc = rocke_b_smax(b, acc, colmax[mw * r + vv]); + V f = rocke_wino::bits_to_f32(b, acc, s->io_dtype); + V e_r = prescale_exp_rt(e, f, ps_scale); + V t_r = rc.t0; + if(r) + { + V kr = ci(e, r); + t_r = rocke_b_add(b, rc.t0, kr); + } + scales[r] = store_e(e, e_rsrc, lane0, tile_ok[r], t_r, e_r); + } + } + + V vbase0; + { + V st = mul_c(e, cb2, n * nwv); + V m1 = rocke_b_mul(b, rc.t0, st); + vbase0 = rocke_b_add(b, m1, c2); + } + input_tiles_transform(e, raw, R, tile_ok, scales, vbase0, v_rsrc, cb2); + (void)err; + return true; +} + +bool emit_input_lds(E& e, std::string* err) +{ + rocke_ir_builder_t* b = e.b; + const Spec* s = e.s; + int n = s->mh + 2, m = s->mh; + int nwv = s->mw + 2, mw = s->mw; + int vec = s->vec; + int bh = s->lds[0], bw = s->lds[1]; + int tpc = s->lds_lanes; + int cc = tpc * vec; + int lv = vec; + int rh = m * bh + n - m; + int rw = mw * bw + nwv - mw; + int block = lds_block(s); + int tpt = s->tiles_per_thread; + int ppx = cc / lv; + int npx = rh * rw; + int loads = (npx * ppx + block - 1) / block; + bool ps = s->prescale; + const rocke_type_t* dt = rocke_wino::ir_dtype(s->io_dtype); + + set_mwgs(e, block); + InputParams ip = input_params(e); + V* g = ip.g; + V nvec = p32(e, "nvec"); + V tiles_h = p32(e, "tiles_h"); + V tiles_w = p32(e, "tiles_w"); + Div cx = div_params(e, "cx"); + Div nbw = div_params(e, "nbw"); + Div nbh = div_params(e, "nbh"); + Div ddw = div_params(e, "ddw"); + Div ddh = div_params(e, "ddh"); + V x_rsrc = rocke_b_buffer_rsrc(b, ip.x, ip.x_bytes); + V v_rsrc = rocke_b_buffer_rsrc(b, ip.v, ip.v_bytes); + V e_rsrc = ps ? rocke_b_buffer_rsrc(b, ip.e_ptr, ip.e_bytes) : nullptr; + V C = g[G_C], Hi = g[G_HI], Wi = g[G_WI]; + V cb2; + { + V k2 = ci(e, 2); + cb2 = rocke_b_mul(b, C, k2); + } + + V tid = rocke_b_thread_id_x(b); + V bid = rocke_b_block_id_x(b); + V r, cxi, bwi, sub, bhi, bsub, img, a; + div_mod(e, cx, bid, &r, &cxi); + div_mod(e, nbw, r, &r, &bwi); + div_mod(e, nbh, r, &sub, &bhi); + div_mod(e, ddw, sub, &r, &bsub); + div_mod(e, ddh, r, &img, &a); + V ti0 = mul_c(e, bhi, bh); + V tj0 = mul_c(e, bwi, bw); + V h0; + { + V s1 = rocke_b_sub(b, a, g[G_PH]); + V m1 = rocke_b_mul(b, ti0, g[G_DHM]); + h0 = rocke_b_add(b, s1, m1); + } + V w0; + { + V s1 = rocke_b_sub(b, bsub, g[G_PW]); + V m1 = rocke_b_mul(b, tj0, g[G_DWM]); + w0 = rocke_b_add(b, s1, m1); + } + V img_row0 = rocke_b_mul(b, img, Hi); + V c0 = mul_c(e, cxi, cc); + + int shape[1] = {loads * block * lv}; + V lds = rocke_b_smem_alloc(b, dt, shape, 1, "wino_xin"); + int ppx_lg = 0; + while((1 << (ppx_lg + 1)) <= ppx) + ++ppx_lg; + for(int k = 0; k < loads; ++k) + { + V q = tid; + if(k) + { + V kb = ci(e, (int64_t)k * block); + q = rocke_b_add(b, tid, kb); + } + V px, pc; + if(ppx > 1) + { + V kl = ci(e, ppx_lg); + px = rocke_b_lshr(b, q, kl); + V km = ci(e, ppx - 1); + pc = rocke_b_land(b, q, km); + } + else + { + px = q; + pc = ci(e, 0); + } + V krw = ci(e, rw); + V rr = rocke_b_div(b, px, krw); + V krw2 = ci(e, rw); + V mr = rocke_b_mul(b, rr, krw2); + V cc_ = rocke_b_sub(b, px, mr); + V mh_ = rocke_b_mul(b, rr, g[G_DH]); + V hi = rocke_b_add(b, h0, mh_); + V mw_ = rocke_b_mul(b, cc_, g[G_DW]); + V wi = rocke_b_add(b, w0, mw_); + V mc = mul_c(e, pc, lv); + V ch = rocke_b_add(b, c0, mc); + V l1 = rocke_b_cmp_ult(b, hi, Hi); + V l2 = rocke_b_cmp_ult(b, wi, Wi); + V ok = rocke_b_land(b, l1, l2); + if((k + 1) * block > npx * ppx) + { + V kn = ci(e, npx); + V lt = rocke_b_cmp_lt(b, px, kn); + ok = rocke_b_land(b, ok, lt); + } + { + V lt = rocke_b_cmp_lt(b, ch, C); + ok = rocke_b_land(b, ok, lt); + } + V a1 = rocke_b_add(b, img_row0, hi); + V m1 = rocke_b_mul(b, a1, Wi); + V pix = rocke_b_add(b, m1, wi); + V m2 = rocke_b_mul(b, pix, cb2); + V m3 = mul_c(e, ch, 2); + V a2 = rocke_b_add(b, m2, m3); + V off = oob(e, ok, a2); + V val = rocke_wino::load_raw(b, x_rsrc, off, s->io_dtype, lv); + V idx = mul_c(e, q, lv); + rocke_b_smem_store_vN(b, lds, &idx, 1, val, lv); + } + V pm_lds = nullptr; + if(ps) + { + V T = g[G_T]; + int npl = (npx + block - 1) / block; + int sh2[1] = {npl * block}; + pm_lds = rocke_b_smem_alloc(b, rocke_i32(), sh2, 1, "wino_pmax"); + for(int k = 0; k < npl; ++k) + { + V px = tid; + if(k) + { + V kb = ci(e, (int64_t)k * block); + px = rocke_b_add(b, tid, kb); + } + V krw = ci(e, rw); + V rr = rocke_b_div(b, px, krw); + V krw2 = ci(e, rw); + V mr = rocke_b_mul(b, rr, krw2); + V cc_ = rocke_b_sub(b, px, mr); + V mh_ = rocke_b_mul(b, rr, g[G_DH]); + V hi = rocke_b_add(b, h0, mh_); + V mw_ = rocke_b_mul(b, cc_, g[G_DW]); + V wi = rocke_b_add(b, w0, mw_); + V l1 = rocke_b_cmp_ult(b, hi, Hi); + V l2 = rocke_b_cmp_ult(b, wi, Wi); + V ok = rocke_b_land(b, l1, l2); + if((k + 1) * block > npx) + { + V kn = ci(e, npx); + V lt = rocke_b_cmp_lt(b, px, kn); + ok = rocke_b_land(b, ok, lt); + } + V a1 = rocke_b_add(b, img_row0, hi); + V m1 = rocke_b_mul(b, a1, Wi); + V pix = rocke_b_add(b, m1, wi); + V a2 = rocke_b_add(b, pix, T); + V k4 = ci(e, 4); + V m2 = rocke_b_mul(b, a2, k4); + V off = oob(e, ok, m2); + V z = ci(e, 0); + V ld = rocke_b_buffer_load(b, e_rsrc, off, z, rocke_i32()); + rocke_b_smem_store_vN(b, pm_lds, &px, 1, ld, 1); + } + } + rocke_b_sync(b); + + int tpc_lg = 0; + while((1 << (tpc_lg + 1)) <= tpc) + ++tpc_lg; + V cl, tb0; + if(tpc > 1) + { + V km = ci(e, tpc - 1); + cl = rocke_b_land(b, tid, km); + V kl = ci(e, tpc_lg); + tb0 = rocke_b_lshr(b, tid, kl); + } + else + { + cl = ci(e, 0); + tb0 = tid; + } + V cv; + { + V m1 = mul_c(e, cxi, tpc); + cv = rocke_b_add(b, m1, cl); + } + V c_ok = rocke_b_cmp_lt(b, cv, nvec); + V ps_scale = ps ? rocke_b_bitcast(b, g[G_PS], rocke_f32()) : nullptr; + for(int rep = 0; rep < tpt; ++rep) + { + V tb = tb0; + if(rep) + { + V k = ci(e, (int64_t)rep * (block / tpc)); + tb = rocke_b_add(b, tb0, k); + } + V bi, bj; + if(bw > 1) + { + V kbw = ci(e, bw); + bi = rocke_b_div(b, tb, kbw); + V mb = mul_c(e, bi, bw); + bj = rocke_b_sub(b, tb, mb); + } + else + { + bi = tb; + bj = ci(e, 0); + } + V ti = rocke_b_add(b, ti0, bi); + V tj = rocke_b_add(b, tj0, bj); + V tile_ok; + { + V l1 = rocke_b_cmp_lt(b, ti, tiles_h); + V l2 = rocke_b_cmp_lt(b, tj, tiles_w); + V la = rocke_b_land(b, l1, l2); + tile_ok = rocke_b_land(b, la, c_ok); + } + V t; + { + V m1 = rocke_b_mul(b, sub, tiles_h); + V a1 = rocke_b_add(b, m1, ti); + V m2 = rocke_b_mul(b, a1, tiles_w); + t = rocke_b_add(b, m2, tj); + } + V pbase; + { + V m1 = mul_c(e, bi, m * rw); + V m2 = mul_c(e, bj, mw); + pbase = rocke_b_add(b, m1, m2); + } + std::vector raw(n); + for(int uu = 0; uu < n; ++uu) + { + for(int vv = 0; vv < nwv; ++vv) + { + V k = ci(e, uu * rw + vv); + V a1 = rocke_b_add(b, pbase, k); + V m1 = mul_c(e, a1, cc); + V m2 = mul_c(e, cl, vec); + V idx = rocke_b_add(b, m1, m2); + V ld = rocke_b_smem_load_vN(b, lds, &idx, 1, dt, vec); + raw[uu].push_back(vec == 1 ? rocke_b_vec_extract(b, ld, 0) : ld); + } + } + VV scales(1, nullptr); + if(ps) + { + V acc = nullptr; + for(int uu = 0; uu < n; ++uu) + { + for(int vv = 0; vv < nwv; ++vv) + { + V k = ci(e, uu * rw + vv); + V ad = rocke_b_add(b, pbase, k); + V ld = rocke_b_smem_load_vN(b, pm_lds, &ad, 1, rocke_i32(), 1); + V pm = rocke_b_vec_extract(b, ld, 0); + acc = acc == nullptr ? pm : rocke_b_smax(b, acc, pm); + } + } + V f = rocke_wino::bits_to_f32(b, acc, s->io_dtype); + V e_t = prescale_exp_rt(e, f, ps_scale); + V ok_e; + { + V c0b = ci(e, 0); + V eq = rocke_b_cmp_eq(b, cv, c0b); + ok_e = rocke_b_land(b, tile_ok, eq); + } + V k4 = ci(e, 4); + V m4 = rocke_b_mul(b, t, k4); + V off = oob(e, ok_e, m4); + V z = ci(e, 0); + V bc = rocke_b_bitcast(b, e_t, rocke_f32()); + rocke_b_buffer_store_f32(b, e_rsrc, off, z, bc); + V z2 = ci(e, 0); + V ne = rocke_b_sub(b, z2, e_t); + scales[0] = rocke_wino::pow2_f32(b, ne); + } + V vbase0; + { + V st = mul_c(e, cb2, n * nwv); + V m1 = rocke_b_mul(b, t, st); + V m2 = mul_c(e, cv, vec * 2); + vbase0 = rocke_b_add(b, m1, m2); + } + VV oks(1, tile_ok); + input_tiles_transform(e, raw, 1, oks, scales, vbase0, v_rsrc, cb2); + } + (void)err; + return true; +} + +/* ------------------------------------------------------------ K4 */ + +const char* const kOutputGeom[] = {"K", "Ho", "Wo", "dH", "dW", "dHm", "dWm"}; +enum +{ + O_K, + O_HO, + O_WO, + O_DH, + O_DW, + O_DHM, + O_DWM, + O_N +}; + +void emit_output(E& e) +{ + rocke_ir_builder_t* b = e.b; + const Spec* s = e.s; + int n = s->mh + 2, m = s->mh; + int nwv = s->mw + 2, mw = s->mw; + int vec = s->vec; + + set_mwgs(e, s->block_size); + V mm = param(e, "mm", gptr(e, rocke_wino::ir_dtype(s->operand_dtype)), NOALIAS | RO); + V e_ptr = param(e, "e", gptr(e, rocke_i32()), NOALIAS | RO); + V y = param(e, "y", gptr(e, rocke_wino::ir_dtype(s->io_dtype)), NOALIAS | WO); + V m_bytes = p32(e, "m_bytes"); + V y_bytes = p32(e, "y_bytes"); + V g[O_N]; + for(int i = 0; i < O_N; ++i) + g[i] = p32(e, kOutputGeom[i]); + V m_rsrc = rocke_b_buffer_rsrc(b, mm, m_bytes); + V y_rsrc = rocke_b_buffer_rsrc(b, y, y_bytes); + Coords rc = run_coords(e, 1); + + V K = g[O_K], Ho = g[O_HO], Wo = g[O_WO]; + V kb2; + { + V k2 = ci(e, 2); + kb2 = rocke_b_mul(b, K, k2); + } + V t = rc.t0; + V ok = rocke_b_land(b, rc.ok, rc.c_ok); + V k = mul_c(e, rc.cv, vec); + V k2v = mul_c(e, k, 2); + V mbase; + { + V st = mul_c(e, kb2, n * nwv); + V m1 = rocke_b_mul(b, t, st); + V a1 = rocke_b_add(b, m1, k2v); + mbase = oob(e, ok, a1); + } + std::vector raw(n); + V off = mbase; + for(int uu = 0; uu < n; ++uu) + { + for(int vv = 0; vv < nwv; ++vv) + { + if(uu || vv) + off = rocke_b_add(b, off, kb2); + raw[uu].push_back(rocke_wino::load_raw(b, m_rsrc, off, s->operand_dtype, vec)); + } + } + V scale = nullptr; + if(s->prescale) + { + V z = ci(e, 0); + V t_safe = rocke_b_select(b, ok, t, z); + V ld = rocke_b_global_load_i32(b, e_ptr, t_safe, 4); + scale = rocke_wino::pow2_f32(b, ld); + } + V ho0; + { + V m1 = rocke_b_mul(b, rc.ti, g[O_DHM]); + ho0 = rocke_b_add(b, rc.a, m1); + } + V wo0; + { + V m1 = rocke_b_mul(b, rc.tj0, g[O_DWM]); + wo0 = rocke_b_add(b, rc.bsub, m1); + } + V img_row0 = rocke_b_mul(b, rc.n, Ho); + VV yrow; + V ho = ho0; + for(int i = 0; i < m; ++i) + { + if(i) + ho = rocke_b_add(b, ho, g[O_DH]); + V lt = rocke_b_cmp_lt(b, ho, Ho); + V okr = rocke_b_land(b, ok, lt); + V a1 = rocke_b_add(b, img_row0, ho); + V m1 = rocke_b_mul(b, a1, Wo); + V m2 = rocke_b_mul(b, m1, kb2); + V a2 = rocke_b_add(b, m2, k2v); + yrow.push_back(oob(e, okr, a2)); + } + VV ycol; + V wo = wo0; + for(int j = 0; j < mw; ++j) + { + if(j) + wo = rocke_b_add(b, wo, g[O_DW]); + V lt = rocke_b_cmp_lt(b, wo, Wo); + V m1 = rocke_b_mul(b, wo, kb2); + ycol.push_back(oob(e, lt, m1)); + } + + rocke_b_sched_barrier(b, 0); + std::vector> colt; + for(int vv = 0; vv < nwv; ++vv) + { + std::vector data; + for(int uu = 0; uu < n; ++uu) + data.push_back(rocke_wino::raw_to_f32(b, raw[uu][vv], vec)); + std::vector per_ch; + for(int ch = 0; ch < vec; ++ch) + { + VV col; + for(int uu = 0; uu < n; ++uu) + col.push_back(data[uu][ch]); + per_ch.push_back(rocke_wino::at_1d(b, m, col)); + } + colt.push_back(per_ch); + } + std::vector> pending; + for(int i = 0; i < m; ++i) + { + std::vector outs; + for(int ch = 0; ch < vec; ++ch) + { + VV row; + for(int vv = 0; vv < nwv; ++vv) + row.push_back(colt[vv][ch][i]); + outs.push_back(rocke_wino::at_1d(b, mw, row)); + } + for(int j = 0; j < mw; ++j) + { + VV vals; + for(int ch = 0; ch < vec; ++ch) + vals.push_back(outs[ch][j]); + if(scale != nullptr) + { + for(auto& x_ : vals) + x_ = rocke_b_fmul(b, x_, scale); + } + V ad = rocke_b_add(b, yrow[i], ycol[j]); + pending.push_back({ad, rocke_wino::pack_vec_f32(b, vals, s->io_dtype)}); + } + } + rocke_b_sched_barrier(b, 0); + int cpol = s->nt_stores ? kCpolNt : 0; + for(auto& pr : pending) + rocke_wino::store_packed(b, y_rsrc, pr.first, pr.second, s->io_dtype, cpol); +} + +bool emit(E& e, std::string* err) +{ + switch(e.s->stage) + { + case ROCKE_WINOGRAD_RT_K1: + emit_filter(e); + return true; + case ROCKE_WINOGRAD_RT_K2A: + emit_prescale(e); + return true; + case ROCKE_WINOGRAD_RT_K2: + return emit_input(e, err); + case ROCKE_WINOGRAD_RT_K4: + emit_output(e); + return true; + } + *err = "unknown stage"; + return false; +} + +} // namespace + +extern "C" { + +rocke_winograd_rt_spec_t rocke_winograd_rt_spec_default(rocke_winograd_rt_stage_t stage) +{ + rocke_winograd_rt_spec_t s; + memset(&s, 0, sizeof s); + s.stage = stage; + s.mh = 4; + s.mw = 4; + s.io_dtype = "fp16"; + s.operand_dtype = "fp16"; + s.vec = 8; + s.filter_rows = 1; + s.block_size = 256; + s.prescale = false; + s.tiles_per_thread = 1; + s.lds[0] = 0; + s.lds[1] = 0; + s.lds_lanes = 16; + s.nt_stores = true; + s.fused = false; + return s; +} + +bool rocke_winograd_rt_is_valid_spec(const rocke_winograd_rt_spec_t* spec, + char* reason, + size_t reason_cap) +{ + std::string err; + if(spec == NULL) + { + set_reason(reason, reason_cap, "null spec"); + return false; + } + if(!validate(spec, &err)) + { + set_reason(reason, reason_cap, err); + return false; + } + set_reason(reason, reason_cap, ""); + return true; +} + +rocke_status_t + rocke_winograd_rt_kernel_name(const rocke_winograd_rt_spec_t* spec, char* out, size_t out_cap) +{ + if(spec == NULL || out == NULL || out_cap == 0) + return ROCKE_ERR_VALUE; + std::string err; + if(!validate(spec, &err)) + return ROCKE_ERR_VALUE; + std::string name = kernel_name(spec); + if(name.size() + 1 > out_cap) + return ROCKE_ERR_VALUE; + memcpy(out, name.c_str(), name.size() + 1); + return ROCKE_OK; +} + +bool rocke_winograd_rt_magic_div(int64_t d, uint32_t* magic, int* shift) +{ + if(d < 1 || d >= ((int64_t)1 << 31) || magic == NULL || shift == NULL) + return false; + if(d == 1) + { + *magic = 0; + *shift = 0; + return true; + } + int lg = 0; + while(((int64_t)1 << lg) < d) + ++lg; + /* ceil(2**(31 + lg) / d) < 2**32 for d >= 2 (see magic_div). */ + unsigned __int128 num = (unsigned __int128)1 << (31 + lg); + unsigned __int128 q = (num + (unsigned __int128)(d - 1)) / (unsigned __int128)d; + *magic = (uint32_t)q; + *shift = lg - 1; + return true; +} + +rocke_kernel_def_t* rocke_build_winograd_rt(rocke_ir_builder_t* b, + const rocke_winograd_rt_spec_t* spec, + const char* arch) +{ + (void)arch; + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + std::string err; + if(spec == NULL) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, "build_winograd_rt: null spec"); + return NULL; + } + if(!validate(spec, &err)) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, err.c_str()); + return NULL; + } + E e{b, spec}; + if(!emit(e, &err)) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, err.c_str()); + return NULL; + } + if(rocke_ir_builder_status(b) != ROCKE_OK) + return NULL; + return rocke_ir_builder_kernel(b); + }); +} + +rocke_kernel_def_t* rocke_build_winograd_rt_new(rocke_ir_builder_t* b, + const rocke_winograd_rt_spec_t* spec, + const char* arch) +{ + if(b != NULL) + memset(b, 0, sizeof(*b)); + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + char name[1024]; + if(b == NULL) + return NULL; + if(spec == NULL) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, "build_winograd_rt: null spec"); + return NULL; + } + std::string err; + if(!validate(spec, &err)) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, err.c_str()); + return NULL; + } + if(rocke_winograd_rt_kernel_name(spec, name, sizeof name) != ROCKE_OK) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, "build_winograd_rt: kernel name too long"); + return NULL; + } + if(rocke_ir_builder_init(b, name) != ROCKE_OK) + return NULL; + return rocke_build_winograd_rt(b, spec, arch); + }); +} + +rocke_status_t rocke_winograd_rt_lower_to_llvm(const rocke_winograd_rt_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap) +{ + rocke_ir_builder_t b = {}; + if(out_ll != NULL) + *out_ll = NULL; + if(spec == NULL || out_ll == NULL) + { + set_reason(err, err_cap, "lower_to_llvm: null spec/out"); + return ROCKE_ERR_VALUE; + } + if(arch == NULL) + arch = "gfx1250"; + rocke_kernel_def_t* kernel = rocke_build_winograd_rt_new(&b, spec, arch); + 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 : "build_winograd_rt failed"); + 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; +} + +} /* extern "C" */ 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 99d8929ed7df..17f733b2ecf0 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 @@ -1319,10 +1319,10 @@ the 4x4, 14x14 and 16x16 patch-embedding stems) and `FWD_SMALLK_KNOBS` over ## Winograd Forward (gfx1250) `library/kernels/common/conv_winograd_*.py` (C++ twins -`instance_conv_winograd_{fwd,gemm_wmma,gen,fused_1d_wmma}.h`). A stride-1, 2-D -forward convolution computed as Winograd `F(m, r)` per axis, with fp16 / bf16 -I/O and fp16 GEMM operands (bf16 operands only for `F(2x2, 3x3)` with bf16 -I/O). Three implementations: +`instance_conv_winograd_{fwd,fwd_rt,gemm_wmma,gen,fused_1d_wmma}.h`). A stride-1, +2-D forward convolution computed as Winograd `F(m, r)` per axis, with fp16 / +bf16 I/O and fp16 GEMM operands (bf16 operands only for `F(2x2, 3x3)` with +bf16 I/O). Four implementations: - **Unfused 3x3** (`conv_winograd_fwd` + `conv_winograd_gemm_wmma`), output tile `mh x mw` with `mh, mw` in {2, 4, 6}: @@ -1345,6 +1345,26 @@ I/O). Three implementations: K1 and the first input-transform stage can run as one merged launch (`WinogradMergedSpec`). Problems past the 32-bit buffer offsets run as batch chunks sharing U. +- **Unfused 3x3, shape-generic (AOT)** (`conv_winograd_fwd_rt` + + `conv_winograd_gemm_wmma`): the same K1 / K2a / K2 / K4 stages, workspace + layouts and arithmetic, with every problem extent a kernel argument, so one + binary serves every problem its knobs fit and the pipeline is compiled + ahead of time like the direct and implicit-GEMM instances. One spec per + stage (`RtFilterSpec`, `RtPrescaleSpec`, `RtInputSpec`, `RtOutputSpec`); + the compile-time knobs are the tile, the dtypes, the channel vector width + (the stage's channel extent must be a multiple of it), the workgroup size, + the K2 register run or LDS region, and the prescale. Divisions by runtime + extents are a multiply-high by a host-computed magic number (`magic_div`, + exact for 31-bit indices). The fused prescale reduces the tile maxima over + a runtime number of channel lanes (up to the workgroup size); wider C runs + the split form (K2a). `stage_launch(spec, RtGeometry(...))` gives a stage's + grid, block and scalar kernargs; `aot_stage_specs()` is the AOT grid and + `default_stage_specs(problem, tile, dtype)` the binaries a pipeline uses by + default (the JIT pipeline's tuned knobs, narrowed to the problem's channel + extents). There is no merged K1 + K2 launch. + `benchmarks/common/benchmark_winograd_aot.py` compiles the grid + (`--compile-all`) and runs cases from it (`--run-from-cache`), choosing each + stage's binary by timing and verifying the assembled pipeline. - **Generalized** (`conv_winograd_gen` + K3 with `taps`): any `F(m, r)` per axis (1-D `1 x r` / `r x 1`, 5x5), and segmented long filters: `S` segments of `F(m, m)` share one input transform and K3 sums @@ -1359,7 +1379,9 @@ The transforms use exact Toom-Cook matrices (`toom_cook_gen`, default interpolation points per `F(m, r)`), applied in fp32. Every transform load and store goes through a buffer resource; padding, dilation holes and tails read zero / drop the store, so the kernels have no bounds branches. The -transform kernels bake the problem extents in; K3 takes them as kernargs. +`conv_winograd_fwd` / `conv_winograd_gen` transform kernels bake the problem +extents in; K3 and the `conv_winograd_fwd_rt` transforms take them as +kernargs. **Plan selection.** `conv_winograd_select.select_winograd(problem, dtype, model=WinogradCostModel(num_cus=...))` returns the safe plan of lowest model @@ -1379,7 +1401,10 @@ K3 GEMM lowers at `llvm23` only, so at `llvm20` / `llvm22` the unfused plans are refused and only the fused 1-D plans remain. `launch_plan_for` returns the plan's `ConvLaunchPlan`: one step per kernel (each built for the problem, kernargs named per step) over the buffers `x`, `w`, `y` and the workspaces -`U`, `V`, `M` (and `e`). +`U`, `V`, `M` (and `e`). A second opt-in candidate, `conv_winograd_fwd_aot` +(`algorithm="winograd_aot"`), runs the selector's 3x3 plans on the +shape-generic binaries (plan kind `unfused2d_rt`): each transform step is +named by its binary's kernel name and carries the problem in its kernargs. ## Direct Wgrad Kernels 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 1c466811947d..7f15fb5de80f 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 @@ -50,7 +50,7 @@ implementation. | `conv_direct_dgrad_wmma.py` | `DirectDgradWmmaSpec` (gfx1250 dgrad) | `instances/convolution.md` | | `conv_direct_fwd_wmma.py` | `DirectFwdWmmaSpec` (gfx1250 forward) | `instances/convolution.md` | | `conv_fwd_small_wmma.py` | `FwdSmallcWmmaSpec`, `FwdSmallkWmmaSpec` (gfx1250 forward, tiny C / K) | `instances/convolution.md` | -| `conv_winograd_fwd.py`, `conv_winograd_gemm_wmma.py`, `conv_winograd_gen.py`, `conv_winograd_fused_1d_wmma.py`, `conv_winograd_select.py` | `WinogradFwdSpec`, `WinogradMergedSpec`, `WinogradGemmWmmaSpec`, `WinogradGemmReduceSpec`, `WinogradGenSpec`, `FusedWino1DSpec`; plan selector `select_winograd` (gfx1250 Winograd forward) | `instances/convolution.md` | +| `conv_winograd_fwd.py`, `conv_winograd_fwd_rt.py`, `conv_winograd_gemm_wmma.py`, `conv_winograd_gen.py`, `conv_winograd_fused_1d_wmma.py`, `conv_winograd_select.py` | `WinogradFwdSpec`, `WinogradMergedSpec`, `RtFilterSpec` / `RtPrescaleSpec` / `RtInputSpec` / `RtOutputSpec` (shape-generic, AOT), `WinogradGemmWmmaSpec`, `WinogradGemmReduceSpec`, `WinogradGenSpec`, `FusedWino1DSpec`; plan selector `select_winograd` (gfx1250 Winograd forward) | `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` | @@ -243,7 +243,7 @@ From `helpers/README.md`: | conv_direct_dgrad_wmma | - | - | - | - | - | yes | - | `emit_direct_params` (AOT kernargs), TDM `tensor_load_to_lds` / `tensor_store_from_lds`, `ds_read_tr16_b128` | | conv_direct_fwd_wmma | - | - | - | - | - | yes | - | `emit_direct_params` (AOT kernargs), TDM `tensor_load_to_lds` / `tensor_store_from_lds`, split barriers, `permlane16` | | conv_fwd_small_wmma (smallc / smallk) | - | - | - | - | - | yes | - | `emit_direct_params` (AOT kernargs), TDM `tensor_load_to_lds` / `tensor_store_from_lds` | -| conv_winograd (fwd / gemm_wmma / gen / fused_1d_wmma) | yes (buffer) | - | - | - | - | yes | - | TDM `tensor_load_to_lds` / `tensor_store_from_lds` (K3), `chiplet_transform_chunked_runtime` (`helpers.grid`), `global_load_async_to_lds`, `ds_read_tr16_b128` (fused 1-D) | +| conv_winograd (fwd / fwd_rt / gemm_wmma / gen / fused_1d_wmma) | yes (buffer) | - | - | - | - | yes | - | TDM `tensor_load_to_lds` / `tensor_store_from_lds` (K3), `chiplet_transform_chunked_runtime` (`helpers.grid`), `global_load_async_to_lds`, `ds_read_tr16_b128` (fused 1-D) | | 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 | - | - | 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 13b84f476ab0..bc4e185baa67 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 @@ -124,7 +124,7 @@ be smaller than 2 GiB. | `conv_direct_dgrad_wmma` | ❌ | ❌ | ❌ | gfx1250 only: wave32 WMMA 16x16x32 + TDM, dgrad with stride 1 or 2 (phase decomposition) or per-axis dilation, dense and grouped (`cpg`, `kpg` multiples of 8, or 2/4/8/16 packed into block-diagonal super-groups with `pack_groups`), `PAD` in 0..(K-1)*dil. Built from `library/kernels`; served by the direct-conv AOT cache, not the dispatcher | | `conv_direct_fwd_wmma` | ❌ | ❌ | ❌ | gfx1250 only: wave32 WMMA 16x16x32 + TDM forward, dense and grouped (`cpg`, `kpg` multiples of 8, or small groups packed block-diagonally), square filters, stride 1 or 2, per-axis dilation. Built from `library/kernels`; served by the direct-conv AOT cache, not the dispatcher | | `conv_direct_fwd_smallc_wmma` / `conv_direct_fwd_smallk_wmma` | ❌ | ❌ | ❌ | gfx1250 only: wave32 WMMA 16x16x32 + TDM forward of tiny channel counts, ungrouped and undilated, square filters: `C <= 4` stems (`K` a multiple of 8) and `K` in (1, 2, 4, 8) layers (`C` a multiple of 8). Built from `library/kernels`; served by the direct-conv AOT cache, not the dispatcher | -| Winograd forward (`conv_winograd_fwd`, `conv_winograd_gemm_wmma`, `conv_winograd_gen`, `conv_winograd_fused_1d_wmma`) | ❌ | ❌ | ❌ | gfx1250 only: fp16 / bf16 I/O, stride 1, 2-D, `C/g` and `K/g` >= 32, dilation up to 35; 3x3 as F(m x n, 3x3) with tiles 2..6 (not 6x6), 1-D / 7x1 / 5x5 and segmented long filters through the generalized pipeline, fused 1-D kernel for `1 x r` with `C` in {32, 64, 128} (fp16). Kernels built per problem (no AOT cache); explicit K2 specs with long register-path runs (`tiles_per_thread`) of tiles >= 4 spill to scratch (see `dsl_docs/instances/convolution.md`); plan from `conv_winograd_select`; opt-in dispatcher route `algorithm="winograd"` (candidate `conv_winograd_fwd`), never picked by `auto` | +| Winograd forward (`conv_winograd_fwd`, `conv_winograd_fwd_rt`, `conv_winograd_gemm_wmma`, `conv_winograd_gen`, `conv_winograd_fused_1d_wmma`) | ❌ | ❌ | ❌ | gfx1250 only: fp16 / bf16 I/O, stride 1, 2-D, `C/g` and `K/g` >= 32, dilation up to 35; 3x3 as F(m x n, 3x3) with tiles 2..6 (not 6x6), 1-D / 7x1 / 5x5 and segmented long filters through the generalized pipeline, fused 1-D kernel for `1 x r` with `C` in {32, 64, 128} (fp16). Kernels built per problem, except the 3x3 pipeline's shape-generic transforms (`conv_winograd_fwd_rt`, compiled ahead of time; opt-in route `algorithm="winograd_aot"`, candidate `conv_winograd_fwd_aot`); explicit K2 specs with long register-path runs (`tiles_per_thread`) of tiles >= 4 spill to scratch (see `dsl_docs/instances/convolution.md`); plan from `conv_winograd_select`; opt-in dispatcher route `algorithm="winograd"` (candidate `conv_winograd_fwd`), never picked by `auto` | --- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json new file mode 100644 index 000000000000..27bec1da0aa6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json @@ -0,0 +1,167 @@ +{ + "aliases": { + "llvm24": "llvm23" + }, + "flavors": { + "llvm20": { + "cases": { + "conv_winograd_fwd_rt/gfx1250/bf16_k2a": { + "arch": "gfx1250", + "bytes": 6097, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_pmax_bf16_v8", + "sha256": "772315f5f68e092013fd61b1ab8b0bf3c63907936d0c354d4820fad0675eeaca", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f2_lds_k2": { + "arch": "gfx1250", + "bytes": 230995, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f2_fp16_opfp16_v8_lds4x8_ll16_t4", + "sha256": "c9dfec3fd03417ec38499f881a69b12fcf41a33e61f6b5f2ef074ff74a4d1aea", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_bf16_fused_prescale_k2": { + "arch": "gfx1250", + "bytes": 138240, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_bf16_opfp16_v4_b256_r1_psf", + "sha256": "f8b2cc767e5634a621f1698aa68c541d8ee3dc2896cd025030e88f3d9581b743", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k1": { + "arch": "gfx1250", + "bytes": 58716, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_filter_f4_fp16_opfp16_v8_fr2", + "sha256": "0dad1effb4857be4a8c842e91c2cc06c0be9543c6bc61fa1b3d2228ae0a4b1d8", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k2": { + "arch": "gfx1250", + "bytes": 102839, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_fp16_opfp16_v4_b256_r1", + "sha256": "e453001a40be48edc66c298568c0d29ea545ba62d397937da9cc5b990c3a52e9", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f6x4_bf16_prescale_k4": { + "arch": "gfx1250", + "bytes": 130889, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_output_f6x4_bf16_opfp16_v4_b256_ps", + "sha256": "01d6e7fc822caf093bd40eff50ac3fb9e33e9877e2933814df87101f8ce6f02a", + "status": "ok" + } + }, + "expected_failures": {} + }, + "llvm22": { + "cases": { + "conv_winograd_fwd_rt/gfx1250/bf16_k2a": { + "arch": "gfx1250", + "bytes": 6189, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_pmax_bf16_v8", + "sha256": "257fda3d9bcf8effe4707d016b0f015aafc73568e1b762641c5c77ea579d452c", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f2_lds_k2": { + "arch": "gfx1250", + "bytes": 231091, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f2_fp16_opfp16_v8_lds4x8_ll16_t4", + "sha256": "9d8effe8ecd4b536ea770b49e476888fb990bd2fce7f09ff7363a8d432d0650e", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_bf16_fused_prescale_k2": { + "arch": "gfx1250", + "bytes": 138379, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_bf16_opfp16_v4_b256_r1_psf", + "sha256": "4a4f41ddbafba07e6615b4c2f458c5ab42f4bfb5dbc3db90d3c43f00d846121e", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k1": { + "arch": "gfx1250", + "bytes": 58812, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_filter_f4_fp16_opfp16_v8_fr2", + "sha256": "71ed2e364305bf6754d2c3fe97810d70380bd51a952e43a9d8896892122bae4b", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k2": { + "arch": "gfx1250", + "bytes": 102935, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_fp16_opfp16_v4_b256_r1", + "sha256": "f7366eb62974f7c6973442ebced385e6d27f59d2e68f0c1e18c49b5ab7a468ca", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f6x4_bf16_prescale_k4": { + "arch": "gfx1250", + "bytes": 130985, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_output_f6x4_bf16_opfp16_v4_b256_ps", + "sha256": "e561b7f691bb9515ca35cd5c19624f69aee5e3b63261495e731a6e0de0b7433c", + "status": "ok" + } + }, + "expected_failures": {} + }, + "llvm23": { + "cases": { + "conv_winograd_fwd_rt/gfx1250/bf16_k2a": { + "arch": "gfx1250", + "bytes": 6253, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_pmax_bf16_v8", + "sha256": "c1ff39de4529279f930d8a298065cae9f62d954445d76fc976a1e36b671367e3", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f2_lds_k2": { + "arch": "gfx1250", + "bytes": 231155, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f2_fp16_opfp16_v8_lds4x8_ll16_t4", + "sha256": "face2b28826e42ecf5ce95fcbe83cb070eb663fdb0a0c39b002420a661474218", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_bf16_fused_prescale_k2": { + "arch": "gfx1250", + "bytes": 138443, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_bf16_opfp16_v4_b256_r1_psf", + "sha256": "455b377ae4439439e12fc3a140e3c25ca968647271ed1f784b1051ae3b83b678", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k1": { + "arch": "gfx1250", + "bytes": 58876, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_filter_f4_fp16_opfp16_v8_fr2", + "sha256": "d86adca253c7e86d3dffe005f0f8291140201e6f4393ac726dbae6cb1ee55a6c", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k2": { + "arch": "gfx1250", + "bytes": 102999, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_fp16_opfp16_v4_b256_r1", + "sha256": "a0c9029f858812440da5b775d95e61acc636fa0c5bdb5aa1b40ac68072fc297c", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f6x4_bf16_prescale_k4": { + "arch": "gfx1250", + "bytes": 131049, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_output_f6x4_bf16_opfp16_v4_b256_ps", + "sha256": "1b87a48cf46fa9ccfcaa822f677bdd0f603e7c99e2dc2df50c7150b704e29ec2", + "status": "ok" + } + }, + "expected_failures": {} + } + }, + "schema": "ck.dsl.ir_golden_sha256/v2" +} 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 95e0244a2d93..b95ecd9b6172 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 @@ -5264,6 +5264,45 @@ def _build(): build_winograd_fwd_case(_stage, **_over), ) + # --- conv_winograd_fwd_rt: shape-generic (AOT) Winograd transforms --- + # Every extent a kernel argument: K1, K2a, K2 (register path with the fused + # and the split prescale, LDS path) and K4. The full matrix (and the C++ + # twin) lives in library/tests/parity/conv_winograd_fwd_rt_emit.*. + def build_winograd_rt_case(cls, **spec_kw): + def _build(): + from kernels.common import conv_winograd_fwd_rt as wr + + return wr.build_stage(getattr(wr, cls)(**spec_kw), arch="gfx1250") + + return _build + + for _case_id, _cls, _over in ( + ("f4_k1", "RtFilterSpec", dict(filter_rows=2)), + ("bf16_k2a", "RtPrescaleSpec", dict(io_dtype="bf16")), + ("f4_k2", "RtInputSpec", dict(vec=4)), + ( + "f4_bf16_fused_prescale_k2", + "RtInputSpec", + dict(io_dtype="bf16", prescale=True, fused=True, vec=4), + ), + ( + "f2_lds_k2", + "RtInputSpec", + dict(mh=2, mw=2, lds=(4, 8), tiles_per_thread=4), + ), + ( + "f6x4_bf16_prescale_k4", + "RtOutputSpec", + dict(mh=6, mw=4, io_dtype="bf16", prescale=True, vec=4), + ), + ): + add( + "conv_winograd_fwd_rt", + f"conv_winograd_fwd_rt/gfx1250/{_case_id}", + "gfx1250", + build_winograd_rt_case(_cls, **_over), + ) + # --- conv_winograd_gemm_wmma: the Winograd batched GEMM (K3, gfx1250) --- # TDM operand ring + WMMA, TDM or per-lane epilogue, split-K slabs and # their reduce. The full matrix (and the C++ twin) lives in From 2b5b27308badfe1a87795f3eb738b171b89eebd1 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Fri, 30 Oct 2026 21:52:27 +0000 Subject: [PATCH 2/5] fix(rocke): tighten the shape-generic Winograd spec and launch checks Review follow-ups for the shape-generic Winograd forward transforms: - Both engines refuse an LDS lane count that is not a positive power of two (the lane mapping is a shift and mask), an LDS region of zero width, a filter_rows below 1 and a K2 / K4 block above 1024 threads. - stage_launch runs RtGeometry.check itself, so a non-3x3, strided or oversized problem is refused there and not only by the callers. - The Winograd launch plans refuse a batch split that is not a positive integer instead of normalizing it. - Six new rejected parity configs; the new CPU, dispatch and GPU tests are in the installed CTest lists. Co-Authored-By: Claude Opus 5.5 --- .../library/dispatch/conv_winograd_plan.py | 19 ++++++++++----- .../kernels/common/conv_winograd_fwd_rt.py | 23 ++++++++++++------ .../test_conv_winograd_aot_dispatch.py | 10 ++++++++ .../tests/parity/conv_winograd_fwd_rt_emit.c | 6 +++++ .../tests/parity/conv_winograd_fwd_rt_emit.py | 12 ++++++++++ .../tests/test_conv_winograd_fwd_rt_cpu.py | 24 +++++++++++++++++++ .../rocke/platform/CMakeLists.txt | 3 +++ .../instances/common/conv_winograd_fwd_rt.cpp | 14 +++++++++-- 8 files changed, 96 insertions(+), 15 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py index 11275ce0e793..1476c462eb83 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/conv_winograd_plan.py @@ -236,6 +236,17 @@ def _build_with_arch(fn, spec, arch): # --------------------------------------------------------------------------- +def _plan_chunks(problem: ConvProblem, plan: WinogradPlan) -> int: + """The plan's batch split: a positive integer dividing N (refused, never + normalized, otherwise).""" + chunks = plan.chunks + if not isinstance(chunks, int) or chunks < 1: + raise ValueError(f"batch split must be a positive integer (got {chunks!r})") + if problem.N % chunks: + raise ValueError(f"batch split {chunks} does not divide N={problem.N}") + return chunks + + def _unfused2d(problem: ConvProblem, dtype: str, plan: WinogradPlan): from kernels.common import conv_winograd_fwd as W from kernels.common.conv_winograd_gemm_wmma import ( @@ -244,9 +255,7 @@ def _unfused2d(problem: ConvProblem, dtype: str, plan: WinogradPlan): ) mh, mw = plan.tile - chunks = max(1, int(plan.chunks)) - if problem.N % chunks: - raise ValueError(f"batch split {chunks} does not divide N={problem.N}") + chunks = _plan_chunks(problem, plan) prescale = dtype == "bf16" and max(mh, mw) >= 4 operand = "bf16" if ((mh, mw) == (2, 2) and dtype == "bf16") else "fp16" xf = W.tuned_transform_options((mh, mw), prescale, problem) @@ -343,9 +352,7 @@ def _unfused2d_rt(problem: ConvProblem, dtype: str, plan: WinogradPlan): ) mh, mw = plan.tile - chunks = max(1, int(plan.chunks)) - if problem.N % chunks: - raise ValueError(f"batch split {chunks} does not divide N={problem.N}") + chunks = _plan_chunks(problem, plan) sub = dataclasses.replace(problem, N=problem.N // chunks) geo = R.RtGeometry(sub, mh, mw) operand, prescale = R.pipeline_dtypes((mh, mw), dtype) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py index afe965dc924b..fa661851facc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py @@ -263,6 +263,11 @@ def _check_common(mh, mw, io, op, vec): raise ValueError(f"vec must be 1/2/4/8 (got {vec})") +def _check_block(block_size): + if block_size < 32 or block_size > 1024 or block_size & (block_size - 1): + raise ValueError("block_size must be a power of two in [32, 1024]") + + def _tile_tag(mh, mw) -> str: return f"f{mh}" if mh == mw else f"f{mh}x{mw}" @@ -280,7 +285,7 @@ class RtFilterSpec: def __post_init__(self): _check_common(self.mh, self.mw, self.io_dtype, self.operand_dtype, self.vec) - if (self.mh + 2) % self.filter_rows: + if self.filter_rows < 1 or (self.mh + 2) % self.filter_rows: raise ValueError("filter_rows must divide n") stage = "K1" @@ -351,14 +356,17 @@ class RtInputSpec: def __post_init__(self): _check_common(self.mh, self.mw, self.io_dtype, self.operand_dtype, self.vec) - if self.block_size < 32 or self.block_size & (self.block_size - 1): - raise ValueError("block_size must be a power of two >= 32") + _check_block(self.block_size) if self.fused and (not self.prescale or self.input_lds): raise ValueError("fused needs the prescale and the register path") if self.tiles_per_thread < 1: raise ValueError("tiles_per_thread must be >= 1") if self.input_lds: bh, bw = self.lds + if bw < 1: + raise ValueError("lds must be (0, 0) or two positive extents") + if self.lds_lanes < 1 or self.lds_lanes & (self.lds_lanes - 1): + raise ValueError("lds_lanes must be a power of two >= 1") if (bh * bw) % self.tiles_per_thread: raise ValueError("tiles_per_thread must divide bh * bw") if self.lds_block > 1024: @@ -429,8 +437,7 @@ class RtOutputSpec: def __post_init__(self): _check_common(self.mh, self.mw, self.io_dtype, self.operand_dtype, self.vec) - if self.block_size < 32 or self.block_size & (self.block_size - 1): - raise ValueError("block_size must be a power of two >= 32") + _check_block(self.block_size) stage = "K4" @@ -1374,8 +1381,10 @@ def stage_launch( spec, geo: RtGeometry ) -> Tuple[Tuple[int, int, int], int, Dict[str, int]]: """``(grid, block, scalar kernarg values)`` of one launch; the buffer - pointers and ``_bytes`` sizes are the caller's.""" - why = spec.reason(geo) + pointers and ``_bytes`` sizes are the caller's. Refuses a problem + outside :meth:`RtGeometry.check` as well as one the spec does not fit.""" + prescale = isinstance(spec, RtPrescaleSpec) or getattr(spec, "prescale", False) + why = geo.check(prescale) or spec.reason(geo) if why: raise ValueError(f"{spec.kernel_name()}: {why}") grid, block, vals = _STAGES[type(spec)][2](spec, geo) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py index b571fc8e9134..1e9214bf80d3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_conv_winograd_aot_dispatch.py @@ -168,6 +168,16 @@ def test_batch_split_relaunches_the_same_binaries(): assert v["y"] == _PTRS["y"] + sz["y"] and v["y_bytes"] == sz["y"] +@pytest.mark.parametrize("kind", ["unfused2d_rt", "unfused2d"]) +@pytest.mark.parametrize("chunks", [0, -2, 1.9]) +def test_invalid_batch_split_is_refused(kind, chunks): + plan = WinogradPlan( + kind=kind, tile=(2, 2), k3="128x128x128_w2x2_f4x4_s2_tdm", chunks=chunks + ) + with pytest.raises(ValueError, match="positive integer"): + winograd_launch_plan(gc._conv_problem(_req(N=4)), "fp16", plan, _ARCH) + + def test_no_3x3_plan_is_refused(): req = _pin(_req(Y=1, X=11, Hi=1, Wi=200, pad_h=0, pad_w=5)) with pytest.raises(ValueError, match="has no AOT kernels"): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c index e109721e53c9..4d014f8b0a71 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.c @@ -77,6 +77,12 @@ static const cfg_t CONFIGS[] = { {"K1", "gfx1250", 4, 4, "fp16", "fp16", 8, 4, 256, 0, 1, 0, 0, 16, 1, 0}, /* 41: rejected: filter_rows does not divide n */ {"K4", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 96, 0, 1, 0, 0, 16, 1, 0}, /* 42: rejected: block not a power of two */ {"K2a", "gfx1250", 4, 4, "fp16", "fp16", 16, 1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 43: rejected: K2a vec 16 */ + {"K1", "gfx1250", 4, 4, "fp16", "fp16", 8, -1, 256, 0, 1, 0, 0, 16, 1, 0}, /* 44: rejected: filter_rows -1 */ + {"K2", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 2048, 0, 1, 0, 0, 16, 1, 0}, /* 45: rejected: K2 block over 1024 */ + {"K4", "gfx1250", 4, 4, "fp16", "fp16", 8, 1, 2048, 0, 1, 0, 0, 16, 1, 0}, /* 46: rejected: K4 block over 1024 */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 4, 4, 8, 3, 1, 0}, /* 47: rejected: LDS lanes not a power of two */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 4, 4, 8, 0, 1, 0}, /* 48: rejected: LDS lanes 0 */ + {"K2", "gfx1250", 2, 2, "fp16", "fp16", 8, 1, 256, 0, 1, 4, 0, 16, 1, 0}, /* 49: rejected: LDS region of zero width */ }; /* clang-format on */ diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py index ef59c469773a..43e60eaed633 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_winograd_fwd_rt_emit.py @@ -55,6 +55,12 @@ # 41 rejected: filter_rows does not divide n [gfx1250] # 42 rejected: block not a power of two [gfx1250] # 43 rejected: K2a vec 16 [gfx1250] +# 44 rejected: filter_rows -1 [gfx1250] +# 45 rejected: K2 block over 1024 [gfx1250] +# 46 rejected: K4 block over 1024 [gfx1250] +# 47 rejected: LDS lanes not a power of two [gfx1250] +# 48 rejected: LDS lanes 0 [gfx1250] +# 49 rejected: LDS region of zero width [gfx1250] from _emit_common import run_emit from kernels.common.conv_winograd_fwd_rt import ( RtFilterSpec, @@ -110,6 +116,12 @@ ('K1', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 4, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 41: rejected: filter_rows does not divide n ('K4', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 96, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 42: rejected: block not a power of two ('K2a', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 16, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 43: rejected: K2a vec 16 + ('K1', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': -1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 44: rejected: filter_rows -1 + ('K2', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 2048, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 45: rejected: K2 block over 1024 + ('K4', 'gfx1250', {'mh': 4, 'mw': 4, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 2048, 'ps': 0, 'tpt': 1, 'bh': 0, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 46: rejected: K4 block over 1024 + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 4, 'bh': 4, 'bw': 8, 'lanes': 3, 'nt': 1, 'fused': 0}), # 47: rejected: LDS lanes not a power of two + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 4, 'bh': 4, 'bw': 8, 'lanes': 0, 'nt': 1, 'fused': 0}), # 48: rejected: LDS lanes 0 + ('K2', 'gfx1250', {'mh': 2, 'mw': 2, 'io': 'fp16', 'op': 'fp16', 'vec': 8, 'fr': 1, 'block': 256, 'ps': 0, 'tpt': 1, 'bh': 4, 'bw': 0, 'lanes': 16, 'nt': 1, 'fused': 0}), # 49: rejected: LDS region of zero width ] # fmt: on diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py index 4d4df5ea40fa..53dc130bb137 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py @@ -59,11 +59,19 @@ def test_magic_div_rejects_zero(): lambda: R.RtFilterSpec(mh=3), lambda: R.RtFilterSpec(io_dtype="bf16", operand_dtype="bf16"), lambda: R.RtFilterSpec(filter_rows=4), + lambda: R.RtFilterSpec(filter_rows=0), + lambda: R.RtFilterSpec(filter_rows=-1), lambda: R.RtOutputSpec(vec=3), lambda: R.RtOutputSpec(block_size=96), + lambda: R.RtOutputSpec(block_size=2048), + lambda: R.RtInputSpec(block_size=2048), lambda: R.RtInputSpec(fused=True), lambda: R.RtInputSpec(mh=2, mw=2, lds=(2, 3), tiles_per_thread=4), lambda: R.RtInputSpec(mh=2, mw=2, lds=(8, 8), lds_lanes=32), + lambda: R.RtInputSpec(mh=2, mw=2, lds=(4, 8), tiles_per_thread=4, lds_lanes=3), + lambda: R.RtInputSpec(mh=2, mw=2, lds=(4, 8), tiles_per_thread=4, lds_lanes=0), + lambda: R.RtInputSpec(mh=2, mw=2, lds=(4, 8), tiles_per_thread=4, lds_lanes=-4), + lambda: R.RtInputSpec(mh=2, mw=2, lds=(4, 0)), lambda: R.RtPrescaleSpec(vec=16), ], ) @@ -161,6 +169,22 @@ def test_stage_launch_refuses_unfit_problems(): ) +def test_stage_launch_refuses_unsupported_geometry(): + """``stage_launch`` runs :meth:`RtGeometry.check` itself: a matching-tile + problem outside the 3x3 stride-1 kernels or the 32-bit offsets is refused.""" + p5 = ConvProblem(N=1, Hi=8, Wi=8, C=8, K=8, Y=5, X=5, pH=2, pW=2) + s2 = ConvProblem(N=1, Hi=8, Wi=8, C=8, K=8, Y=3, X=3, pH=1, pW=1, sH=2, sW=2) + big = ConvProblem(N=64, Hi=512, Wi=512, C=256, K=256, Y=3, X=3, pH=1, pW=1) + for p in (p5, s2): + for spec in (R.RtFilterSpec(), R.RtInputSpec(), R.RtOutputSpec()): + with pytest.raises(ValueError, match="3x3 stride-1"): + R.stage_launch(spec, R.RtGeometry(p, 4, 4)) + with pytest.raises(ValueError, match="32-bit"): + R.stage_launch(R.RtInputSpec(), R.RtGeometry(big, 4, 4)) + with pytest.raises(ValueError, match="32-bit"): + R.stage_launch(R.RtPrescaleSpec(), R.RtGeometry(big, 4, 4)) + + def test_geometry_refuses_non_3x3_and_huge_buffers(): p5 = ConvProblem(N=1, Hi=8, Wi=8, C=8, K=8, Y=5, X=5, pH=2, pW=2) assert R.RtGeometry(p5, 4, 4).check(False) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index ec47264b8e2a..8ff6eb6e175b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -611,10 +611,12 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_direct_fwd_wmma.py" "tests/library/tests/test_conv_fwd_split_k.py" "tests/library/tests/test_conv_winograd_fwd_cpu.py" + "tests/library/tests/test_conv_winograd_fwd_rt_cpu.py" "tests/library/tests/test_conv_winograd_gemm_wmma.py" "tests/library/tests/test_conv_winograd_gen_cpu.py" "tests/library/tests/test_conv_winograd_select.py" "tests/library/tests/dispatch/test_conv_winograd_dispatch.py" + "tests/library/tests/dispatch/test_conv_winograd_aot_dispatch.py" "tests/library/tests/test_direct_kernel_cache.py" "tests/library/tests/test_direct_depthwise_fwd_cpl.py" "tests/library/tests/test_conv_gradient_contract.py" @@ -864,6 +866,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_conv_dgrad_fused_weights.py" "tests/library/tests/test_direct_mfma_dgrad_correctness.py" "tests/library/tests/test_conv_winograd_fwd_gpu.py" + "tests/library/tests/test_conv_winograd_fwd_rt_gpu.py" "tests/library/tests/test_conv_winograd_gemm_wmma.py" "tests/library/tests/test_conv_winograd_gen_gpu.py" "tests/library/tests/test_conv_winograd_dispatch_gpu.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp index f77db1cb495e..2cf560b11ee4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_winograd_fwd_rt.cpp @@ -91,9 +91,9 @@ bool check_common(const Spec* s, std::string* err) bool check_block(const Spec* s, std::string* err) { - if(s->block_size < 32 || (s->block_size & (s->block_size - 1))) + if(s->block_size < 32 || s->block_size > 1024 || (s->block_size & (s->block_size - 1))) { - *err = "block_size must be a power of two >= 32"; + *err = "block_size must be a power of two in [32, 1024]"; return false; } return true; @@ -144,6 +144,16 @@ bool validate(const Spec* s, std::string* err) } if(input_lds(s)) { + if(s->lds[1] < 1) + { + *err = "lds must be (0, 0) or two positive extents"; + return false; + } + if(s->lds_lanes < 1 || (s->lds_lanes & (s->lds_lanes - 1))) + { + *err = "lds_lanes must be a power of two >= 1"; + return false; + } if((s->lds[0] * s->lds[1]) % s->tiles_per_thread) { *err = "tiles_per_thread must divide bh * bw"; From 0f78ac00c1a2609c77dfdc714d378d31ea750f7c Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sat, 31 Oct 2026 11:18:54 +0000 Subject: [PATCH 3/5] fix(rocke): size the fused-prescale Winograd K2 to its channel lanes The default shape-generic K2 with the fused prescale ran at block 1024 for anything over 256 channel lanes. A 1024-thread workgroup caps a wave at 128 VGPRs, so the wider binaries spilled, and at 512 lanes the block-1024 K2 was slower than the block-512 one the JIT pipeline uses (block = max(block, lanes)). - The AOT grid adds fused K2 binaries at block 512. - The default is the smallest block of 256 / 512 / 1024 that holds all lanes; tiles over 36 bins (F(6x4), F(4x6)) take the split prescale past 512 lanes, where their block-1024 K2 spills most and the split is faster. - CPU test of the block choice; GPU cases for bf16 C = 256, 2048 and 4096 (the fused max over more than one wave, blocks 512 / 1024, and the split). Co-Authored-By: Claude Opus 5.5 --- .../kernels/common/conv_winograd_fwd_rt.py | 23 +++++++++++------ .../tests/test_conv_winograd_fwd_rt_cpu.py | 25 +++++++++++++++++++ .../tests/test_conv_winograd_fwd_rt_gpu.py | 5 ++++ .../dsl_docs/instances/convolution.md | 6 +++-- 4 files changed, 50 insertions(+), 9 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py index fa661851facc..6cdbb4cd7f1f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_winograd_fwd_rt.py @@ -1418,6 +1418,14 @@ def stage_launch( _LDS_TPT = 4 _LDS_LANES = 16 _LDS_MIN_FILL = 0.75 +#: Workgroup sizes of the fused-prescale K2 binaries: the smallest that holds +#: all of a problem's channel lanes is the default (as the JIT pipeline's +#: ``block = max(block, lanes)``; a larger block caps the registers per lane +#: and spills). +_FUSED_BLOCKS = (256, 512, 1024) +#: Bins above which a fused K2 at block 1024 (128 VGPRs per lane) spills so +#: much that the split prescale is faster. +_FUSED_1024_MAX_BINS = 36 def pipeline_dtypes(tile, io_dtype: str) -> Tuple[str, bool]: @@ -1437,7 +1445,7 @@ def aot_stage_specs() -> list: """Every stage spec of the AOT grid: per I/O dtype, K2a at each vector width; per tile and operand dtype (bf16 F(2,3) also with fp16 operands), K1, K2 and K4 at each vector width -- K2 on the register path, with the - fused prescale at blocks 256 and 1024 when the tile takes the prescale, + fused prescale at blocks 256, 512 and 1024 when the tile takes the prescale, and for tiles under 36 bins also with 2-tile runs and through the LDS region.""" out = [] @@ -1456,7 +1464,7 @@ def aot_stage_specs() -> list: out.append(RtOutputSpec(mh, mw, io, op, v, 256, prescale=ps)) out.append(RtInputSpec(mh, mw, io, op, v, 256, prescale=ps)) if ps: - for blk in (256, 1024): + for blk in _FUSED_BLOCKS: out.append( RtInputSpec( mh, mw, io, op, v, blk, prescale=True, fused=True @@ -1488,9 +1496,9 @@ def default_stage_specs(problem: ConvProblem, tile, io_dtype: str) -> dict: """The AOT binaries a pipeline uses by default, ``{"K1", ("K2a",) "K2", "K4"}`` -> spec: the JIT pipeline's tuned knobs (8-wide channel vectors for F(2,3)-sized tiles, 4-wide for larger ones, narrowed to divide the - channel extent; the fused prescale while a workgroup can hold all - channels, else the split one; the LDS region for small tiles whose tile - grid fills it). A pure function of the problem; every spec is in + channel extent; the fused prescale on the smallest block that holds all + channels -- up to 1024 lanes, 512 for tiles over 36 bins -- else the split + one; the LDS region for small tiles whose tile grid fills it). A pure function of the problem; every spec is in :func:`aot_stage_specs`.""" mh, mw = _split_tile(tile) operand, ps = pipeline_dtypes((mh, mw), io_dtype) @@ -1507,8 +1515,9 @@ def default_stage_specs(problem: ConvProblem, tile, io_dtype: str) -> dict: } if ps: lanes = _pow2_ceil(p.C // vc) - if lanes <= 1024: - blk = 256 if lanes <= 256 else 1024 + cap = 1024 if (mh + 2) * (mw + 2) <= _FUSED_1024_MAX_BINS else 512 + if lanes <= cap: + blk = next(b for b in _FUSED_BLOCKS if b >= lanes) out["K2"] = RtInputSpec( mh, mw, io_dtype, operand, vc, blk, prescale=True, fused=True ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py index 53dc130bb137..50c71ecbdb33 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_cpu.py @@ -105,6 +105,8 @@ def test_aot_grid_builds(): ConvProblem(N=2, Hi=16, Wi=16, C=128, K=128, Y=3, X=3, pH=1, pW=1, groups=4), ConvProblem(N=1, Hi=8, Wi=8, C=16384, K=64, Y=3, X=3, pH=1, pW=1), ConvProblem(N=1, Hi=5, Wi=7, C=3, K=5, Y=3, X=3, pH=1, pW=1), + ConvProblem(N=1, Hi=8, Wi=8, C=2048, K=64, Y=3, X=3, pH=1, pW=1), + ConvProblem(N=1, Hi=8, Wi=8, C=4096, K=64, Y=3, X=3, pH=1, pW=1), ] @@ -130,6 +132,29 @@ def test_fused_prescale_falls_back_to_split_for_wide_c(): assert "K2a" not in narrow and narrow["K2"].fused +@pytest.mark.parametrize( + "C,tile,block", + [ + (256, (4, 4), 256), + (1024, (4, 4), 256), + (2048, (4, 4), 512), + (2048, (6, 4), 512), + (4096, (4, 4), 1024), + (4096, (4, 2), 1024), + (4096, (6, 4), None), # 48 bins: block 1024 spills, the split is faster + (8192, (4, 4), None), + ], +) +def test_fused_prescale_block_is_the_smallest_that_holds_c(C, tile, block): + p = ConvProblem(N=1, Hi=8, Wi=8, C=C, K=64, Y=3, X=3, pH=1, pW=1) + specs = R.default_stage_specs(p, tile, "bf16") + if block is None: + assert "K2a" in specs and not specs["K2"].fused + else: + assert "K2a" not in specs and specs["K2"].fused + assert specs["K2"].block_size == block + + @pytest.mark.parametrize("problem", _PROBLEMS[:4]) def test_stage_launch_spells_every_scalar(problem): for tile in ((4, 4), (2, 2), (6, 4)): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py index ef6b13a82d38..240bc1fbb0bc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_winograd_fwd_rt_gpu.py @@ -157,6 +157,11 @@ def _plans(req, tiles=None): ), # fmt: skip ("bf16", dict(N=2, C=128, K=128, Hi=16, Wi=16, Y=3, X=3, pad_h=1, pad_w=1, G=4)), ("fp16", dict(N=1, C=64, K=64, Hi=17, Wi=23, Y=3, X=3, pad_h=0, pad_w=2)), + # bf16 wide C: the fused prescale's max over more than one wave of channel + # lanes, at blocks 256 / 512 / 1024, and the split prescale past them. + ("bf16", dict(N=1, C=256, K=64, Hi=12, Wi=12, Y=3, X=3, pad_h=1, pad_w=1)), + ("bf16", dict(N=1, C=2048, K=64, Hi=8, Wi=8, Y=3, X=3, pad_h=1, pad_w=1)), + ("bf16", dict(N=1, C=4096, K=32, Hi=8, Wi=8, Y=3, X=3, pad_h=1, pad_w=1)), ] 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 17f733b2ecf0..6acb17960efb 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 @@ -1356,8 +1356,10 @@ bf16 I/O). Four implementations: the K2 register run or LDS region, and the prescale. Divisions by runtime extents are a multiply-high by a host-computed magic number (`magic_div`, exact for 31-bit indices). The fused prescale reduces the tile maxima over - a runtime number of channel lanes (up to the workgroup size); wider C runs - the split form (K2a). `stage_launch(spec, RtGeometry(...))` gives a stage's + a runtime number of channel lanes (up to the workgroup size); the default + is the smallest block (256, 512 or 1024) that holds all lanes, since a + larger block caps the registers per lane and spills. Wider C (and tiles over + 36 bins past 512 lanes) runs the split form (K2a). `stage_launch(spec, RtGeometry(...))` gives a stage's grid, block and scalar kernargs; `aot_stage_specs()` is the AOT grid and `default_stage_specs(problem, tile, dtype)` the binaries a pipeline uses by default (the JIT pipeline's tuned knobs, narrowed to the problem's channel From 7fe0c229a3c4fbf3380e3929d73033546d07d488 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 16:52:54 +0000 Subject: [PATCH 4/5] test(rocke): re-bless the shape-generic Winograd golden for LLVM 24 develop now records LLVM 24 as its own golden flavor instead of an alias of LLVM 23. Re-record the conv_winograd_fwd_rt shard after merging develop: drop the llvm24 -> llvm23 alias and add the six llvm24 cases. The llvm20/22/23 hashes are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../conv_winograd_fwd_rt.json | 57 ++++++++++++++++++- 1 file changed, 54 insertions(+), 3 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json index 27bec1da0aa6..0046899575d6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_winograd_fwd_rt.json @@ -1,7 +1,5 @@ { - "aliases": { - "llvm24": "llvm23" - }, + "aliases": {}, "flavors": { "llvm20": { "cases": { @@ -161,6 +159,59 @@ } }, "expected_failures": {} + }, + "llvm24": { + "cases": { + "conv_winograd_fwd_rt/gfx1250/bf16_k2a": { + "arch": "gfx1250", + "bytes": 6253, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_pmax_bf16_v8", + "sha256": "c1ff39de4529279f930d8a298065cae9f62d954445d76fc976a1e36b671367e3", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f2_lds_k2": { + "arch": "gfx1250", + "bytes": 231155, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f2_fp16_opfp16_v8_lds4x8_ll16_t4", + "sha256": "face2b28826e42ecf5ce95fcbe83cb070eb663fdb0a0c39b002420a661474218", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_bf16_fused_prescale_k2": { + "arch": "gfx1250", + "bytes": 138443, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_bf16_opfp16_v4_b256_r1_psf", + "sha256": "455b377ae4439439e12fc3a140e3c25ca968647271ed1f784b1051ae3b83b678", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k1": { + "arch": "gfx1250", + "bytes": 58876, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_filter_f4_fp16_opfp16_v8_fr2", + "sha256": "d86adca253c7e86d3dffe005f0f8291140201e6f4393ac726dbae6cb1ee55a6c", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f4_k2": { + "arch": "gfx1250", + "bytes": 102999, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_input_f4_fp16_opfp16_v4_b256_r1", + "sha256": "a0c9029f858812440da5b775d95e61acc636fa0c5bdb5aa1b40ac68072fc297c", + "status": "ok" + }, + "conv_winograd_fwd_rt/gfx1250/f6x4_bf16_prescale_k4": { + "arch": "gfx1250", + "bytes": 131049, + "family": "conv_winograd_fwd_rt", + "kernel_name": "wino_rt_output_f6x4_bf16_opfp16_v4_b256_ps", + "sha256": "1b87a48cf46fa9ccfcaa822f677bdd0f603e7c99e2dc2df50c7150b704e29ec2", + "status": "ok" + } + }, + "expected_failures": {} } }, "schema": "ck.dsl.ir_golden_sha256/v2" From 01cf5889cae4875fbd62ac3477a10fae9b7f4367 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 18:24:23 +0000 Subject: [PATCH 5/5] feat(rocke): run the AOT Winograd sweep in benchmark_conv_compare benchmark_conv_compare.py now runs benchmark_winograd_aot.py as a third engine beside the direct-conv and implicit-GEMM sweeps. It runs by default on a --run-from-cache comparison whose directory holds the Winograd cache, and --skip-winograd leaves it out. A JIT comparison, a non-fwd --direction or a missing cache skip it with an info line. - case_results: "winograd" is an engine, with its own merged and CSV columns. The winner is the fastest engine that measured the case, and the speedup is over the next best. - benchmark_winograd_aot.py reads the compare inputs: --miopen-cmd, the shape flags, and --dtype / --direction overrides. It skips the cases the winograd_aot route refuses before it loads the cache. It times the assembled pipeline with --timing staged|block. --results writes case_results rows, --skip-measured resumes (it replaces --resume), and it prints the standard Best: TFLOPS line. A case with no fitting pipeline exits 0; one whose pipelines all fail verification exits 1. - run_conv_shards.py --script winograd shards the sweep. Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_conv_compare.py | 172 ++++++++-- .../common/benchmark_winograd_aot.py | 314 ++++++++++++++---- .../library/benchmarks/common/case_results.py | 78 +++-- .../benchmarks/common/run_conv_shards.py | 9 +- .../library/tests/test_run_conv_shards.py | 312 +++++++++++++++++ .../dsl_docs/instances/convolution.md | 3 + 6 files changed, 761 insertions(+), 127 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py index cfa5c526a9b2..cf77e907ceb1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py @@ -1,12 +1,13 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Side-by-side comparison of direct-conv and implicit-GEMM conv benchmarks. +"""Side-by-side comparison of direct-conv, implicit-GEMM and Winograd conv benchmarks. Runs ``benchmark_direct_conv.py`` and ``benchmark_implicit_gemm_conv.py`` on the -same convolution shape, then prints a compact summary of the best kernel found -by each approach. +same convolution shape, and on a cache run also ``benchmark_winograd_aot.py`` +(the AOT Winograd ``F(m x m, 3 x 3)`` forward), then prints a compact summary +of the best kernel found by each approach. -The two benchmark scripts are invoked as subprocesses so this script carries no +The benchmark scripts are invoked as subprocesses so this script carries no benchmark logic of its own — it only forwards arguments and aggregates results. Usage @@ -34,11 +35,18 @@ --N 8 --Hi 56 --Wi 56 --C 64 --K 64 --groups 4 The two scripts' entries have distinct identities, so one directory can hold -both caches. +both caches. The Winograd sweep runs out of the same directory when its own +cache is there too (``benchmark_winograd_aot.py --compile-all --cache-dir +kernel_cache``), on the 3x3 stride-1 forward cases; without that cache, on a +JIT comparison or with ``--skip-winograd`` it is left out. The ``--top`` / ``--warmup`` / ``--iters`` / ``--jobs`` / ``--verify`` / ``--dtype`` / ``--direction`` / ``--run-from-cache`` / ``--early-stop`` / ``--early-stop-after`` / ``--fast-rank`` flags are forwarded to both scripts. +The Winograd sweep takes the shape, ``--warmup`` / ``--iters`` / ``--verify`` / +``--dtype`` / ``--direction`` / ``--timing`` and the results files; it times +each stage's binaries and the assembled pipelines rather than a list of +kernels, so it has no early stopping or fast ranking to forward to. Early stopping carries over from one script to the other: direct conv runs first and records each case's best time (a kernel that fails ``--verify`` is @@ -191,11 +199,9 @@ def _cpg_valid_for_direct(C: int, K: int, groups: int) -> tuple[bool, str]: return True, "" -def _build_shape_args(args) -> list[str]: - """Shared shape args forwarded to both scripts.""" - a: list[str] = [ - "--arch", - args.arch, +def _shape_flags(args) -> list[str]: + """The explicit shape, as every script spells it.""" + return [ "--N", str(args.N), "--Hi", @@ -220,6 +226,15 @@ def _build_shape_args(args) -> list[str]: str(args.pW), "--groups", str(args.groups), + ] + + +def _build_shape_args(args) -> list[str]: + """Shared shape args forwarded to both scripts.""" + a: list[str] = [ + "--arch", + args.arch, + *_shape_flags(args), "--top", str(args.top), "--warmup", @@ -301,6 +316,54 @@ def _build_implicit_only_args(args) -> list[str]: return a +def _build_winograd_args(args) -> list[str]: + """Args of the Winograd AOT sweep: its case input, the measurement flags + and the cache. It has no kernel list to cut short, so none of the early + stopping or fast-ranking flags apply to it.""" + a: list[str] = [ + "--arch", + args.arch, + "--warmup", + str(args.warmup), + "--iters", + str(args.iters), + "--run-from-cache", + args.run_from_cache, + "--timing", + args.timing, + ] + if args.verify: + a.append("--verify") + if args.dtype is not None: + a += ["--dtype", args.dtype] + if args.direction is not None: + a += ["--direction", args.direction] + if args.miopen_cmd: + a += ["--miopen-cmd", args.miopen_cmd] + elif args.miopen_file: + a += ["--miopen-file", args.miopen_file] + else: + a += _shape_flags(args) + return a + + +def _winograd_skip_reason(args) -> str: + """Why the Winograd sweep sits this comparison out (``""``: it runs).""" + if args.run_from_cache is None: + return "the Winograd sweep runs from an AOT cache (--run-from-cache)" + if args.direction not in (None, "fwd"): + return "Winograd runs the forward direction only" + from benchmarks.common.benchmark_winograd_aot import cache_index + + index = cache_index(args.run_from_cache, args.arch) + if not index.is_file(): + return ( + f"no Winograd cache ({index} is missing; build it with " + f"benchmark_winograd_aot.py --compile-all)" + ) + return "" + + def _build_miopen_args(args) -> list[str]: """MIOpen input args forwarded to both scripts.""" a: list[str] = [ @@ -327,30 +390,36 @@ def _build_miopen_args(args) -> list[str]: # The line an EarlyStop prints when a seeded sweep skipped every kernel. _ALL_STOPPED_RE = re.compile(r"all \d+ early-stopped: .*") +# The line the Winograd sweep prints for a case it does not run. +_NOT_WINOGRAD_RE = re.compile(r"not a Winograd case: (.*)") + def _write_results( path: str, rows_direct: str, rows_implicit: str, + rows_winograd: str | None = None, *, miopen_file: str | None = None, miopen_cmd: str | None = None, timing: str | None = None, ) -> None: - """Join the two sweeps' per-case rows into ``path`` (and a ``.csv``). + """Join the sweeps' per-case rows into ``path`` (and a ``.csv``). - The two files are what the children wrote under ``--results``; a sweep that - was skipped (``--skip-direct`` / ``--skip-implicit``) or that reached no - case leaves none, and :func:`case_results.read_rows` reads that as no rows, - so the merged table is simply one-sided. + The files are what the children wrote under ``--results``; a sweep that was + skipped (``--skip-direct`` / ``--skip-implicit`` / ``--skip-winograd``) or + that reached no case leaves none, and :func:`case_results.read_rows` reads + that as no rows, so the merged table simply lacks that side. The children key their rows by geometry, not by the line they read, so the MIOpenDriver command goes back on here, from the input this run was given. - ``timing`` is how both children timed their kernels; the merge takes only + ``timing`` is how the children timed their kernels; the merge takes only rows measured that way (``None``: the rows must agree among themselves). """ rows = case_results.merge( - case_results.read_rows(rows_direct) + case_results.read_rows(rows_implicit), + case_results.read_rows(rows_direct) + + case_results.read_rows(rows_implicit) + + (case_results.read_rows(rows_winograd) if rows_winograd else []), case_results.miopen_cmds_by_key(miopen_file=miopen_file, miopen_cmd=miopen_cmd), timing=timing, ) @@ -366,6 +435,8 @@ def _print_summary( direct: tuple[float | None, str], implicit: tuple[float | None, str], implicit_note: str = "", + winograd: "tuple[float | None, str] | None" = None, + winograd_note: str = "", ) -> None: direct_tflops, direct_name = direct implicit_tflops, implicit_name = implicit @@ -383,6 +454,8 @@ def _fmt(label, tflops, name, note=""): _fmt("direct-conv", direct_tflops, direct_name) _fmt("implicit-GEMM", implicit_tflops, implicit_name, implicit_note) + if winograd is not None: + _fmt("Winograd", winograd[0], winograd[1], winograd_note) if direct_tflops is not None and implicit_tflops is not None: ratio = direct_tflops / implicit_tflops @@ -391,15 +464,32 @@ def _fmt(label, tflops, name, note=""): f"\n Speedup direct/implicit: {ratio:.3f}x → {winner} wins", flush=True, ) + others = [ + (tflops, label) + for tflops, label in ( + (direct_tflops, "direct-conv"), + (implicit_tflops, "implicit-GEMM"), + ) + if tflops is not None + ] + if winograd is not None and winograd[0] is not None and others: + best_tflops, best_label = max(others) + ratio = winograd[0] / best_tflops + winner = "Winograd" if ratio >= 1.0 else best_label + print( + f" Speedup Winograd/{best_label}: {ratio:.3f}x → {winner} wins", + flush=True, + ) print(f"{'='*72}", flush=True) def main() -> int: parser = argparse.ArgumentParser( description=( - "Compare direct-conv and implicit-GEMM benchmarks for the same conv shape. " - "Both benchmark scripts are run as subprocesses; their output is echoed and " - "their best TFLOPS results are printed side-by-side." + "Compare direct-conv, implicit-GEMM and (on a cache run, 3x3 stride-1 " + "forward) Winograd benchmarks for the same conv shape. The benchmark " + "scripts are run as subprocesses; their output is echoed and their best " + "TFLOPS results are printed side-by-side." ) ) parser.add_argument( @@ -486,6 +576,13 @@ def main() -> int: dest="skip_implicit", help="skip the implicit-GEMM benchmark (direct-conv only)", ) + parser.add_argument( + "--skip-winograd", + action="store_true", + dest="skip_winograd", + help="skip the Winograd AOT benchmark (it runs by default on a cache " + "run whose directory holds the Winograd cache)", + ) rank_grp = parser.add_argument_group( "Fast ranking", @@ -718,8 +815,16 @@ def main() -> int: ) args.skip_direct = True + winograd_note = "" + if not args.skip_winograd: + winograd_note = _winograd_skip_reason(args) + if winograd_note: + print(f"[info] Winograd skipped: {winograd_note}", file=sys.stderr) + args.skip_winograd = True + direct_result: tuple[float | None, str] = (None, "") implicit_result: tuple[float | None, str] = (None, "") + winograd_result: tuple[float | None, str] = (None, "") implicit_note = "" with tempfile.TemporaryDirectory() as tmp: @@ -747,18 +852,25 @@ def main() -> int: rows_implicit = str( results_path.with_name(results_path.stem + ".implicit.jsonl") ) + rows_winograd = str( + results_path.with_name(results_path.stem + ".winograd.jsonl") + ) else: rows_direct = str(Path(tmp) / "direct_rows.jsonl") rows_implicit = str(Path(tmp) / "implicit_rows.jsonl") + rows_winograd = str(Path(tmp) / "winograd_rows.jsonl") + winograd_args = [] if args.skip_winograd else _build_winograd_args(args) if args.results: direct_args += ["--results", rows_direct] implicit_args += ["--results", rows_implicit] + winograd_args += ["--results", rows_winograd] if args.skip_measured: # Each sweep skips what *it* already recorded: a run killed - # between the two has the case on one side only, and resuming has - # to finish the other rather than skip it. + # between two of them has the case on some sides only, and + # resuming has to finish the others rather than skip it. direct_args += ["--skip-measured", args.skip_measured] implicit_args += ["--skip-measured", args.skip_measured] + winograd_args += ["--skip-measured", args.skip_measured] overall_rc = 0 @@ -788,17 +900,31 @@ def main() -> int: if tflops is None and m: implicit_note = m.group(0) + if not args.skip_winograd: + script = _SCRIPT_DIR / "benchmark_winograd_aot.py" + tflops, name, output, rc = _run_script(script, winograd_args) + winograd_result = (tflops, name) + if rc != 0: + overall_rc = 1 + print(_sweep_failure_message(script, rc), file=sys.stderr) + m = _NOT_WINOGRAD_RE.search(output) + if tflops is None and m: + winograd_note = f"not a Winograd case: {m.group(1)}" + if args.results: _write_results( args.results, rows_direct, rows_implicit, + rows_winograd, miopen_file=args.miopen_file, miopen_cmd=args.miopen_cmd, timing=args.timing, ) - _print_summary(direct_result, implicit_result, implicit_note) + _print_summary( + direct_result, implicit_result, implicit_note, winograd_result, winograd_note + ) return overall_rc diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py index e29746337837..2359122948b0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_winograd_aot.py @@ -10,13 +10,21 @@ * ``--compile-all --cache-dir DIR``: build every binary of the grid below into ``DIR//wino_rt`` (CPU only, ``--jobs`` workers); -* ``--run-from-cache DIR --miopen-file F``: for each 3x3 stride-1 case and - each tile (all of :data:`conv_winograd_fwd_rt.AOT_TILES`), pick each stage's - fastest fitting binary by timing it alone (K3: every plan, as - ``benchmark_winograd_conv.py`` sweeps them), verify the assembled pipeline - against the reference conv, and time the whole chain (min over - ``--repeats``). One result row per case (``--results``, JSON Lines) with - every tile's pipeline, its time and its binaries in ``top``. +* ``--run-from-cache DIR``: for each 3x3 stride-1 forward case (from + ``--miopen-file``, ``--miopen-cmd`` or the shape flags; the other cases are + skipped with the reason) and each tile (all of + :data:`conv_winograd_fwd_rt.AOT_TILES`), pick each stage's fastest fitting + binary by timing it alone (K3: every plan, as ``benchmark_winograd_conv.py`` + sweeps them), verify the assembled pipeline against the reference conv, and + time the whole chain as one call (``--timing``, best of ``--repeats``). + +The cache run takes the inputs and writes the rows of the other AOT conv +sweeps: ``--results`` writes one :mod:`benchmarks.common.case_results` row per +case (engine ``winograd``, the fastest tile's pipeline, the next tiles as +runners-up), ``--skip-measured`` resumes, and the ``Best:`` line is the one +``benchmark_conv_compare.py`` reads -- the comparison runs this sweep beside +the direct-conv and implicit-GEMM ones, and ``run_conv_shards.py --script +winograd`` shards it. The prescale and operand dtype follow ``benchmark_winograd_conv.py``'s ``auto`` rules (bf16 operands only for F(2,3) with bf16 I/O, also tried with @@ -151,11 +159,16 @@ def _compile_one(item): return name, meta, bytes(art.hsaco), art.kernel_name +def cache_index(cache_dir: str, arch: str) -> Path: + """The index ``--compile-all`` writes last: a cache run needs it.""" + return Path(cache_dir) / arch / "wino_rt" / "index.json" + + def compile_all(cache_dir: str, arch: str, jobs: int) -> int: import multiprocessing from concurrent.futures import ProcessPoolExecutor - out = Path(cache_dir) / arch / "wino_rt" + out = cache_index(cache_dir, arch).parent out.mkdir(parents=True, exist_ok=True) work = [(k, p, arch) for k, p in _job_list(arch)] print(f"compile-all: {len(work)} binaries -> {out}", flush=True) @@ -182,7 +195,7 @@ def compile_all(cache_dir: str, arch: str, jobs: int) -> int: index[name] = meta if i % 200 == 0: print(f" {i}/{len(work)} ({time.time() - t0:.0f}s)", flush=True) - (out / "index.json").write_text(json.dumps(index, indent=1, sort_keys=True)) + cache_index(cache_dir, arch).write_text(json.dumps(index, indent=1, sort_keys=True)) print( f"compile-all done: {len(index)} binaries, {n_bad} failed " f"({time.time() - t0:.0f}s)", @@ -213,8 +226,9 @@ def art(self): def load_cache(cache_dir: str, arch: str) -> List[Binary]: from kernels.common import conv_winograd_fwd_rt as W - d = Path(cache_dir) / arch / "wino_rt" - index = json.loads((d / "index.json").read_text()) + path = cache_index(cache_dir, arch) + d = path.parent + index = json.loads(path.read_text()) out = [] for name, meta in sorted(index.items()): spec = None @@ -458,10 +472,28 @@ def _fits(binary: Binary, tp: TilePipe) -> bool: return s.reason(tp.geo) is None +def _time_us(run, args) -> float: + """One call of ``run`` (a whole pipeline) in microseconds, measured the + ``--timing`` way: ``staged`` is the median of single stalled calls, as + hipDNN times a convolution; ``block`` the average of back-to-back calls.""" + if args.timing == "staged": + from statistics import median + + from rocke.runtime.stall_gate import time_launches_staged + + samples = time_launches_staged(run, warmup=args.warmup, iters=args.iters) + return median(samples) * 1e3 + from rocke.runtime import time_launches + + return time_launches(run, warmup=args.warmup, iters=args.iters) * 1e3 + + def run_case( case: Case, tile, op: str, cache: List[Binary], args, log ) -> Optional[dict]: - """The best verified pipeline of one tile: ``{label, ms, binaries, stages}``.""" + """The best verified pipeline of one tile: ``{label, ms, binaries, err}``; + ``{"failed": True}`` when no K3 finalist verifies, ``None`` when no + pipeline fits.""" from benchmarks.common import winograd_k3_wmma as K3W from benchmarks.common.benchmark_winograd_conv import verify_error from rocke.runtime import time_launches @@ -474,6 +506,8 @@ def run_case( tp = TilePipe(case, tile, op, prescale, chunks) try: + # Ranking a stage's binaries times each alone, back to back; only the + # assembled pipeline is timed the --timing way. def us(fn): return time_launches(fn, warmup=args.warmup, iters=args.iters) * 1e3 @@ -576,18 +610,20 @@ def us(fn): best = (run, k3label, k3bins, err) break if best is None: - return None + return dict(failed=True) run, k3label, k3bins, err = best total = None for _ in range(args.repeats): if args.cooldown > 0: time.sleep(args.cooldown) - t = us(run) + t = _time_us(run, args) total = t if total is None else min(total, t) bins = [pick[st].name for st in stages] + list(k3bins) + k2a = f"-k2a={pick['K2a'].name}" if "K2a" in pick else "" label = ( f"winograd-{case.io}-F{tile[0]}x{tile[1]}-op={op}-ps={int(prescale)}" - f"-k1={pick['K1'].name}-k2={pick['K2'].name}-k3={k3label}-k4={pick['K4'].name}" + f"-k1={pick['K1'].name}{k2a}-k2={pick['K2'].name}-k3={k3label}" + f"-k4={pick['K4'].name}" ) log( f" F{tile[0]}x{tile[1]} op={op} ps={int(prescale)} chunks={chunks}: " @@ -607,44 +643,138 @@ def _pow2(n: int, cap: int = 8) -> int: return v +def winograd_case_reason(problem, dtype: str, direction: str) -> Optional[str]: + """Why the AOT pipeline does not run a case (``None``: it may). + + The ``winograd_aot`` route's rule (``winograd_ineligible_reason``), + narrowed to the 3x3 filters the shape-generic transforms build.""" + from kernels.common.conv_winograd_select import winograd_ineligible_reason + + if direction != "fwd": + return "Winograd runs the forward direction only" + why = winograd_ineligible_reason(problem, dtype) + if why is None and (problem.Y, problem.X) != (3, 3): + why = f"{problem.Y}x{problem.X} filter (the AOT transforms are 3x3)" + return why + + +def cases_of(args) -> list: + """``(problem, dtype, direction)`` of every case the run was given. + + As the other conv sweeps read them: an MIOpenDriver line runs the + directions of its ``-F`` and its own dtype, ``--direction`` / ``--dtype`` + override both, and the shape flags are one case (fp16 fwd by default).""" + from benchmarks.common.benchmark_implicit_gemm_conv import ( + _FORW_TO_DIRECTIONS, + parse_miopen_cmd, + ) + + if args.miopen_file or args.miopen_cmd: + if args.miopen_cmd: + lines = [args.miopen_cmd] + else: + lines = Path(args.miopen_file).read_text().splitlines() + out = [] + for line in lines: + line = line.strip() + if not line or line.startswith("#"): + continue + problem, dtype, forw = parse_miopen_cmd(line) + directions = ( + [args.direction] + if args.direction + else _FORW_TO_DIRECTIONS.get(forw, ["fwd"]) + ) + out += [(problem, args.dtype or dtype, d) for d in directions] + return out + from kernels.common._conv_implicit_gemm_common import ConvProblem + + problem = ConvProblem( + N=args.N, + Hi=args.Hi, + Wi=args.Wi, + C=args.C, + K=args.K, + Y=args.Y, + X=args.X, + sH=args.sH, + sW=args.sW, + pH=args.pH, + pW=args.pW, + groups=args.groups, + ) + return [(problem, args.dtype or "fp16", args.direction or "fwd")] + + def run_from_cache(args) -> int: + from benchmarks.common.case_results import ( + RUNNERS_UP, + ResultsWriter, + measured_keys, + same_results_file, + ) + from benchmarks.common.early_stop import case_key + + # Filtered before anything is loaded: a comparison hands this sweep every + # case, and the ones it cannot run should cost a line, not a cache load. + cases = [] + for problem, dtype, direction in cases_of(args): + why = winograd_case_reason(problem, dtype, direction) + if why is not None: + print( + f"[skip] {problem.short()} {dtype} {direction}: " + f"not a Winograd case: {why}", + flush=True, + ) + continue + cases.append((problem, dtype)) + if not cases: + return 0 + # Read before the writer opens --results: the two may name one file. + done = measured_keys(args.skip_measured, "winograd", args.timing) + import torch # noqa: F401 -- claim the HIP context first - from benchmarks.common.benchmark_implicit_gemm_conv import parse_miopen_cmd - from benchmarks.common.early_stop import case_key from rocke.runtime.hip_module import Runtime + if not args.cus: + from benchmarks.common.winograd_k3_wmma import device_cus + + args.cus = device_cus() + if args.timing == "staged": + # Checked once, as batched_timing.time_batch_staged does: a device + # without the stall gate would otherwise fail every pipeline alike. + from rocke.runtime.hip_module import HipError + from rocke.runtime.stall_gate import default_gate + + try: + default_gate() + except HipError as e: + raise HipError( + f"{e}; this device cannot be timed with --timing staged, rerun " + f"with --timing block" + ) from e + cache = load_cache(args.run_from_cache, args.arch) print(f"cache: {len(cache)} binaries", flush=True) - cmds = [ - ln.strip() - for ln in open(args.miopen_file) - if ln.strip() and not ln.startswith("#") - ] - out = open(args.results, "a" if args.resume else "w") if args.results else None - done = set() - if args.resume and args.results and Path(args.results).exists(): - for ln in open(args.results): - try: - done.add(json.loads(ln)["miopen_cmd"]) - except (ValueError, KeyError): - pass + results_out = ResultsWriter( + args.results, + "winograd", + args.timing, + append=same_results_file(args.results, args.skip_measured), + ) rt = Runtime() rc = 0 - for ci, cmd in enumerate(cmds, 1): - if cmd in done: + n_skipped = 0 + for ci, (problem, dtype) in enumerate(cases, 1): + key = case_key(problem, dtype, "fwd") + if key in done: + n_skipped += 1 continue - problem, dtype, forw = parse_miopen_cmd(cmd) - if ( - forw != 1 - or (problem.Y, problem.X) != (3, 3) - or (problem.sH, problem.sW) != (1, 1) - ): - print(f"[skip] not a 3x3 stride-1 forward: {cmd}", flush=True) - continue - print(f"\n== [{ci}/{len(cmds)}] {problem.short()} {dtype}", flush=True) + print(f"\n== [{ci}/{len(cases)}] {problem.short()} {dtype}", flush=True) case = Case(rt, problem, dtype, args.seed, args.verify) res = [] + failed = False try: for tile in _tiles(): for op in _operands(tile, dtype): @@ -655,47 +785,85 @@ def run_from_cache(args) -> int: except (ValueError, RuntimeError) as e: print(f" [skip] F{tile} op={op}: {e}", flush=True) r = None - if r: - res.append(r) + if r is None: + continue + if r.get("failed"): + failed = True + continue + res.append(r) finally: case.free() if not res: - print(" no Winograd pipeline passed", flush=True) - rc = 1 + if failed: + # A pipeline that fits and computes the wrong answer is a bug; + # one that does not fit is coverage. + print(" no Winograd pipeline verified", flush=True) + rc = 1 + else: + print(" no Winograd pipeline fits this case", flush=True) continue res.sort(key=lambda r: r["ms"]) best = res[0] - print(f"Best: {best['ms'] * 1e3:.2f} us — {best['label']}", flush=True) - if out: - row = dict( - key=case_key(problem, dtype, "fwd"), - miopen_cmd=cmd, - engine="winograd", - shape=problem.short(), - dtype=dtype, - direction="fwd", - ms=round(best["ms"], 6), - kernel=best["label"], - timing="wino_aot", - top=[[r["label"], round(r["ms"], 6), r["binaries"]] for r in res], - ) - out.write(json.dumps(row, sort_keys=True) + "\n") - out.flush() - if out: - out.close() + sec = best["ms"] * 1e-3 + tflops = problem.flops / sec * 1e-12 + gbps = (case.xb + case.wb + case.yb) / sec * 1e-9 + # The line benchmark_conv_compare.py reads, as the other sweeps print. + print( + f"\nBest: {tflops:.1f} TFLOPS — {best['label']} {best['ms']:.3f} ms", + flush=True, + ) + results_out.row( + key=key, + case=ci, + shape=problem.short(), + dtype=dtype, + direction="fwd", + tflops=tflops, + ms=best["ms"], + gbps=gbps, + kernel=best["label"], + candidates=len(res), + # Every pipeline reported was verified, when --verify asked. + passed=True if args.verify else None, + runners_up=[r["label"] for r in res[1 : 1 + RUNNERS_UP]], + ) + results_out.close() + if n_skipped: + print( + f"\n{n_skipped} case(s) skipped: already in {args.skip_measured}", + flush=True, + ) return rc def main(argv=None) -> int: + from benchmarks.common import batched_timing, case_results + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--arch", default="gfx1250") ap.add_argument("--compile-all", action="store_true") ap.add_argument("--cache-dir", default="./kernel_cache") ap.add_argument("--jobs", type=int, default=os.cpu_count() or 8) ap.add_argument("--run-from-cache", default=None, metavar="DIR") - ap.add_argument("--miopen-file", default=None) - ap.add_argument("--results", default=None) - ap.add_argument("--resume", action="store_true") + ap.add_argument( + "--miopen-file", default=None, help="MIOpenDriver commands, one per line" + ) + ap.add_argument("--miopen-cmd", default=None, help="one MIOpenDriver command") + ap.add_argument( + "--dtype", + default=None, + choices=["fp16", "bf16"], + help="I/O dtype (default: fp16, or each MIOpenDriver command's)", + ) + ap.add_argument( + "--direction", + default=None, + choices=["fwd", "dgrad", "wgrad"], + help="direction (default: fwd, or each command's -F); Winograd runs fwd only", + ) + case_results.add_results_arg(ap, merged=False) + case_results.add_skip_measured_arg(ap) + batched_timing.add_timing_arg(ap) ap.add_argument("--verify", action="store_true") ap.add_argument("--warmup", type=int, default=3) ap.add_argument("--iters", type=int, default=10) @@ -704,15 +872,17 @@ def main(argv=None) -> int: ap.add_argument("--seed", type=int, default=42) ap.add_argument("--k3-finalists", type=int, default=4) ap.add_argument("--cus", type=int, default=0, help="0: query the device") + shape = ap.add_argument_group("Shape", "one case, without an MIOpenDriver input") + for name, default in ( + ("N", 8), ("Hi", 56), ("Wi", 56), ("C", 64), ("K", 64), ("Y", 3), ("X", 3), + ("sH", 1), ("sW", 1), ("pH", 1), ("pW", 1), ("groups", 1), + ): # fmt: skip + shape.add_argument(f"--{name}", type=int, default=default) args = ap.parse_args(argv) if args.compile_all: return compile_all(args.cache_dir, args.arch, args.jobs) - if not args.run_from_cache or not args.miopen_file: - ap.error("give --compile-all, or --run-from-cache with --miopen-file") - if not args.cus: - from benchmarks.common.winograd_k3_wmma import device_cus - - args.cus = device_cus() + if not args.run_from_cache: + ap.error("give --compile-all or --run-from-cache") return run_from_cache(args) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py index 6bb4b9b2a573..cc61bcae0232 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py @@ -12,7 +12,9 @@ So each sweep writes one JSON object per case as the case finishes -- :class:`ResultsWriter`, one call site per script, next to the ``Best:`` print -- -and :func:`merge` joins the two engines' rows into one table. +and :func:`merge` joins the engines' rows into one table. The engines are the +direct-conv and implicit-GEMM sweeps, and the AOT Winograd pipeline +(``benchmark_winograd_aot.py``) on the 3x3 stride-1 forward cases it runs. The join key is :func:`benchmarks.common.early_stop.case_key`: the normalised geometry plus dtype and direction, the key ``--early-stop-seed`` already uses to @@ -49,8 +51,8 @@ from benchmarks.common.batched_timing import BLOCK_TIMING, TIMINGS -#: The two sweeps, by the name their rows carry. -ENGINES = ("direct", "implicit") +#: The sweeps, by the name their rows carry. +ENGINES = ("direct", "implicit", "winograd") #: How many kernels below the winner each engine records, per case. #: @@ -81,6 +83,7 @@ def _engine_fields(engine: str) -> tuple: "direction", *_engine_fields("direct"), *_engine_fields("implicit"), + *_engine_fields("winograd"), "winner", "winner_ms", "winner_kernel", @@ -97,15 +100,17 @@ def _engine_fields(engine: str) -> tuple: #: #: The implicit side lists its whole ranking because that is where near-ties #: live -- the three names are what re-picking the winner would cost you. The -#: direct side lists only its winner: the two tables are joined per case, and -#: a row wide enough to hold both rankings is one nobody reads. +#: direct and Winograd sides list only their winner: the tables are joined per +#: case, and a row wide enough to hold every ranking is one nobody reads. CSV_FIELDS = ( "miopen_cmd", "direct_kernel", "implicit_kernel", *(f"implicit_kernel_{rank}" for rank in range(2, RUNNERS_UP + 2)), + "winograd_kernel", "direct_ms", "implicit_ms", + "winograd_ms", "winner_ms", "winner_kernel", "timing", @@ -121,19 +126,19 @@ def add_results_arg(parser, *, merged: bool) -> None: """ if merged: what = ( - "carrying the best direct-conv and implicit-GEMM kernel side by " - f"side (plus the next {RUNNERS_UP} kernels of each by name, which " - "is what a near-tie costs you), and a .csv of the same rows next " - "to it -- the .csv narrowed to the reading columns: the " + "carrying the best direct-conv, implicit-GEMM and Winograd kernel " + f"side by side (plus the next {RUNNERS_UP} kernels of each by name, " + "which is what a near-tie costs you), and a .csv of the same rows " + "next to it -- the .csv narrowed to the reading columns: the " "MIOpenDriver command, each engine's pick, their times, and the " - "winner. Written once both sweeps have finished, from the rows " + "winner. Written once every sweep has finished, from the rows " "each of them recorded." ) else: what = ( f"carrying the best kernel, its time, the next {RUNNERS_UP} " "kernels by name, and the key benchmark_conv_compare.py joins the " - "two sweeps on. Rows are appended as each case finishes, so an " + "sweeps on. Rows are appended as each case finishes, so an " "interrupted run keeps what it measured." ) parser.add_argument( @@ -200,7 +205,7 @@ def measured_keys(path, engine: str, timing: str = BLOCK_TIMING) -> set: def expand_merged_for_skip(rows: Iterable[dict]) -> List[dict]: """Per-engine skip rows from already-merged (:data:`MERGED_FIELDS`) rows. - A merged row carries both engines' times under ``{engine}_ms`` and no + A merged row carries every engine's time under ``{engine}_ms`` and no ``engine`` field of its own, so it is not what :func:`measured_keys` (and the sweeps' own ``--skip-measured`` reader) looks for. One output row per engine side the merged row actually has a time for -- :func:`merge` leaves @@ -460,8 +465,8 @@ def merge( ) -> List[dict]: """Join per-engine rows into one row per case. - Rows may come from both engines and from several shards. A key only one - engine reached still gets a row with the other side left ``None`` -- that + Rows may come from every engine and from several shards. A key only one + engine reached still gets a row with the other sides left ``None`` -- that is the common and interesting case (a shape the direct kernels do not cover), not an error. Within one engine and key, the faster row wins, so merging a case measured twice is idempotent. @@ -535,32 +540,38 @@ def merge( for field in MERGED_FIELDS: merged.setdefault(field, None) merged["winner"], merged["speedup"] = _verdict( - merged["direct_ms"], merged["implicit_ms"] + {engine: merged[f"{engine}_ms"] for engine in ENGINES} ) if merged["winner"] is not None: merged["winner_ms"] = merged[f"{merged['winner']}_ms"] merged["winner_kernel"] = merged[f"{merged['winner']}_kernel"] - # A row the other engine never reached keeps both empty, as ``winner`` - # does: one time is not a comparison, and the engine that did run it is + # A row only one engine reached keeps both empty, as ``winner`` does: + # one time is not a comparison, and the engine that did run it is # already on the row under its own name. merged["miopen_cmd"] = (miopen_cmds or {}).get(key) or miopen_cmd_from_key(key) out.append({field: merged[field] for field in MERGED_FIELDS}) return out -def _verdict(direct_ms, implicit_ms): - """``(winner, speedup)``: which engine was faster and by how much. +def _verdict(times: Dict[str, Optional[float]]): + """``(winner, speedup)``: which engine was fastest and by how much. - ``speedup`` is the loser's time over the winner's, so it reads as "the - winner is N times faster" whichever side won. Both ``None`` when only one + ``times`` is each engine's time, ``None`` where it did not reach the case. + ``speedup`` is the runner-up's time over the winner's, so it reads as "the + winner is N times faster than the next best" whichever side won; a tie goes + to the engine listed first in :data:`ENGINES`. Both ``None`` when only one engine reached the case -- a one-sided row is a fact about coverage, not a comparison. """ - if direct_ms is None or implicit_ms is None: + timed = sorted( + (ms, ENGINES.index(engine), engine) + for engine, ms in times.items() + if ms is not None + ) + if len(timed) < 2: return None, None - if direct_ms <= implicit_ms: - return "direct", round(implicit_ms / direct_ms, 3) if direct_ms else None - return "implicit", round(direct_ms / implicit_ms, 3) if implicit_ms else None + (best_ms, _, winner), (next_ms, _, _) = timed[:2] + return winner, round(next_ms / best_ms, 3) if best_ms else None def write(rows: Sequence[dict], path) -> Path: @@ -587,7 +598,7 @@ def write(rows: Sequence[dict], path) -> Path: def summary(rows: Sequence[dict]) -> str: - """One line: how many cases, and how the two engines split them.""" + """One line: how many cases, and how the engines split them.""" both = sum(1 for r in rows if r["direct_ms"] is not None and r["implicit_ms"]) direct_only = sum( 1 for r in rows if r["direct_ms"] is not None and r["implicit_ms"] is None @@ -595,10 +606,21 @@ def summary(rows: Sequence[dict]) -> str: implicit_only = sum( 1 for r in rows if r["direct_ms"] is None and r["implicit_ms"] is not None ) - wins = sum(1 for r in rows if r["winner"] == "direct") + wins = sum( + 1 + for r in rows + if r["direct_ms"] is not None + and r["implicit_ms"] is not None + and r["direct_ms"] <= r["implicit_ms"] + ) timing = f" ({rows[0].get('timing')} timing)" if rows else "" - return ( + line = ( f"{len(rows)} case(s){timing}: {both} measured by both " f"(direct faster on {wins}), {direct_only} direct-only, " f"{implicit_only} implicit-only" ) + winograd = sum(1 for r in rows if r.get("winograd_ms") is not None) + if winograd: + fastest = sum(1 for r in rows if r["winner"] == "winograd") + line += f"; Winograd measured {winograd} (fastest on {fastest})" + return line diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py index 8c4ce2c5c3b4..a2ebcff5b327 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py @@ -83,6 +83,7 @@ "compare": "benchmark_conv_compare.py", "direct": "benchmark_direct_conv.py", "implicit": "benchmark_implicit_gemm_conv.py", + "winograd": "benchmark_winograd_aot.py", } #: Shards per card when ``--shards`` is not given. See the module docstring: @@ -376,7 +377,7 @@ def make_parser() -> argparse.ArgumentParser: "--script", choices=sorted(SCRIPTS), default="compare", - help="benchmark to run in each shard (default: compare, both engines)", + help="benchmark to run in each shard (default: compare, every engine)", ) ap.add_argument( "--log-dir", @@ -569,9 +570,9 @@ def _progress(shards_done: int, busy: int) -> None: ] if args.script == "compare": # Each shard ran benchmark_conv_compare.py, whose own --results is - # already a merged row (both engines, winner, MIOpenDriver command): - # case_results.merge() expects the raw per-engine rows the "direct" - # and "implicit" scripts write, so handing it these would see no + # already a merged row (every engine, winner, MIOpenDriver command): + # case_results.merge() expects the raw per-engine rows the "direct", + # "implicit" and "winograd" scripts write, so handing it these would see no # "engine" field on any row and silently drop them all. A fresh row # replaces a carried one of the same key -- resume only reaches this # case again because --skip-measured did not cover it. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py index 1ee3262446f9..60b6423410e6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py @@ -947,3 +947,315 @@ def test_a_merge_of_one_timing_needs_no_argument(): ) def test_the_runner_reads_the_timing_it_forwards(forwarded, timing): assert run_conv_shards.forwarded_timing(forwarded) == timing + + +# --------------------------------------------------------------------------- winograd +def test_winograd_is_a_third_engine_on_the_row(): + rows = case_results.merge( + [ + _row("direct", "k", 1.0, "d"), + _row("implicit", "k", 0.8, "i"), + _row("winograd", "k", 0.4, "w", runners_up=["w-2"]), + ] + ) + assert len(rows) == 1 + assert rows[0]["winograd_kernel"] == "w" and rows[0]["winograd_kernel_2"] == "w-2" + assert rows[0]["winner"] == "winograd" + assert rows[0]["winner_kernel"] == "w" and rows[0]["winner_ms"] == 0.4 + # Over the next best engine, not over the slowest. + assert rows[0]["speedup"] == 2.0 + + +def test_winograd_against_one_other_engine_is_a_comparison(): + rows = case_results.merge( + [_row("implicit", "k", 0.5, "i"), _row("winograd", "k", 1.0, "w")] + ) + assert rows[0]["direct_ms"] is None + assert rows[0]["winner"] == "implicit" and rows[0]["speedup"] == 2.0 + + +def test_a_winograd_only_case_names_no_winner(): + rows = case_results.merge([_row("winograd", "k", 1.0, "w")]) + assert rows[0]["winner"] is None and rows[0]["speedup"] is None + + +def test_the_summary_counts_winograd_when_it_ran(): + two = case_results.merge( + [_row("direct", "k", 0.5, "d"), _row("implicit", "k", 1.0, "i")] + ) + assert "Winograd" not in case_results.summary(two) + three = case_results.merge( + [ + _row("direct", "k", 0.5, "d"), + _row("implicit", "k", 1.0, "i"), + _row("winograd", "k", 0.25, "w"), + ] + ) + line = case_results.summary(three) + # direct still beat implicit on the case, though Winograd won it. + assert "direct faster on 1" in line + assert "Winograd measured 1 (fastest on 1)" in line + + +def test_the_csv_carries_the_winograd_pick(tmp_path): + rows = case_results.merge( + [_row("direct", "k", 0.5, "d"), _row("winograd", "k", 0.25, "w")] + ) + csv_path = case_results.write(rows, tmp_path / "results.jsonl") + header, body = csv_path.read_text().splitlines()[:2] + fields = dict(zip(header.split(","), body.split(","))) + assert fields["winograd_kernel"] == "w" and fields["winograd_ms"] == "0.25" + + +def test_merged_rows_resume_the_winograd_side_too(): + merged = case_results.merge( + [_row("implicit", "k", 1.0, "i"), _row("winograd", "k", 0.5, "w")] + ) + skip = case_results.expand_merged_for_skip(merged) + assert sorted(r["engine"] for r in skip) == ["implicit", "winograd"] + + +def test_compare_merges_the_winograd_rows(tmp_path): + from benchmarks.common import benchmark_conv_compare as compare + + implicit = tmp_path / "implicit.jsonl" + implicit.write_text(json.dumps(_row("implicit", "k", 1.0, "i")) + "\n") + winograd = tmp_path / "winograd.jsonl" + winograd.write_text(json.dumps(_row("winograd", "k", 0.5, "w")) + "\n") + out = tmp_path / "merged.jsonl" + + compare._write_results( + str(out), str(tmp_path / "absent.jsonl"), str(implicit), str(winograd) + ) + + rows = case_results.read_rows(out) + assert rows[0]["winograd_kernel"] == "w" and rows[0]["winner"] == "winograd" + + +def _record_scripts(monkeypatch, compare) -> list: + """Stub the sweeps; return the ``(script name, argv)`` of each launch.""" + ran = [] + + def fake(script, extra, timeout=None): + ran.append((script.name, list(extra))) + return 1.0, "kern", "", 0 + + monkeypatch.setattr(compare, "_run_script", fake) + return ran + + +def _wino_cache(tmp_path, arch="gfx1250"): + from benchmarks.common.benchmark_winograd_aot import cache_index + + index = cache_index(str(tmp_path), arch) + index.parent.mkdir(parents=True) + index.write_text("{}") + return index + + +def test_compare_runs_winograd_from_a_cache_that_holds_it(tmp_path, monkeypatch): + from benchmarks.common import benchmark_conv_compare as compare + + _wino_cache(tmp_path) + ran = _record_scripts(monkeypatch, compare) + monkeypatch.setattr( + sys, + "argv", + ["x", "--arch", "gfx1250", "--run-from-cache", str(tmp_path), "--fast-rank", + "--early-stop", "1.5", "--results", str(tmp_path / "r.jsonl"), + "--skip-measured", str(tmp_path / "m.jsonl")], + ) # fmt: skip + + assert compare.main() == 0 + assert [name for name, _ in ran] == [ + "benchmark_direct_conv.py", + "benchmark_implicit_gemm_conv.py", + "benchmark_winograd_aot.py", + ] + argv = ran[-1][1] + flags = {a for a in argv if a.startswith("--")} + # Its own inputs, the timing and the files; none of the kernel-list flags. + assert {"--run-from-cache", "--timing", "--N", "--C", "--dtype"} <= flags + assert {"--results", "--skip-measured"} <= flags + assert not flags & {"--fast-rank", "--early-stop", "--top", "--jobs"} + assert argv[argv.index("--results") + 1].endswith("r.winograd.jsonl") + assert argv[argv.index("--dtype") + 1] == "fp16" + + +def test_compare_forwards_an_miopen_command_to_winograd(tmp_path, monkeypatch): + from benchmarks.common import benchmark_conv_compare as compare + + _wino_cache(tmp_path) + ran = _record_scripts(monkeypatch, compare) + monkeypatch.setattr( + sys, + "argv", + ["x", "--arch", "gfx1250", "--run-from-cache", str(tmp_path), + "--miopen-cmd", _MIOPEN_LINE], + ) # fmt: skip + + assert compare.main() == 0 + argv = ran[-1][1] + assert argv[argv.index("--miopen-cmd") + 1] == _MIOPEN_LINE + assert "--N" not in argv + + +@pytest.mark.parametrize( + "extra, why", + [ + ([], "no Winograd cache"), + (["--direction", "wgrad"], "forward direction only"), + (["--skip-winograd"], None), + ], +) +def test_compare_leaves_winograd_out(tmp_path, monkeypatch, capsys, extra, why): + from benchmarks.common import benchmark_conv_compare as compare + + if why != "no Winograd cache": + _wino_cache(tmp_path) + ran = _record_scripts(monkeypatch, compare) + monkeypatch.setattr( + sys, + "argv", + ["x", "--arch", "gfx1250", "--run-from-cache", str(tmp_path)] + extra, + ) + + assert compare.main() == 0 + assert "benchmark_winograd_aot.py" not in [name for name, _ in ran] + err = capsys.readouterr().err + if why: + assert "Winograd skipped: " in err and why in err + else: + assert "Winograd skipped" not in err + + +def test_a_jit_comparison_leaves_winograd_out(monkeypatch, capsys): + from benchmarks.common import benchmark_conv_compare as compare + + ran = _record_scripts(monkeypatch, compare) + monkeypatch.setattr(sys, "argv", ["x", "--skip-direct"]) + + assert compare.main() == 0 + assert [name for name, _ in ran] == ["benchmark_implicit_gemm_conv.py"] + assert "AOT cache" in capsys.readouterr().err + + +def test_compare_fails_when_the_winograd_sweep_dies(tmp_path, monkeypatch, capsys): + from benchmarks.common import benchmark_conv_compare as compare + + _wino_cache(tmp_path) + + def fake(script, extra, timeout=None): + if script.name == "benchmark_winograd_aot.py": + return None, "", "", 1 + return 1.0, "kern", "", 0 + + monkeypatch.setattr(compare, "_run_script", fake) + monkeypatch.setattr( + sys, "argv", ["x", "--arch", "gfx1250", "--run-from-cache", str(tmp_path)] + ) + + assert compare.main() == 1 + assert "benchmark_winograd_aot.py exited 1" in capsys.readouterr().err + + +def test_compare_says_why_a_shape_is_not_a_winograd_case(tmp_path, monkeypatch, capsys): + from benchmarks.common import benchmark_conv_compare as compare + + _wino_cache(tmp_path) + + def fake(script, extra, timeout=None): + if script.name == "benchmark_winograd_aot.py": + return None, "", "[skip] x fp16 fwd: not a Winograd case: stride != 1\n", 0 + return 1.0, "kern", "", 0 + + monkeypatch.setattr(compare, "_run_script", fake) + monkeypatch.setattr( + sys, "argv", ["x", "--arch", "gfx1250", "--run-from-cache", str(tmp_path)] + ) + + assert compare.main() == 0 + assert "no result — not a Winograd case: stride != 1" in capsys.readouterr().out + + +def test_the_shard_runner_can_shard_the_winograd_sweep(): + assert run_conv_shards.SCRIPTS["winograd"] == "benchmark_winograd_aot.py" + assert (run_conv_shards._SCRIPT_DIR / run_conv_shards.SCRIPTS["winograd"]).is_file() + + +# The Winograd sweep's own reading of its cases: host-side, before the cache or +# the device is touched. +def _wino_args(*argv): + from benchmarks.common import benchmark_winograd_aot as W + + captured = {} + + def fake_run(args): + captured["args"] = args + return 0 + + import unittest.mock + + with unittest.mock.patch.object(W, "run_from_cache", fake_run): + assert W.main(["--run-from-cache", "c", *argv]) == 0 + return captured["args"] + + +def test_winograd_reads_the_directions_of_an_miopen_line(): + from benchmarks.common import benchmark_winograd_aot as W + + line = _MIOPEN_LINE.replace("-F 1", "-F 0").replace("-u 2 -v 2", "-u 1 -v 1") + cases = W.cases_of(_wino_args("--miopen-cmd", line)) + assert [(d, dt) for _p, dt, d in cases] == [ + ("fwd", "bf16"), + ("dgrad", "bf16"), + ("wgrad", "bf16"), + ] + reasons = [W.winograd_case_reason(p, dt, d) for p, dt, d in cases] + assert reasons[0] is None + assert all("forward direction only" in r for r in reasons[1:]) + # --dtype / --direction override the line, as in the other sweeps. + over = W.cases_of(_wino_args("--miopen-cmd", line, "--dtype", "fp16", + "--direction", "fwd")) # fmt: skip + assert [(d, dt) for _p, dt, d in over] == [("fwd", "fp16")] + + +def test_winograd_reads_the_shape_flags(): + from benchmarks.common import benchmark_winograd_aot as W + + ((p, dtype, direction),) = W.cases_of( + _wino_args("--N", "2", "--C", "32", "--K", "48", "--groups", "2") + ) + assert (p.N, p.C, p.K, p.groups, p.Y, p.X) == (2, 32, 48, 2, 3, 3) + assert (dtype, direction) == ("fp16", "fwd") + assert _wino_args().timing == "staged" + + +@pytest.mark.parametrize( + "shape, why", + [ + (dict(Y=1, X=1, pH=0, pW=0), "1x1"), + (dict(sH=2, sW=2), "stride"), + (dict(Y=5, X=5, pH=2, pW=2), "3x3"), + (dict(C=2, K=2), "C/g"), + ], +) +def test_winograd_refuses_what_the_aot_route_refuses(shape, why): + from benchmarks.common import benchmark_winograd_aot as W + + base = dict(N=1, Hi=16, Wi=16, C=64, K=64, Y=3, X=3, pH=1, pW=1) + assert W.winograd_case_reason(ConvProblem(**base), "fp16", "fwd") is None + p = ConvProblem(**{**base, **shape}) + assert why in W.winograd_case_reason(p, "fp16", "fwd") + + +def test_a_winograd_run_with_no_winograd_case_loads_nothing(tmp_path, capsys): + """The comparison hands the sweep every case; one it cannot run is a skip + line and exit 0 -- no cache (there is none here), no device.""" + from benchmarks.common import benchmark_winograd_aot as W + + rc = W.main( + ["--run-from-cache", str(tmp_path / "absent"), "--sH", "2", "--sW", "2"] + ) + assert rc == 0 + assert "not a Winograd case: stride != 1" in capsys.readouterr().out 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 6acb17960efb..a3766a422a9a 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 @@ -1367,6 +1367,9 @@ bf16 I/O). Four implementations: `benchmarks/common/benchmark_winograd_aot.py` compiles the grid (`--compile-all`) and runs cases from it (`--run-from-cache`), choosing each stage's binary by timing and verifying the assembled pipeline. + `benchmark_conv_compare.py --run-from-cache` runs it beside the direct-conv + and implicit-GEMM sweeps when the cache directory holds its binaries, and + `--results` merges the three engines' rows. - **Generalized** (`conv_winograd_gen` + K3 with `taps`): any `F(m, r)` per axis (1-D `1 x r` / `r x 1`, 5x5), and segmented long filters: `S` segments of `F(m, m)` share one input transform and K3 sums