From e66ded8f26faf91e9eaae4efa014e76ff9c06293 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Mon, 5 Oct 2026 19:17:52 +0000 Subject: [PATCH 1/7] fix(rocke): fix fwd perf regressions after switch to aot --- .../common/benchmark_conv_compare.py | 160 +++++++--- .../common/benchmark_direct_conv.py | 26 +- .../common/benchmark_implicit_gemm_conv.py | 23 +- .../benchmarks/common/direct_kernel_sweep.py | 38 ++- .../library/benchmarks/common/early_stop.py | 115 ++++++- .../library/benchmarks/common/kernel_sweep.py | 24 +- .../kernels/common/conv_direct_grouped.py | 280 ++++++++++++------ .../rocke/library/tests/test_conv_abi.py | 13 +- .../rocke/library/tests/test_early_stop.py | 64 ++++ .../conv_direct_grouped_build_depthwise.cpp | 197 +++++++++--- ...direct_grouped_build_depthwise_spatial.cpp | 233 +++++++++++---- .../differential/golden/llvm_gfx_all.json | 12 +- 12 files changed, 900 insertions(+), 285 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 9c12a5dab865..eb3f85e83a1f 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 @@ -37,8 +37,14 @@ both caches. The ``--top`` / ``--warmup`` / ``--iters`` / ``--jobs`` / ``--verify`` / -``--dtype`` / ``--direction`` / ``--run-from-cache`` flags are forwarded to both -scripts. ``--dtype`` / ``--direction`` default to fp16 / fwd for explicit shapes; +``--dtype`` / ``--direction`` / ``--run-from-cache`` / ``--early-stop`` / +``--early-stop-after`` flags are forwarded to both scripts. + +Early stopping carries over from one script to the other: direct conv runs +first and records each case's best time, and the implicit-GEMM sweep starts +each case from that time instead of from nothing -- so from its first kernel it +skips any that cannot come within ``--early-stop`` times the direct result. +``--no-seed-from-direct`` turns that off. ``--dtype`` / ``--direction`` default to fp16 / fwd for explicit shapes; with a MIOpenDriver command they are taken from the command unless given. Notes @@ -57,6 +63,7 @@ import re import subprocess import sys +import tempfile from pathlib import Path _SCRIPT_DIR = Path(__file__).parent @@ -65,56 +72,75 @@ _BEST_RE = re.compile(r"Best:\s*([\d.]+)\s*TFLOPS\s*[—\-]\s*(.+)") +def _tee(stream, sink, lines: list) -> None: + """Echo *stream* to *sink* line by line as it arrives, keeping a copy.""" + for line in iter(stream.readline, ""): + sink.write(line) + sink.flush() + lines.append(line) + stream.close() + + def _run_script( script: Path, extra_args: list[str], timeout: "float | None" = None ) -> tuple[float | None, str, str]: - """Run *script* with *extra_args*; return (best_tflops, kernel_name, stdout).""" - cmd = [sys.executable, str(script)] + extra_args - print(f"\n{'='*72}", flush=True) - print(f"Running: {script.name} {' '.join(extra_args)}", flush=True) - print(f"{'='*72}", flush=True) + """Run *script* with *extra_args*; return (best_tflops, kernel_name, output). + The script's stdout and stderr are echoed live, as it prints them; ``output`` + is the captured stdout followed by stderr.""" import os import signal + import threading + # -u: a Python child block-buffers a piped stdout, which would hold its + # progress back until it exits. + cmd = [sys.executable, "-u", str(script)] + extra_args + print(f"\n{'='*72}", flush=True) + print(f"Running: {script.name} {' '.join(extra_args)}", flush=True) + print(f"{'='*72}", flush=True) + + out_lines: list = [] + err_lines: list = [] with subprocess.Popen( cmd, text=True, + bufsize=1, stdout=subprocess.PIPE, stderr=subprocess.PIPE, start_new_session=True, ) as popen: + readers = [ + threading.Thread( + target=_tee, args=(popen.stdout, sys.stdout, out_lines), daemon=True + ), + threading.Thread( + target=_tee, args=(popen.stderr, sys.stderr, err_lines), daemon=True + ), + ] + for t in readers: + t.start() try: - raw_out, raw_err = popen.communicate(timeout=timeout) + popen.wait(timeout=timeout) except subprocess.TimeoutExpired: try: os.killpg(os.getpgid(popen.pid), signal.SIGKILL) except ProcessLookupError: pass - raw_out, raw_err = popen.communicate() - elapsed = timeout + popen.wait() print( - f"\n[timeout] {script.name} exceeded {elapsed:.0f}s limit — killed.", + f"\n[timeout] {script.name} exceeded {timeout:.0f}s limit — killed.", file=sys.stderr, flush=True, ) - if raw_err: - print(raw_err, end="", file=sys.stderr, flush=True) - print(raw_out, end="", flush=True) - m = _BEST_RE.search(raw_out) - if m: - return float(m.group(1)), m.group(2).strip(), raw_out - return None, "", raw_out - - # Echo output to the terminal so the user can see the sweep progress. - print(raw_out, end="", flush=True) - if raw_err: - print(raw_err, end="", file=sys.stderr, flush=True) + for t in readers: + t.join() + raw_out = "".join(out_lines) + output = raw_out + "".join(err_lines) m = _BEST_RE.search(raw_out) if m: - return float(m.group(1)), m.group(2).strip(), raw_out - return None, "", raw_out + return float(m.group(1)), m.group(2).strip(), output + return None, "", output def _cpg_valid_for_direct(C: int, K: int, groups: int) -> tuple[bool, str]: @@ -185,6 +211,10 @@ def _build_mode_args(args) -> list[str]: a += ["--direction", args.direction] if args.run_from_cache is not None: a += ["--run-from-cache", args.run_from_cache] + if args.early_stop is not None: + a += ["--early-stop", str(args.early_stop)] + if args.early_stop_after is not None: + a += ["--early-stop-after", str(args.early_stop_after)] return a @@ -225,9 +255,14 @@ def _build_miopen_args(args) -> list[str]: return a + _build_mode_args(args) +# The line an EarlyStop prints when a seeded sweep skipped every kernel. +_ALL_STOPPED_RE = re.compile(r"all \d+ early-stopped: .*") + + def _print_summary( direct: tuple[float | None, str], implicit: tuple[float | None, str], + implicit_note: str = "", ) -> None: direct_tflops, direct_name = direct implicit_tflops, implicit_name = implicit @@ -236,14 +271,15 @@ def _print_summary( print("COMPARISON SUMMARY", flush=True) print(f"{'='*72}", flush=True) - def _fmt(label, tflops, name): + def _fmt(label, tflops, name, note=""): if tflops is None: - print(f" {label:<20} (no result — skipped or failed)", flush=True) + why = note or "skipped or failed" + print(f" {label:<20} (no result — {why})", flush=True) else: print(f" {label:<20} {tflops:7.1f} TFLOPS {name}", flush=True) _fmt("direct-conv", direct_tflops, direct_name) - _fmt("implicit-GEMM", implicit_tflops, implicit_name) + _fmt("implicit-GEMM", implicit_tflops, implicit_name, implicit_note) if direct_tflops is not None and implicit_tflops is not None: ratio = direct_tflops / implicit_tflops @@ -292,6 +328,26 @@ def main() -> int: action="store_true", help="forward --verify to both benchmark scripts", ) + parser.add_argument( + "--early-stop", + type=float, + default=None, + metavar="FACTOR", + help="forwarded to both scripts (default: theirs)", + ) + parser.add_argument( + "--early-stop-after", + type=int, + default=None, + metavar="N", + help="forwarded to both scripts (default: theirs)", + ) + parser.add_argument( + "--no-seed-from-direct", + action="store_true", + help="do not start the implicit-GEMM early stopping from the direct-conv " + "best times", + ) parser.add_argument( "--dtype", default=None, @@ -445,24 +501,38 @@ def main() -> int: direct_result: tuple[float | None, str] = (None, "") implicit_result: tuple[float | None, str] = (None, "") + implicit_note = "" + + with tempfile.TemporaryDirectory() as tmp: + # Direct conv's best time per case, handed to the implicit-GEMM sweep + # as its starting early-stop bound. + seed_file = str(Path(tmp) / "direct_best.json") + seed = not args.skip_direct and not args.no_seed_from_direct + if seed: + direct_args += ["--early-stop-record", seed_file] + + if not args.skip_direct: + tflops, name, _ = _run_script( + _SCRIPT_DIR / "benchmark_direct_conv.py", + direct_args, + # The JIT sweep compiles; a cache run only launches. + timeout=None if args.run_from_cache else 240, + ) + direct_result = (tflops, name) + + if not args.skip_implicit: + if seed and Path(seed_file).is_file(): + implicit_args += ["--early-stop-seed", seed_file] + tflops, name, output = _run_script( + _SCRIPT_DIR / "benchmark_implicit_gemm_conv.py", + implicit_args, + ) + implicit_result = (tflops, name) + m = _ALL_STOPPED_RE.search(output) + if tflops is None and m: + implicit_note = m.group(0) - if not args.skip_direct: - tflops, name, _ = _run_script( - _SCRIPT_DIR / "benchmark_direct_conv.py", - direct_args, - # The JIT sweep compiles; a cache run only launches. - timeout=None if args.run_from_cache else 240, - ) - direct_result = (tflops, name) - - if not args.skip_implicit: - tflops, name, _ = _run_script( - _SCRIPT_DIR / "benchmark_implicit_gemm_conv.py", - implicit_args, - ) - implicit_result = (tflops, name) - - _print_summary(direct_result, implicit_result) + _print_summary(direct_result, implicit_result, implicit_note) return 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py index 2e7bba889562..f960ddce589b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py @@ -588,7 +588,7 @@ def _run_depthwise_sweep( flush=True, ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "fwd") n_run = 0 for combo, spec, kernel in pending: block_w, block_waves = combo @@ -684,6 +684,8 @@ def _run_depthwise_sweep( if not results: print("No valid depthwise configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) @@ -796,7 +798,7 @@ def _run_sweep( flush=True, ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "fwd") n_run = 0 for combo, spec, kernel in pending: block_q, block_groups, double_buffer = combo @@ -894,6 +896,8 @@ def _run_sweep( if not results: print("No valid configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) @@ -1038,7 +1042,7 @@ def _run_wgrad_sweep( flush=True, ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "wgrad") n_run = 0 for combo, spec, kernel in pending: waves_k, waves_c, waves_q, hpb, mk = combo @@ -1136,6 +1140,8 @@ def _run_wgrad_sweep( if not results: print("No valid wgrad configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r["tflops"], reverse=True) @@ -1458,7 +1464,7 @@ def _run_dgrad_sweep( flush=True, ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "dgrad") n_run = 0 for combo, spec, kernel_or_pair in pending: is_mfma_pair = isinstance(kernel_or_pair, tuple) @@ -1717,6 +1723,8 @@ def run_mfma_dgrad(): if not results: print("No valid dgrad configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r["tflops"], reverse=True) @@ -1797,7 +1805,11 @@ def _cache_dispatch(args, arch: str, cases: list) -> int: return rc cache = KernelCache(Path(args.run_from_cache), arch) - rc = describe_cache(cache) + # Only the direct kernels: the implicit-GEMM entries sharing the cache + # are far more numerous and this run never launches them. + rc = describe_cache( + cache, directions=("direct_fwd", "direct_dgrad", "direct_dgrad_helper") + ) if rc: return rc return _run_from_cache(args, arch, cases, cache, directions) @@ -1894,7 +1906,7 @@ def _u8(t): flop = float(p.flops) bytes_xfer = float(In_t.nbytes + W_t.nbytes + Out_t.nbytes) results = [] - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, p, dtype, direction) for plan in plans: try: launches = [ @@ -1956,6 +1968,8 @@ def run(launches=launches): rt.free(ptr) if not results: print(" no cached kernel launched successfully", flush=True) + if _stop.summary(): + print(f" {_stop.summary()}", flush=True) overall_rc = overall_rc or 1 continue results.sort(key=lambda r: r[0], reverse=True) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index 35262a2fc1d3..9b338776ea04 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -2286,7 +2286,7 @@ def _make(*shape): flush=True, ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "fwd") n_run = 0 for combo, spec, artifact in pending: tile_m, tile_n, tile_k, warp_m, warp_n, warp_tile_mn, pipeline, epilogue = combo @@ -2414,6 +2414,8 @@ def _make(*shape): if not results: print("No valid configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) @@ -2758,7 +2760,7 @@ def _degrees(pipeline: str, async_dma: bool) -> tuple: # Set of config keys that have been pruned (skip remaining split-K degrees). _pruned_configs: set = set() - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "wgrad") n_run = 0 for combo, spec, resolved_split_k, artifact in pending: ( @@ -3195,6 +3197,8 @@ def _launch_two_stage( if not results: print("No valid wgrad configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) @@ -3400,7 +3404,7 @@ def _make(*shape): flush=True, ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, dtype, "dgrad") n_measured = 0 for _combo, spec, resolved_split_k, artifact in pending: sub_gemms = spec.compute_sub_gemms() @@ -3534,6 +3538,8 @@ def _launch_atomic(): if not results: print("No valid dgrad configurations found.", file=sys.stderr) + if _stop.summary(): + print(_stop.summary(), file=sys.stderr) return 1, [] results.sort(key=lambda r: r.tflops, reverse=True) @@ -3608,7 +3614,12 @@ def _cache_dispatch(args, arch, target, cases) -> int: return rc cache = KernelCache(Path(args.run_from_cache), arch) - rc = describe_cache(cache) + # Only the directions this run launches (see _run_from_cache): describing + # the whole cache would parse every other direction's entries for nothing. + used = directions or sorted( + {case[2] if len(case) > 2 else args.direction for case in cases} + ) + rc = describe_cache(cache, directions=used) if rc: return rc return _run_from_cache(args, arch, target, cases, cache, directions) @@ -3754,7 +3765,7 @@ def _u8(t): if args.verify or args.dump_fail else None ) - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop.for_case(args, problem, case_dtype, direction) n_failed = 0 n_cand = len(candidates) @@ -3920,6 +3931,8 @@ def run( ) if not results: print(" No successful launches.", flush=True) + if _stop.summary(): + print(f" {_stop.summary()}", flush=True) overall_rc = 1 continue diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py index 1e09a47aedb4..c623e6540f8d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py @@ -130,6 +130,16 @@ def _all_pads(kh: int) -> range: for k, s in itertools.product(_DEPTHWISE_FILTERS, _STRIDES) ), ), + # Depthwise forward for groups < wave_size: lanes split over channels and + # output columns. Same capabilities; the spec validator offers it only to + # problems with fewer groups than a wave. + "direct_depthwise_spatial": ( + "fwd", + tuple( + DirectCaps(KH=k, PAD=_same_pad(k), stride=s, cpg=1, kpg=1) + for k, s in itertools.product(_DEPTHWISE_FILTERS, _STRIDES) + ), + ), # Depthwise dgrad (row-streaming kernel): same filters and strides. "direct_depthwise_dgrad": ( "dgrad", @@ -173,6 +183,9 @@ def _knob_grid(variant: str) -> Iterator[dict]: elif variant in ("direct_depthwise", "direct_depthwise_dgrad"): for bw, waves in itertools.product(DW_BLOCK_W, DW_BLOCK_WAVES): yield dict(block_w=bw, block_waves=waves) + elif variant == "direct_depthwise_spatial": + for waves in DW_BLOCK_WAVES: + yield dict(block_waves=waves) elif variant == "direct_grouped_dgrad": for bq, bg in itertools.product(DGRAD_BLOCK_Q, BLOCK_GROUPS): yield dict(block_q=bq, block_groups=bg) @@ -215,6 +228,10 @@ def make_spec(variant: str, problem, knobs: dict): return dc.DirectDepthwiseSpec( problem=problem, name="rocke_direct_depthwise", **knobs ) + if variant == "direct_depthwise_spatial": + return dc.DirectDepthwiseSpatialSpec( + problem=problem, name="rocke_direct_depthwise_spatial", **knobs + ) if variant == "direct_depthwise_dgrad": return dc.DirectDepthwiseDgradStreamSpec( problem=problem, name="rocke_direct_dw_dgrad", **knobs @@ -273,6 +290,8 @@ def validate_spec(variant: str, spec, arch: str) -> Tuple[bool, str]: return dc.is_valid_spec(spec, arch=arch) if variant == "direct_depthwise": return dc.is_valid_depthwise_spec(spec, arch=arch) + if variant == "direct_depthwise_spatial": + return dc.is_valid_depthwise_spatial_spec(spec, arch=arch) if variant == "direct_depthwise_dgrad": return dc.is_valid_depthwise_dgrad_stream_spec(spec, arch=arch) if variant == "direct_grouped_dgrad": @@ -291,6 +310,7 @@ def _build_kernel(variant: str, spec, arch: str): "direct_grouped": dc.build_direct_conv, "direct_grouped_dgrad_mfma": dc.build_direct_conv, "direct_depthwise": dc.build_direct_depthwise, + "direct_depthwise_spatial": dc.build_direct_depthwise_spatial, "direct_depthwise_dgrad": dc.build_direct_depthwise_dgrad_streaming, "direct_grouped_dgrad": dc.build_direct_conv_dgrad, _TRANSPOSE: dc.build_direct_transpose_weights_dgrad, @@ -381,18 +401,24 @@ def _caps_of(identity: KernelIdentity) -> DirectCaps: ) -def probe_problem(caps: DirectCaps, dtype: str = DIRECT_DTYPE): +def probe_problem( + caps: DirectCaps, dtype: str = DIRECT_DTYPE, variant: Optional[str] = None +): """A problem with exactly ``caps`` that every knob in the grid fits. The emitted IR does not depend on N, H, W or groups (the shape-invariance cases in test_conv_abi.py assert it), so the probe only has to satisfy the spec validators: 64 groups divide every block_groups, 256 depthwise - channels fill the widest depthwise block, and 64x64 leaves every stride-2 - output wider than the widest W tile. + channels fill the widest depthwise block (the spatial depthwise kernel + takes fewer groups than a wave, so it gets 8), and 64x64 leaves every + stride-2 output wider than the widest W tile. """ from kernels.common.conv_direct_grouped import DirectConvProblem - groups = 256 if caps.cpg == 1 else 64 + if variant == "direct_depthwise_spatial": + groups = 8 + else: + groups = 256 if caps.cpg == 1 else 64 return DirectConvProblem( N=1, H=64, @@ -435,7 +461,7 @@ def _caps_jobs( """ out: List[Tuple[str, BuildJob]] = [] n_raw = 0 - probe = probe_problem(caps, dtype) + probe = probe_problem(caps, dtype, variant) for knobs in _knob_grid(variant): n_raw += 1 if variant == "direct_grouped_dgrad_mfma" and _mfma_knobs_reason(probe, knobs): @@ -598,7 +624,7 @@ def build_direct_job(job: BuildJob, arch: str, dtype: str): raise ValueError(f"direct conv builds {DIRECT_DTYPES}, not {dtype!r}") caps = DirectCaps(**job.caps) variant = job.identity.algorithm - spec = make_spec(variant, probe_problem(caps, dtype), job.spec_kwargs) + spec = make_spec(variant, probe_problem(caps, dtype, variant), job.spec_kwargs) ok, why = validate_spec(variant, spec, arch) if not ok: raise ValueError(why) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py index 2cc8221de52d..bbd6f85d1f3f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py @@ -7,11 +7,22 @@ launches run anyway, so they are timed too: a kernel whose warmup is already ``factor`` times slower than the best kernel measured so far is reported and skipped instead of being timed. + +The best time can also come from another run: ``--early-stop-record FILE`` +keeps each case's best time in FILE (rewritten whenever one improves, so a run +that is killed still leaves what it measured), and +``--early-stop-seed FILE`` starts each case's sweep from the times in such a +file. ``benchmark_conv_compare.py`` uses that to seed the implicit-GEMM sweep +with the direct-conv results. Cases are matched by their normalized geometry, +dtype and direction, so the two scripts' problem types meet on one key. """ from __future__ import annotations -from typing import Callable, Optional +import json +from dataclasses import asdict +from pathlib import Path +from typing import Callable, Dict, Optional DEFAULT_FACTOR = 5.0 # Kernels timed in full before early stopping engages: until then the best @@ -38,10 +49,69 @@ def add_early_stop_arg(parser) -> None: dest="early_stop_after", metavar="N", help="time the first N kernels of a sweep in full before --early-stop " - f"starts skipping (default: {DEFAULT_AFTER})", + f"starts skipping (default: {DEFAULT_AFTER}); a seeded case skips from " + "the first kernel", + ) + parser.add_argument( + "--early-stop-seed", + default=None, + dest="early_stop_seed", + metavar="FILE", + help="JSON of per-case best times (as written by --early-stop-record) " + "to start each case's early stopping from", + ) + parser.add_argument( + "--early-stop-record", + default=None, + dest="early_stop_record", + metavar="FILE", + help="keep each case's best measured time in FILE (JSON)", ) +def case_key(problem: object, dtype: str, direction: str) -> str: + """One key for a conv case, whichever script's problem type describes it.""" + from kernels.common.conv_args import ConvGeometry + + if hasattr(problem, "KH"): # DirectConvProblem + geom = ConvGeometry.from_direct_problem(problem) + else: + geom = ConvGeometry.from_problem(problem) + fields = ",".join(f"{k}={v}" for k, v in sorted(asdict(geom).items())) + return f"{direction}/{dtype}/{fields}" + + +# Per-process state behind --early-stop-seed / --early-stop-record, set up on +# the first EarlyStop.for_case call. +_seeds: Optional[Dict[str, float]] = None +_records: Dict[str, float] = {} +_record_path: Optional[Path] = None + + +def _load(args) -> Dict[str, float]: + global _seeds, _record_path + if _seeds is None: + seed_path = getattr(args, "early_stop_seed", None) + _seeds = {} + if seed_path and Path(seed_path).is_file(): + text = Path(seed_path).read_text() + _seeds = {k: float(v) for k, v in json.loads(text).items()} + record_path = getattr(args, "early_stop_record", None) + _record_path = Path(record_path) if record_path else None + return _seeds + + +def _record(key: str, ms: float) -> None: + prev = _records.get(key) + if prev is not None and ms >= prev: + return + _records[key] = ms + if _record_path is not None: + tmp = _record_path.with_name(_record_path.name + ".tmp") + tmp.write_text(json.dumps(_records, indent=1, sort_keys=True) + "\n") + tmp.replace(_record_path) + + class EarlyStop: """Times kernels for one sweep, skipping the hopeless ones. @@ -53,14 +123,35 @@ class EarlyStop: """ def __init__( - self, factor: float = DEFAULT_FACTOR, after: int = DEFAULT_AFTER + self, + factor: float = DEFAULT_FACTOR, + after: int = DEFAULT_AFTER, + *, + seed_ms: Optional[float] = None, + record_key: Optional[str] = None, ) -> None: self.factor = factor self.after = after self.n_measured = 0 - self.best_ms: Optional[float] = None + # A seed is a time another run measured on this case: it is trusted + # from the first kernel, so ``after`` does not apply to it. + self.seed_ms = seed_ms + self.best_ms: Optional[float] = seed_ms self.n_stopped = 0 self.last_warmup_ms: Optional[float] = None + self.record_key = record_key + + @classmethod + def for_case(cls, args, problem: object, dtype: str, direction: str) -> "EarlyStop": + """The early stopping for one case: seeded and recorded per the flags.""" + seeds = _load(args) + key = case_key(problem, dtype, direction) + return cls( + args.early_stop, + args.early_stop_after, + seed_ms=seeds.get(key), + record_key=key if getattr(args, "early_stop_record", None) else None, + ) def measure( self, @@ -75,7 +166,7 @@ def measure( engaged = ( self.factor > 0 and self.best_ms is not None - and self.n_measured >= self.after + and (self.seed_ms is not None or self.n_measured >= self.after) and warmup > 0 ) if engaged: @@ -91,15 +182,25 @@ def measure( self.n_measured += 1 if self.best_ms is None or ms < self.best_ms: self.best_ms = ms + if self.record_key is not None: + _record(self.record_key, ms) return ms def report(self, label: str) -> None: """Print the line for a kernel :meth:`measure` just skipped.""" + seeded = " (seed)" if self.best_ms == self.seed_ms else "" print( f" [early-stop] {label}: warmup {self.last_warmup_ms:.3f} ms > " - f"{self.factor:g} x best {self.best_ms:.3f} ms", + f"{self.factor:g} x best {self.best_ms:.3f} ms{seeded}", flush=True, ) def summary(self) -> str: - return f"{self.n_stopped} early-stopped" if self.n_stopped else "" + if not self.n_stopped: + return "" + if not self.n_measured and self.seed_ms is not None: + return ( + f"all {self.n_stopped} early-stopped: none came within " + f"{self.factor:g} x the seeded best {self.seed_ms:.3f} ms" + ) + return f"{self.n_stopped} early-stopped" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py index 932f94a628fe..2bf43e5f15e8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py @@ -1521,13 +1521,29 @@ def _launch_values_for(direction, problem, identity, ptrs, sizes, extras): ) -def describe_cache(cache: KernelCache, log=print) -> int: - """Print what is in the cache, grouped by direction.""" +def describe_cache( + cache: KernelCache, log=print, directions: Optional[Sequence[str]] = None +) -> int: + """Print what is in the cache, grouped by direction. + + ``directions`` limits the listing to the directions a run will use: every + entry is a metadata file to parse, and a full implicit-GEMM cache holds + hundreds of thousands of them -- reading all of them costs tens of seconds + before a run that needs a few hundred. + """ by_direction: Dict[str, int] = {} - for identity, _ in cache.list_all(): + entries = ( + cache.list_all() + if directions is None + else [e for d in directions for e in cache.list_all(d)] + ) + for identity, _ in entries: by_direction[identity.direction] = by_direction.get(identity.direction, 0) + 1 if not by_direction: - log("AOT cache is empty.") + if directions is None: + log("AOT cache is empty.") + else: + log(f"AOT cache has no {', '.join(directions)} kernels.") return 2 stale_by_direction = {} for direction, count in sorted(by_direction.items()): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index 9edcf442ba00..ccf855076b89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -5450,10 +5450,10 @@ def build_direct_depthwise_dgrad_streaming( ) -> KernelDef: """Build the ho-streaming depthwise dgrad kernel. - Streams dY rows (ho) in a Python-unrolled loop. For each ho, all r-taps - are evaluated (KH iterations, fully unrolled), contributing to hi values - via circular accumulator slots. W is preloaded into registers once before - the loop. The flush condition is computed at Python build time for each ho. + Streams dY rows (ho) in a runtime loop, KH rows per iteration, with the + next iteration's rows prefetched. For each ho, all r-taps are evaluated + (KH iterations, fully unrolled), contributing to hi values via circular + accumulator slots. W is preloaded into registers once before the loop. Tensor roles: A param — dY: output gradient, shape [N, Ho, Wo, groups], NHWK @@ -5608,52 +5608,81 @@ def _slot_index(slot: int, j: int) -> int: wi_vals = [b.add(wi_tile_start, b.const_i32(j)) for j in range(BLOCK_W)] wi_oks = [b.cmp_lt(wi, c_Wi) for wi in wi_vals] - for_op = b.scf_for_iter(c0, y_end_v, c_KH_v, slot_args, iv_name="dgs_ho") - with for_op as (y_base, loop_slots): - slots_flat = list(loop_slots) + # Tap (s, j) reads dY column wo = (wi + PAD - s) / stride, which depends on + # j - s only: one load per distinct d = j - s serves every tap (and every + # r) that reads it. The column bound, the stride divisibility and the + # channel guard do not depend on the row, so they are hoisted too. + cols = [] + for d in range(1 - p.KW, BLOCK_W): + wi_p_s = b.add(wi_tile_start, b.const_i32(p.PAD + d)) + if stride > 1: + div_ok = b.cmp_eq(b.mod(wi_p_s, c_stride_v), c0) + wo = b.div(wi_p_s, c_stride_v) + else: + div_ok = None + wo = wi_p_s + col_ok = b.land(b.land(b.cmp_ge(wi_p_s, c0), b.cmp_lt(wo, c_Wo)), ch_ok) + if stride > 1: + col_ok = b.land(col_ok, div_ok) + cols.append((wo, col_ok)) + + def load_rows(y0: Value) -> List[Value]: + """Raw dY values of rows y0 .. y0 + KH - 1, row-major by column. + + Rows past Ho, columns outside [0, Wo) or off the stride and lanes past + the channel count load from the OOB sentinel, which the buffer + descriptor returns as zero. + """ + vals: List[Value] = [] + for jj in range(p.KH): + y_v = b.add(y0, b.const_i32(jj)) + ho_ok = b.cmp_lt(y_v, c_Ho) + for wo, col_ok in cols: + dy_off, _ = dy_desc.offset(b, n=n, ho=y_v, wo=wo, ch=ch) + safe_dy = b.select( + b.land(col_ok, ho_ok), b.mul(dy_off, c_half_bytes), oob_sentinel + ) + vals.append( + b.buffer_load_bf16(a_rsrc, safe_dy, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_dy, c0) + ) + return vals + + # The dY rows are software-pipelined one iteration ahead: each iteration + # issues the loads of the next KH rows and carries them to the next + # iteration, then computes from the rows the previous one loaded. Loaded + # and used in the same iteration, the loads are scheduled next to their + # FMAs and each waits out the full memory latency before the next issues. + n_cols = len(cols) + row_names = [ + f"dgs_dy_r{jj}_d{d}" for jj in range(p.KH) for d in range(1 - p.KW, BLOCK_W) + ] + row_args = list(zip(row_names, load_rows(c0))) + + for_op = b.scf_for_iter(c0, y_end_v, c_KH_v, slot_args + row_args, iv_name="dgs_ho") + with for_op as (y_base, loop_vals): + slots_flat = list(loop_vals[: len(slot_args)]) + cur_rows = loop_vals[len(slot_args) :] + next_rows = load_rows(b.add(y_base, c_KH_v)) for jj in range(p.KH): y_v = b.add(y_base, b.const_i32(jj)) - ho_ok = b.cmp_lt(y_v, c_Ho) + dy_row = [ + b.cast_to_f32(dy_h) + for dy_h in cur_rows[jj * n_cols : (jj + 1) * n_cols] + ] + # Every tap accumulates, also into the slots of dX rows outside + # [0, H): the flush below never stores those, and resets each slot + # before a later row reuses it. for r_const in range(p.KH): # hi = y*stride + r - PAD. Its slot is static because # y_base*stride is a multiple of n_slots = stride*KH. slot = (jj * stride + r_const - p.PAD) % n_slots - hi_v = b.add(b.mul(y_v, c_stride_v), b.const_i32(r_const - p.PAD)) - hi_ok = b.land(b.cmp_ge(hi_v, c0), b.cmp_lt(hi_v, p_Hi)) - row_ok = b.land(hi_ok, ho_ok) - for s_const in range(p.KW): for j in range(BLOCK_W): - wi = wi_vals[j] - # wo = (wi + PAD - s) / stride, only where it divides. - wi_p_s = b.add(wi, b.const_i32(p.PAD - s_const)) - if stride > 1: - div_ok = b.cmp_eq(b.mod(wi_p_s, c_stride_v), c0) - wo = b.div(wi_p_s, c_stride_v) - else: - div_ok = None - wo = wi_p_s - - wo_ok = b.land(b.cmp_ge(wi_p_s, c0), b.cmp_lt(wo, c_Wo)) - tap_valid = b.land( - b.land(wo_ok, row_ok), b.land(ch_ok, wi_oks[j]) - ) - if stride > 1: - tap_valid = b.land(tap_valid, div_ok) - - dy_off, _ = dy_desc.offset(b, n=n, ho=y_v, wo=wo, ch=ch) - safe_dy = b.select( - tap_valid, b.mul(dy_off, c_half_bytes), oob_sentinel - ) - dy_h = ( - b.buffer_load_bf16(a_rsrc, safe_dy, c0) - if p.dtype == "bf16" - else b.buffer_load_f16(a_rsrc, safe_dy, c0) - ) - dy_f32 = b.select(tap_valid, b.cast_to_f32(dy_h), zero_f32) - + dy_f32 = dy_row[j - s_const + p.KW - 1] flat = _slot_index(slot, j) slots_flat[flat] = b.fma( weights_f32[r_const][s_const], dy_f32, slots_flat[flat] @@ -5684,7 +5713,7 @@ def _slot_index(slot: int, j: int) -> int: # Reset the slot for the next wrap. slots_flat[flat] = zero_f32 - b.scf_yield(*slots_flat) + b.scf_yield(*slots_flat, *next_rows) return b.kernel @@ -5708,7 +5737,8 @@ class DirectDepthwiseSpec: ``by * block_ch + w * 64 + l``. - Weights (``KH × KW`` values per channel) are preloaded into registers before the H-streaming loop. - - H-streaming loop (Python-level unroll, ``H + KH - 1`` iterations): + - H-streaming loop (runtime loop over ``H + KH - 1`` rows, ``KH`` rows + per iteration, the next iteration's input rows prefetched): circular ``KH``-slot accumulators per output W position, flushed to global memory one output row at a time. For stride > 1 only input rows that align to an output position produce a write. @@ -5907,6 +5937,14 @@ def build_direct_depthwise( # n_groups = ceil(n_iters / KH) times and the inner KH steps are unrolled # with STATIC slot indices, so the preloaded weights are still referenced # directly (no scatter, no runtime weight loads). + # + # The input rows are software-pipelined one group ahead: iteration g + # first issues the loads of group g + 1 and carries them to the next + # iteration, then computes group g from the rows the previous iteration + # loaded. Loaded and used in the same iteration, the rows are pulled back + # next to their FMAs by the scheduler, and every load then waits out the + # full memory latency before the next one issues (the unrolled JIT body + # had no such problem; in a loop it costs 2-6x). c1 = b.const_i32(1) c_KH = b.const_i32(p.KH) c_stride_rv = b.const_i32(c_stride_dw) @@ -5914,9 +5952,49 @@ def build_direct_depthwise( n_iters_v = b.add(p_Hi, b.const_i32(p.KH - 1)) n_groups_v = b.div(b.add(n_iters_v, b.const_i32(p.KH - 1)), c_KH) + # One load per distinct input column of a row: tap (w_out, s) reads + # column w_out * stride + s. ``row_cols`` maps each column to the first + # tap that reads it, in tap order. + row_cols: Dict[int, Tuple[int, int]] = {} + for w_out in range(BLOCK_W): + for s_const in range(p.KW): + row_cols.setdefault(w_out * c_stride_dw + s_const, (w_out, s_const)) + + def load_group(y0: Value) -> List[Value]: + """Raw input values of rows y0 .. y0 + KH - 1, row-major by column. + + Rows past the last iteration, halo rows and columns, and lanes past + the channel count all load from the OOB sentinel, which the buffer + descriptor returns as zero. + """ + vals: List[Value] = [] + for j in range(p.KH): + y_j = b.add(y0, b.const_i32(j)) + j_valid = b.cmp_lt(y_j, n_iters_v) + for w_out, s_const in row_cols.values(): + a_off, valid = a_desc.offset( + b, + n=n, + y_iter=y_j, + wo=b.add(q_tile_start, b.const_i32(w_out)), + s_off=b.const_i32(s_const), + c=ch, + ) + ok = b.land(b.land(valid, j_valid), ch_in_range) + safe_off = b.select(ok, b.mul(a_off, c_half_bytes), oob_sentinel) + vals.append( + b.buffer_load_bf16(a_rsrc, safe_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_off, c0) + ) + return vals + + n_acc = p.KH * BLOCK_W dw_iter_args = [ (f"dw_acc_kh{kh}_w{w}", zero_f32) for kh in range(p.KH) for w in range(BLOCK_W) ] + row_names = [f"dw_a_r{j}_c{col}" for j in range(p.KH) for col in row_cols] + dw_iter_args += list(zip(row_names, load_group(c0))) group_loop = b.scf_for_iter( c0, n_groups_v, @@ -5925,32 +6003,23 @@ def build_direct_depthwise( iv_name="dw_grp", elide_trailing_barrier=False, ) - with group_loop as (grp_iv, loop_accs): - new_accs = list(loop_accs) + with group_loop as (grp_iv, loop_vals): + new_accs = list(loop_vals[:n_acc]) + cur_rows = loop_vals[n_acc:] + next_rows = load_group(b.mul(b.add(grp_iv, c1), c_KH)) + col_idx = {col: i for i, col in enumerate(row_cols)} for j in range(p.KH): y_j = b.add(b.mul(grp_iv, c_KH), b.const_i32(j)) j_valid = b.cmp_lt(y_j, n_iters_v) + a_row = [ + b.cast_to_f32(a_h) + for a_h in cur_rows[j * len(row_cols) : (j + 1) * len(row_cols)] + ] for w_out in range(BLOCK_W): - w_pos = b.add(q_tile_start, b.const_i32(w_out)) for s_const in range(p.KW): - a_off, valid = a_desc.offset( - b, - n=n, - y_iter=y_j, - wo=w_pos, - s_off=b.const_i32(s_const), - c=ch, - ) - ok = b.land(b.land(valid, j_valid), ch_in_range) - safe_off = b.select(ok, b.mul(a_off, c_half_bytes), oob_sentinel) - a_h = ( - b.buffer_load_bf16(a_rsrc, safe_off, c0) - if p.dtype == "bf16" - else b.buffer_load_f16(a_rsrc, safe_off, c0) - ) - a_f32 = b.select(ok, b.cast_to_f32(a_h), zero_f32) + a_f32 = a_row[col_idx[w_out * c_stride_dw + s_const]] for r_const in range(p.KH): p_idx = (j - r_const + p.KH) % p.KH # STATIC slot idx = p_idx * BLOCK_W + w_out @@ -5991,7 +6060,7 @@ def build_direct_depthwise( for w_out in range(BLOCK_W): new_accs[P_FLUSH_j * BLOCK_W + w_out] = zero_f32 - b.scf_yield(*new_accs) + b.scf_yield(*new_accs, *next_rows) return b.kernel @@ -6118,9 +6187,7 @@ def build_direct_depthwise_spatial( p = spec.problem WAVE = spec.wave_size THREADS = spec.threads_per_block - n_w = spec.n_w_per_wave - BLOCK_W = spec.block_w - Wo = p.Wo + BLOCK_WAVES = spec.block_waves c_stride_dw = p.stride # AOT: the row count follows the runtime height, so only the @@ -6158,22 +6225,23 @@ def build_direct_depthwise_spatial( bx = b.block_id_x() n = b.block_id_z() + # AOT: groups is a kernarg, so the W positions per wave (and with them the + # block's W tile) are derived at run time; the host sizes the grid with + # the same n_w = wave_size // groups (see direct_launch_geometry). + n_w = b.div(c_wave, c_groups) q_out = b.add( - b.mul(bx, b.const_i32(BLOCK_W)), - b.add(b.mul(wave_id, b.const_i32(n_w)), w_in_wave), + b.mul(bx, b.mul(b.const_i32(BLOCK_WAVES), n_w)), + b.add(b.mul(wave_id, n_w), w_in_wave), ) # Guard: wasted threads when groups * n_w < wave_size - w_valid = b.cmp_lt(w_in_wave, b.const_i32(n_w)) + w_valid = b.cmp_lt(w_in_wave, n_w) q_ok = b.land(b.cmp_lt(q_out, c_Wo), w_valid) a_rsrc = b.buffer_rsrc(A, A_bytes) b_rsrc = b.buffer_rsrc(Bp, B_bytes) d_rsrc = b.buffer_rsrc(D, D_bytes) - a_desc = direct_a_descriptor_dynamic( - b, params, pad=p.PAD, stride=p.stride, w_upper=("wo", "s_off") - ) b_desc = TensorDescriptor.naive( "B", lengths=[p.total_k, p.KH, p.KW, 1], coord_names=("k", "r", "s", "c") ) @@ -6203,7 +6271,56 @@ def build_direct_depthwise_spatial( n_iters_v = b.add(p_Hi, b.const_i32(p.KH - 1)) n_groups_v = b.div(b.add(n_iters_v, b.const_i32(p.KH - 1)), c_KH) + # Split each input offset into a per-lane column part and a wave-uniform + # row part. The column part -- and with it the column bounds and the lane + # guard -- does not change down the image, so it is computed once here: + # one VGPR per filter column, the OOB sentinel where the tap falls + # outside [0, W) or the lane is idle. The row part (batch and input row) + # is the same for every lane and travels in the load's scalar offset, so + # a load in the loop needs no vector address arithmetic at all. + col_offs: List[Value] = [] + for s_const in range(p.KW): + wi = b.add(b.mul(q_out, c_stride_rv), b.const_i32(s_const - p.PAD)) + col_ok = b.land(b.land(b.cmp_ge(wi, c0), b.cmp_lt(wi, params["p_Wi"])), q_ok) + col_off = b.mul(b.add(b.mul(wi, params["p_A_stride_wi"]), ch), c_half_bytes) + col_offs.append(b.select(col_ok, col_off, oob_sentinel)) + + def load_row(y: Value) -> List[Value]: + """Raw input values of row ``y``, one per filter column. + + Halo rows and columns, rows past the image and idle lanes load from + the OOB sentinel, which the buffer descriptor returns as zero; an + invalid row also drops its scalar offset so the sentinel stays the + whole offset. + """ + hi = b.add(y, b.const_i32(-p.PAD)) + row_ok = b.land(b.cmp_ge(hi, c0), b.cmp_lt(hi, p_Hi)) + row_off = b.mul( + b.add(b.mul(n, params["p_A_stride_n"]), b.mul(hi, params["p_A_stride_hi"])), + c_half_bytes, + ) + row_off = b.select(row_ok, row_off, c0) + vals: List[Value] = [] + for col_off in col_offs: + safe_off = b.select(row_ok, col_off, oob_sentinel) + vals.append( + b.buffer_load_bf16(a_rsrc, safe_off, row_off) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, safe_off, row_off) + ) + return vals + + # The input rows are software-pipelined PF rows ahead: step y issues the + # loads of row y + PF before computing row y, and the PF rows in flight + # at the end of an iteration are carried into the next one. Loaded and + # used in the same step, every load would wait out the full memory + # latency before the next one issues. PF keeps ~16 loads in flight + # without carrying a whole KH x KW window for the large filters. + PF = min(p.KH, -(-16 // p.KW)) iter_args = [(f"sp_acc_{kh}", zero_f32) for kh in range(p.KH)] + for k in range(PF): + for s_const, a_h in enumerate(load_row(b.const_i32(k))): + iter_args.append((f"sp_a_r{k}_s{s_const}", a_h)) group_loop = b.scf_for_iter( c0, n_groups_v, @@ -6212,25 +6329,18 @@ def build_direct_depthwise_spatial( iv_name="sp_grp", elide_trailing_barrier=False, ) - with group_loop as (grp_iv, loop_accs): - new_accs = list(loop_accs) + with group_loop as (grp_iv, loop_vals): + new_accs = list(loop_vals[: p.KH]) + window = [loop_vals[p.KH + k * p.KW : p.KH + (k + 1) * p.KW] for k in range(PF)] for j in range(p.KH): y_j = b.add(b.mul(grp_iv, c_KH), b.const_i32(j)) j_valid = b.cmp_lt(y_j, n_iters_v) + window.append(load_row(b.add(y_j, b.const_i32(PF)))) + a_row = [b.cast_to_f32(a_h) for a_h in window.pop(0)] for s_const in range(p.KW): - a_off, valid = a_desc.offset( - b, n=n, y_iter=y_j, wo=q_out, s_off=b.const_i32(s_const), c=ch - ) - ok = b.land(b.land(valid, j_valid), q_ok) - safe_off = b.select(ok, b.mul(a_off, c_half_bytes), oob_sentinel) - a_h = ( - b.buffer_load_bf16(a_rsrc, safe_off, c0) - if p.dtype == "bf16" - else b.buffer_load_f16(a_rsrc, safe_off, c0) - ) - a_f32 = b.select(ok, b.cast_to_f32(a_h), zero_f32) + a_f32 = a_row[s_const] for r_const in range(p.KH): p_idx = (j - r_const + p.KH) % p.KH # STATIC new_accs[p_idx] = b.fma( @@ -6264,7 +6374,7 @@ def build_direct_depthwise_spatial( new_accs[P_FLUSH_j] = zero_f32 # unconditional static reset - b.scf_yield(*new_accs) + b.scf_yield(*new_accs, *[a_h for row in window for a_h in row]) return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py index f7a996216c39..bf0ef2ec5ac9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py @@ -431,9 +431,11 @@ def _direct_builders(): dc.DirectDepthwiseSpec(problem=pr), arch=_ARCH ), ), + # Fewer groups than a wave; the multiples below keep it that way while + # changing the W positions each wave covers. ( "depthwise_spatial", - dict(cpg=1, kpg=1, groups=32), + dict(cpg=1, kpg=1, groups=3), "fwd", lambda pr: dc.build_direct_depthwise_spatial( dc.DirectDepthwiseSpatialSpec(problem=pr), arch=_ARCH @@ -494,12 +496,6 @@ def _direct_builders(): ] -# The spatial depthwise kernel packs every channel of an image into one -# wavefront, so its lane layout -- and with it the emitted IR -- follows the -# group count. It is the one direct kernel whose groups are not a kernarg. -_GROUPS_BAKED = {"depthwise_spatial"} - - @pytest.mark.parametrize( "label,channels,direction,build", _direct_builders(), @@ -525,8 +521,7 @@ def test_direct_conv_abi_and_shape_invariance(label, channels, direction, build) [(2, 14, 14), (3, 29, 37), (1, 64, 64)], (1, 2, 3) ): shape = dict(channels) - if label not in _GROUPS_BAKED: - shape["groups"] *= group_mult + shape["groups"] *= group_mult problem = DirectConvProblem(N=N, H=H, W=W, KH=3, KW=3, PAD=1, stride=1, **shape) kernel = build(problem) fingerprints.append(_op_signature(kernel)) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py index d066c7cf7969..70669335a9d3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py @@ -54,3 +54,67 @@ def test_waits_for_the_first_kernels(monkeypatch): assert stop.measure(slow, warmup=3, iters=10) == 100.0 # ...from the fourth on early stopping applies. assert stop.measure(slow, warmup=3, iters=10) is None + + +def _reset_handoff(monkeypatch): + import benchmarks.common.early_stop as es + + monkeypatch.setattr(es, "_seeds", None) + monkeypatch.setattr(es, "_records", {}) + monkeypatch.setattr(es, "_record_path", None) + return es + + +def test_direct_and_implicit_problems_share_a_case_key(): + """The two scripts' problem types meet on one key, so a seed carries over.""" + from kernels.common._conv_implicit_gemm_common import ConvProblem + from kernels.common.conv_direct_grouped import DirectConvProblem + + from benchmarks.common.early_stop import case_key + + direct = DirectConvProblem(N=2, H=16, W=16, groups=4, cpg=8, kpg=8) + igemm = ConvProblem(N=2, Hi=16, Wi=16, C=32, K=32, Y=3, X=3, pH=1, pW=1, groups=4) + assert case_key(direct, "fp16", "fwd") == case_key(igemm, "fp16", "fwd") + assert case_key(direct, "fp16", "fwd") != case_key(igemm, "bf16", "fwd") + assert case_key(direct, "fp16", "fwd") != case_key(igemm, "fp16", "dgrad") + + +def test_record_then_seed_hands_the_best_over(monkeypatch, tmp_path): + """One run records its best; the next starts from it, from the first kernel.""" + from argparse import Namespace + + from kernels.common._conv_implicit_gemm_common import ConvProblem + + fast, slow, ok = object(), object(), object() + _fake_timer(monkeypatch, {fast: 1.0, slow: 9.0, ok: 3.0}) + problem = ConvProblem(N=2, Hi=16, Wi=16, C=32, K=32, Y=3, X=3, pH=1, pW=1) + record = tmp_path / "best.json" + + es = _reset_handoff(monkeypatch) + args = Namespace( + early_stop=5.0, early_stop_after=100, early_stop_record=str(record) + ) + first = es.EarlyStop.for_case(args, problem, "fp16", "fwd") + first.measure(slow, warmup=3, iters=10) + first.measure(fast, warmup=3, iters=10) + assert record.is_file() + + es = _reset_handoff(monkeypatch) + args = Namespace(early_stop=5.0, early_stop_after=100, early_stop_seed=str(record)) + second = es.EarlyStop.for_case(args, problem, "fp16", "fwd") + assert second.best_ms == 1.0 + # Seeded: `after` does not hold the first kernels back. + assert second.measure(slow, warmup=3, iters=10) is None + assert second.measure(ok, warmup=3, iters=10) == 3.0 + # A case the seed does not know starts from nothing. + other = es.EarlyStop.for_case(args, problem, "fp16", "dgrad") + assert other.best_ms is None + + +def test_all_stopped_against_a_seed_says_so(monkeypatch): + slow = object() + _fake_timer(monkeypatch, {slow: 9.0}) + stop = EarlyStop(5.0, after=100, seed_ms=1.0) + assert stop.measure(slow, warmup=3, iters=10) is None + assert "all 1 early-stopped" in stop.summary() + assert "1.000 ms" in stop.summary() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp index b5b68bcf8565..07fd4c134237 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp @@ -220,7 +220,7 @@ void rocke_dconv_dw_load_weights(rocke_dconv_dw_ctx_t* ctx) } /* ===================================================================== * - * H-row streaming loop (stride-aware, Python lines 2664-2712 updated) + * H-row streaming loop (stride-aware) * * Accumulator: acc[w_out][slot] — BLOCK_W × KH scalar f32 values. * The flush uses stride-aware ho_row = p_flush_val / c_stride_dw and only @@ -230,7 +230,66 @@ void rocke_dconv_dw_load_weights(rocke_dconv_dw_ctx_t* ctx) * The rows stream through a scf_for_iter group loop over * n_groups = ceil(n_iters / KH) groups of KH iterations each; the trip count * follows the runtime height (AOT), so there is no build-time-unrolled form. + * The input rows are software-pipelined one group ahead: each iteration + * issues the loads of the next group and carries them to the next iteration + * (see the Python source for why). * ===================================================================== */ + +/* Raw input values of rows y0 .. y0 + KH - 1, row-major by column: one load + * per distinct input column, taken at the first tap (col_w, col_s) that + * reads it. Out-of-range rows, columns and lanes load from the OOB + * sentinel, which the buffer descriptor returns as zero. */ +static void dw_load_group(rocke_dconv_dw_ctx_t* ctx, + rocke_value_t* y0, + rocke_value_t* n_iters_v, + const int* col_w, + const int* col_s, + int n_cols, + rocke_value_t** out) +{ + rocke_ir_builder_t* b = ctx->b; + int j, c; + + for(j = 0; j < ctx->p.KH; ++j) + { + rocke_value_t* y_j = rocke_b_add(b, y0, rocke_b_const_i32(b, j)); + rocke_value_t* j_valid = rocke_b_cmp_lt(b, y_j, n_iters_v); + + for(c = 0; c < n_cols; ++c) + { + rocke_value_t* a_off = NULL; + rocke_value_t* valid = NULL; + rocke_value_t* ok; + rocke_value_t* safe_off; + const char* off_names[5]; + rocke_value_t* off_vals[5]; + + off_names[0] = "n"; + off_vals[0] = ctx->n; + off_names[1] = "y_iter"; + off_vals[1] = y_j; + off_names[2] = "wo"; + { + rocke_value_t* cwo = rocke_b_const_i32(b, col_w[c]); + off_vals[2] = rocke_b_add(b, ctx->q_tile_start, cwo); + } + off_names[3] = "s_off"; + off_vals[3] = rocke_b_const_i32(b, col_s[c]); + off_names[4] = "c"; + off_vals[4] = ctx->ch; + rocke_transforms_descriptor_offset( + b, ctx->a_desc, off_names, off_vals, 5, &a_off, &valid); + + ok = rocke_b_land(b, rocke_b_land(b, valid, j_valid), ctx->ch_in_range); + safe_off = rocke_b_select( + b, ok, rocke_b_mul(b, a_off, ctx->c_half_bytes), ctx->oob_sentinel); + out[j * n_cols + c] = ctx->is_bf16 + ? rocke_b_buffer_load_bf16(b, ctx->a_rsrc, safe_off, ctx->c0) + : rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); + } + } +} + rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; @@ -238,6 +297,7 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) int KH = p->KH; int KW = p->KW; int BLOCK_W = ctx->BLOCK_W; + int stride = ctx->c_stride_dw; int n_iters = ctx->n_iters; /* AOT: the row count follows the runtime height, so only the * grouped-period scf.for form is usable -- the build-time-unrolled variant @@ -247,10 +307,26 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) * directly (no scatter, no runtime weight loads). */ (void)n_iters; { - int num_iargs = KH * BLOCK_W; + /* Distinct input columns of a row in first-tap order: tap (w_out, s) + * reads column w_out * stride + s; col_idx maps a column to its + * position in that order. */ + int max_col = (BLOCK_W - 1) * stride + KW; + int* col_w = (int*)alloca((size_t)max_col * sizeof(int)); + int* col_s = (int*)alloca((size_t)max_col * sizeof(int)); + int* col_num = (int*)alloca((size_t)max_col * sizeof(int)); + int* col_idx = (int*)alloca((size_t)max_col * sizeof(int)); + int n_cols = 0; + int num_acc = KH * BLOCK_W; + int num_rows; + int num_iargs; rocke_iter_arg_t* iargs; rocke_for_t group_loop; rocke_value_t** new_accs; + rocke_value_t** first_rows; + rocke_value_t** next_rows; + rocke_value_t** yields; + rocke_value_t* const* cur_rows; + rocke_value_t** a_row; rocke_value_t* c1; rocke_value_t* c_KH; rocke_value_t* c_stride_rv; @@ -258,13 +334,46 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) rocke_value_t* n_groups_v; int j; - /* Name storage: KH*BLOCK_W names, each at most 32 chars. */ + { + int w, s; + for(w = 0; w < max_col; ++w) + col_idx[w] = -1; + for(w = 0; w < BLOCK_W; ++w) + { + for(s = 0; s < KW; ++s) + { + int col = w * stride + s; + if(col_idx[col] >= 0) + continue; + col_idx[col] = n_cols; + col_w[n_cols] = w; + col_s[n_cols] = s; + col_num[n_cols] = col; + ++n_cols; + } + } + } + num_rows = KH * n_cols; + num_iargs = num_acc + num_rows; + first_rows = (rocke_value_t**)alloca((size_t)num_rows * sizeof(rocke_value_t*)); + next_rows = (rocke_value_t**)alloca((size_t)num_rows * sizeof(rocke_value_t*)); + yields = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); + a_row = (rocke_value_t**)alloca((size_t)n_cols * sizeof(rocke_value_t*)); + + c1 = rocke_b_const_i32(b, 1); + c_KH = rocke_b_const_i32(b, KH); + c_stride_rv = rocke_b_const_i32(b, stride); + /* n_iters = Hi + KH - 1 rows; the loop walks them KH at a time. */ + n_iters_v = rocke_b_add(b, ctx->params.p_Hi, rocke_b_const_i32(b, KH - 1)); + n_groups_v = rocke_b_div(b, rocke_b_add(b, n_iters_v, rocke_b_const_i32(b, KH - 1)), c_KH); + + /* Name storage: accumulators then carried rows, each at most 32 chars. */ { char(*name_store)[32] = (char(*)[32])alloca((size_t)num_iargs * 32 * sizeof(char)); - int kh, w, idx = 0; + int kh, w, c, idx = 0; iargs = (rocke_iter_arg_t*)alloca((size_t)num_iargs * sizeof(rocke_iter_arg_t)); - new_accs = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); + new_accs = (rocke_value_t**)alloca((size_t)num_acc * sizeof(rocke_value_t*)); for(kh = 0; kh < KH; ++kh) { @@ -276,15 +385,20 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) ++idx; } } + /* Group 0's rows: the loop's first iteration computes them. */ + dw_load_group(ctx, ctx->c0, n_iters_v, col_w, col_s, n_cols, first_rows); + for(kh = 0; kh < KH; ++kh) + { + for(c = 0; c < n_cols; ++c) + { + snprintf(name_store[idx], 32, "dw_a_r%d_c%d", kh, col_num[c]); + iargs[idx].name = name_store[idx]; + iargs[idx].init = first_rows[kh * n_cols + c]; + ++idx; + } + } } - c1 = rocke_b_const_i32(b, 1); - c_KH = rocke_b_const_i32(b, KH); - c_stride_rv = rocke_b_const_i32(b, ctx->c_stride_dw); - /* n_iters = Hi + KH - 1 rows; the loop walks them KH at a time. */ - n_iters_v = rocke_b_add(b, ctx->params.p_Hi, rocke_b_const_i32(b, KH - 1)); - n_groups_v = rocke_b_div(b, rocke_b_add(b, n_iters_v, rocke_b_const_i32(b, KH - 1)), c_KH); - group_loop = rocke_b_scf_for_iter(b, ctx->c0, n_groups_v, @@ -299,15 +413,24 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) /* Seed new_accs from loop-carried values */ { int i; - for(i = 0; i < num_iargs; ++i) + for(i = 0; i < num_acc; ++i) new_accs[i] = group_loop.iter_vars[i]; } + cur_rows = group_loop.iter_vars + num_acc; + /* Prefetch the next group before computing this one. */ + dw_load_group(ctx, + rocke_b_mul(b, rocke_b_add(b, group_loop.iv, c1), c_KH), + n_iters_v, + col_w, + col_s, + n_cols, + next_rows); for(j = 0; j < KH; ++j) { rocke_value_t* y_j; rocke_value_t* j_valid; - int w_out, s_const, r_const; + int w_out, s_const, r_const, c; int P_FLUSH_j = (j + 1) % KH; rocke_value_t* p_flush_rv; rocke_value_t* flush_ge; @@ -321,46 +444,15 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) y_j = rocke_b_add(b, mul_gk, cj); } j_valid = rocke_b_cmp_lt(b, y_j, n_iters_v); + for(c = 0; c < n_cols; ++c) + a_row[c] = rocke_b_cast_to_f32(b, cur_rows[j * n_cols + c]); /* Python order: outer w_out, inner s_const, innermost r_const. */ for(w_out = 0; w_out < BLOCK_W; ++w_out) { - rocke_value_t* w_pos; - { - rocke_value_t* cwo = rocke_b_const_i32(b, w_out); - w_pos = rocke_b_add(b, ctx->q_tile_start, cwo); - } - for(s_const = 0; s_const < KW; ++s_const) { - rocke_value_t* a_off = NULL; - rocke_value_t* valid = NULL; - rocke_value_t* ok; - rocke_value_t* safe_off; - rocke_value_t* a_h; - rocke_value_t* a_f32; - const char* off_names[5]; - rocke_value_t* off_vals[5]; - - off_names[0] = "n"; - off_vals[0] = ctx->n; - off_names[1] = "y_iter"; - off_vals[1] = y_j; - off_names[2] = "wo"; - off_vals[2] = w_pos; - off_names[3] = "s_off"; - off_vals[3] = rocke_b_const_i32(b, s_const); - off_names[4] = "c"; - off_vals[4] = ctx->ch; - rocke_transforms_descriptor_offset( - b, ctx->a_desc, off_names, off_vals, 5, &a_off, &valid); - - ok = rocke_b_land(b, rocke_b_land(b, valid, j_valid), ctx->ch_in_range); - safe_off = rocke_b_select( - b, ok, rocke_b_mul(b, a_off, ctx->c_half_bytes), ctx->oob_sentinel); - a_h = ctx->is_bf16 ? rocke_b_buffer_load_bf16(b, ctx->a_rsrc, safe_off, ctx->c0) - : rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); - a_f32 = rocke_b_select(b, ok, rocke_b_cast_to_f32(b, a_h), ctx->zero_f32); + rocke_value_t* a_f32 = a_row[col_idx[w_out * stride + s_const]]; for(r_const = 0; r_const < KH; ++r_const) { @@ -435,7 +527,14 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) } } - rocke_b_scf_yield(b, new_accs, num_iargs); + { + int i; + for(i = 0; i < num_acc; ++i) + yields[i] = new_accs[i]; + for(i = 0; i < num_rows; ++i) + yields[num_acc + i] = next_rows[i]; + } + rocke_b_scf_yield(b, yields, num_iargs); rocke_b_region_leave(b); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp index 1f7b43b46b0c..e9a1bea5fb54 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp @@ -61,6 +61,54 @@ static int sp_wo(const rocke_direct_conv_problem_t* p) return (p->W + 2 * p->PAD - p->KW) / s + 1; } +/* Raw input values of row y, one per filter column (Python: load_row). + * The per-lane column offsets come precomputed (OOB sentinel where the tap + * falls outside the image or the lane is idle); the wave-uniform row offset + * travels in the load's scalar offset. Halo rows load from the sentinel with + * a zero scalar offset, which the buffer descriptor returns as zero. */ +static void sp_load_row(rocke_ir_builder_t* bld, + rocke_value_t* y, + int PAD, + int KW, + const rocke_dconv_params_t* params, + rocke_value_t* n, + rocke_value_t* const* col_offs, + rocke_value_t* a_rsrc, + rocke_value_t* c_half_bytes, + rocke_value_t* oob_sentinel, + rocke_value_t* c0, + int is_bf16, + rocke_value_t** out) +{ + rocke_value_t* hi; + rocke_value_t* row_ok; + rocke_value_t* row_off; + int s_const; + + { + rocke_value_t* c_pad = rocke_b_const_i32(bld, -PAD); + hi = rocke_b_add(bld, y, c_pad); + } + { + rocke_value_t* ge = rocke_b_cmp_ge(bld, hi, c0); + rocke_value_t* lt = rocke_b_cmp_lt(bld, hi, params->p_Hi); + row_ok = rocke_b_land(bld, ge, lt); + } + { + rocke_value_t* off_n = rocke_b_mul(bld, n, params->p_A_stride_n); + rocke_value_t* off_h = rocke_b_mul(bld, hi, params->p_A_stride_hi); + row_off = rocke_b_mul(bld, rocke_b_add(bld, off_n, off_h), c_half_bytes); + } + row_off = rocke_b_select(bld, row_ok, row_off, c0); + + for(s_const = 0; s_const < KW; ++s_const) + { + rocke_value_t* safe_off = rocke_b_select(bld, row_ok, col_offs[s_const], oob_sentinel); + out[s_const] = is_bf16 ? rocke_b_buffer_load_bf16(bld, a_rsrc, safe_off, row_off) + : rocke_b_buffer_load_f16(bld, a_rsrc, safe_off, row_off); + } +} + /* ===================================================================== * * rocke_build_direct_depthwise_spatial * ===================================================================== */ @@ -69,7 +117,7 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( { rocke_ir_builder_t* bld = b; const rocke_direct_conv_problem_t* p; - int WAVE, BLOCK_WAVES, THREADS, n_w, BLOCK_W; + int WAVE, BLOCK_WAVES, THREADS; int Wo, c_stride_dw; int total_c, total_k; int is_bf16; @@ -96,6 +144,7 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( rocke_value_t* w_in_wave; rocke_value_t* bx; rocke_value_t* n; + rocke_value_t* n_w; rocke_value_t* q_out; rocke_value_t* w_valid; rocke_value_t* q_ok; @@ -104,7 +153,6 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( rocke_value_t* b_rsrc; rocke_value_t* d_rsrc; - const rocke_tensor_descriptor_t* a_desc; const rocke_tensor_descriptor_t* b_desc_v; const rocke_tensor_descriptor_t* d_desc; @@ -135,10 +183,7 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( p = &spec->problem; WAVE = spec->wave_size; BLOCK_WAVES = spec->block_waves; - (void)BLOCK_WAVES; THREADS = rocke_direct_depthwise_spatial_threads_per_block(spec); - n_w = rocke_direct_depthwise_spatial_n_w_per_wave(spec); - BLOCK_W = rocke_direct_depthwise_spatial_block_w(spec); c_stride_dw = p->stride > 0 ? p->stride : 1; Wo = sp_wo(p); total_c = rocke_direct_conv_problem_total_c(p); @@ -186,19 +231,23 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( bx = rocke_b_block_id_x(bld); n = rocke_b_block_id_z(bld); - /* q_out = bx*BLOCK_W + wave_id*n_w + w_in_wave - * Force Python left-to-right: const(BLOCK_W), mul, const(n_w), mul, add, add. */ + /* AOT: groups is a kernarg, so the W positions per wave (and with them the + * block's W tile) are derived at run time; the host sizes the grid with + * the same n_w = wave_size / groups. + * q_out = bx*(BLOCK_WAVES*n_w) + wave_id*n_w + w_in_wave + * Force Python left-to-right: const(BLOCK_WAVES), mul, mul, mul, add, add. */ + n_w = rocke_b_div(bld, c_wave, params.p_groups); { - rocke_value_t* c_bw = rocke_b_const_i32(bld, BLOCK_W); - rocke_value_t* mul_bx = rocke_b_mul(bld, bx, c_bw); - rocke_value_t* c_nw = rocke_b_const_i32(bld, n_w); - rocke_value_t* mul_wv = rocke_b_mul(bld, wave_id, c_nw); + rocke_value_t* c_bwv = rocke_b_const_i32(bld, BLOCK_WAVES); + rocke_value_t* block_w = rocke_b_mul(bld, c_bwv, n_w); + rocke_value_t* mul_bx = rocke_b_mul(bld, bx, block_w); + rocke_value_t* mul_wv = rocke_b_mul(bld, wave_id, n_w); rocke_value_t* inner = rocke_b_add(bld, mul_wv, w_in_wave); q_out = rocke_b_add(bld, mul_bx, inner); } /* Guard: wasted threads when groups * n_w < wave_size */ - w_valid = rocke_b_cmp_lt(bld, w_in_wave, rocke_b_const_i32(bld, n_w)); + w_valid = rocke_b_cmp_lt(bld, w_in_wave, n_w); q_ok = rocke_b_land(bld, rocke_b_cmp_lt(bld, q_out, c_Wo), w_valid); /* ---- buffer resources ---- */ @@ -207,15 +256,6 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( d_rsrc = rocke_b_buffer_rsrc(bld, D, D_bytes); /* ---- descriptors ---- */ - { - /* a_desc with runtime extents; PAD and stride stay build-time. */ - rocke_dynamic_tensor_descriptor_t* a_dyn - = rocke_dconv_a_descriptor_dynamic(bld, ¶ms, p->PAD, c_stride_dw, "wo", "s_off"); - if(!a_dyn) - return NULL; - a_desc = &a_dyn->base; - } - { /* b_desc_v = naive("B", [total_k, KH, KW, 1]) */ static const char* const b_coords[4] = {"k", "r", "s", "c"}; @@ -270,14 +310,22 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( * ================================================================== */ /* AOT: the row count follows the runtime height, so the rows stream * through an scf.for over KH-row periods; a build-time unroll would - * bake Hi into the trip count. */ + * bake Hi into the trip count. The input rows are software-pipelined PF + * rows ahead (see the Python source for why). */ { /* ---- scf_for_iter group loop path ---- */ int KH = p->KH; - int num_iargs = KH; + int KW = p->KW; + int PF = (16 + KW - 1) / KW < KH ? (16 + KW - 1) / KW : KH; + int num_iargs = KH + PF * KW; rocke_iter_arg_t* iargs; rocke_for_t group_loop; rocke_value_t** new_accs; + /* window[k * KW + s]: rows in flight, oldest first; one spare row + * slot for the row a step appends before it drops the oldest. */ + rocke_value_t** window; + rocke_value_t* a_row[ROCKE_DCONV_DW_MAX_KW]; + rocke_value_t* col_offs[ROCKE_DCONV_DW_MAX_KW]; rocke_value_t* c1; rocke_value_t* c_KH; rocke_value_t* c_stride_rv; @@ -285,12 +333,48 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( rocke_value_t* n_groups_v; int j; + window = (rocke_value_t**)alloca((size_t)(PF + 1) * KW * sizeof(rocke_value_t*)); + + c1 = rocke_b_const_i32(bld, 1); + c_KH = rocke_b_const_i32(bld, KH); + c_stride_rv = rocke_b_const_i32(bld, c_stride_dw); + /* n_iters = Hi + KH - 1 rows; the loop walks them KH at a time. */ + n_iters_v = rocke_b_add(bld, params.p_Hi, rocke_b_const_i32(bld, KH - 1)); + n_groups_v + = rocke_b_div(bld, rocke_b_add(bld, n_iters_v, rocke_b_const_i32(bld, KH - 1)), c_KH); + + /* Per-lane column offsets, hoisted out of the loop (Python: col_offs). */ + { + int s_const; + for(s_const = 0; s_const < KW; ++s_const) + { + rocke_value_t* wi; + rocke_value_t* col_ok; + rocke_value_t* col_off; + { + rocke_value_t* q_s = rocke_b_mul(bld, q_out, c_stride_rv); + rocke_value_t* c_s = rocke_b_const_i32(bld, s_const - p->PAD); + wi = rocke_b_add(bld, q_s, c_s); + } + { + rocke_value_t* ge = rocke_b_cmp_ge(bld, wi, c0); + rocke_value_t* lt = rocke_b_cmp_lt(bld, wi, params.p_Wi); + col_ok = rocke_b_land(bld, rocke_b_land(bld, ge, lt), q_ok); + } + { + rocke_value_t* w_off = rocke_b_mul(bld, wi, params.p_A_stride_wi); + col_off = rocke_b_mul(bld, rocke_b_add(bld, w_off, ch), c_half_bytes); + } + col_offs[s_const] = rocke_b_select(bld, col_ok, col_off, oob_sentinel); + } + } + { char(*name_store)[32] = (char(*)[32])alloca((size_t)num_iargs * 32 * sizeof(char)); - int kh; + int kh, k, s_const, idx; iargs = (rocke_iter_arg_t*)alloca((size_t)num_iargs * sizeof(rocke_iter_arg_t)); - new_accs = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); + new_accs = (rocke_value_t**)alloca((size_t)KH * sizeof(rocke_value_t*)); for(kh = 0; kh < KH; ++kh) { @@ -298,16 +382,33 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( iargs[kh].name = name_store[kh]; iargs[kh].init = zero_f32; } + /* Rows 0 .. PF-1: in flight when the loop starts. */ + idx = KH; + for(k = 0; k < PF; ++k) + { + sp_load_row(bld, + rocke_b_const_i32(bld, k), + p->PAD, + KW, + ¶ms, + n, + col_offs, + a_rsrc, + c_half_bytes, + oob_sentinel, + c0, + is_bf16, + window + k * KW); + for(s_const = 0; s_const < KW; ++s_const) + { + snprintf(name_store[idx], 32, "sp_a_r%d_s%d", k, s_const); + iargs[idx].name = name_store[idx]; + iargs[idx].init = window[k * KW + s_const]; + ++idx; + } + } } - c1 = rocke_b_const_i32(bld, 1); - c_KH = rocke_b_const_i32(bld, KH); - c_stride_rv = rocke_b_const_i32(bld, c_stride_dw); - /* n_iters = Hi + KH - 1 rows; the loop walks them KH at a time. */ - n_iters_v = rocke_b_add(bld, params.p_Hi, rocke_b_const_i32(bld, KH - 1)); - n_groups_v - = rocke_b_div(bld, rocke_b_add(bld, n_iters_v, rocke_b_const_i32(bld, KH - 1)), c_KH); - group_loop = rocke_b_scf_for_iter(bld, c0, n_groups_v, @@ -321,8 +422,10 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( { int i; - for(i = 0; i < num_iargs; ++i) + for(i = 0; i < KH; ++i) new_accs[i] = group_loop.iter_vars[i]; + for(i = 0; i < PF * KW; ++i) + window[i] = group_loop.iter_vars[KH + i]; } for(j = 0; j < KH; ++j) @@ -353,35 +456,30 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( } j_valid = rocke_b_cmp_lt(bld, y_j, n_iters_v); - for(s_const = 0; s_const < p->KW; ++s_const) + /* Issue row y_j + PF, then compute row y_j (the oldest in flight). */ + { + rocke_value_t* cpf = rocke_b_const_i32(bld, PF); + sp_load_row(bld, + rocke_b_add(bld, y_j, cpf), + p->PAD, + KW, + ¶ms, + n, + col_offs, + a_rsrc, + c_half_bytes, + oob_sentinel, + c0, + is_bf16, + window + PF * KW); + } + for(s_const = 0; s_const < KW; ++s_const) + a_row[s_const] = rocke_b_cast_to_f32(bld, window[s_const]); + memmove(window, window + KW, (size_t)PF * KW * sizeof(rocke_value_t*)); + + for(s_const = 0; s_const < KW; ++s_const) { - rocke_value_t* a_off = NULL; - rocke_value_t* valid = NULL; - rocke_value_t* ok; - rocke_value_t* safe_off; - rocke_value_t* a_h; - rocke_value_t* a_f32; - const char* on[5]; - rocke_value_t* ov[5]; - - on[0] = "n"; - ov[0] = n; - on[1] = "y_iter"; - ov[1] = y_j; - on[2] = "wo"; - ov[2] = q_out; - on[3] = "s_off"; - ov[3] = rocke_b_const_i32(bld, s_const); - on[4] = "c"; - ov[4] = ch; - rocke_transforms_descriptor_offset(bld, a_desc, on, ov, 5, &a_off, &valid); - - ok = rocke_b_land(bld, rocke_b_land(bld, valid, j_valid), q_ok); - safe_off - = rocke_b_select(bld, ok, rocke_b_mul(bld, a_off, c_half_bytes), oob_sentinel); - a_h = is_bf16 ? rocke_b_buffer_load_bf16(bld, a_rsrc, safe_off, c0) - : rocke_b_buffer_load_f16(bld, a_rsrc, safe_off, c0); - a_f32 = rocke_b_select(bld, ok, rocke_b_cast_to_f32(bld, a_h), zero_f32); + rocke_value_t* a_f32 = a_row[s_const]; for(r_const = 0; r_const < KH; ++r_const) { @@ -443,7 +541,16 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( new_accs[P_FLUSH_j] = zero_f32; } - rocke_b_scf_yield(bld, new_accs, num_iargs); + { + rocke_value_t** yields + = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); + int i; + for(i = 0; i < KH; ++i) + yields[i] = new_accs[i]; + for(i = 0; i < PF * KW; ++i) + yields[KH + i] = window[i]; + rocke_b_scf_yield(bld, yields, num_iargs); + } rocke_b_region_leave(bld); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 14975fbfb60c..443ee6995dd3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -125,8 +125,8 @@ "conv_direct_grouped": { "0": "bba7ec55a99618a311f5a632fe0a76ceb8904356f18f1d68468ee4e35d39de07", "1": "41caa5bc1a1e21b8ac280942fe658bb2e1b7813d1780de124edc1a97a82d9d48", - "10": "b7476209198e9097633fa1552eba421df9a26522ba3b0f47dec2e6deabc639c3", - "11": "9efcaeb7394974bc45ba675e34ab961efec270e6aa50b2e6d2ea8a1cda422449", + "10": "6ad18c90e3246ea1e62294e65caa5e8c247aad5a1b0244f6b46e28ba2ed2c9e7", + "11": "6bb7a73fa1daa80de9a4094b50f7166068710b810dd9e8e9ea12271e3c0cba29", "12": "18585f97b7f64e5e8a44df4134eab70f9f3b18ccb766db4de74e5c113b5e1095", "13": "62e24ebde4896d81a6789d283a8667c2937ac3e1a197df8ee7ab3e6bfc28f0c7", "14": "66545c341ebc8c3a0dec020a8103559bebb8099841463b396a3f07e3d6bbdcc1", @@ -138,8 +138,8 @@ "2": "5d93548857f2edf16e8fc1e839ffe6a5ed829952dc75b0edb438ab905db26de6", "20": "4ca258b9d883e4326a8b2dbe1d8aa252ad1ddd1c1c0e84ab4b5a5e2e8a087388", "21": "e3a7414cf1fcd93b058b9d5e267f8ff3f23b8e7aec494355c5e721b7371727a9", - "22": "86cf2e891598673a3c2b45f5cca8df8785a2ea63192d4d44fbd4a608b3520b58", - "23": "49c25e3043fac5c886adcfa92818f066c8e538a2d84a6f89c36eb79c732e965a", + "22": "6d277e8ae2810023f6283bb245c3943b90dfea789e91c31f5217d027229a60f5", + "23": "15a4768eb2f30649e82771216fa905993a46c3616b27460b829e6e8b45c1a87a", "24": "bd0167063cda9af5fc38d9cbb3b10fca05cd9a2fba943e46153dd2b66f06ea52", "25": "f4ad024735075c18f83ffa6507069a378ab3297f0a111c404b2bac2811d3e57a", "26": "a53aba5a8e55b8f07ccc843e701c78dfd42b85f85aa6daa9ab65bc1a6d0875a3", @@ -151,8 +151,8 @@ "5": "abb1d9ba0ef3a0e97f89d5b033e6ed1626aeb25b80defcda6c36e6151b397ad1", "6": "2ef16c7aea9fbbe3e2f4bc9c13cea6e2d5e94692d83396757058822e9d9d3cb6", "7": "7fd61c571db4c684397cf9afcc063d2b6295c9b52ab13cfbfe84588bcbf366a1", - "8": "d00572540d3e237b85de49ad1b5183235802dd6b68f77b0a3eb31bec792c963a", - "9": "f475f511e38477e0231a6b2efa276745236bd4627498272182d54b537f8e00b9" + "8": "73be0cf06c11b24b28277fcf835dcaa1d78294c8a518e157cc0d60bcb6ff8cca", + "9": "c80751c9816a2847c82823222162161ac1783caa608e28dcb730e78cba8a3099" }, "conv_implicit_gemm": { "0": "a09ed5bdeb9da980fee688bd046209d9ce126fc79754020e56c720034409f926", From a5a275cffc971428888250e912bc6c3a1e7fd4ba Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 6 Oct 2026 10:10:46 +0000 Subject: [PATCH 2/7] fixes --- .../common/benchmark_conv_compare.py | 2 +- .../common/benchmark_direct_conv.py | 35 +- .../common/benchmark_implicit_gemm_conv.py | 11 +- .../kernels/common/conv_direct_grouped.py | 254 +++++--- .../tests/parity/conv_direct_grouped_emit.c | 21 + .../tests/parity/conv_direct_grouped_emit.py | 21 + .../instance_conv_direct_grouped_internal.h | 4 +- .../conv_direct_grouped_build_depthwise.cpp | 616 ++++++++++-------- ...direct_grouped_build_depthwise_spatial.cpp | 28 + .../python/rocke/runtime/hip_module.py | 55 +- .../platform/python/rocke/runtime/launcher.py | 31 + .../differential/golden/llvm_gfx_all.json | 13 +- .../rocke/platform/tests/test_rocke.py | 39 ++ 13 files changed, 735 insertions(+), 395 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 eb3f85e83a1f..e2380aa965b9 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 @@ -307,7 +307,7 @@ def main() -> int: parser.add_argument( "--top", type=int, - default=5, + default=10, help="top-N results to show per benchmark (default: 5)", ) parser.add_argument( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py index f960ddce589b..b3a17475bcf8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py @@ -231,6 +231,13 @@ def parse_miopen_cmd_direct(cmd: str): return problem, dtype, miopen_args.forw +def _dw_config_label(block_w, block_waves) -> str: + """A depthwise sweep entry's knobs; the spatial kernel has no block_w.""" + if block_w is None: + return f"spatial bwv={block_waves}" + return f"bw={block_w:3d} bwv={block_waves}" + + def _sample_combos(combos: list, frac: float, seed: int) -> list: import random @@ -443,7 +450,7 @@ def _print_depthwise_results( print(hdr) print("-" * width) for rank, r in enumerate(results[:top_n], 1): - cfg = f"bw={r.block_w:3d} bwv={r.block_waves}" + cfg = _dw_config_label(r.block_w, r.block_waves) if show_verify: v = "PASS" if r.passed else "FAIL" print( @@ -645,7 +652,7 @@ def _run_depthwise_sweep( rt.memset(D_dev, 0, D_t.nbytes) ms = _stop.measure( - lambda: launcher(values, config=cfg), + launcher.bind(values, config=cfg), warmup=args.warmup, iters=args.iters, stream=stream, @@ -671,7 +678,7 @@ def _run_depthwise_sweep( ) ) print( - f"[{n_run:4d}] bw={block_w:3d} bwv={block_waves}" + f"[{n_run:4d}] {_dw_config_label(block_w, block_waves)}" f" {cur_tflops:6.1f} TFLOPS {ms:.3f} ms", flush=True, ) @@ -855,7 +862,7 @@ def _run_sweep( rt.memset(D_dev, 0, D_t.nbytes) ms = _stop.measure( - lambda: launcher(values, config=cfg), + launcher.bind(values, config=cfg), warmup=args.warmup, iters=args.iters, stream=stream, @@ -1095,7 +1102,7 @@ def _run_wgrad_sweep( cfg_wg = LaunchConfig(grid=grid, block=block_dim) ms = _stop.measure( - lambda: launcher(values, config=cfg_wg), + launcher.bind(values, config=cfg_wg), warmup=args.warmup, iters=args.iters, stream=0, @@ -1681,7 +1688,7 @@ def run_mfma_dgrad(): else: cfg = LaunchConfig(grid=grid, block=block_dim) ms = _stop.measure( - lambda: launcher(values, config=cfg), + launcher.bind(values, config=cfg), warmup=args.warmup, iters=args.iters, stream=0, @@ -1925,9 +1932,16 @@ def _u8(t): print(f" [skip] {plan.label}: {e}", flush=True) continue - def run(launches=launches): - for launcher, values, cfg in launches: - launcher(values, config=cfg) + # Each step packed once: a timed loop that packs per launch + # times the host, not a short kernel. + bound = [ + launcher.bind(values, config=cfg) + for launcher, values, cfg in launches + ] + + def run(bound=bound): + for launch in bound: + launch() passed = None if ref_out is not None: @@ -1983,6 +1997,9 @@ def run(launches=launches): print( f" {rank:>4} {tflops:>7.1f} {ms:>8.3f} {gbps:>7.1f} {verdict:>6} {label}" ) + # The line benchmark_conv_compare.py reads, as the JIT sweep prints. + tflops, _, ms, _, label = results[0] + print(f"\nBest: {tflops:.1f} TFLOPS — {label} {ms:.3f} ms", flush=True) return overall_rc diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index 9b338776ea04..73a151f7efb2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -2354,7 +2354,7 @@ def _make(*shape): rt.memset(D_dev, 0, D_t.nbytes) ms = _stop.measure( - lambda: launcher(values, config=cfg), + launcher.bind(values, config=cfg), warmup=args.warmup, iters=args.iters, stream=stream, @@ -2886,7 +2886,7 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): else: _vals_snap = dict(values) _cfg_snap = cfg - timed_fn = lambda _v=_vals_snap, _c=_cfg_snap: launcher(_v, config=_c) + timed_fn = launcher.bind(_vals_snap, config=_cfg_snap) ms = _stop.measure( timed_fn, @@ -3479,7 +3479,7 @@ def _launch_atomic(): timed_fn = _launch_atomic else: - timed_fn = lambda: launcher(values, config=cfg) + timed_fn = launcher.bind(values, config=cfg) ms = _stop.measure( timed_fn, warmup=args.warmup, iters=args.iters, stream=stream @@ -3829,7 +3829,7 @@ def _u8(t): ), stream=0, ) - run = lambda: launcher(values, config=cfg) # noqa: E731 + run = launcher.bind(values, config=cfg) if direction == "wgrad" and ident.two_stage: # Two-stage is three steps and all of them are the # algorithm's cost: zero the scratch (Stage 1 @@ -3945,6 +3945,9 @@ def run( f"{ms:>8.3f} ms {label}", flush=True, ) + # The line benchmark_conv_compare.py reads, as the JIT sweep prints. + tflops, _, ms, label = results[0] + print(f"\nBest: {tflops:.1f} TFLOPS — {label} {ms:.3f} ms", flush=True) return overall_rc diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index ccf855076b89..9f9c25baab0c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -5737,8 +5737,9 @@ class DirectDepthwiseSpec: ``by * block_ch + w * 64 + l``. - Weights (``KH × KW`` values per channel) are preloaded into registers before the H-streaming loop. - - H-streaming loop (runtime loop over ``H + KH - 1`` rows, ``KH`` rows - per iteration, the next iteration's input rows prefetched): + - H-streaming loop (runtime loop over the ``H`` input rows, ``KH`` rows + per iteration, input rows prefetched a few rows ahead, then a + load-free loop over the rows that only complete output rows): circular ``KH``-slot accumulators per output W position, flushed to global memory one output row at a time. For stride > 1 only input rows that align to an output position produce a write. @@ -5891,11 +5892,6 @@ def build_direct_depthwise( b_rsrc = b.buffer_rsrc(Bp, B_bytes) d_rsrc = b.buffer_rsrc(D, D_bytes) - # A[N, H, W, C] NHWC descriptor with h and w boundary embeds. - a_desc = direct_a_descriptor_dynamic( - b, params, pad=p.PAD, stride=p.stride, w_upper=("wo", "s_off") - ) - # B[total_k, KH, KW, 1] KRSC descriptor (cpg=1: last dim is always 0). b_desc = TensorDescriptor.naive( "B", @@ -5933,24 +5929,26 @@ def build_direct_depthwise( # ---- H-streaming loop ---- # AOT: the row count follows the runtime height, so only the # grouped-period scf.for form is usable -- the Python-unrolled variant - # would bake H into the trip count. The outer loop runs - # n_groups = ceil(n_iters / KH) times and the inner KH steps are unrolled - # with STATIC slot indices, so the preloaded weights are still referenced - # directly (no scatter, no runtime weight loads). + # would bake H into the trip count. The loops below walk the rows KH at + # a time and unroll the KH steps of a group with STATIC slot indices, so + # the preloaded weights are still referenced directly (no scatter, no + # runtime weight loads). # - # The input rows are software-pipelined one group ahead: iteration g - # first issues the loads of group g + 1 and carries them to the next - # iteration, then computes group g from the rows the previous iteration - # loaded. Loaded and used in the same iteration, the rows are pulled back - # next to their FMAs by the scheduler, and every load then waits out the - # full memory latency before the next one issues (the unrolled JIT body - # had no such problem; in a loop it costs 2-6x). + # Streaming row y holds input row hi = y - PAD and completes output row + # (y - (KH-1)) / stride. Of the H + KH - 1 rows only hi in [0, H) has + # data: the first PAD rows are halo and are skipped outright (the stream + # starts at y = PAD), and the last PAD rows only complete output rows. + # The FMA loop therefore covers ceil(H / KH) groups and a second loop, + # with no loads and no FMAs, flushes what is left. Without the split the + # loop iterates (H + KH - 1) / H times per output row -- substantially + # more than necessary for small H -- and the kernel becomes FMA-bound. c1 = b.const_i32(1) c_KH = b.const_i32(p.KH) c_stride_rv = b.const_i32(c_stride_dw) - # n_iters = H + KH - 1 rows; the loop walks them KH at a time. + # n_iters = H + KH - 1 rows, the last one completing the last output row. n_iters_v = b.add(p_Hi, b.const_i32(p.KH - 1)) - n_groups_v = b.div(b.add(n_iters_v, b.const_i32(p.KH - 1)), c_KH) + n_fma_groups_v = b.div(b.add(p_Hi, b.const_i32(p.KH - 1)), c_KH) + n_groups_v = b.div(b.add(p_Hi, b.const_i32(p.PAD + p.KH - 1)), c_KH) # One load per distinct input column of a row: tap (w_out, s) reads # column w_out * stride + s. ``row_cols`` maps each column to the first @@ -5959,108 +5957,164 @@ def build_direct_depthwise( for w_out in range(BLOCK_W): for s_const in range(p.KW): row_cols.setdefault(w_out * c_stride_dw + s_const, (w_out, s_const)) + n_cols = len(row_cols) + + # Every lane of a wave reads the same input row and column, each at its + # own channel, so an input offset is the lane's channel plus a + # wave-uniform term for batch, row and column -- computed on the SALU, + # one VALU add per load. The uniform term is the OOB sentinel for a tap + # outside the image, which keeps the sum past the end of the buffer (the + # descriptor returns zero) as long as the tensor is under 2 GiB, the + # limit of the i32 offsets anyway. Bounds travel as integers -- one + # term per column, the sentinel where the column is outside [0, W) -- + # rather than as booleans: a uniform boolean feeding a load offset is a + # 64-bit lane mask, and a mask per column held across the loop spills + # the SGPRs. Lanes past the channel count need no guard here: they load + # whatever lies there and their sums are never stored. + ch_off = b.mul(ch, c_half_bytes) + wi_tile = b.add(b.mul(q_tile_start, c_stride_rv), b.const_i32(-p.PAD)) + col_terms: List[Value] = [] + for col in row_cols: + wi = b.add(wi_tile, b.const_i32(col)) + col_ok = b.land(b.cmp_ge(wi, c0), b.cmp_lt(wi, params["p_Wi"])) + col_off = b.mul(b.mul(wi, params["p_A_stride_wi"]), c_half_bytes) + col_terms.append(b.select(col_ok, col_off, oob_sentinel)) + n_off = b.mul(n, params["p_A_stride_n"]) - def load_group(y0: Value) -> List[Value]: - """Raw input values of rows y0 .. y0 + KH - 1, row-major by column. - - Rows past the last iteration, halo rows and columns, and lanes past - the channel count all load from the OOB sentinel, which the buffer - descriptor returns as zero. - """ + def load_row(y: Value) -> List[Value]: + """Raw input values of streaming row ``y``, one per distinct column.""" + hi = b.add(y, b.const_i32(-p.PAD)) + row_ok = b.cmp_lt(hi, p_Hi) + row_off = b.mul(b.add(n_off, b.mul(hi, params["p_A_stride_hi"])), c_half_bytes) vals: List[Value] = [] - for j in range(p.KH): - y_j = b.add(y0, b.const_i32(j)) - j_valid = b.cmp_lt(y_j, n_iters_v) - for w_out, s_const in row_cols.values(): - a_off, valid = a_desc.offset( - b, - n=n, - y_iter=y_j, - wo=b.add(q_tile_start, b.const_i32(w_out)), - s_off=b.const_i32(s_const), - c=ch, - ) - ok = b.land(b.land(valid, j_valid), ch_in_range) - safe_off = b.select(ok, b.mul(a_off, c_half_bytes), oob_sentinel) - vals.append( - b.buffer_load_bf16(a_rsrc, safe_off, c0) - if p.dtype == "bf16" - else b.buffer_load_f16(a_rsrc, safe_off, c0) - ) + for col_term in col_terms: + s_term = b.select(row_ok, b.add(row_off, col_term), oob_sentinel) + v_off = b.add(ch_off, s_term) + vals.append( + b.buffer_load_bf16(a_rsrc, v_off, c0) + if p.dtype == "bf16" + else b.buffer_load_f16(a_rsrc, v_off, c0) + ) return vals + def row_y(grp_iv: Value, j: int) -> Value: + """Streaming row of step ``j`` of group ``grp_iv``.""" + return b.add(b.mul(grp_iv, c_KH), b.const_i32(p.PAD + j)) + + def flush(grp_iv: Value, j: int, y_j: Value, accs: List[Value]) -> None: + """Store the output row streaming row ``y_j`` completes; reset its slot.""" + j_valid = b.cmp_lt(y_j, n_iters_v) + P_FLUSH_j = (p.PAD + j + 1) % p.KH # STATIC + p_flush_rv = b.add(y_j, b.const_i32(-(p.KH - 1))) + + # y_j >= KH - 1 holds from group 1 on; in group 0 from step PAD on. + if j >= p.PAD: + flush_ge = j_valid + else: + flush_ge = b.land(b.cmp_lt(c0, grp_iv), j_valid) + + if c_stride_dw == 1: + should_flush = flush_ge + else: + flush_stride = b.cmp_eq(b.mod(p_flush_rv, c_stride_rv), c0) + should_flush = b.land(flush_ge, flush_stride) + + ho_row_j = b.div(p_flush_rv, c_stride_rv) + + for w_out in range(BLOCK_W): + out_q = b.add(q_tile_start, b.const_i32(w_out)) + out_q_ok = b.land(b.cmp_lt(out_q, c_W), ch_in_range) + store_ok = b.land(out_q_ok, should_flush) + acc_val = accs[P_FLUSH_j * BLOCK_W + w_out] # STATIC index + d_off, _ = d_desc.offset(b, n=n, h=ho_row_j, w=out_q, k=ch) + safe_d = b.select(store_ok, b.mul(d_off, c_half_bytes), oob_sentinel) + if p.dtype == "bf16": + b.buffer_store_bf16(d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc_val)) + else: + b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_val)) + + for w_out in range(BLOCK_W): + accs[P_FLUSH_j * BLOCK_W + w_out] = zero_f32 + + # The input rows are software-pipelined PF rows ahead: step y issues the + # loads of row y + PF before computing row y, and the PF rows in flight + # at the end of an iteration are carried into the next one. Loaded and + # used in the same step, every load would wait out the full memory + # latency before the next one issues. Carrying a whole KH-row group + # instead holds KH x n_cols loads in registers on top of the KH x BLOCK_W + # accumulators and the KH x KW weights; a few rows hide the latency as + # well. + PF = min(p.KH, -(-16 // n_cols)) n_acc = p.KH * BLOCK_W dw_iter_args = [ (f"dw_acc_kh{kh}_w{w}", zero_f32) for kh in range(p.KH) for w in range(BLOCK_W) ] - row_names = [f"dw_a_r{j}_c{col}" for j in range(p.KH) for col in row_cols] - dw_iter_args += list(zip(row_names, load_group(c0))) - group_loop = b.scf_for_iter( + for k in range(PF): + for col, a_h in zip(row_cols, load_row(b.const_i32(p.PAD + k))): + dw_iter_args.append((f"dw_a_r{k}_c{col}", a_h)) + fma_loop = b.scf_for_iter( c0, - n_groups_v, + n_fma_groups_v, c1, dw_iter_args, iv_name="dw_grp", elide_trailing_barrier=False, ) - with group_loop as (grp_iv, loop_vals): + with fma_loop as (grp_iv, loop_vals): new_accs = list(loop_vals[:n_acc]) - cur_rows = loop_vals[n_acc:] - next_rows = load_group(b.mul(b.add(grp_iv, c1), c_KH)) + window = [ + loop_vals[n_acc + k * n_cols : n_acc + (k + 1) * n_cols] for k in range(PF) + ] col_idx = {col: i for i, col in enumerate(row_cols)} for j in range(p.KH): - y_j = b.add(b.mul(grp_iv, c_KH), b.const_i32(j)) - j_valid = b.cmp_lt(y_j, n_iters_v) - a_row = [ - b.cast_to_f32(a_h) - for a_h in cur_rows[j * len(row_cols) : (j + 1) * len(row_cols)] - ] + y_j = row_y(grp_iv, j) + # Issue row y_j + PF, then compute row y_j (the oldest in flight). + window.append(load_row(b.add(y_j, b.const_i32(PF)))) + a_row = [b.cast_to_f32(a_h) for a_h in window.pop(0)] for w_out in range(BLOCK_W): for s_const in range(p.KW): a_f32 = a_row[col_idx[w_out * c_stride_dw + s_const]] for r_const in range(p.KH): - p_idx = (j - r_const + p.KH) % p.KH # STATIC slot + # STATIC slot of output row (y_j - r) / stride. + p_idx = (p.PAD + j - r_const) % p.KH idx = p_idx * BLOCK_W + w_out new_accs[idx] = b.fma( weights_f32[r_const][s_const], a_f32, new_accs[idx] ) - P_FLUSH_j = (j + 1) % p.KH # STATIC - p_flush_rv = b.add(y_j, b.const_i32(-(p.KH - 1))) - - if j >= p.KH - 1: - flush_ge = j_valid - else: - flush_ge = b.land(b.cmp_lt(c0, grp_iv), j_valid) - - if c_stride_dw == 1: - should_flush = flush_ge - else: - flush_stride = b.cmp_eq(b.mod(p_flush_rv, c_stride_rv), c0) - should_flush = b.land(flush_ge, flush_stride) - - ho_row_j = b.div(p_flush_rv, c_stride_rv) - - for w_out in range(BLOCK_W): - out_q = b.add(q_tile_start, b.const_i32(w_out)) - out_q_ok = b.land(b.cmp_lt(out_q, c_W), ch_in_range) - store_ok = b.land(out_q_ok, should_flush) - acc_val = new_accs[P_FLUSH_j * BLOCK_W + w_out] # STATIC index - d_off, _ = d_desc.offset(b, n=n, h=ho_row_j, w=out_q, k=ch) - safe_d = b.select(store_ok, b.mul(d_off, c_half_bytes), oob_sentinel) - if p.dtype == "bf16": - b.buffer_store_bf16( - d_rsrc, safe_d, c0, b.trunc_f32_to_bf16(acc_val) - ) - else: - b.buffer_store_f16(d_rsrc, safe_d, c0, b.trunc_f32_to_f16(acc_val)) - - for w_out in range(BLOCK_W): - new_accs[P_FLUSH_j * BLOCK_W + w_out] = zero_f32 - - b.scf_yield(*new_accs, *next_rows) + flush(grp_iv, j, y_j, new_accs) + + # Pass the accumulators through an empty asm on their way into the + # next iteration. Without it LLVM's SLP vectorizer pairs the loop- + # carried FMA chains into v_pk_fma_f32, which requires their operands + # in adjacent registers; the register copies needed to arrange that + # consume enough VGPRs to halve occupancy. The asm emits no instruction. + new_accs = [ + b.inline_asm("", "=v,0", [a], result_type=F32, sideeffect=False) + for a in new_accs + ] + b.scf_yield(*new_accs, *(a_h for row in window for a_h in row)) + + # The rows past the image: no input, only output rows to complete. + flush_loop = b.scf_for_iter( + n_fma_groups_v, + n_groups_v, + c1, + [ + (f"dw_tail_acc_kh{kh}_w{w}", fma_loop.results[kh * BLOCK_W + w]) + for kh in range(p.KH) + for w in range(BLOCK_W) + ], + iv_name="dw_tail", + elide_trailing_barrier=False, + ) + with flush_loop as (grp_iv, loop_vals): + new_accs = list(loop_vals) + for j in range(p.KH): + flush(grp_iv, j, row_y(grp_iv, j), new_accs) + b.scf_yield(*new_accs) return b.kernel @@ -6195,6 +6249,11 @@ def build_direct_depthwise_spatial( # bake H into the trip count. b = IRBuilder(spec.kernel_name()) b.kernel.attrs["max_workgroup_size"] = THREADS + # Two waves per SIMD: the KH x KW weights and the prefetch window push + # larger filters just over the threshold for one wave, and the two-wave + # budget spills a few loop-invariant weights while being faster overall. + # Small filters fit either way. + b.kernel.attrs["waves_per_eu"] = 2 io_type = _io_type(p.dtype) @@ -6374,6 +6433,13 @@ def load_row(y: Value) -> List[Value]: new_accs[P_FLUSH_j] = zero_f32 # unconditional static reset + # The accumulators pass through an empty asm into the next iteration, + # which keeps the SLP vectorizer from pairing the loop-carried FMA + # chains into v_pk_fma_f32 (see build_direct_depthwise). + new_accs = [ + b.inline_asm("", "=v,0", [a], result_type=F32, sideeffect=False) + for a in new_accs + ] b.scf_yield(*new_accs, *[a_h for row in window for a_h in row]) return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c index 062dd26c9195..adb748d6dd23 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c @@ -543,6 +543,27 @@ static int make_cfg(int idx, *kind = KIND_WGRAD; *arch = "gfx950"; return 0; + case 32: + /* depthwise 7x7 stride 2 bf16: PAD > 1 shifts the row stream's static + * slots, and stride 2 gates its flushes on the row parity. */ + p.N = 2; + p.H = 14; + p.W = 14; + p.groups = 128; + p.cpg = 1; + p.kpg = 1; + p.KH = 7; + p.KW = 7; + p.PAD = 3; + p.stride = 2; + p.dtype = "bf16"; + *sdw = rocke_direct_depthwise_spec_default(); + sdw->problem = p; + sdw->block_w = 8; + sdw->block_waves = 2; + *kind = KIND_DW; + *arch = "gfx950"; + return 0; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py index bb40a420e0fc..8e0d3e7446c2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py @@ -439,6 +439,27 @@ def _spec(idx: int): DirectConvWgradSpec(problem=p, mfma_k=16, ho_per_block=2), "gfx950", ) + if idx == 32: + # depthwise 7x7 stride 2 bf16: PAD > 1 shifts the row stream's static + # slots, and stride 2 gates its flushes on the row parity. + p = DirectConvProblem( + N=2, + H=14, + W=14, + groups=128, + cpg=1, + kpg=1, + KH=7, + KW=7, + PAD=3, + stride=2, + dtype="bf16", + ) + return ( + "dw", + DirectDepthwiseSpec(problem=p, block_w=8, block_waves=2), + "gfx950", + ) raise SystemExit(f"unknown config index {idx}") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h index b2938dcfd727..3ea225e89ede 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h @@ -653,8 +653,8 @@ typedef struct rocke_dconv_32c_ctx * KH*KW scalar FMA products into a BLOCK_W * KH circular register array. * ===================================================================== */ #define ROCKE_DCONV_DW_MAX_BLOCK_W 64 -#define ROCKE_DCONV_DW_MAX_KH 8 -#define ROCKE_DCONV_DW_MAX_KW 8 +#define ROCKE_DCONV_DW_MAX_KH 32 +#define ROCKE_DCONV_DW_MAX_KW 32 typedef struct rocke_dconv_dw_ctx { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp index 07fd4c134237..249a82935dcf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp @@ -63,6 +63,15 @@ bool rocke_dconv_dw_prologue(rocke_dconv_dw_ctx_t* ctx) } ctx->p = spec->problem; + /* The weight table is sized for filters up to the depthwise limit. */ + if(ctx->p.KH > ROCKE_DCONV_DW_MAX_KH || ctx->p.KW > ROCKE_DCONV_DW_MAX_KW) + { + if(b->status == ROCKE_OK) + { + b->status = ROCKE_ERR_VALUE; + } + return false; + } ctx->BLOCK_W = spec->block_w; ctx->BLOCK_WAVES = spec->block_waves; @@ -129,7 +138,6 @@ bool rocke_dconv_dw_prologue(rocke_dconv_dw_ctx_t* ctx) /* ===================================================================== * * Descriptor phase (Python lines 2614-2646) * - * a_desc: A[N,H,W,total_c] with two embeds (y boundary, w boundary). * b_desc: B[total_k,KH,KW,1] naive. * d_desc: D[N,H,W,total_k] naive. * @@ -148,15 +156,6 @@ void rocke_dconv_dw_build_descriptors(rocke_dconv_dw_ctx_t* ctx) rocke_ir_builder_t* b = ctx->b; int total_k = rocke_direct_conv_problem_total_k(&ctx->p); - /* a_desc with runtime extents; PAD and stride stay build-time. */ - { - rocke_dynamic_tensor_descriptor_t* a_dyn = rocke_dconv_a_descriptor_dynamic( - b, &ctx->params, ctx->p.PAD, ctx->c_stride_dw, "wo", "s_off"); - if(!a_dyn) - return; - ctx->a_desc = &a_dyn->base; - } - /* b_desc = naive("B", [total_k, KH, KW, 1], coord_names=(k,r,s,c)) */ { static const char* const b_coords[4] = {"k", "r", "s", "c"}; @@ -222,72 +221,151 @@ void rocke_dconv_dw_load_weights(rocke_dconv_dw_ctx_t* ctx) /* ===================================================================== * * H-row streaming loop (stride-aware) * - * Accumulator: acc[w_out][slot] — BLOCK_W × KH scalar f32 values. + * Accumulator: acc[slot][w_out] — KH × BLOCK_W scalar f32 values. * The flush uses stride-aware ho_row = p_flush_val / c_stride_dw and only - * stores when p_flush_val % c_stride_dw == 0 and ho_row < Ho. - * ch_in_range guards loads and stores for partial channel tiles. + * stores when p_flush_val % c_stride_dw == 0. + * ch_in_range guards the stores for partial channel tiles. * - * The rows stream through a scf_for_iter group loop over - * n_groups = ceil(n_iters / KH) groups of KH iterations each; the trip count - * follows the runtime height (AOT), so there is no build-time-unrolled form. - * The input rows are software-pipelined one group ahead: each iteration - * issues the loads of the next group and carries them to the next iteration - * (see the Python source for why). + * Streaming row y holds input row hi = y - PAD. The stream starts at + * y = PAD (the halo rows above the image hold no data), an FMA loop covers + * ceil(Hi / KH) groups of KH rows with the input rows prefetched PF rows + * ahead, and a load-free loop flushes the output rows the rows past the + * image complete. The trip counts follow the runtime height (AOT), so there + * is no build-time-unrolled form. See the Python source for why. * ===================================================================== */ -/* Raw input values of rows y0 .. y0 + KH - 1, row-major by column: one load - * per distinct input column, taken at the first tap (col_w, col_s) that - * reads it. Out-of-range rows, columns and lanes load from the OOB - * sentinel, which the buffer descriptor returns as zero. */ -static void dw_load_group(rocke_dconv_dw_ctx_t* ctx, - rocke_value_t* y0, - rocke_value_t* n_iters_v, - const int* col_w, - const int* col_s, - int n_cols, - rocke_value_t** out) +/* Per-kernel values the streaming phases share. */ +typedef struct dw_stream +{ + int KH; + int KW; + int BLOCK_W; + int stride; + int n_cols; + rocke_value_t* c_KH; + rocke_value_t* c_stride_rv; + rocke_value_t* n_iters_v; + rocke_value_t* ch_off; + rocke_value_t* n_off; + rocke_value_t** col_terms; /* per distinct column; the OOB sentinel off-image */ +} dw_stream_t; + +/* Raw input values of streaming row y, one per distinct column: the lane's + * channel offset plus a wave-uniform term, the OOB sentinel off the image. */ +static void dw_load_row(rocke_dconv_dw_ctx_t* ctx, + const dw_stream_t* st, + rocke_value_t* y, + rocke_value_t** out) { rocke_ir_builder_t* b = ctx->b; - int j, c; + rocke_value_t* hi = rocke_b_add(b, y, rocke_b_const_i32(b, -ctx->p.PAD)); + rocke_value_t* row_ok = rocke_b_cmp_lt(b, hi, ctx->params.p_Hi); + rocke_value_t* row_off; + int c; - for(j = 0; j < ctx->p.KH; ++j) { - rocke_value_t* y_j = rocke_b_add(b, y0, rocke_b_const_i32(b, j)); - rocke_value_t* j_valid = rocke_b_cmp_lt(b, y_j, n_iters_v); + rocke_value_t* h_off = rocke_b_mul(b, hi, ctx->params.p_A_stride_hi); + row_off = rocke_b_mul(b, rocke_b_add(b, st->n_off, h_off), ctx->c_half_bytes); + } + for(c = 0; c < st->n_cols; ++c) + { + rocke_value_t* rc = rocke_b_add(b, row_off, st->col_terms[c]); + rocke_value_t* s_term = rocke_b_select(b, row_ok, rc, ctx->oob_sentinel); + rocke_value_t* v_off = rocke_b_add(b, st->ch_off, s_term); + out[c] = ctx->is_bf16 ? rocke_b_buffer_load_bf16(b, ctx->a_rsrc, v_off, ctx->c0) + : rocke_b_buffer_load_f16(b, ctx->a_rsrc, v_off, ctx->c0); + } +} - for(c = 0; c < n_cols; ++c) - { - rocke_value_t* a_off = NULL; - rocke_value_t* valid = NULL; - rocke_value_t* ok; - rocke_value_t* safe_off; - const char* off_names[5]; - rocke_value_t* off_vals[5]; - - off_names[0] = "n"; - off_vals[0] = ctx->n; - off_names[1] = "y_iter"; - off_vals[1] = y_j; - off_names[2] = "wo"; - { - rocke_value_t* cwo = rocke_b_const_i32(b, col_w[c]); - off_vals[2] = rocke_b_add(b, ctx->q_tile_start, cwo); - } - off_names[3] = "s_off"; - off_vals[3] = rocke_b_const_i32(b, col_s[c]); - off_names[4] = "c"; - off_vals[4] = ctx->ch; - rocke_transforms_descriptor_offset( - b, ctx->a_desc, off_names, off_vals, 5, &a_off, &valid); - - ok = rocke_b_land(b, rocke_b_land(b, valid, j_valid), ctx->ch_in_range); - safe_off = rocke_b_select( - b, ok, rocke_b_mul(b, a_off, ctx->c_half_bytes), ctx->oob_sentinel); - out[j * n_cols + c] = ctx->is_bf16 - ? rocke_b_buffer_load_bf16(b, ctx->a_rsrc, safe_off, ctx->c0) - : rocke_b_buffer_load_f16(b, ctx->a_rsrc, safe_off, ctx->c0); - } +/* Streaming row of step j of group grp_iv: grp_iv*KH + (PAD + j). Emits the + * mul first, then the const, matching Python's left-to-right evaluation. */ +static rocke_value_t* + dw_row_y(rocke_dconv_dw_ctx_t* ctx, const dw_stream_t* st, rocke_value_t* grp_iv, int j) +{ + rocke_ir_builder_t* b = ctx->b; + rocke_value_t* mul_gk = rocke_b_mul(b, grp_iv, st->c_KH); + rocke_value_t* cj = rocke_b_const_i32(b, ctx->p.PAD + j); + return rocke_b_add(b, mul_gk, cj); +} + +/* Store the output row streaming row y_j completes; reset its slot. */ +static void dw_flush(rocke_dconv_dw_ctx_t* ctx, + const dw_stream_t* st, + rocke_value_t* grp_iv, + int j, + rocke_value_t* y_j, + rocke_value_t** accs) +{ + rocke_ir_builder_t* b = ctx->b; + int KH = st->KH; + int BLOCK_W = st->BLOCK_W; + int P_FLUSH_j = (ctx->p.PAD + j + 1) % KH; /* STATIC */ + rocke_value_t* j_valid = rocke_b_cmp_lt(b, y_j, st->n_iters_v); + rocke_value_t* p_flush_rv = rocke_b_add(b, y_j, rocke_b_const_i32(b, -(KH - 1))); + rocke_value_t* flush_ge; + rocke_value_t* should_flush; + rocke_value_t* ho_row_j; + int w_out; + + /* y_j >= KH - 1 holds from group 1 on; in group 0 from step PAD on. */ + if(j >= ctx->p.PAD) + { + flush_ge = j_valid; } + else + { + flush_ge = rocke_b_land(b, rocke_b_cmp_lt(b, ctx->c0, grp_iv), j_valid); + } + + if(st->stride == 1) + { + should_flush = flush_ge; + } + else + { + rocke_value_t* flush_stride + = rocke_b_cmp_eq(b, rocke_b_mod(b, p_flush_rv, st->c_stride_rv), ctx->c0); + should_flush = rocke_b_land(b, flush_ge, flush_stride); + } + + ho_row_j = rocke_b_div(b, p_flush_rv, st->c_stride_rv); + + for(w_out = 0; w_out < BLOCK_W; ++w_out) + { + rocke_value_t* out_q = rocke_b_add(b, ctx->q_tile_start, rocke_b_const_i32(b, w_out)); + rocke_value_t* out_q_ok + = rocke_b_land(b, rocke_b_cmp_lt(b, out_q, ctx->c_W), ctx->ch_in_range); + rocke_value_t* store_ok = rocke_b_land(b, out_q_ok, should_flush); + rocke_value_t* acc_val = accs[P_FLUSH_j * BLOCK_W + w_out]; + rocke_value_t* d_off = NULL; + rocke_value_t* d_valid = NULL; + rocke_value_t* safe_d; + const char* off_names[4]; + rocke_value_t* off_vals[4]; + + off_names[0] = "n"; + off_vals[0] = ctx->n; + off_names[1] = "h"; + off_vals[1] = ho_row_j; + off_names[2] = "w"; + off_vals[2] = out_q; + off_names[3] = "k"; + off_vals[3] = ctx->ch; + rocke_transforms_descriptor_offset( + b, ctx->d_desc, off_names, off_vals, 4, &d_off, &d_valid); + + safe_d = rocke_b_select( + b, store_ok, rocke_b_mul(b, d_off, ctx->c_half_bytes), ctx->oob_sentinel); + if(ctx->is_bf16) + rocke_b_buffer_store_bf16( + b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_bf16(b, acc_val)); + else + rocke_b_buffer_store_f16( + b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_f16(b, acc_val)); + } + + for(w_out = 0; w_out < BLOCK_W; ++w_out) + accs[P_FLUSH_j * BLOCK_W + w_out] = ctx->zero_f32; } rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) @@ -298,245 +376,227 @@ rocke_kernel_def_t* rocke_dconv_dw_stream_h_loop(rocke_dconv_dw_ctx_t* ctx) int KW = p->KW; int BLOCK_W = ctx->BLOCK_W; int stride = ctx->c_stride_dw; - int n_iters = ctx->n_iters; - /* AOT: the row count follows the runtime height, so only the - * grouped-period scf.for form is usable -- the build-time-unrolled variant - * would bake Hi into the trip count. The outer loop runs - * n_groups = ceil(n_iters / KH) times and the inner KH steps are unrolled - * with STATIC slot indices, so the preloaded weights are still referenced - * directly (no scatter, no runtime weight loads). */ - (void)n_iters; - { - /* Distinct input columns of a row in first-tap order: tap (w_out, s) - * reads column w_out * stride + s; col_idx maps a column to its - * position in that order. */ - int max_col = (BLOCK_W - 1) * stride + KW; - int* col_w = (int*)alloca((size_t)max_col * sizeof(int)); - int* col_s = (int*)alloca((size_t)max_col * sizeof(int)); - int* col_num = (int*)alloca((size_t)max_col * sizeof(int)); - int* col_idx = (int*)alloca((size_t)max_col * sizeof(int)); - int n_cols = 0; - int num_acc = KH * BLOCK_W; - int num_rows; - int num_iargs; - rocke_iter_arg_t* iargs; - rocke_for_t group_loop; - rocke_value_t** new_accs; - rocke_value_t** first_rows; - rocke_value_t** next_rows; - rocke_value_t** yields; - rocke_value_t* const* cur_rows; - rocke_value_t** a_row; - rocke_value_t* c1; - rocke_value_t* c_KH; - rocke_value_t* c_stride_rv; - rocke_value_t* n_iters_v; - rocke_value_t* n_groups_v; - int j; + /* Distinct input columns of a row in first-tap order: tap (w_out, s) + * reads column w_out * stride + s; col_idx maps a column to its + * position in that order. */ + int max_col = (BLOCK_W - 1) * stride + KW; + int* col_num = (int*)alloca((size_t)max_col * sizeof(int)); + int* col_idx = (int*)alloca((size_t)max_col * sizeof(int)); + int n_cols = 0; + int num_acc = KH * BLOCK_W; + int PF; + int num_win; + int num_iargs; + dw_stream_t st; + rocke_iter_arg_t* iargs; + rocke_for_t fma_loop; + rocke_for_t flush_loop; + rocke_value_t** new_accs; + rocke_value_t** window; + rocke_value_t** yields; + rocke_value_t** a_row; + rocke_value_t* c1; + rocke_value_t* n_fma_groups_v; + rocke_value_t* n_groups_v; + char(*name_store)[48]; + int j, i; + { + int w, s; + for(w = 0; w < max_col; ++w) + col_idx[w] = -1; + for(w = 0; w < BLOCK_W; ++w) { - int w, s; - for(w = 0; w < max_col; ++w) - col_idx[w] = -1; - for(w = 0; w < BLOCK_W; ++w) + for(s = 0; s < KW; ++s) { - for(s = 0; s < KW; ++s) - { - int col = w * stride + s; - if(col_idx[col] >= 0) - continue; - col_idx[col] = n_cols; - col_w[n_cols] = w; - col_s[n_cols] = s; - col_num[n_cols] = col; - ++n_cols; - } + int col = w * stride + s; + if(col_idx[col] >= 0) + continue; + col_idx[col] = n_cols; + col_num[n_cols] = col; + ++n_cols; } } - num_rows = KH * n_cols; - num_iargs = num_acc + num_rows; - first_rows = (rocke_value_t**)alloca((size_t)num_rows * sizeof(rocke_value_t*)); - next_rows = (rocke_value_t**)alloca((size_t)num_rows * sizeof(rocke_value_t*)); - yields = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); - a_row = (rocke_value_t**)alloca((size_t)n_cols * sizeof(rocke_value_t*)); - - c1 = rocke_b_const_i32(b, 1); - c_KH = rocke_b_const_i32(b, KH); - c_stride_rv = rocke_b_const_i32(b, stride); - /* n_iters = Hi + KH - 1 rows; the loop walks them KH at a time. */ - n_iters_v = rocke_b_add(b, ctx->params.p_Hi, rocke_b_const_i32(b, KH - 1)); - n_groups_v = rocke_b_div(b, rocke_b_add(b, n_iters_v, rocke_b_const_i32(b, KH - 1)), c_KH); - - /* Name storage: accumulators then carried rows, each at most 32 chars. */ + } + /* Rows in flight: enough for ~16 loads, at most a group. */ + PF = (16 + n_cols - 1) / n_cols; + if(PF > KH) + PF = KH; + num_win = PF * n_cols; + num_iargs = num_acc + num_win; + + st.KH = KH; + st.KW = KW; + st.BLOCK_W = BLOCK_W; + st.stride = stride; + st.n_cols = n_cols; + st.col_terms = (rocke_value_t**)alloca((size_t)n_cols * sizeof(rocke_value_t*)); + /* window holds PF rows plus the one being issued. */ + window = (rocke_value_t**)alloca((size_t)(num_win + n_cols) * sizeof(rocke_value_t*)); + yields = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); + a_row = (rocke_value_t**)alloca((size_t)n_cols * sizeof(rocke_value_t*)); + new_accs = (rocke_value_t**)alloca((size_t)num_acc * sizeof(rocke_value_t*)); + iargs = (rocke_iter_arg_t*)alloca((size_t)num_iargs * sizeof(rocke_iter_arg_t)); + name_store = (char(*)[48])alloca((size_t)num_iargs * 48 * sizeof(char)); + + c1 = rocke_b_const_i32(b, 1); + st.c_KH = rocke_b_const_i32(b, KH); + st.c_stride_rv = rocke_b_const_i32(b, stride); + /* n_iters = Hi + KH - 1 rows, the last one completing the last output row. */ + st.n_iters_v = rocke_b_add(b, ctx->params.p_Hi, rocke_b_const_i32(b, KH - 1)); + n_fma_groups_v + = rocke_b_div(b, rocke_b_add(b, ctx->params.p_Hi, rocke_b_const_i32(b, KH - 1)), st.c_KH); + n_groups_v = rocke_b_div( + b, rocke_b_add(b, ctx->params.p_Hi, rocke_b_const_i32(b, p->PAD + KH - 1)), st.c_KH); + + /* Input offset = the lane's channel + a wave-uniform term per column. */ + st.ch_off = rocke_b_mul(b, ctx->ch, ctx->c_half_bytes); + { + rocke_value_t* q_s = rocke_b_mul(b, ctx->q_tile_start, st.c_stride_rv); + rocke_value_t* c_mpad = rocke_b_const_i32(b, -p->PAD); + rocke_value_t* wi_tile = rocke_b_add(b, q_s, c_mpad); + int c; + for(c = 0; c < n_cols; ++c) { - char(*name_store)[32] = (char(*)[32])alloca((size_t)num_iargs * 32 * sizeof(char)); - int kh, w, c, idx = 0; - - iargs = (rocke_iter_arg_t*)alloca((size_t)num_iargs * sizeof(rocke_iter_arg_t)); - new_accs = (rocke_value_t**)alloca((size_t)num_acc * sizeof(rocke_value_t*)); + rocke_value_t* wi = rocke_b_add(b, wi_tile, rocke_b_const_i32(b, col_num[c])); + rocke_value_t* ge = rocke_b_cmp_ge(b, wi, ctx->c0); + rocke_value_t* lt = rocke_b_cmp_lt(b, wi, ctx->params.p_Wi); + rocke_value_t* col_ok = rocke_b_land(b, ge, lt); + rocke_value_t* w_off = rocke_b_mul(b, wi, ctx->params.p_A_stride_wi); + rocke_value_t* col_off = rocke_b_mul(b, w_off, ctx->c_half_bytes); + st.col_terms[c] = rocke_b_select(b, col_ok, col_off, ctx->oob_sentinel); + } + } + st.n_off = rocke_b_mul(b, ctx->n, ctx->params.p_A_stride_n); - for(kh = 0; kh < KH; ++kh) + /* Iter args: accumulators, then rows PAD .. PAD+PF-1 in flight. */ + { + int kh, w, k, c, idx = 0; + for(kh = 0; kh < KH; ++kh) + { + for(w = 0; w < BLOCK_W; ++w) { - for(w = 0; w < BLOCK_W; ++w) - { - snprintf(name_store[idx], 32, "dw_acc_kh%d_w%d", kh, w); - iargs[idx].name = name_store[idx]; - iargs[idx].init = ctx->zero_f32; - ++idx; - } + snprintf(name_store[idx], 48, "dw_acc_kh%d_w%d", kh, w); + iargs[idx].name = name_store[idx]; + iargs[idx].init = ctx->zero_f32; + ++idx; } - /* Group 0's rows: the loop's first iteration computes them. */ - dw_load_group(ctx, ctx->c0, n_iters_v, col_w, col_s, n_cols, first_rows); - for(kh = 0; kh < KH; ++kh) + } + for(k = 0; k < PF; ++k) + { + dw_load_row(ctx, &st, rocke_b_const_i32(b, p->PAD + k), window + k * n_cols); + for(c = 0; c < n_cols; ++c) { - for(c = 0; c < n_cols; ++c) - { - snprintf(name_store[idx], 32, "dw_a_r%d_c%d", kh, col_num[c]); - iargs[idx].name = name_store[idx]; - iargs[idx].init = first_rows[kh * n_cols + c]; - ++idx; - } + snprintf(name_store[idx], 48, "dw_a_r%d_c%d", k, col_num[c]); + iargs[idx].name = name_store[idx]; + iargs[idx].init = window[k * n_cols + c]; + ++idx; } } + } + + fma_loop = rocke_b_scf_for_iter(b, + ctx->c0, + n_fma_groups_v, + c1, + iargs, + num_iargs, + "dw_grp", + /*unroll=*/false, + /*elide_trailing_barrier=*/false); + rocke_b_region_enter(b, fma_loop.body); + + for(i = 0; i < num_acc; ++i) + new_accs[i] = fma_loop.iter_vars[i]; + for(i = 0; i < num_win; ++i) + window[i] = fma_loop.iter_vars[num_acc + i]; + + for(j = 0; j < KH; ++j) + { + rocke_value_t* y_j = dw_row_y(ctx, &st, fma_loop.iv, j); + int w_out, s_const, r_const, c; - group_loop = rocke_b_scf_for_iter(b, - ctx->c0, - n_groups_v, - c1, - iargs, - num_iargs, - "dw_grp", - /*unroll=*/false, - /*elide_trailing_barrier=*/false); - rocke_b_region_enter(b, group_loop.body); - - /* Seed new_accs from loop-carried values */ + /* Issue row y_j + PF, then compute row y_j (the oldest in flight). */ { - int i; - for(i = 0; i < num_acc; ++i) - new_accs[i] = group_loop.iter_vars[i]; + rocke_value_t* cpf = rocke_b_const_i32(b, PF); + dw_load_row(ctx, &st, rocke_b_add(b, y_j, cpf), window + num_win); } - cur_rows = group_loop.iter_vars + num_acc; - /* Prefetch the next group before computing this one. */ - dw_load_group(ctx, - rocke_b_mul(b, rocke_b_add(b, group_loop.iv, c1), c_KH), - n_iters_v, - col_w, - col_s, - n_cols, - next_rows); - - for(j = 0; j < KH; ++j) + for(c = 0; c < n_cols; ++c) + a_row[c] = rocke_b_cast_to_f32(b, window[c]); + memmove(window, window + n_cols, (size_t)num_win * sizeof(rocke_value_t*)); + + /* Python order: outer w_out, inner s_const, innermost r_const. */ + for(w_out = 0; w_out < BLOCK_W; ++w_out) { - rocke_value_t* y_j; - rocke_value_t* j_valid; - int w_out, s_const, r_const, c; - int P_FLUSH_j = (j + 1) % KH; - rocke_value_t* p_flush_rv; - rocke_value_t* flush_ge; - rocke_value_t* should_flush; - rocke_value_t* ho_row_j; - - /* y_j = grp_iv*c_KH + j -- emit mul first, then const(j), matching Python. */ + for(s_const = 0; s_const < KW; ++s_const) { - rocke_value_t* mul_gk = rocke_b_mul(b, group_loop.iv, c_KH); - rocke_value_t* cj = rocke_b_const_i32(b, j); - y_j = rocke_b_add(b, mul_gk, cj); - } - j_valid = rocke_b_cmp_lt(b, y_j, n_iters_v); - for(c = 0; c < n_cols; ++c) - a_row[c] = rocke_b_cast_to_f32(b, cur_rows[j * n_cols + c]); + rocke_value_t* a_f32 = a_row[col_idx[w_out * stride + s_const]]; - /* Python order: outer w_out, inner s_const, innermost r_const. */ - for(w_out = 0; w_out < BLOCK_W; ++w_out) - { - for(s_const = 0; s_const < KW; ++s_const) + for(r_const = 0; r_const < KH; ++r_const) { - rocke_value_t* a_f32 = a_row[col_idx[w_out * stride + s_const]]; - - for(r_const = 0; r_const < KH; ++r_const) - { - int p_idx = ((j - r_const + KH) % KH) * BLOCK_W + w_out; - new_accs[p_idx] = rocke_b_fma( - b, ctx->weights_f32[r_const][s_const], a_f32, new_accs[p_idx]); - } + /* STATIC slot of output row (y_j - r) / stride. */ + int p_idx = ((p->PAD + j - r_const + KH) % KH) * BLOCK_W + w_out; + new_accs[p_idx] = rocke_b_fma( + b, ctx->weights_f32[r_const][s_const], a_f32, new_accs[p_idx]); } } + } - /* Flush / store for slot P_FLUSH_j */ - p_flush_rv = rocke_b_add(b, y_j, rocke_b_const_i32(b, -(KH - 1))); - - if(j >= KH - 1) - { - flush_ge = j_valid; - } - else - { - flush_ge = rocke_b_land(b, rocke_b_cmp_lt(b, ctx->c0, group_loop.iv), j_valid); - } - - if(ctx->c_stride_dw == 1) - { - should_flush = flush_ge; - } - else - { - rocke_value_t* flush_stride - = rocke_b_cmp_eq(b, rocke_b_mod(b, p_flush_rv, c_stride_rv), ctx->c0); - should_flush = rocke_b_land(b, flush_ge, flush_stride); - } - - ho_row_j = rocke_b_div(b, p_flush_rv, c_stride_rv); + dw_flush(ctx, &st, fma_loop.iv, j, y_j, new_accs); + } - for(w_out = 0; w_out < BLOCK_W; ++w_out) - { - rocke_value_t* out_q - = rocke_b_add(b, ctx->q_tile_start, rocke_b_const_i32(b, w_out)); - rocke_value_t* out_q_ok - = rocke_b_land(b, rocke_b_cmp_lt(b, out_q, ctx->c_W), ctx->ch_in_range); - rocke_value_t* store_ok = rocke_b_land(b, out_q_ok, should_flush); - rocke_value_t* acc_val = new_accs[P_FLUSH_j * BLOCK_W + w_out]; - rocke_value_t* d_off = NULL; - rocke_value_t* d_valid = NULL; - rocke_value_t* safe_d; - const char* off_names[4]; - rocke_value_t* off_vals[4]; - - off_names[0] = "n"; - off_vals[0] = ctx->n; - off_names[1] = "h"; - off_vals[1] = ho_row_j; - off_names[2] = "w"; - off_vals[2] = out_q; - off_names[3] = "k"; - off_vals[3] = ctx->ch; - rocke_transforms_descriptor_offset( - b, ctx->d_desc, off_names, off_vals, 4, &d_off, &d_valid); - - safe_d = rocke_b_select( - b, store_ok, rocke_b_mul(b, d_off, ctx->c_half_bytes), ctx->oob_sentinel); - if(ctx->is_bf16) - rocke_b_buffer_store_bf16( - b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_bf16(b, acc_val)); - else - rocke_b_buffer_store_f16( - b, ctx->d_rsrc, safe_d, ctx->c0, rocke_b_trunc_f32_to_f16(b, acc_val)); - - /* Unconditional static reset */ - new_accs[P_FLUSH_j * BLOCK_W + w_out] = ctx->zero_f32; - } + /* An empty asm on each accumulator keeps LLVM's SLP vectorizer from + * pairing the loop-carried FMA chains into v_pk_fma_f32 (see Python). */ + { + const rocke_type_t* f32_ty = rocke_f32(); + rocke_inline_asm_opts_t opts; + memset(&opts, 0, sizeof opts); + opts.sideeffect = false; + opts.sideeffect_set = true; + for(i = 0; i < num_acc; ++i) + { + rocke_op_t* asm_op + = rocke_b_inline_asm(b, "", "=v,0", &new_accs[i], 1, &f32_ty, 1, &opts); + new_accs[i] = asm_op ? asm_op->results[0] : NULL; } - + } + for(i = 0; i < num_acc; ++i) + yields[i] = new_accs[i]; + for(i = 0; i < num_win; ++i) + yields[num_acc + i] = window[i]; + rocke_b_scf_yield(b, yields, num_iargs); + rocke_b_region_leave(b); + + /* The rows past the image: no input, only output rows to complete. */ + { + int kh, w, idx = 0; + for(kh = 0; kh < KH; ++kh) { - int i; - for(i = 0; i < num_acc; ++i) - yields[i] = new_accs[i]; - for(i = 0; i < num_rows; ++i) - yields[num_acc + i] = next_rows[i]; + for(w = 0; w < BLOCK_W; ++w) + { + snprintf(name_store[idx], 48, "dw_tail_acc_kh%d_w%d", kh, w); + iargs[idx].name = name_store[idx]; + iargs[idx].init = fma_loop.op->results[idx]; + ++idx; + } } - rocke_b_scf_yield(b, yields, num_iargs); - rocke_b_region_leave(b); } + flush_loop = rocke_b_scf_for_iter(b, + n_fma_groups_v, + n_groups_v, + c1, + iargs, + num_acc, + "dw_tail", + /*unroll=*/false, + /*elide_trailing_barrier=*/false); + rocke_b_region_enter(b, flush_loop.body); + for(i = 0; i < num_acc; ++i) + new_accs[i] = flush_loop.iter_vars[i]; + for(j = 0; j < KH; ++j) + dw_flush(ctx, &st, flush_loop.iv, j, dw_row_y(ctx, &st, flush_loop.iv, j), new_accs); + rocke_b_scf_yield(b, new_accs, num_acc); + rocke_b_region_leave(b); return rocke_ir_builder_kernel(b); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp index e9a1bea5fb54..66b621586cc5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp @@ -181,6 +181,15 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( } p = &spec->problem; + /* The weight table is sized for filters up to the depthwise limit. */ + if(p->KH > ROCKE_DCONV_DW_MAX_KH || p->KW > ROCKE_DCONV_DW_MAX_KW) + { + if(bld->status == ROCKE_OK) + { + bld->status = ROCKE_ERR_VALUE; + } + return NULL; + } WAVE = spec->wave_size; BLOCK_WAVES = spec->block_waves; THREADS = rocke_direct_depthwise_spatial_threads_per_block(spec); @@ -196,6 +205,9 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( is_bf16 = (p->dtype && strcmp(p->dtype, "bf16") == 0) ? 1 : 0; rocke_attr_set_int(bld, &bld->kernel->attrs, "max_workgroup_size", THREADS); + /* Two waves per SIMD: the KH x KW weights and the prefetch window push + * larger filters just over the one-wave threshold; two waves are faster. */ + rocke_attr_set_int(bld, &bld->kernel->attrs, "waves_per_eu", 2); /* ---- parameters ---- */ { @@ -541,6 +553,22 @@ rocke_kernel_def_t* rocke_build_direct_depthwise_spatial( new_accs[P_FLUSH_j] = zero_f32; } + /* An empty asm on each accumulator keeps LLVM's SLP vectorizer from + * pairing the loop-carried FMA chains into v_pk_fma_f32 (see Python). */ + { + const rocke_type_t* f32_ty = rocke_f32(); + rocke_inline_asm_opts_t opts; + int i; + memset(&opts, 0, sizeof opts); + opts.sideeffect = false; + opts.sideeffect_set = true; + for(i = 0; i < KH; ++i) + { + rocke_op_t* asm_op + = rocke_b_inline_asm(bld, "", "=v,0", &new_accs[i], 1, &f32_ty, 1, &opts); + new_accs[i] = asm_op ? asm_op->results[0] : NULL; + } + } { rocke_value_t** yields = (rocke_value_t**)alloca((size_t)num_iargs * sizeof(rocke_value_t*)); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py index a067d231f26b..9c704addf2c3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py @@ -21,7 +21,7 @@ import ctypes from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from ..core.arch import base_arch_from_target_id from ._ctypes_bind import _LazyFn @@ -633,6 +633,59 @@ def launch( bucket.append(((args_buf, size_buf, extra), evt)) return evt + def prepare_launch( + self, + fn: _HipFunctionHandle, + grid: Tuple[int, int, int], + block: Tuple[int, int, int], + args_packed: bytes, + *, + shared_bytes: int = 0, + stream: int = 0, + ) -> "Callable[[], None]": + """Build every ctypes argument of one launch once; return a + zero-argument callable that only enqueues it. + + For loops that launch the same kernel with the same arguments + over and over -- the timed loop of a benchmark. :meth:`launch` + rebuilds the args buffer, the ``extra`` array and the scalar + ctypes on each call, which costs more host time than a small + kernel runs, so a timed loop of those measures the host and not + the GPU. The callable owns its args buffer, so it stays valid + for as long as the callable is alive and is never written to + again: the ``extra`` path's late host-buffer read (see + :func:`rocke.runtime.packing.pack_args_kernelparams`) always + sees the right bytes. + """ + args_buf = (ctypes.c_ubyte * len(args_packed)).from_buffer_copy(args_packed) + size_buf = ctypes.c_size_t(len(args_packed)) + extra = (ctypes.c_void_p * 5)( + HIP_LAUNCH_PARAM_BUFFER_POINTER, + ctypes.cast(args_buf, ctypes.c_void_p), + HIP_LAUNCH_PARAM_BUFFER_SIZE, + ctypes.cast(ctypes.pointer(size_buf), ctypes.c_void_p), + HIP_LAUNCH_PARAM_END, + ) + call_args = ( + fn, + ctypes.c_uint(grid[0]), + ctypes.c_uint(grid[1]), + ctypes.c_uint(grid[2]), + ctypes.c_uint(block[0]), + ctypes.c_uint(block[1]), + ctypes.c_uint(block[2]), + ctypes.c_uint(shared_bytes), + ctypes.c_void_p(int(stream)), + None, + extra, + ) + keep = (args_buf, size_buf, extra) + + def enqueue(_keep=keep) -> None: + _check(_hipModuleLaunchKernel(*call_args), "hipModuleLaunchKernel") + + return enqueue + def launch_blocking( self, fn: _HipFunctionHandle, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py index b718bbbd9d52..4a99359b0e2d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py @@ -490,6 +490,37 @@ def __call__( rt.retain_for_stream(stream, *values.values()) return LaunchSummary(launches=1) + def bind( + self, + values: Mapping[str, Any], + *, + config: LaunchConfig, + ) -> Callable[[], None]: + """Pack ``values`` once; return a zero-argument asynchronous launch. + + The callable enqueues the kernel with exactly these arguments on + ``config``'s stream and returns -- no packing, no stream lookup, + no fence. Hand it to :func:`time_launches`: a timed loop of + ``launcher(values, config=...)`` pays the Python packing on every + launch, and for a kernel shorter than that the loop times the + host. The callable keeps ``values`` alive for as long as it + lives; drain the stream before dropping it. + """ + launch = _runtime().prepare_launch( + self._fn, + config.grid, + config.block, + self._packer(values), + shared_bytes=config.shared_bytes, + stream=resolve_stream(config.stream), + ) + keep = tuple(values.values()) + + def run(_keep=keep) -> None: + launch() + + return run + def __repr__(self) -> str: key_str = f", cache_key={self._cache_key!r}" if self._cache_key else "" return f"KernelLauncher({self._kernel_name!r}{key_str})" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 443ee6995dd3..45cf6fc2e516 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -125,8 +125,8 @@ "conv_direct_grouped": { "0": "bba7ec55a99618a311f5a632fe0a76ceb8904356f18f1d68468ee4e35d39de07", "1": "41caa5bc1a1e21b8ac280942fe658bb2e1b7813d1780de124edc1a97a82d9d48", - "10": "6ad18c90e3246ea1e62294e65caa5e8c247aad5a1b0244f6b46e28ba2ed2c9e7", - "11": "6bb7a73fa1daa80de9a4094b50f7166068710b810dd9e8e9ea12271e3c0cba29", + "10": "937610c7bd042724ac0ae9c12ce6a4f5381e63090973b2f92ff4ba375dced941", + "11": "7f667791c50d4f1ce0271a0782db3e54d6c00952fbc46b4d2172cab8f08ad4d6", "12": "18585f97b7f64e5e8a44df4134eab70f9f3b18ccb766db4de74e5c113b5e1095", "13": "62e24ebde4896d81a6789d283a8667c2937ac3e1a197df8ee7ab3e6bfc28f0c7", "14": "66545c341ebc8c3a0dec020a8103559bebb8099841463b396a3f07e3d6bbdcc1", @@ -138,8 +138,8 @@ "2": "5d93548857f2edf16e8fc1e839ffe6a5ed829952dc75b0edb438ab905db26de6", "20": "4ca258b9d883e4326a8b2dbe1d8aa252ad1ddd1c1c0e84ab4b5a5e2e8a087388", "21": "e3a7414cf1fcd93b058b9d5e267f8ff3f23b8e7aec494355c5e721b7371727a9", - "22": "6d277e8ae2810023f6283bb245c3943b90dfea789e91c31f5217d027229a60f5", - "23": "15a4768eb2f30649e82771216fa905993a46c3616b27460b829e6e8b45c1a87a", + "22": "cf2c5c499ff2d7890a5f47affa07d5280231b2b8737c44fb7cfedb305c6635b0", + "23": "211168b7f2ca03564a941640d6adff5bf785ee325bc663e103f27a1f8d5b020c", "24": "bd0167063cda9af5fc38d9cbb3b10fca05cd9a2fba943e46153dd2b66f06ea52", "25": "f4ad024735075c18f83ffa6507069a378ab3297f0a111c404b2bac2811d3e57a", "26": "a53aba5a8e55b8f07ccc843e701c78dfd42b85f85aa6daa9ab65bc1a6d0875a3", @@ -147,12 +147,13 @@ "3": "c21cbf7a6bab52d6a5c8f9085ff3fbfb5b63bc5c05ee9c1ab8c4a2d056ca5a32", "30": "16075e7b7f50652db14d4342e79d0d0b576e44945a2272aa01781afa13a42c82", "31": "da67c021330cbc92f3d0c96fc9bc1d6efcffc53ebbcb19c5d93b4b0b2ab19854", + "32": "166f8859b616f4233433fc481225d40b485eac1465efdb3aaae49b06224644ee", "4": "9c15c589f1d5e6f896e170a3c484b86490a564f3fd9e71a39c39d705c0a40b65", "5": "abb1d9ba0ef3a0e97f89d5b033e6ed1626aeb25b80defcda6c36e6151b397ad1", "6": "2ef16c7aea9fbbe3e2f4bc9c13cea6e2d5e94692d83396757058822e9d9d3cb6", "7": "7fd61c571db4c684397cf9afcc063d2b6295c9b52ab13cfbfe84588bcbf366a1", - "8": "73be0cf06c11b24b28277fcf835dcaa1d78294c8a518e157cc0d60bcb6ff8cca", - "9": "c80751c9816a2847c82823222162161ac1783caa608e28dcb730e78cba8a3099" + "8": "993ca72cc55ad6a7d2aa146ef0daa374464637427696294c7f69587ded651ce9", + "9": "12791f90bfbe4b4d84b154a30a67c77910ccbec6a779d70c81ca4ce3b34defee" }, "conv_implicit_gemm": { "0": "a09ed5bdeb9da980fee688bd046209d9ce126fc79754020e56c720034409f926", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index 13d540571328..7a0d224846f1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -5811,6 +5811,45 @@ def test_lower_to_hip_include_prologue_can_be_disabled(self): self.assertIn("void bare(", out) +class TestLauncherBind(unittest.TestCase): + """:meth:`KernelLauncher.bind` packs once and then only enqueues. + + No GPU is required: the launcher is built without loading a module and + the runtime is replaced by a fake that records ``prepare_launch``. + """ + + def test_bind_packs_once_and_enqueues_per_call(self): + from unittest import mock + + import rocke.runtime.launcher as L + + packed = [] + enqueued = [] + + def packer(values): + packed.append(dict(values)) + return b"\x01\x02" + + class FakeRuntime: + def prepare_launch(self, fn, grid, block, args, *, shared_bytes, stream): + self.prepared = (fn, grid, block, args, shared_bytes, stream) + return lambda: enqueued.append(1) + + launcher = L.KernelLauncher.__new__(L.KernelLauncher) + launcher._fn = "fn" + launcher._packer = packer + rt = FakeRuntime() + cfg = L.LaunchConfig(grid=(2, 1, 1), block=(64, 1, 1), stream=7) + with mock.patch.object(L, "_runtime", return_value=rt): + run = launcher.bind({"a": 1}, config=cfg) + for _ in range(3): + run() + + self.assertEqual(packed, [{"a": 1}]) + self.assertEqual(len(enqueued), 3) + self.assertEqual(rt.prepared, ("fn", (2, 1, 1), (64, 1, 1), b"\x01\x02", 0, 7)) + + class TestLauncherFenceContract(unittest.TestCase): """Per-launch event-fence policy in :mod:`rocke.runtime.launcher`. From 92292b1777e717cb58920cd87c0c92afd0ecdb1e Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 6 Oct 2026 10:12:15 +0000 Subject: [PATCH 3/7] fixes --- .../rocke_representative_ir_sha256.json | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index e092ab90b4ab..204224eda222 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -596,10 +596,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 259692, + "bytes": 222395, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "44d9bc3067fe38cc7e0db3e37a6f12b7f8b8e19fc6af5d296a8afd0ba8d33532", + "sha256": "2d65018c095d527989d0722d8d28fcec9b3a592774f0c284ccb092f74747fd07", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -652,10 +652,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 259692, + "bytes": 222395, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "0633f9155298ce6742d9f1ae2d220ea80e25440bd4477f90b83c8f74baad61ea", + "sha256": "e47cdf2c3e8ee8019db0d4f5636a9ceef7aa335d6446d45c8a7993c0ea97bdb4", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { @@ -1849,10 +1849,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 259831, + "bytes": 222534, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "c02bf0d0736d0927e96c9b46b9562d22aa9fa3c802b6ecab2d0394dda4d138ab", + "sha256": "8c4dfca822180b78acc6ac483576e2a7b831203b2c99e759fc3d9fdfe3c27d5e", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -1905,10 +1905,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 259831, + "bytes": 222534, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "98a6b0065f72219d125a9f3c09f7efa74a202af2669b5c7ef612f50a9690031a", + "sha256": "f91c2bb16c7a1674add9d7f88f24cd46da2cef12751b87eb7ac58ca3502cc0d1", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { @@ -3102,10 +3102,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 259835, + "bytes": 222538, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "ba7233457c6390326196ce2428d0d01b5b6397a7fc4dccfc2bae56b2fd88ea40", + "sha256": "3eadcea0a9fa5da060583ce3358edcaa27de443ca87c815173b1b70d46846d9f", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -3158,10 +3158,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 259835, + "bytes": 222538, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "f6a98e9101c9ee49ab7349e47f066df6a99a0bcacabd876135cad67aea9ea150", + "sha256": "2a109170be3ab1df4e720d864a39e9516de84897912deac46662d55c2d39cb72", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { From 8e6751d22b1554f097bf46acf02fde1a22952760 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 6 Oct 2026 06:01:07 -0500 Subject: [PATCH 4/7] ir --- .../rocke_representative_ir_sha256.json | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 204224eda222..5de9b23296ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -596,10 +596,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 222395, + "bytes": 180680, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "2d65018c095d527989d0722d8d28fcec9b3a592774f0c284ccb092f74747fd07", + "sha256": "20403a77e59bfdd92900a814591df87d3c72caf67843d04789c4b04bd19f9fcd", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -652,10 +652,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 222395, + "bytes": 180680, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "e47cdf2c3e8ee8019db0d4f5636a9ceef7aa335d6446d45c8a7993c0ea97bdb4", + "sha256": "f23796300bfedb7099fa740dd78d6ccadc3d28b11017a34b49f92e3bcf5c0d41", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { @@ -1849,10 +1849,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 222534, + "bytes": 180819, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "8c4dfca822180b78acc6ac483576e2a7b831203b2c99e759fc3d9fdfe3c27d5e", + "sha256": "eff70c8198c20978aa43529b0f86d1f3d9f6d5347f3ee33d0293da24e209bd7d", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -1905,10 +1905,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 222534, + "bytes": 180819, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "f91c2bb16c7a1674add9d7f88f24cd46da2cef12751b87eb7ac58ca3502cc0d1", + "sha256": "601265a8868379a961a843cbdd14d4ab32058c73cf731de8d4085a30b3a202e9", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { @@ -3102,10 +3102,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 222538, + "bytes": 180823, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "3eadcea0a9fa5da060583ce3358edcaa27de443ca87c815173b1b70d46846d9f", + "sha256": "a930c98b2e0e1d4eb6dc4eb059af54ddd58fdd9450731be6d993a76a668d8ea5", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -3158,10 +3158,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 222538, + "bytes": 180823, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "2a109170be3ab1df4e720d864a39e9516de84897912deac46662d55c2d39cb72", + "sha256": "b0aff29da91dbf33235cd2ad25f3e27df7892f72f864ecfa699c297033c379f0", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { From deaacc05d2c8ff0b19f155016e1930f391b7c952 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 6 Oct 2026 06:12:31 -0500 Subject: [PATCH 5/7] feat(rocke): non-grouped direct conv fwd Co-Authored-By: Claude Opus 5.5 (1M context) --- .../common/benchmark_direct_conv.py | 259 ++- .../benchmarks/common/direct_kernel_sweep.py | 93 +- .../library/benchmarks/common/kernel_cache.py | 3 + .../common/README_conv_direct_grouped.md | 4 + .../common/README_conv_direct_nongrouped.md | 318 ++++ .../rocke/library/kernels/common/conv_abi.py | 5 + .../kernels/common/conv_direct_nongrouped.py | 1015 ++++++++++++ .../tests/parity/conv_direct_grouped_emit.c | 160 +- .../tests/parity/conv_direct_grouped_emit.py | 186 ++- .../tests/test_builder_signature_contract.py | 1 + .../rocke/library/tests/test_conv_abi.py | 43 + .../tests/test_direct_conv_correctness.py | 461 +++++- .../library/tests/test_direct_kernel_cache.py | 45 +- .../platform/cpp/bindings/rocke_engine.cpp | 94 ++ .../rocke/platform/cpp/core/arch/query.cpp | 55 + .../platform/cpp/include/rocke/arch_target.h | 6 + .../rocke/instance_conv_direct_nongrouped.h | 161 ++ .../conv_direct_grouped_spec_and_validate.cpp | 75 +- .../common/conv_direct_nongrouped.cpp | 1415 +++++++++++++++++ .../dsl_docs/architecture/kernel_taxonomy.md | 1 + .../dsl_docs/instances/convolution.md | 80 + .../platform/dsl_docs/instances/index.md | 2 + .../platform/dsl_docs/reference/api_index.md | 1 + .../platform/dsl_docs/reference/file_index.md | 1 + .../platform/python/rocke/core/backend.py | 75 + .../python/rocke/instances/SUPPORT_MATRIX.md | 1 + .../rocke_representative_ir_sha256.json | 141 +- .../instances/rocke_ir_parity_harness.py | 93 ++ 28 files changed, 4691 insertions(+), 103 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_nongrouped.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_nongrouped.cpp diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py index b3a17475bcf8..f56aa41d57a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py @@ -2,20 +2,22 @@ # SPDX-License-Identifier: MIT """Tile sweep benchmark for the parametric direct convolution kernels. -Two kernel families are covered: - cpg == 1 (groups == C == K) — depthwise: ``DirectDepthwiseSpec``, scalar fma. - cpg >= 4, cpg % 4 == 0 — grouped: ``DirectConvSpec``, mfma_f32_16x16x16_f16. +Three kernel families are covered: + cpg == 1 (groups == C == K) — depthwise: ``DirectDepthwiseSpec``, scalar fma. + groups == 1 — non-grouped: ``DirectNongroupedConvSpec``, LDS halo + reuse + mfma_f32_32x32x16 / 16x16x32. + cpg >= 4, cpg % 4 == 0 — grouped: ``DirectConvSpec``, mfma_f32_16x16x16. The variant is selected automatically from C / groups. Run examples: python benchmark_direct_conv.py --N 8 --Hi 56 --Wi 56 --C 64 --K 64 --groups 64 # depthwise - python benchmark_direct_conv.py --N 8 --Hi 56 --Wi 56 --C 64 --K 64 --groups 1 # grouped cpg=64 + python benchmark_direct_conv.py --N 8 --Hi 56 --Wi 56 --C 64 --K 64 --groups 1 # non-grouped python benchmark_direct_conv.py --N 8 --Hi 56 --Wi 56 --C 1024 --K 1024 --groups 64 --verify Kernel cache. Without cache flags every run builds its own kernels. The -kernels take N, H, W and groups as kernel arguments, so they can instead be -compiled once and reused for any matching shape: +kernels take N, H, W and groups as kernel arguments (the non-grouped one also C +and K), so they can instead be compiled once and reused for any matching shape: python benchmark_direct_conv.py --compile-all --cache-dir ./kernel_cache --jobs 0 python benchmark_direct_conv.py --run-from-cache ./kernel_cache --N 8 --Hi 56 --Wi 56 \ --C 1024 --K 1024 --groups 64 --verify @@ -54,6 +56,10 @@ DW_BLOCK_WAVES as _DW_BLOCK_WAVES, ) +# The non-grouped (groups == 1) geometry sweep lives next to the kernel, in +# ``kernels.common.conv_direct_nongrouped.nongrouped_specs``, because the useful +# tile widths depend on Wo. + # --------------------------------------------------------------------------- # Result records @@ -72,6 +78,16 @@ class Result: passed: "bool | None" = None +@dataclass +class NonGroupedResult: + kernel_name: str + label: str + ms: float + tflops: float + gbps: float + passed: "bool | None" = None + + @dataclass class DepthwiseResult: kernel_name: str @@ -912,6 +928,230 @@ def _run_sweep( return 0, results +# --------------------------------------------------------------------------- +# Non-grouped (groups == 1) sweep +# --------------------------------------------------------------------------- + + +def _run_nongrouped_sweep( + *, + args, + problem, + dtype: str = "fp16", + arch: str, + compile_kernel, + jobs: int, + synchronize_and_release, + time_launches, + Runtime, + KernelLauncher, + LaunchConfig, + u8, +) -> "tuple[int, List[NonGroupedResult]]": + """Sweep :class:`DirectNongroupedConvSpec` over tile / wave / swizzle geometry. + + The grouped ``DirectConvSpec`` draws its parallelism from the ``groups`` + axis and collapses to a single wave per output row when ``groups == 1``, + so it is not a usable fallback here. + """ + import torch + + from kernels.common.conv_abi import conv_direct_args_signature + from kernels.common.conv_args import ConvArgs + from kernels.common.conv_direct_nongrouped import ( + build_direct_conv_nongrouped, + nongrouped_specs, + ) + from rocke.runtime.hip_module import HipError + + p = problem + + torch.manual_seed(42) + _torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16 + A_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=_torch_dtype).uniform_(-1.0, 1.0) + B_t = torch.empty(p.total_k, p.KH, p.KW, p.cpg, dtype=_torch_dtype).uniform_( + -1.0, 1.0 + ) + D_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=_torch_dtype) + + bytes_xfer = float(A_t.nbytes + B_t.nbytes + D_t.nbytes) + flop = float(p.flops) + # Direct conv is AOT: the whole shape travels as kernargs. + sig = conv_direct_args_signature(dtype) + + specs = nongrouped_specs(p, arch=arch, name="rocke_bench_direct_conv_nongrouped") + if args.sample is not None: + total = len(specs) + specs = _sample_combos(specs, args.sample, args.seed) + print( + f"Sampling {len(specs)}/{total} combinations " + f"({args.sample*100:.0f}%, seed={args.seed}).", + flush=True, + ) + + print( + f"Sweeping {len(specs)} non-grouped combinations for {arch} {dtype} " + f"{p.short()} (C={p.cpg}, K={p.kpg}) ...", + flush=True, + ) + + n_skipped = 0 + pending = [] + for spec in specs: + try: + kernel = build_direct_conv_nongrouped(spec, arch=arch) + except ValueError: + n_skipped += 1 + continue + pending.append((spec, kernel)) + + if not pending: + print("No valid non-grouped configurations for this shape.", file=sys.stderr) + return 1, [] + + artifact_map = _compile_kernels_parallel( + [k for _, k in pending], compile_kernel, arch, jobs + ) + n_built = len(artifact_map) + + rt = Runtime() + A_dev = rt.alloc(A_t.nbytes) + B_dev = rt.alloc(B_t.nbytes) + D_dev = rt.alloc(D_t.nbytes) + rt.memcpy_h2d(A_dev, u8(A_t), A_t.nbytes) + rt.memcpy_h2d(B_dev, u8(B_t), B_t.nbytes) + rt.memset(D_dev, 0, D_t.nbytes) + + ref_out = None + if args.verify or args.dump_fail: + ref_out = _conv_reference(A_t, B_t, _DirectConvProblemAdapter(p)) + print( + f"Reference computed via torch ({tuple(ref_out.shape)}, {ref_out.dtype}).", + flush=True, + ) + + values = ConvArgs.from_problem(p).to_launch_values( + int(A_dev), + int(B_dev), + int(D_dev), + A_t.nbytes, + B_t.nbytes, + D_t.nbytes, + ) + + results: List[NonGroupedResult] = [] + _stop = EarlyStop(args.early_stop, args.early_stop_after) + n_run = 0 + for spec, kernel in pending: + artifact = artifact_map.get(kernel.name) + if artifact is None: + n_skipped += 1 + continue + try: + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=sig, + ) + except HipError as e: + n_skipped += 1 + print( + f"[skip] kernel load failed for {artifact.kernel_name}: {e}", + file=sys.stderr, + flush=True, + ) + continue + + grid = spec.grid() + block = (spec.threads_per_block, 1, 1) + stream = 0 + cfg = LaunchConfig(grid=grid, block=block, stream=stream) + + kernel_passed = None + if args.verify or args.dump_fail: + rt.memset(D_dev, 0, D_t.nbytes) + stopped, kernel_passed = _verify_kernel( + rt=rt, + launcher=launcher, + values=values, + grid=grid, + block=block, + out_dev=D_dev, + out_t=D_t, + ref_out=ref_out, + kernel_name=artifact.kernel_name, + dump_fail=args.dump_fail, + u8=u8, + ) + if stopped: + rt.free(A_dev) + rt.free(B_dev) + rt.free(D_dev) + return 1, [] + rt.memset(D_dev, 0, D_t.nbytes) + + ms = _stop.measure( + lambda: launcher(values, config=cfg), + warmup=args.warmup, + iters=args.iters, + stream=stream, + ) + if ms is None: + _stop.report(artifact.kernel_name) + continue + synchronize_and_release(stream) + + label = ( + f"th{spec.tile_h} tw{spec.tile_w} tk{spec.tile_k} ck{spec.ck} " + f"w{spec.waves_m}x{spec.waves_n} a{spec.atom} " + f"wgm{spec.swizzle_wgm} iglp{spec.iglp} we{spec.waves_per_eu}" + ) + n_run += 1 + results.append( + NonGroupedResult( + kernel_name=artifact.kernel_name, + label=label, + ms=ms, + tflops=(flop / ms) * 1e-9, + gbps=(bytes_xfer / ms) * 1e-6, + passed=kernel_passed, + ) + ) + print( + f"[{n_run:4d}] {label} {results[-1].tflops:6.1f} TFLOPS {ms:.3f} ms", + flush=True, + ) + + rt.free(A_dev) + rt.free(B_dev) + rt.free(D_dev) + print(f"\nSweep done: {n_built} compiled, {n_skipped} skipped.", flush=True) + + if not results: + print("No valid configurations found.", file=sys.stderr) + return 1, [] + + results.sort(key=lambda r: r.tflops, reverse=True) + top_n = min(args.top, len(results)) + width = 110 + print(f"\n{'='*width}") + print(f"Top {top_n} non-grouped configurations for {arch} {dtype} {p.short()}") + print(f"{'='*width}") + print( + f"{'rank':>4} {'TFLOPS':>7} {'ms':>8} {'GBps':>7} " + + (f"{'verify':>6} " if args.verify else "") + + "config" + ) + print("-" * width) + for rank, r in enumerate(results[:top_n], 1): + v = f"{'PASS' if r.passed else 'FAIL':>6} " if args.verify else "" + print( + f"{rank:>4} {r.tflops:>7.1f} {r.ms:>8.3f} {r.gbps:>7.1f} {v}{r.label}" + ) + print(f"\nBest: {results[0].tflops:.1f} TFLOPS — {results[0].kernel_name}") + return 0, results + + # --------------------------------------------------------------------------- # Wgrad sweep # --------------------------------------------------------------------------- @@ -2267,8 +2507,9 @@ def main() -> int: cpg = args.C // args.groups kpg = args.K // args.groups - # For fwd direction, cpg == kpg is required. For dgrad it need not hold. - if args.direction == "fwd" and cpg != kpg: + # The grouped/depthwise fwd kernels require cpg == kpg; the non-grouped + # (groups == 1) fwd kernel takes any C/K, and dgrad need not hold it. + if args.direction == "fwd" and args.groups != 1 and cpg != kpg: print( f"error: cpg={cpg} != kpg={kpg}; forward direct grouped conv requires C/groups == K/groups", file=sys.stderr, @@ -2359,6 +2600,8 @@ def _u8(t): rc, _ = _run_wgrad_sweep(problem=problem, dtype=dtype, **_common) elif cpg == 1: rc, _ = _run_depthwise_sweep(problem=problem, dtype=dtype, **_common) + elif problem.groups == 1: + rc, _ = _run_nongrouped_sweep(problem=problem, dtype=dtype, **_common) else: rc, _ = _run_sweep(problem=problem, dtype=dtype, **_common) all_rc = all_rc or rc diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py index c623e6540f8d..edf89358c8a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py @@ -15,6 +15,10 @@ follows the set the reference direct-conv library covers on CDNA and drops what rocke's kernels cannot compute (see the table's comments). +The non-grouped (``groups == 1``) forward kernel bakes less: its channel counts +are kernargs as well, so its capabilities record ``cpg = kpg = 0`` (runtime) +and one binary serves every ``C``/``K`` its tile divides. + The run side is split in two. :func:`direct_plans` is pure CPU: it picks the cached kernels that fit a problem, re-validates each spec against the real shape, and returns a launch plan -- binaries, signatures, grids and the @@ -75,7 +79,11 @@ @dataclass(frozen=True) class DirectCaps: - """What one binary is baked for; everything else about a shape is runtime.""" + """What one binary is baked for; everything else about a shape is runtime. + + ``cpg = kpg = 0`` means the channel counts are runtime too (the non-grouped + kernel). + """ KH: int PAD: int @@ -96,6 +104,9 @@ def _grouped_channels() -> Tuple[int, ...]: _DEPTHWISE_FILTERS = (3, 5, 7, 9, 11) _STRIDES = (1, 2) +# The non-grouped kernel exists for 3x3: its LDS halo reuse is what lets it beat +# implicit GEMM there. A 1x1 has no halo to share and is a plain GEMM. +_NONGROUPED_FILTERS = (3,) def _same_pad(kh: int) -> int: @@ -122,6 +133,17 @@ def _all_pads(kh: int) -> range: DirectCaps(KH=3, PAD=1, stride=1, cpg=c, kpg=c) for c in _grouped_channels() ), ), + # Non-grouped forward (groups == 1), C and K runtime. The kernel bounds + # rows and columns against the runtime output extents, so it is not tied + # to "same" padding the way the row-streaming kernels are; the cache bakes + # the "same" padding the forward rows here share. + "direct_nongrouped": ( + "fwd", + tuple( + DirectCaps(KH=k, PAD=_same_pad(k), stride=s, cpg=0, kpg=0) + for k, s in itertools.product(_NONGROUPED_FILTERS, _STRIDES) + ), + ), # Depthwise forward: square odd filters 3..11, stride 1 and 2. "direct_depthwise": ( "fwd", @@ -176,8 +198,13 @@ def _all_pads(kh: int) -> range: _REORGANIZE = "direct_dgrad_reorganize" -def _knob_grid(variant: str) -> Iterator[dict]: - if variant == "direct_grouped": +def _knob_grid(variant: str, arch: str, dtype: str) -> Iterator[dict]: + if variant == "direct_nongrouped": + from kernels.common.conv_direct_nongrouped import nongrouped_knobs + + # Every width (no Wo to pick from): plan_for keeps those that fit. + yield from nongrouped_knobs(arch, dtype) + elif variant == "direct_grouped": for bq, bg, db in itertools.product(BLOCK_Q, BLOCK_GROUPS, DOUBLE_BUFFER): yield dict(block_q=bq, block_groups=bg, double_buffer=db) elif variant in ("direct_depthwise", "direct_depthwise_dgrad"): @@ -221,9 +248,14 @@ def make_spec(variant: str, problem, knobs: dict): derived from the original problem separately. """ from kernels.common import conv_direct_grouped as dc + from kernels.common import conv_direct_nongrouped as dn if variant == "direct_grouped": return dc.DirectConvSpec(problem=problem, name="rocke_direct_conv", **knobs) + if variant == "direct_nongrouped": + return dn.DirectNongroupedConvSpec( + problem=problem, name="rocke_direct_conv_nongrouped", **knobs + ) if variant == "direct_depthwise": return dc.DirectDepthwiseSpec( problem=problem, name="rocke_direct_depthwise", **knobs @@ -283,8 +315,11 @@ def _uses_coalesced_weights(knobs: dict) -> bool: def validate_spec(variant: str, spec, arch: str) -> Tuple[bool, str]: """``(ok, reason)`` for ``spec`` on ``arch``, from the kernel's own validator.""" from kernels.common import conv_direct_grouped as dc + from kernels.common import conv_direct_nongrouped as dn try: + if variant == "direct_nongrouped": + return dn.is_valid_nongrouped_spec(spec, arch=arch) if variant in ("direct_grouped", "direct_grouped_dgrad_mfma"): spec.validate() return dc.is_valid_spec(spec, arch=arch) @@ -305,9 +340,11 @@ def validate_spec(variant: str, spec, arch: str) -> Tuple[bool, str]: def _build_kernel(variant: str, spec, arch: str): from kernels.common import conv_direct_grouped as dc + from kernels.common import conv_direct_nongrouped as dn builders = { "direct_grouped": dc.build_direct_conv, + "direct_nongrouped": dn.build_direct_conv_nongrouped, "direct_grouped_dgrad_mfma": dc.build_direct_conv, "direct_depthwise": dc.build_direct_depthwise, "direct_depthwise_spatial": dc.build_direct_depthwise_spatial, @@ -411,10 +448,26 @@ def probe_problem( spec validators: 64 groups divide every block_groups, 256 depthwise channels fill the widest depthwise block (the spatial depthwise kernel takes fewer groups than a wave, so it gets 8), and 64x64 leaves every - stride-2 output wider than the widest W tile. + stride-2 output wider than the widest W tile. Runtime channel counts + (``cpg == 0``, the non-grouped kernel) take one group of 256: a multiple of + every channel chunk and of every channel tile. """ from kernels.common.conv_direct_grouped import DirectConvProblem + if caps.cpg == 0: + return DirectConvProblem( + N=1, + H=64, + W=64, + groups=1, + cpg=256, + kpg=256, + KH=caps.KH, + KW=caps.KW, + PAD=caps.PAD, + stride=caps.stride, + dtype=dtype, + ) if variant == "direct_depthwise_spatial": groups = 8 else: @@ -462,7 +515,7 @@ def _caps_jobs( out: List[Tuple[str, BuildJob]] = [] n_raw = 0 probe = probe_problem(caps, dtype, variant) - for knobs in _knob_grid(variant): + for knobs in _knob_grid(variant, arch, dtype): n_raw += 1 if variant == "direct_grouped_dgrad_mfma" and _mfma_knobs_reason(probe, knobs): continue @@ -505,11 +558,13 @@ def _caps_worker(payload): return _caps_jobs(*payload) -def count_direct_jobs(directions: Sequence[str]) -> int: +def count_direct_jobs( + directions: Sequence[str], arch: str, dtype: str = DIRECT_DTYPE +) -> int: """Raw candidate count :func:`enumerate_direct_jobs` walks (the progress denominator): capabilities times tuning knobs, before validation.""" return sum( - len(caps_list) * sum(1 for _ in _knob_grid(variant)) + len(caps_list) * sum(1 for _ in _knob_grid(variant, arch, dtype)) for variant, (direction, caps_list) in DIRECT_CAPABILITIES.items() if direction in directions ) @@ -552,7 +607,7 @@ def enumerate_direct_jobs( total = 0 if log is not None: - total = count_direct_jobs(directions) + total = count_direct_jobs(directions, arch, dtype) log(f" {total} candidates to check with {jobs} process(es)") n_raw = n_valid = 0 @@ -724,8 +779,12 @@ def _conv_step(identity, hsaco_path, meta, spec, problem, direction, buffers): from kernels.common.conv_abi import conv_direct_args_signature from kernels.common.conv_args import ConvArgs from kernels.common.conv_direct_grouped import direct_launch_geometry + from kernels.common.conv_direct_nongrouped import DirectNongroupedConvSpec - grid, block = direct_launch_geometry(spec) + if isinstance(spec, DirectNongroupedConvSpec): + grid, block = spec.grid(), (spec.threads_per_block, 1, 1) + else: + grid, block = direct_launch_geometry(spec) return DirectStep( kernel_name=meta["kernel_name"], hsaco_path=hsaco_path, @@ -768,6 +827,18 @@ def plan_for( ok, why = validate_spec(variant, spec, arch) if not ok: return None, why + if variant == "direct_nongrouped": + # The cache holds every width; offer only those a sweep for this + # shape would try, so no binary spends MFMA work past the image edge + # when one that does not is cached. + from kernels.common.conv_direct_nongrouped import tile_w_candidates + + widths = tile_w_candidates(problem.Wo, spec.atom_tile) + if spec.tile_w not in widths: + return None, ( + f"tile_w={spec.tile_w} does not fit Wo={problem.Wo} " + f"(widths tried for it: {widths})" + ) if identity.direction == "direct_fwd": step = _conv_step( @@ -845,6 +916,10 @@ def direct_plans( # A binary only serves the operand dtype it was built for. if identity.dtype_a != getattr(problem, "dtype", DIRECT_DTYPE): continue + # The non-grouped family only exists for one group; offering it to a + # grouped problem would bury the grouped kernels' reasons in noise. + if identity.algorithm == "direct_nongrouped" and problem.groups != 1: + continue plan, why = plan_for(cache, identity, hsaco_path, meta, problem, arch) if plan is None: rejected.append((identity, why)) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py index ad2d32c69eaa..3648d6146d03 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py @@ -748,6 +748,9 @@ def supports_problem( if actual is not None and int(actual) != baked: return False, f"{field}={baked} but problem has {actual}" + # 0 means runtime here, for direct conv too: the non-grouped direct + # kernel takes its channel counts as kernargs (cpg == 0 is never a + # real shape, unlike PAD == 0 above). if identity.cpg and identity.cpg != cpg: return False, f"kernel baked cpg={identity.cpg}, problem has {cpg}" if identity.kpg and identity.kpg != kpg: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_grouped.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_grouped.md index 231b180a7ebf..c7e4b869f3db 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_grouped.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_grouped.md @@ -59,6 +59,9 @@ The fixed-`cpg` variants in the table require `kpg = cpg`. The generic positive multiple of 4. If neither fits, use the implicit-GEMM kernel ([`conv_implicit_gemm.py`](conv_implicit_gemm.py)). +For `groups == 1` these variants degenerate to one wave per output row; use the +non-grouped direct conv instead ([`README_conv_direct_nongrouped.md`](README_conv_direct_nongrouped.md)). + ### cpg = 1 — Depthwise (`DirectDepthwiseSpec`) Each output channel is computed independently using scalar FMA. The kernel @@ -429,6 +432,7 @@ For a backward-weights pass on gfx942 use the implicit-GEMM wgrad kernel | Depthwise (`cpg = 1`, `stride = 1`) | Direct-conv depthwise | | Backward-weights, gfx950, `cpg` and `kpg` ≥ 16, `stride = 1` | Direct-conv wgrad | | Backward-weights, anything else | Implicit-GEMM wgrad | +| Non-grouped (`groups = 1`), `stride ≤ 2` | Sweep `conv_direct_nongrouped` and implicit-GEMM, keep the faster ([`README_conv_direct_nongrouped.md`](README_conv_direct_nongrouped.md)) | --- diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md new file mode 100644 index 000000000000..6ebf086b3add --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md @@ -0,0 +1,318 @@ +# Direct Non-grouped Convolution (`groups == 1`) + +Source: [`conv_direct_nongrouped.py`](conv_direct_nongrouped.py) — +`DirectNongroupedConvSpec`, `is_valid_nongrouped_spec`, `build_direct_conv_nongrouped`, +`nongrouped_specs`, `nongrouped_knobs`, `tile_w_candidates`. + +C++ engine twin: `platform/cpp/instances/common/conv_direct_nongrouped.cpp` +(public header `rocke/instance_conv_direct_nongrouped.h`), byte-identical `.ll`. + +## Why a separate family + +The grouped direct-conv kernels ([`README_conv_direct_grouped.md`](README_conv_direct_grouped.md)) +take all their parallelism from the `groups` axis: one wave owns one group. At +`groups == 1` that degenerates to a single wave per output row holding +`ceil(K/16)` accumulator tiles, which leaves the machine almost idle. This family +is the `groups == 1` counterpart. It keeps the property that makes a direct conv +worth having next to implicit GEMM on 3x3, stride ≤ 2 shapes: + +> the input tile is staged in LDS **once per channel chunk, with its halo**, and +> all `KH*KW` filter taps read shifted sub-tiles of that single staged copy. + +An implicit-GEMM formulation re-reads the activations once per tap, so its +activation-side traffic is roughly `KH*KW` times larger and the L2 has to absorb +the redundancy. + +## Operands + +| Role | Tensor | Layout | Notes | +|------|--------|--------|-------| +| A | input activations | NHWC `[N, H, W, C]` | `C % 8 == 0` (dwordx4 staging loads) | +| B | weights | KRSC `[K, KH, KW, C]` | | +| D | output | NHWK `[N, Ho, Wo, K]` | `K % atom_tile == 0` | + +fp16 and bf16 I/O, f32 accumulation. `stride ∈ {1, 2}`, any `KH`/`KW`/`PAD`. + +## AOT kernel arguments + +The kernel is AOT: it declares the direct-conv kernarg block of +[`conv_abi.conv_direct_arg_names(direction="fwd")`](conv_abi.py), the same +block every direct-conv kernel takes: + +```text +A, B, D, A_bytes, B_bytes, D_bytes, +p_N, p_Hi, p_Wi, p_Ho, p_Wo, p_groups, p_total_c, p_total_k, +p_A_stride_{n,hi,wi}, p_D_stride_{n,ho,wo} +``` + +| Runtime (kernarg) | Baked (capability / knob) | +|-------------------|---------------------------| +| `N`, `H`, `W`, `Ho`, `Wo`, NHWC / NHWK strides | `KH`, `KW`, `stride`, `PAD`, dtype | +| `C` (`p_total_c`), `K` (`p_total_k`) | tile, waves, atom, `ck`, `lds_pad`, swizzle, schedule knobs | + +Unlike the grouped kernels, the channel counts are runtime too: they only set +the channel loop's trip count and the staging / store masks. `p_groups` is +ignored (the family is `groups == 1`). One binary therefore serves every shape +whose `C` is a multiple of `ck` (and of 8) and whose `K` is a multiple of the +atom tile; `spec.problem` only has to pass the validator, and the emitted IR +does not depend on it (`library/tests/test_conv_abi.py` asserts this). + +Launch with the shared helpers: + +```python +from kernels.common.conv_abi import conv_direct_args_signature +from kernels.common.conv_args import ConvArgs + +signature = conv_direct_args_signature(dtype) +values = ConvArgs.from_problem(problem).to_launch_values( + a_ptr, b_ptr, d_ptr, a_bytes, b_bytes, d_bytes) +``` + +Two host-side contracts move from build time to launch time, so run +`validate()` (or `is_valid_nongrouped_spec`) on the spec rebuilt for the +problem actually launched — the benchmark sweep and the AOT cache's +`plan_for` both do: + +* every operand tensor must end at or below the `0x7F000000` masked-access + offset (see optimization 5); +* the flattened grid must fit `gridDim.x` (2³¹ − 1). + +## Computation + +Per workgroup, with M = output channels and N = output pixels: + +```text +tile_k output channels (M; multiple of the atom tile) +tile_h x tile_w output pixels (N; tile_w a multiple of the atom tile) +ck input channels reduced per LDS stage + +for c0 in range(0, C, ck): # runtime scf.for; accumulators carried + stage X[(tile_h-1)*s+KH, (tile_w-1)*s+KW, ck] -> LDS (halo included) + stage W[tile_k, KH, KW, ck] -> LDS (MFMA fragment order) + for k_atom, s: # Python-unrolled + read activation fragments for every input row of the wave window + for r: + read weight fragments for tap (r, s) + for row, col-block, m-tile: acc += mfma(W_frag, X_frag[row*stride + r]) +``` + +Waves split the block `waves_m x waves_n`: `waves_m` along output channels, +`waves_n` along output rows; every wave keeps all `tile_w` columns. + +## The optimizations + +1. **LDS halo reuse.** The activation tile is staged with its `KH-1` / `KW-1` + halo once per channel chunk; all taps read shifted sub-tiles. This is the + reason the family exists. + +2. **Tap-shared activation fragments**. + Output row `row` at tap row `r` reads input row `row*stride + r`. Indexing + activation fragments by *input* row lets the `KH` taps of one filter column + share them: `(ROWS_W-1)*stride + KH` distinct fragments per `(s, k_atom)` + instead of `ROWS_W*KH`. Filter column `s` is the outer loop so every fragment + loaded for one `(s, k_atom)` is consumed by all `KH` tap rows and all + channel tiles of the wave before it dies. + +3. **Weights pre-swizzled into MFMA fragment order.** Weight slot + `((tap*M_TILES + m_tile)*KATOMS + k_atom)*64 + lane` holds exactly the + `FRAG` halves that `lane` feeds to the MFMA. Both the staging store and the + consuming read are linear in `lane`, so neither side can bank-conflict and no + padding is needed. The staging thread's global read is `FRAG` contiguous + channels of one `k_out` row — contiguous in KRSC. + +4. **Conflict-free activation reads via `lds_pad`.** The activation tile is + stored `[pos][c]` with channel stride `ck + lds_pad`. The default pad of 8 + halves spreads a `ds_read_b128` wave across all banks on gfx950 (64 banks, see + [`arch/gfx950.md`](../../../platform/dsl_docs/optimization/arch/gfx950.md)). + Neighbouring pads (4, 16) conflict; treat the value as tuned, not arbitrary. + +5. **Hoisted staging predication**. + The validity masks of the staging loads (image border, `K` tail, pass tail) + are loop-invariant. They are folded once, before the channel loop, into the + *base* byte offset: `select(valid, base, 0x7F000000)`. A buffer load whose + offset lies past `num_records` returns zero — exactly the value a padded pixel + or an out-of-range filter row needs — so the loop body carries no compare, + `and` or `select` at all. The epilogue redirects its masked stores to the + same offset (one sentinel for every masked access), where the hardware drops + them. `validate()` guarantees the trick holds: every operand tensor must end + at or below `0x7F000000` bytes, and `0x7F000000` plus the largest per-chunk + channel offset must still fit in i32. The tensor sizes are kernargs, so this + holds only for problems the spec was validated against — see + [AOT kernel arguments](#aot-kernel-arguments). + +6. **Software prefetch through loop-carried registers** + . The global loads + for chunk `i+1` are issued right after the barrier that publishes chunk `i` + and travel into the next iteration as `scf.for` iter-args, so chunk `i`'s + MFMAs cover the DRAM latency of chunk `i+1`. The tail iteration loads one + chunk past `C`; buffer loads clamp rather than fault, so no masking is needed. + +7. **Square 16-wide atoms for widths that are not a multiple of 32.** With a + 32-wide pixel tile the last tile of each row computes columns past `Wo`, and + that wasted MFMA work is spread over every block of the row. + `16x16x32` / `16x16x16` atoms plus `tile_w_candidates(Wo, atom_tile)` — which + prefers widths that divide `Wo` exactly — remove the waste. All four atoms + are square (`M == N == tile`), which is what lets one code path serve them. + +8. **`iglp_opt(0)` on the channel loop**. + The canned MFMA/memory interleave was the largest single scheduling win; + level 1 was worse than no hint at all. `iglp=None` leaves the scheduler alone. + +9. **Chiplet-aware grid decode.** The 1-D grid of (spatial cell × channel tile) + is remapped with `chiplet_aware_super_tile_dynamic` (the tile counts follow + the runtime extents) so the workgroups one XCD receives share operands; `swizzle_wgm` is the number of + channel tiles walked back-to-back per spatial cell. Ordering the channel tile + as the fast-varying axis is what matters; `swizzle_wgm=1` (spatial fastest) + is clearly worse, while the chiplet remap on top contributes little. + +10. **A clamped runtime trip count.** The channel loop runs `max(C/ck, 1)` + times. `validate()` makes the clamp a no-op (`C` is a positive multiple of + `ck`), but it lets LLVM prove the loop runs: with a bare `C/ck` it keeps a + zero-trip guard and schedules the prefetch loads late in the body. + +11. **Guards that turn hangs and silent corruption into errors.** + * `validate()` rejects tiles needing more than 256 accumulator registers per + lane: past that the backend scheduler does not just slow down, it hangs. + * `validate()` rejects a flattened grid past 2³¹ − 1 workgroups (the kernel + computes the workgroup count in i32), and `is_valid_nongrouped_spec` + never raises — the arch-table queries sit inside its error handling too. + * Staging loops are sized in whole thread-passes, so the last pass can own + slots past the tile. Both LDS arrays carry a scratch tail those slots are + redirected to; without it they overwrite the neighbouring allocation (symptom: row 0 correct, row 1 + wrong). + +### Levers measured and rejected + +Kept as knobs, off by default, because the sweep found them unhelpful on the +target shapes: + +| Lever | Knob | Why it did not help | +|-------|------|---------------------| +| LDS ping-pong (1 barrier per chunk instead of 2) | `double_buffer=True` | Doubles LDS, which drops the CU from two resident blocks to one — the lost overlap costs more than the saved barrier. | +| Occupancy hint | `waves_per_eu` | Mattered only without `iglp_opt(0)`; negligible with it. | +| Wider tiles on the shapes that still trail | sweep axes | No change: the residual is barrier / `waitcnt` exposure at two waves per SIMD, not bandwidth. | + +## Spec + +```python +@dataclass(frozen=True) +class DirectNongroupedConvSpec: + problem: DirectConvProblem # groups must be 1 + name: str = "direct_conv_nongrouped" + tile_h: int = 16 # output rows per workgroup + tile_w: int = 32 # output cols per workgroup + tile_k: int = 128 # output channels per workgroup + ck: int = 16 # input channels per LDS chunk + waves_m: int = 2 # waves along tile_k + waves_n: int = 4 # waves along tile_h + atom: str = "32x32x16" # 32x32x16 | 32x32x8 | 16x16x32 | 16x16x16 + wave_size: int = 64 + lds_pad: int = 8 + chiplet_swizzle: bool = True + swizzle_wgm: int = 8 + chiplet_chunk: int = 64 + num_xcds: int = 8 + double_buffer: bool = False + iglp: int | None = None # nongrouped_specs sweeps (0,) + waves_per_eu: int | None = None # nongrouped_specs sweeps (None, 3) +``` + +Constraints (`validate()` then `is_valid_nongrouped_spec(spec, arch)`): + +* `groups == 1`, dtype fp16/bf16, atom known and present on the target + (`32x32x16` and `16x16x32` need gfx950; `32x32x8` / `16x16x16` also gfx942); +* tile, wave, `ck`, stride, `C` and `K` values positive (checked before any + derived geometry, so a bad spec gets a reason, not a `ZeroDivisionError`); +* `tile_w`, `tile_k`, `K` multiples of the atom tile; `tile_h % waves_n == 0`; + `(tile_k / atom_tile) % waves_m == 0`; +* `ck % atom_k == 0`, `ck % 8 == 0`, `C % ck == 0`, `C % 8 == 0`; +* `lds_pad` a non-negative multiple of 8 (a negative pad overlaps neighbouring + pixels in LDS; the vec8 LDS access is emitted `align 16`, so the pixel stride + `ck + lds_pad` must stay a multiple of 8 halves); `swizzle_wgm >= 1`; `iglp` None or ≥ 0; `waves_per_eu` None or ≥ 1 + (the C++ port encodes None as -1 / 0, so those values are reserved in both + engines); flattened grid ≤ 2³¹ − 1 workgroups; A, B and D each ≤ + `0x7F000000` bytes; accumulators ≤ 256 registers per lane; +* `stride ∈ {1, 2}`; threads, wave size and `lds_bytes` within the arch limits. + +Usage: + +```python +from kernels.common.conv_direct_grouped import DirectConvProblem +from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, build_direct_conv_nongrouped, is_valid_nongrouped_spec, +) + +p = DirectConvProblem(N=4, H=64, W=64, groups=1, cpg=640, kpg=640, dtype="bf16") +spec = DirectNongroupedConvSpec(problem=p, tile_h=16, tile_w=64, tile_k=128, ck=16, + waves_m=2, waves_n=4, iglp=0) +ok, why = is_valid_nongrouped_spec(spec, arch="gfx950") +kernel = build_direct_conv_nongrouped(spec, arch="gfx950") +``` + +## Launch grid + +```text +grid = spec.grid() = (n_w_tiles * n_h_tiles * N * n_k_tiles, 1, 1) +block = (spec.threads_per_block, 1, 1) # waves_m * waves_n * 64 +``` + +The grid is flat; the kernel recomputes the tile counts from `p_Wo`, `p_Ho`, +`p_total_k` and `p_N` and decodes `(k_tile, spatial cell)` from the block id +(through the chiplet swizzle when enabled). The host computes the grid from the +problem it launches, so one binary is launched with whatever grid the shape +needs. + +## Candidate generation and benchmarking + +`nongrouped_specs(problem, arch=...)` returns every valid, name-deduplicated +spec over `tile_h ∈ {8, 16}`, `tile_k ∈ {32, 64, 128, 256}`, +`ck ∈ {16, 32, 64}`, five wave layouts, and per MFMA tile size the widest-K +atom the target supports (`32x32x16` / `16x16x32` on gfx950, `32x32x8` / +`16x16x16` on gfx942) with the `tile_w_candidates` of each atom. +`benchmarks/common/benchmark_direct_conv.py` routes `groups == 1` shapes to +this sweep: + +```bash +python benchmarks/common/benchmark_direct_conv.py --N 4 --Hi 64 --Wi 64 --C 640 --K 640 --groups 1 --verify +``` + +Every candidate a sweep can pick must be correct, not just the hand-picked +geometries, so the correctness suite runs an evenly spaced sample of +`nongrouped_specs` on silicon. + +`nongrouped_specs` filters `nongrouped_knobs(arch, dtype, Wo=...)`; without +`Wo` the same generator yields every width up to six atoms, which is the grid +the AOT kernel cache builds (variant `direct_nongrouped` in +`benchmarks/common/direct_kernel_sweep.py`, 3x3 filters at stride 1 and 2, +`cpg = kpg = 0` in the identity, i.e. runtime channels). At run time the cache +offers a problem the binaries whose width `tile_w_candidates(Wo, atom_tile)` +would pick and whose spec validates on the real shape. + +## Tests and gates + +| Gate | What it proves | Command | +|------|----------------|---------| +| C++ ↔ Python byte-identity | the C++ port emits identical `.ll` for 15 configs covering every atom, dtype, stride, filter size, partial tile, uneven staging pass, DB/SB, swizzle on/off, iglp, waves_per_eu, gfx942/gfx950 | `python platform/tools/check_byte_identity.py --only conv_direct_grouped` (also with `ROCKE_LLVM_FLAVOR=llvm22`) | +| Emitters | the two sides of that gate; shared with the grouped direct-conv family, non-grouped configs are indices 33+ | `tests/parity/conv_direct_grouped_emit.{py,c}` | +| IR golden | Python lowering is byte-stable (5 cases, all llvm flavors) | `conv_direct_nongrouped/*` in `platform/tests/instances/rocke_ir_parity_harness.py` | +| On-silicon numerics | output vs `torch.nn.functional.conv2d` (fp32 reference), rel. tol 5e-2 fp16 / 1e-1 bf16; includes one binary launched on four unrelated shapes | `pytest tests/test_direct_conv_correctness.py -k Nongrouped` | +| AOT ABI / shape invariance | kernarg order matches `conv_direct_arg_names("fwd")`; emitted IR independent of `N`/`H`/`W`/`C`/`K` | `pytest tests/test_conv_abi.py -k nongrouped` | +| AOT cache | runtime channels, width filtering, `groups == 1` only, plans pack against the signature | `pytest tests/test_direct_kernel_cache.py` | + +`rocke.core.backend.lower_conv_direct_nongrouped(spec, backend="python"|"cpp"|"both")` +reaches the C++ builder directly through the `rocke_engine.conv_direct_nongrouped_*` +binding. + +When changing the builder, change `conv_direct_nongrouped.cpp` in the same change and +keep one IR op per C++ statement in Python order: C++ leaves argument evaluation +order unspecified, so folding two emitting calls into one argument list +reorders `arith.constant`s and breaks byte-identity. + +## Open items + +* No dispatch wiring: `library/dispatch/` does not select this family yet. +* The remaining gap on the shapes where it trails looks like barrier + + `waitcnt` exposure at two waves per SIMD. The untried lever is async + DRAM→LDS (`raw_ptr_buffer_load_lds`), which would drop the prefetch + registers; it needs lane-contiguous LDS writes, which conflicts with the + `lds_pad` that keeps the activation reads conflict-free. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py index 5f204e83c07d..a214c9528c89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py @@ -317,6 +317,11 @@ def conv_direct_arg_names(*, direction: str = "fwd") -> List[ArgSpec]: The weight tensor needs no stride args: its strides are products of ``KH``/``KW``/``cpg`` alone, all build-time. + The non-grouped forward kernel (``conv_direct_nongrouped``, ``groups == + 1``) takes the same block but bakes less: its channel counts only bound + its channel loop and masks, so it reads them from ``p_total_c`` / + ``p_total_k`` (and the weight row length from ``p_total_c``) instead. + ``direction="wgrad"`` binds dY (NHWK) to ``A`` and X (NHWC) to ``B``; ``D`` is the fp32 dW the kernel atomically accumulates into, which is filter-shaped and so needs no stride args either. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py new file mode 100644 index 000000000000..5655d52b253d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py @@ -0,0 +1,1015 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Direct non-grouped (``groups == 1``) convolution, NHWC layout. + +Why a second direct-conv family +------------------------------- +:mod:`kernels.common.conv_direct_grouped` draws all of its parallelism from the +``groups`` axis: one wave owns one convolution group, and the accumulator tile is +``ceil(kpg/16)`` M-tiles deep. + +This module is the ``groups == 1`` counterpart. It keeps the property that +makes a direct conv beat implicit GEMM on 3x3 stride<=2 shapes: the input tile +is staged in LDS **once per channel chunk, with its halo**, and all ``KH*KW`` +filter taps read shifted sub-tiles out of that one staged copy. An +implicit-GEMM formulation re-reads the activations once per tap, so its A-side +traffic is ~``KH*KW`` times larger, and the L2 has to absorb that redundancy. + +Shape of the computation +------------------------ +Per workgroup, with M = output channels and N = output pixels:: + + tile_k output channels (M, multiple of 32) + tile_h x tile_w output pixels (N, tile_w a multiple of 32) + ck input channels reduced per LDS stage + + for c0 in range(0, C, ck): # runtime scf.for, accumulators carried + stage X[(tile_h-1)*s+KH, (tile_w-1)*s+KW, ck] -> LDS (halo included) + stage W[tile_k, KH, KW, ck] -> LDS (fragment order) + for (r, s) in taps: # Python-unrolled + for m_tile, n_tile, k_atom: + acc = mfma_f32_32x32xK(W_lds, X_lds, acc) + +LDS layouts +----------- +``X`` is stored ``[pos][c]`` with ``pos = ih_lds * LDS_IN_W + iw_lds`` and a +channel stride of ``ck + lds_pad``. The pad is what keeps the 32 lanes of a +``ds_read_b128`` on distinct bank quads: with ``ck=16, pad=8`` the stride is +24 halves = 12 dwords, so lanes 0..7 land on bank quads 0,12,24,4,16,28,8,20 -- +eight distinct quads covering the full 128-byte LDS service width. + +``W`` is stored **pre-swizzled into MFMA fragment order**: fragment slot +``((tap * M_TILES + m_tile) * KATOMS + k_atom) * 64 + lane`` holds exactly the +``FRAG`` halves that ``lane`` feeds to the MFMA. Both the staging store and the +consuming read are then linear in ``lane``, so no padding is needed and neither +side can bank-conflict. The staging thread's global read is ``FRAG`` contiguous +channels of one ``k_out`` row, which is what the KRSC layout makes contiguous. + +AOT kernel arguments +-------------------- +The kernel takes the direct-conv kernarg block of +:func:`kernels.common.conv_abi.conv_direct_arg_names` (``direction="fwd"``): +batch, input and output extents, ``C`` (``p_total_c``), ``K`` (``p_total_k``) +and the NHWC / NHWK strides are runtime values, so one binary serves any image +size and any channel counts its tile divides. The filter (``KH``/``KW``), +``stride`` and ``PAD`` shape the staged halo and the unrolled tap loop, so they +stay build-time capabilities, like the tile geometry. Unlike the grouped +kernels this family does not bake ``C``/``K`` either: they only set the channel +loop's trip count and the store / staging masks. + +Software pipeline +----------------- +The global loads for chunk ``i+1`` are issued immediately after the barrier that +publishes chunk ``i``, and travel to the next iteration as loop-carried registers. +The MFMAs of chunk ``i`` therefore cover the DRAM latency of chunk ``i+1``. The +tail iteration loads one chunk past ``C``; the result is discarded, and buffer +loads clamp rather than fault, so no masking is needed on that path. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterator, List, Tuple + +from rocke.core.ir import ( + F32, + IRBuilder, + KernelDef, + Value, +) + +from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + _buf_load_vN, + _buf_store_vN, + _io_type, + _mfma, + _trunc_f32, + emit_direct_params, +) + +# Global load width for the activation staging path, in halves (dwordx4). +_X_LOAD_VEC = 8 + +# Byte offset every masked-out buffer access is redirected to (see +# build_direct_conv_nongrouped): the staging loads that fall outside the image +# or past K, and the epilogue stores that fall outside the output. A buffer +# access whose offset is at or past num_records is dropped (a load returns +# zero), so every operand tensor must end at or below this offset, and +# ``_OOB_BASE + chunk offset`` must still fit in i32 -- +# DirectNongroupedConvSpec.validate() enforces both. The tensor sizes are +# runtime values in the AOT kernel, so this is a launch-time contract: the +# spec has to be validated against the problem it is launched on. +_OOB_BASE = 0x7F000000 +_I32_MAX = (1 << 31) - 1 + +# Atom name -> (tile, K, frag). All supported atoms are square (M == N == tile), +# which is what lets one code path serve them: a lane's operand row/column is +# ``lane % tile`` and its K slice is ``(lane // tile) * frag``. ``frag`` is the +# per-lane operand width in halves, and ``tile * tile / 64`` accumulator floats +# per lane follow from the shape. +# +# The 16-wide atoms exist for output widths that are not a multiple of 32: with +# a 32-wide pixel tile the last tile in each row computes columns past ``Wo``, +# and that wasted MFMA work is spread over every block of the row. +_ATOMS = { + "32x32x16": (32, 16, 8), + "32x32x8": (32, 8, 4), + "16x16x32": (16, 32, 8), + "16x16x16": (16, 16, 4), +} + + +@dataclass(frozen=True) +class DirectNongroupedConvSpec: + """One concrete non-grouped direct-conv kernel configuration. + + ``problem.groups`` must be 1; the grouped families in + :mod:`kernels.common.conv_direct_grouped` cover ``groups > 1``. + """ + + problem: DirectConvProblem + name: str = "direct_conv_nongrouped" + + tile_h: int = 16 # output rows per workgroup + tile_w: int = 32 # output cols per workgroup (multiple of 32) + tile_k: int = 128 # output channels per workgroup (multiple of 32) + ck: int = 16 # input channels staged per LDS chunk + + waves_m: int = 2 # waves splitting tile_k + waves_n: int = 4 # waves splitting tile_h (each keeps all tile_w columns) + + atom: str = "32x32x16" + wave_size: int = 64 + lds_pad: int = 8 # extra halves on the LDS activation channel stride + + # Grid swizzle. The launch grid is flattened to 1-D and remapped so that the + # workgroups a single XCD receives share operands. ``swizzle_wgm`` is the + # number of *channel* tiles grouped per spatial cell: wgm = n_k_tiles walks + # every output-channel tile of one spatial cell back to back (activation tile + # loaded into that XCD's L2 once), wgm = 1 walks every spatial cell of one + # channel tile (weight slab loaded once). Intermediate values trade between. + chiplet_swizzle: bool = True + swizzle_wgm: int = 8 + chiplet_chunk: int = 64 + num_xcds: int = 8 + + # Ping-pong the staged tiles across two LDS buffers. Single-buffered, the + # channel loop needs two barriers per chunk with nothing but ds_writes + # between them; ping-ponged, chunk i+1 is written while chunk i is still + # feeding MFMAs and one barrier per chunk suffices. Costs 2x LDS. + double_buffer: bool = False + # ``iglp_opt`` level for the channel loop, or None to leave the backend + # scheduler alone. 0 is the canned GEMM MFMA/memory interleave. + iglp: "int | None" = None + # ``amdgpu-waves-per-eu`` occupancy hint. The prefetch registers carried + # across the channel loop can push register use high enough to leave one + # wave per SIMD, which leaves nothing to cover the barriers per chunk. + waves_per_eu: "int | None" = None + + # ---- derived geometry ------------------------------------------------- + + @property + def atom_tile(self) -> int: + """MFMA M == N (channels per M-tile, pixels per N-tile).""" + return _ATOMS[self.atom][0] + + @property + def atom_k(self) -> int: + return _ATOMS[self.atom][1] + + @property + def frag(self) -> int: + """Halves per lane in an A/B operand fragment.""" + return _ATOMS[self.atom][2] + + @property + def acc_per_lane(self) -> int: + t = self.atom_tile + return t * t // self.wave_size + + @property + def n_col_blocks(self) -> int: + return self.tile_w // self.atom_tile + + @property + def m_tiles_total(self) -> int: + return self.tile_k // self.atom_tile + + @property + def k_atoms(self) -> int: + return self.ck // self.atom_k + + @property + def rows_per_wave(self) -> int: + return self.tile_h // self.waves_n + + @property + def m_tiles_per_wave(self) -> int: + return self.m_tiles_total // self.waves_m + + @property + def n_tiles_per_wave(self) -> int: + return self.rows_per_wave * self.n_col_blocks + + @property + def threads_per_block(self) -> int: + return self.waves_m * self.waves_n * self.wave_size + + @property + def lds_in_h(self) -> int: + p = self.problem + return (self.tile_h - 1) * p.stride + p.KH + + @property + def lds_in_w(self) -> int: + p = self.problem + return (self.tile_w - 1) * p.stride + p.KW + + @property + def c_stride(self) -> int: + """LDS channel stride of the staged activation tile, in halves.""" + return self.ck + self.lds_pad + + @property + def lds_x_halves(self) -> int: + return self.lds_in_h * self.lds_in_w * self.c_stride + + @property + def lds_w_halves(self) -> int: + p = self.problem + return self.tile_k * p.KH * p.KW * self.ck + + @property + def lds_bytes(self) -> int: + # + one scratch fragment per array (see build_direct_conv_nongrouped). + one = self.lds_x_halves + _X_LOAD_VEC + self.lds_w_halves + self.frag + return 2 * one * (2 if self.double_buffer else 1) + + @property + def acc_vgprs(self) -> int: + """f32 accumulator registers per lane.""" + return self.m_tiles_per_wave * self.n_tiles_per_wave * self.acc_per_lane + + @property + def tile_counts(self) -> Tuple[int, int, int]: + """``(n_w_tiles, n_h_tiles, n_k_tiles)``.""" + p = self.problem + return ( + (p.Wo + self.tile_w - 1) // self.tile_w, + (p.Ho + self.tile_h - 1) // self.tile_h, + (p.kpg + self.tile_k - 1) // self.tile_k, + ) + + def grid(self) -> Tuple[int, int, int]: + n_wt, n_ht, n_kt = self.tile_counts + return (n_wt * n_ht * self.problem.N * n_kt, 1, 1) + + def kernel_name(self) -> str: + from rocke.helpers.spec import kernel_name_join + + p = self.problem + return kernel_name_join( + self.name, + p.short(), + f"t{self.tile_h}x{self.tile_w}x{self.tile_k}", + f"ck{self.ck}", + f"w{self.waves_m}x{self.waves_n}", + f"a{self.atom}", + f"g{self.swizzle_wgm}" if self.chiplet_swizzle else "gnone", + "db" if self.double_buffer else "", + f"iglp{self.iglp}" if self.iglp is not None else "", + f"we{self.waves_per_eu}" if self.waves_per_eu is not None else "", + "bf16" if p.dtype == "bf16" else "", + ) + + def validate(self) -> None: + p = self.problem + if p.dtype not in ("fp16", "bf16"): + raise ValueError(f"DirectNongroupedConvSpec: unsupported dtype {p.dtype!r}") + if p.groups != 1: + raise ValueError( + f"DirectNongroupedConvSpec is the groups==1 family (got groups={p.groups}); " + f"use conv_direct_grouped for grouped shapes" + ) + if self.atom not in _ATOMS: + raise ValueError( + f"unknown atom {self.atom!r}; expected one of {list(_ATOMS)}" + ) + if ( + self.tile_h <= 0 + or self.tile_w <= 0 + or self.tile_k <= 0 + or self.ck <= 0 + or self.waves_m <= 0 + or self.waves_n <= 0 + or self.wave_size <= 0 + or p.stride <= 0 + or p.cpg <= 0 + or p.kpg <= 0 + ): + # Checked before any derived geometry: those divide by these. A + # positive C is also what lets the kernel assume its channel loop + # runs at least once. + raise ValueError( + "DirectNongroupedConvSpec: tile, wave, stride and channel " + "parameters must be positive" + ) + t = self.atom_tile + if self.tile_w % t != 0: + raise ValueError(f"tile_w must be a multiple of {t} (got {self.tile_w})") + if self.tile_k % t != 0: + raise ValueError(f"tile_k must be a multiple of {t} (got {self.tile_k})") + if self.tile_h % self.waves_n != 0: + raise ValueError( + f"tile_h {self.tile_h} not divisible by waves_n {self.waves_n}" + ) + if self.m_tiles_total % self.waves_m != 0: + raise ValueError( + f"tile_k/{self.atom_tile} = {self.m_tiles_total} not divisible by " + f"waves_m {self.waves_m}" + ) + if self.ck % self.atom_k != 0: + raise ValueError(f"ck {self.ck} not divisible by atom K {self.atom_k}") + if self.ck % _X_LOAD_VEC != 0: + raise ValueError(f"ck {self.ck} not divisible by {_X_LOAD_VEC}") + if p.cpg % self.ck != 0: + raise ValueError(f"C {p.cpg} not divisible by ck {self.ck}") + if p.kpg % self.atom_tile != 0: + raise ValueError( + f"K {p.kpg} must be a multiple of {self.atom_tile} (got {p.kpg})" + ) + if p.cpg % _X_LOAD_VEC != 0: + raise ValueError(f"C {p.cpg} must be a multiple of {_X_LOAD_VEC}") + if self.threads_per_block > 1024: + raise ValueError(f"threads_per_block {self.threads_per_block} > 1024") + if self.lds_pad < 0 or self.lds_pad % 8 != 0: + # A negative pad shrinks the pixel stride below the ck halves each + # pixel stages, so neighbouring pixels overlap in LDS. The vec8 LDS + # store/load is emitted with align 16, so the pixel stride + # (ck + lds_pad halves) must be a multiple of 8 halves. + raise ValueError( + f"lds_pad must be a non-negative multiple of 8 to keep " + f"ds_read_b128 aligned (got {self.lds_pad})" + ) + if self.swizzle_wgm < 1: + raise ValueError(f"swizzle_wgm must be >= 1 (got {self.swizzle_wgm})") + if self.iglp is not None and self.iglp < 0: + raise ValueError(f"iglp must be None or >= 0 (got {self.iglp})") + if self.waves_per_eu is not None and self.waves_per_eu < 1: + raise ValueError( + f"waves_per_eu must be None or >= 1 (got {self.waves_per_eu})" + ) + if self.grid()[0] > _I32_MAX: + # The kernel decodes its tile from a flat 1-D block id and computes + # the workgroup count in i32; past gridDim.x's 2**31 - 1 limit both + # the launch and that product would overflow. (The D-size bound + # below implies this one; it is checked on its own so the reason + # names the real limit.) + raise ValueError( + f"flattened grid of {self.grid()[0]} workgroups exceeds " + f"{_I32_MAX}: grow the tile or split the batch" + ) + a_bytes = 2 * p.N * p.H * p.W * p.cpg + b_bytes = 2 * p.kpg * p.KH * p.KW * p.cpg + d_bytes = 2 * p.N * p.Ho * p.Wo * p.kpg + if max(a_bytes, b_bytes, d_bytes) > _OOB_BASE: + # Masked staging loads and epilogue stores go to _OOB_BASE and rely + # on it being past num_records; a larger tensor would hand the + # loads real data and take the stores. + raise ValueError( + f"tensors must fit below the {_OOB_BASE:#x}-byte masked-access " + f"offset (A {a_bytes} B, B {b_bytes} B, D {d_bytes} B)" + ) + if _OOB_BASE + 2 * (p.cpg + self.ck) > _I32_MAX: + # The last prefetch runs one chunk (two when double-buffered) past C. + raise ValueError( + f"C {p.cpg} too large: the masked-load offset plus the channel " + f"offset would overflow i32" + ) + if self.acc_vgprs > 256: + # A lane cannot hold more than the 256-entry accumulator file; past + # that the config is not merely slow, it makes the backend + # scheduler blow up (a 2048-accumulator tile hung the compiler). + raise ValueError( + f"accumulator tile needs {self.acc_vgprs} registers per lane " + f"(max 256): shrink tile_k/tile_h/tile_w or add waves" + ) + + +def is_valid_nongrouped_spec( + spec: DirectNongroupedConvSpec, arch: str = "gfx950" +) -> Tuple[bool, str]: + """Return ``(ok, reason)`` for ``spec`` on ``arch`` without raising.""" + from rocke.core.arch import ArchTarget + + try: + target = ArchTarget.from_gfx(arch) + except KeyError as e: # pragma: no cover - arch table lookup + return False, str(e) + + # Everything past the arch lookup stays inside the try: the arch-table + # queries are part of the no-raise contract too, not just validate(). + try: + spec.validate() + p = spec.problem + if p.stride not in (1, 2): + return False, f"stride {p.stride} is not supported (expected 1 or 2)" + ab = "bf16" if p.dtype == "bf16" else "f16" + t, k = spec.atom_tile, spec.atom_k + if not target.mma.has_shape( + a_dtype=ab, b_dtype=ab, c_dtype="fp32", m=t, n=t, k=k + ): + return False, f"missing mfma_f32_{spec.atom}_{ab} on {arch}" + if spec.wave_size != target.wave_size: + return ( + False, + f"wave_size {spec.wave_size} != {arch} wave {target.wave_size}", + ) + if spec.threads_per_block > target.max_threads_per_block: + return ( + False, + f"threads_per_block {spec.threads_per_block} exceeds arch limit", + ) + if not target.fits_lds(spec.lds_bytes): + return ( + False, + f"LDS {spec.lds_bytes} B exceeds {target.lds_capacity_bytes} B", + ) + except (ValueError, KeyError) as e: + return False, str(e) + return True, "ok" + + +# --------------------------------------------------------------------------- +# Candidate generation +# --------------------------------------------------------------------------- + +# Swept geometry. ``tile_w`` is not in this table because the useful values +# depend on ``Wo``: when ``Wo`` is not a multiple of ``tile_w``, the last tile in +# each row spends MFMA work on pixels past the image edge. +# ``tile_h`` stops at 8 because a shorter tile re-reads a larger share of halo +# rows per output row (``KH - 1`` extra rows amortised over ``tile_h``). +_SWEEP_TILE_H = (8, 16) +_SWEEP_TILE_K = (32, 64, 128, 256) +_SWEEP_CK = (16, 32, 64) +_SWEEP_WAVES = ((2, 2), (1, 4), (2, 4), (4, 2), (1, 8)) +# Per MFMA tile size, the atoms to sweep in order of preference: the widest-K +# atom the target has wins (fewer MFMAs per chunk), so gfx950 sweeps +# 32x32x16 / 16x16x32 and gfx942 -- which has neither -- 32x32x8 / 16x16x16. +_SWEEP_ATOMS = (("32x32x16", "32x32x8"), ("16x16x32", "16x16x16")) + + +def _sweep_atoms(arch: str, dtype: str) -> "list[str]": + """The preferred supported atom of each tile size on ``arch``.""" + from rocke.core.arch import ArchTarget + + target = ArchTarget.from_gfx(arch) + ab = "bf16" if dtype == "bf16" else "f16" + out = [] + for prefs in _SWEEP_ATOMS: + for atom in prefs: + t, k, _ = _ATOMS[atom] + if target.mma.has_shape( + a_dtype=ab, b_dtype=ab, c_dtype="fp32", m=t, n=t, k=k + ): + out.append(atom) + break + return out + + +def tile_w_candidates(Wo: int, atom_tile: int, max_mult: int = 6) -> "list[int]": + """Widths worth trying for ``Wo`` with an ``atom_tile``-wide MFMA N dimension. + + Prefers widths that tile ``Wo`` with no wasted columns; if none do, falls + back to the single width with the least waste so the shape is still covered. + """ + cands = [atom_tile * m for m in range(1, max_mult + 1)] + exact = [w for w in cands if Wo % w == 0] + if exact: + return exact + return [min(cands, key=lambda w: -(-Wo // w) * w)] + + +def nongrouped_knobs( + arch: str, + dtype: str, + *, + Wo: "int | None" = None, + iglp: "tuple[int | None, ...]" = (0,), + waves_per_eu: "tuple[int | None, ...]" = (None, 3), + swizzle_wgm: "tuple[int, ...]" = (8,), +) -> "Iterator[dict]": + """The swept :class:`DirectNongroupedConvSpec` geometry, as spec kwargs. + + With ``Wo`` the widths are the :func:`tile_w_candidates` of that output + width -- what a sweep for one shape needs. Without it every width up to + six atoms is produced: the shape-independent grid an AOT cache builds, since + a cached binary serves any ``Wo`` and the run side picks the widths that fit + the shape. Unvalidated: :func:`is_valid_nongrouped_spec` drops what a given + problem or arch cannot run. + """ + import itertools + + for atom in _sweep_atoms(arch, dtype): + at = _ATOMS[atom][0] + widths = ( + tile_w_candidates(Wo, at) + if Wo is not None + else [at * m for m in range(1, 7)] + ) + for tw in widths: + for th, tk, ck, (wm, wn), wgm, ig, we in itertools.product( + _SWEEP_TILE_H, + _SWEEP_TILE_K, + _SWEEP_CK, + _SWEEP_WAVES, + swizzle_wgm, + iglp, + waves_per_eu, + ): + yield dict( + tile_h=th, + tile_w=tw, + tile_k=tk, + ck=ck, + waves_m=wm, + waves_n=wn, + atom=atom, + swizzle_wgm=wgm, + iglp=ig, + waves_per_eu=we, + ) + + +def nongrouped_specs( + problem: DirectConvProblem, + *, + arch: str = "gfx950", + name: str = "direct_conv_nongrouped", + iglp: "tuple[int | None, ...]" = (0,), + waves_per_eu: "tuple[int | None, ...]" = (None, 3), + swizzle_wgm: "tuple[int, ...]" = (8,), +) -> "list[DirectNongroupedConvSpec]": + """Every valid :class:`DirectNongroupedConvSpec` worth benchmarking for ``problem``. + + Deduplicated by kernel name, so callers can compile the list directly. + """ + out: "list[DirectNongroupedConvSpec]" = [] + seen = set() + for knobs in nongrouped_knobs( + arch, + problem.dtype, + Wo=problem.Wo, + iglp=iglp, + waves_per_eu=waves_per_eu, + swizzle_wgm=swizzle_wgm, + ): + spec = DirectNongroupedConvSpec(problem=problem, name=name, **knobs) + ok, _ = is_valid_nongrouped_spec(spec, arch=arch) + if not ok: + continue + key = spec.kernel_name() + if key in seen: + continue + seen.add(key) + out.append(spec) + return out + + +def build_direct_conv_nongrouped( + spec: DirectNongroupedConvSpec, arch: str = "gfx950" +) -> KernelDef: + """Build the IR for one non-grouped NHWC direct convolution kernel. + + AOT: the binary bakes the filter, ``stride``, ``PAD``, dtype and the tile + geometry of ``spec``; batch, extents and channel counts are kernargs (see + the module docstring). ``spec.problem`` only has to pass the validator -- + the emitted IR does not depend on its ``N``/``H``/``W``/``C``/``K``. + """ + ok, why = is_valid_nongrouped_spec(spec, arch=arch) + if not ok: + raise ValueError(f"invalid DirectNongroupedConvSpec for {arch}: {why}") + + p = spec.problem + io_type = _io_type(p.dtype) + dtype = p.dtype + + KH, KW, S, PAD = p.KH, p.KW, p.stride, p.PAD + N_TAPS = KH * KW + + TH, TW, TK, CK = spec.tile_h, spec.tile_w, spec.tile_k, spec.ck + THREADS = spec.threads_per_block + WAVE = spec.wave_size + FRAG = spec.frag + AK = spec.atom_k + AT = spec.atom_tile # MFMA M == N + ACC = spec.acc_per_lane + QUADS = ACC // 4 # accumulator slots come in quads of consecutive channels + NCB = spec.n_col_blocks + M_TILES = spec.m_tiles_total + KATOMS = spec.k_atoms + ROWS_W = spec.rows_per_wave + MT_W = spec.m_tiles_per_wave + NT_W = spec.n_tiles_per_wave + CSTRIDE = spec.c_stride + LDS_IN_H, LDS_IN_W = spec.lds_in_h, spec.lds_in_w + + # Staging pass counts. + X_CV = CK // _X_LOAD_VEC # channel vectors per staged pixel + X_VECS = LDS_IN_H * LDS_IN_W * X_CV + X_PASSES = (X_VECS + THREADS - 1) // THREADS + W_SLOTS = N_TAPS * M_TILES * KATOMS * WAVE + W_PASSES = (W_SLOTS + THREADS - 1) // THREADS + + b = IRBuilder(spec.kernel_name()) + b.kernel.attrs["max_workgroup_size"] = THREADS + if spec.waves_per_eu is not None: + b.kernel.attrs["waves_per_eu"] = spec.waves_per_eu + + params = emit_direct_params(b, io_type=io_type) + A = params["A"] + Bp = params["B"] + D = params["D"] + A_bytes = params["A_bytes"] + B_bytes = params["B_bytes"] + D_bytes = params["D_bytes"] + p_N = params["p_N"] + p_Hi = params["p_Hi"] + p_Wi = params["p_Wi"] + p_Ho = params["p_Ho"] + p_Wo = params["p_Wo"] + # groups == 1, so the totals are the full channel counts. + C = params["p_total_c"] + K = params["p_total_k"] + + c0 = b.const_i32(0) + c1 = b.const_i32(1) + c_half = b.const_i32(2) + # Predication for the staging loads is loop-invariant, so it is folded into + # the *base* offset once, before the channel loop, rather than re-selected + # every iteration. validate() guarantees ``_OOB_BASE`` is at or past the + # end of every operand tensor and leaves headroom for ``+ c_off`` without + # wrapping i32, so ``base + c_off`` stays out of range for the whole loop. + # A buffer load whose voffset exceeds num_records returns zero, which is + # exactly the value a padded pixel or an out-of-range filter row needs -- + # so the masked *value* select disappears from the loop as well. The + # epilogue redirects its masked stores to the same offset. + oob_base = b.const_i32(_OOB_BASE) + + a_rsrc = b.buffer_rsrc(A, A_bytes) + b_rsrc = b.buffer_rsrc(Bp, B_bytes) + d_rsrc = b.buffer_rsrc(D, D_bytes) + + # The staging loops are sized in whole thread-passes, so the last pass can + # own slots past the end of the tile. Those slots write into a scratch tail + # instead of running off the array (and over the neighbouring allocation). + DB = spec.double_buffer + x_stage = spec.lds_x_halves + _X_LOAD_VEC + w_stage = spec.lds_w_halves + FRAG + x_dump = spec.lds_x_halves + w_dump = spec.lds_w_halves + nbuf = 2 if DB else 1 + X_smem = b.smem_alloc(io_type, [1, x_stage * nbuf], name_hint="lds_x") + W_smem = b.smem_alloc(io_type, [1, w_stage * nbuf], name_hint="lds_w") + + # ---- thread / wave decomposition ------------------------------------- + tid = b.thread_id_x() + lane = b.mod(tid, b.const_i32(WAVE)) + wave_id = b.div(tid, b.const_i32(WAVE)) + wave_m = b.div(wave_id, b.const_i32(spec.waves_n)) + wave_n = b.mod(wave_id, b.const_i32(spec.waves_n)) + lane_lo = b.mod(lane, b.const_i32(AT)) # M row / N column inside the atom + lane_hi = b.div(lane, b.const_i32(AT)) # K slice inside the atom + + # ---- grid decode ------------------------------------------------------ + # Flat 1-D grid of (spatial cell x channel tile). The tile *sizes* are + # build-time, but how many tiles the problem needs follows the runtime + # extents, so the counts are a handful of scalar divides once per + # workgroup. The swizzle decides which of the two axes a single XCD walks + # contiguously; see DirectNongroupedConvSpec. + n_wt = b.div(b.add(p_Wo, b.const_i32(TW - 1)), b.const_i32(TW)) + n_ht = b.div(b.add(p_Ho, b.const_i32(TH - 1)), b.const_i32(TH)) + n_kt = b.div(b.add(K, b.const_i32(TK - 1)), b.const_i32(TK)) + n_hw = b.mul(n_wt, n_ht) + n_cells = b.mul(n_hw, p_N) + + wgid = b.block_id_x() + if spec.chiplet_swizzle: + from rocke.helpers.grid import chiplet_aware_super_tile_dynamic + + # A swizzle_wgm above the runtime channel-tile count needs no clamp: + # the swizzle sizes its last group as min(wgm, tiles left). + sw = chiplet_aware_super_tile_dynamic( + b, + wgid, + num_pid_m=n_kt, + num_pid_n=n_cells, + wgm=spec.swizzle_wgm, + num_xcds=spec.num_xcds, + chunk_size=spec.chiplet_chunk, + ) + k_tile, cell = sw.row, sw.col + else: + k_tile = b.mod(wgid, n_kt) + cell = b.div(wgid, n_kt) + + n_img = b.div(cell, n_hw) + st = b.mod(cell, n_hw) + w_tile = b.mod(st, n_wt) + h_tile = b.div(st, n_wt) + + # Output-tile origins. + out_h0 = b.mul(h_tile, b.const_i32(TH)) + out_w0 = b.mul(w_tile, b.const_i32(TW)) + k_base = b.mul(k_tile, b.const_i32(TK)) + + # Input origin of the staged (halo-inclusive) activation tile. + in_h0 = b.sub(b.mul(out_h0, b.const_i32(S)), b.const_i32(PAD)) + in_w0 = b.sub(b.mul(out_w0, b.const_i32(S)), b.const_i32(PAD)) + + # ---- staging metadata (channel-independent; hoisted out of the C loop) -- + + def _x_pass_meta(): + """Per-pass (global byte offset base, LDS half index) for X.""" + meta = [] + img_base = b.mul(n_img, params["p_A_stride_n"]) + for j in range(X_PASSES): + v = b.add(tid, b.const_i32(j * THREADS)) + cv = b.mod(v, b.const_i32(X_CV)) + pos = b.div(v, b.const_i32(X_CV)) + ih_l = b.div(pos, b.const_i32(LDS_IN_W)) + iw_l = b.mod(pos, b.const_i32(LDS_IN_W)) + ih = b.add(in_h0, ih_l) + iw = b.add(in_w0, iw_l) + ok_h = b.land(b.cmp_ge(ih, c0), b.cmp_lt(ih, p_Hi)) + ok_w = b.land(b.cmp_ge(iw, c0), b.cmp_lt(iw, p_Wi)) + valid = b.land(ok_h, ok_w) + if X_VECS % THREADS != 0: + in_tile = b.cmp_lt(v, b.const_i32(X_VECS)) + valid = b.land(valid, in_tile) + else: + in_tile = None + # base = n*stride_n + ih*stride_hi + iw*stride_wi + cv*VEC; the + # chunk offset is added per iteration. + pix_off = b.add( + b.mul(ih, params["p_A_stride_hi"]), b.mul(iw, params["p_A_stride_wi"]) + ) + elems = b.add(b.add(img_base, pix_off), b.mul(cv, b.const_i32(_X_LOAD_VEC))) + lds_idx = b.add( + b.mul(pos, b.const_i32(CSTRIDE)), b.mul(cv, b.const_i32(_X_LOAD_VEC)) + ) + if in_tile is not None: + lds_idx = b.select(in_tile, lds_idx, b.const_i32(x_dump)) + base = b.select(valid, b.mul(elems, c_half), oob_base) + meta.append((base, lds_idx)) + return meta + + def _w_pass_meta(): + """Per-pass (global byte offset base, LDS half index) for W. + + Slot ``g`` decodes to the MFMA fragment that lane ``g % 64`` consumes for + ``(tap, m_tile, k_atom)``, so the LDS index is simply ``g * FRAG``. + """ + meta = [] + for j in range(W_PASSES): + g = b.add(tid, b.const_i32(j * THREADS)) + g_lane = b.mod(g, b.const_i32(WAVE)) + rest = b.div(g, b.const_i32(WAVE)) + katom = b.mod(rest, b.const_i32(KATOMS)) + rest2 = b.div(rest, b.const_i32(KATOMS)) + m_tile = b.mod(rest2, b.const_i32(M_TILES)) + tap = b.div(rest2, b.const_i32(M_TILES)) + r = b.div(tap, b.const_i32(KW)) + s = b.mod(tap, b.const_i32(KW)) + + k_out = b.add( + k_base, + b.add( + b.mul(m_tile, b.const_i32(AT)), + b.mod(g_lane, b.const_i32(AT)), + ), + ) + c_in_chunk = b.add( + b.mul(katom, b.const_i32(AK)), + b.mul(b.div(g_lane, b.const_i32(AT)), b.const_i32(FRAG)), + ) + valid = b.cmp_lt(k_out, K) + if W_SLOTS % THREADS != 0: + in_tile = b.cmp_lt(g, b.const_i32(W_SLOTS)) + valid = b.land(valid, in_tile) + else: + in_tile = None + # base = ((k_out*KH + r)*KW + s)*C + c_in_chunk (KRSC) + krs_h = b.mul(k_out, b.const_i32(KH)) + krs_hr = b.add(krs_h, r) + krs_w = b.mul(krs_hr, b.const_i32(KW)) + krs = b.add(krs_w, s) + elems = b.add(b.mul(krs, C), c_in_chunk) + lds_idx = b.mul(g, b.const_i32(FRAG)) + if in_tile is not None: + lds_idx = b.select(in_tile, lds_idx, b.const_i32(w_dump)) + base = b.select(valid, b.mul(elems, c_half), oob_base) + meta.append((base, lds_idx)) + return meta + + x_meta = _x_pass_meta() + w_meta = _w_pass_meta() + + x_dwords = _X_LOAD_VEC // 2 + w_dwords = FRAG // 2 + + def issue_stage_loads(c_off: Value): + """Issue the global loads for one channel chunk; returns raw registers.""" + xs = [ + _buf_load_vN(b, dtype, a_rsrc, b.add(base, c_off), c0, x_dwords) + for base, _ in x_meta + ] + ws = [ + _buf_load_vN(b, dtype, b_rsrc, b.add(base, c_off), c0, w_dwords) + for base, _ in w_meta + ] + return xs, ws + + def commit_stage(xs, ws, x_buf=None, w_buf=None): + for (_, lds_idx), val in zip(x_meta, xs): + idx = lds_idx if x_buf is None else b.add(lds_idx, x_buf) + b.smem_store_vN(X_smem, [c0, idx], val, _X_LOAD_VEC) + for (_, lds_idx), val in zip(w_meta, ws): + idx = lds_idx if w_buf is None else b.add(lds_idx, w_buf) + b.smem_store_vN(W_smem, [c0, idx], val, FRAG) + + # ---- per-lane LDS read bases ----------------------------------------- + # X: idx = wave row term + lane term + compile-time term. + x_wave_term = b.mul( + b.mul(wave_n, b.const_i32(ROWS_W * S)), b.const_i32(LDS_IN_W * CSTRIDE) + ) + x_lane_term = b.add( + b.mul(lane_lo, b.const_i32(S * CSTRIDE)), b.mul(lane_hi, b.const_i32(FRAG)) + ) + x_read_base = b.add(x_wave_term, x_lane_term) + # W: idx = wave_m block + lane*FRAG + compile-time term. + w_read_base = b.add( + b.mul(wave_m, b.const_i32(MT_W * KATOMS * WAVE * FRAG)), + b.mul(lane, b.const_i32(FRAG)), + ) + + def read_a_frag(base: Value, tap: int, mt: int, katom: int) -> Value: + off = ((tap * M_TILES + mt) * KATOMS + katom) * WAVE * FRAG + idx = b.add(base, b.const_i32(off)) + return b.smem_load_vN(W_smem, c0, idx, dtype=io_type, n=FRAG) + + def read_b_frag(base: Value, in_row: int, cb: int, s: int, katom: int) -> Value: + """One activation fragment, addressed by *input* row inside the wave window. + + Taps share activations: output row ``row`` at tap row ``r`` reads input + row ``row * stride + r``, so the ``KH`` taps of a column only need + ``(ROWS_W - 1) * stride + KH`` distinct fragments instead of + ``ROWS_W * KH``. Indexing by input row is what makes that reuse + expressible. + """ + pos = in_row * LDS_IN_W + cb * AT * S + s + off = pos * CSTRIDE + katom * AK + idx = b.add(base, b.const_i32(off)) + return b.smem_load_vN(X_smem, c0, idx, dtype=io_type, n=FRAG) + + # ---- main channel loop ------------------------------------------------ + zero_acc = b.zero_vec_f32(ACC) + n_acc = MT_W * NT_W + n_in_rows = (ROWS_W - 1) * S + KH + + def emit_mfmas(accs, x_base, w_base): + """One channel chunk of MFMAs against the staged tiles. + + Filter-column (``s``) outermost: every activation fragment loaded for + one ``(s, k_atom)`` is consumed by all ``KH`` tap rows and all ``MT_W`` + channel tiles. + """ + for katom in range(KATOMS): + for s_c in range(KW): + b_frags = [ + [read_b_frag(x_base, ir, cb, s_c, katom) for cb in range(NCB)] + for ir in range(n_in_rows) + ] + for r_c in range(KH): + tap = r_c * KW + s_c + a_frags = [ + read_a_frag(w_base, tap, mt, katom) for mt in range(MT_W) + ] + for row in range(ROWS_W): + bf_row = b_frags[row * S + r_c] + for cb in range(NCB): + nt = row * NCB + cb + for mt in range(MT_W): + idx = mt * NT_W + nt + accs[idx] = _mfma( + b, + dtype, + spec.atom, + a_frags[mt], + bf_row[cb], + accs[idx], + ) + + c_ck_bytes = b.const_i32(CK * 2) + xs0, ws0 = issue_stage_loads(c0) + + if DB: + # Prologue: publish chunk 0 into buffer 0 and have chunk 1 in flight. + commit_stage(xs0, ws0) + xs0, ws0 = issue_stage_loads(c_ck_bytes) + + iter_args: List[Tuple[str, Value]] = [(f"acc{i}", zero_acc) for i in range(n_acc)] + iter_args += [(f"xs{i}", v) for i, v in enumerate(xs0)] + iter_args += [(f"ws{i}", v) for i, v in enumerate(ws0)] + + # Runtime trip count. validate() guarantees C is a positive whole number of + # chunks, so the clamp to one chunk never changes it. It is there for LLVM: + # unable to prove the loop runs at all, it keeps a zero-trip guard and + # schedules the prefetch loads late in the body. + n_chunks = b.smax(b.div(C, b.const_i32(CK)), c1) + loop = b.scf_for_iter( + c0, + n_chunks, + c1, + iter_args, + iv_name="c_iter", + elide_trailing_barrier=False, + ) + with loop as (c_iv, carried): + accs = list(carried[:n_acc]) + xs = list(carried[n_acc : n_acc + X_PASSES]) + ws = list(carried[n_acc + X_PASSES :]) + + if spec.iglp is not None: + b.iglp_opt(spec.iglp) + + if DB: + # Parity of the induction variable picks the live buffer; every + # fragment offset stays a compile-time constant off these bases. + par = b.mod(c_iv, b.const_i32(2)) + cur_x = b.mul(par, b.const_i32(x_stage)) + cur_w = b.mul(par, b.const_i32(w_stage)) + nxt_x = b.sub(b.const_i32(x_stage), cur_x) + nxt_w = b.sub(b.const_i32(w_stage), cur_w) + + b.sync() + # Stage chunk i+1 into the idle buffer while chunk i still feeds MFMAs. + commit_stage(xs, ws, nxt_x, nxt_w) + c_off = b.mul(b.add(c_iv, b.const_i32(2)), c_ck_bytes) + xs_n, ws_n = issue_stage_loads(c_off) + emit_mfmas(accs, b.add(x_read_base, cur_x), b.add(w_read_base, cur_w)) + else: + # Publish the chunk prefetched during the previous iteration. + b.sync() + commit_stage(xs, ws) + b.sync() + c_off = b.mul(b.add(c_iv, c1), c_ck_bytes) + xs_n, ws_n = issue_stage_loads(c_off) + emit_mfmas(accs, x_read_base, w_read_base) + + b.scf_yield(*accs, *xs_n, *ws_n) + + accs_out = list(loop.results[:n_acc]) + + # ---- epilogue --------------------------------------------------------- + # Square-atom accumulator: slot i -> row = (i//4)*(AT//4) + lane_hi*4 + (i%4), + # col = lane_lo. Slots 4q..4q+3 are four consecutive output channels, so a + # quad packs into one dwordx2 store. (AT=32 -> 4 quads spaced 8 channels + # apart; AT=16 -> a single quad.) + k_lane_base = b.add( + b.add(k_base, b.mul(wave_m, b.const_i32(MT_W * AT))), + b.mul(lane_hi, b.const_i32(4)), + ) + row_base = b.add(out_h0, b.mul(wave_n, b.const_i32(ROWS_W))) + img_out_base = b.mul(n_img, params["p_D_stride_n"]) + + for mt in range(MT_W): + k_mt = b.add(k_lane_base, b.const_i32(mt * AT)) + for row in range(ROWS_W): + out_h = b.add(row_base, b.const_i32(row)) + h_ok = b.cmp_lt(out_h, p_Ho) + row_off = b.add(img_out_base, b.mul(out_h, params["p_D_stride_ho"])) + for cb in range(NCB): + nt = row * NCB + cb + acc = accs_out[mt * NT_W + nt] + out_w = b.add(b.add(out_w0, b.const_i32(cb * AT)), lane_lo) + hw_ok = b.land(h_ok, b.cmp_lt(out_w, p_Wo)) + base_off = b.add(row_off, b.mul(out_w, params["p_D_stride_wo"])) + for q in range(QUADS): + k_out = b.add(k_mt, b.const_i32(q * (AT // 4))) + valid = b.land(hw_ok, b.cmp_lt(k_out, K)) + d_off = b.mul(b.add(base_off, k_out), c_half) + # Same sentinel as the staging loads: validate() keeps D + # at or below _OOB_BASE, so the redirected store lies + # past num_records and the hardware drops it. + safe = b.select(valid, d_off, oob_base) + quad = b.vec_pack( + [b.vec_extract(acc, 4 * q + j) for j in range(4)], F32 + ) + _buf_store_vN( + b, dtype, d_rsrc, safe, c0, _trunc_f32(b, dtype, quad), 2 + ) + + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c index adb748d6dd23..1cd46c6459ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c @@ -6,7 +6,8 @@ * argv[1] (the config index), builds the rocke_direct_conv_16c_spec_t / * rocke_direct_conv_4c_spec_t / rocke_direct_conv_8c_spec_t / * rocke_direct_conv_32c_spec_t / rocke_direct_depthwise_spec_t / - * rocke_direct_conv_wgrad_spec_t identically to + * rocke_direct_conv_wgrad_spec_t (and, from index 32, the non-grouped + * rocke_direct_conv_nongrouped_spec_t) identically to * the Python emitter conv_direct_grouped_emit.py, builds the kernel via the * matching rocke_build_direct_conv_*_new function and lowers via * rocke_lower_kernel_to_llvm (per-config arch, flavor AUTO) and prints the .ll @@ -17,6 +18,7 @@ #include #include "rocke/instance_conv_direct_grouped.h" +#include "rocke/instance_conv_direct_nongrouped.h" #include "rocke/ir.h" #include "rocke/ir_serialize.h" #include "rocke/lower_llvm.h" @@ -32,9 +34,139 @@ enum KIND_SPATIAL = 5, KIND_DGRAD = 6, KIND_DW_DGRAD = 7, - KIND_WGRAD = 8 + KIND_WGRAD = 8, + KIND_NONGROUPED = 9 }; +#define NONGROUPED_CFG_BASE 33 + +static rocke_direct_conv_problem_t mk_nongrouped_problem( + int N, int H, int W, int C, int K, int KH, int KW, int PAD, int stride, const char* dtype) +{ + rocke_direct_conv_problem_t p = rocke_direct_conv_problem_default(); + p.N = N; + p.H = H; + p.W = W; + p.groups = 1; + p.cpg = C; + p.kpg = K; + p.KH = KH; + p.KW = KW; + p.PAD = PAD; + p.stride = stride; + p.dtype = dtype; + return p; +} + +/* Python _BASE: t8x32x64, ck32, 2x2 waves, iglp 0 (everything else default). */ +static rocke_direct_conv_nongrouped_spec_t mk_nongrouped_spec(rocke_direct_conv_problem_t p) +{ + rocke_direct_conv_nongrouped_spec_t s = rocke_direct_conv_nongrouped_spec_default(); + s.problem = p; + s.tile_h = 8; + s.tile_w = 32; + s.tile_k = 64; + s.ck = 32; + s.waves_m = 2; + s.waves_n = 2; + s.iglp = 0; + return s; +} + +/* Non-grouped (groups == 1) DirectNongroupedConvSpec configs, emitted as indices + * NONGROUPED_CFG_BASE + idx. Returns 0 on success, -1 if unknown. */ +static int make_nongrouped_cfg(int idx, rocke_direct_conv_nongrouped_spec_t* s, const char** arch) +{ + *arch = "gfx950"; + switch(idx) + { + case 0: + *s = mk_nongrouped_spec(mk_nongrouped_problem(2, 16, 32, 64, 128, 3, 3, 1, 1, "bf16")); + return 0; + case 1: + *s = mk_nongrouped_spec(mk_nongrouped_problem(2, 16, 32, 64, 128, 3, 3, 1, 1, "fp16")); + return 0; + case 2: + *s = mk_nongrouped_spec(mk_nongrouped_problem(2, 16, 32, 64, 128, 3, 3, 1, 1, "bf16")); + s->double_buffer = true; + return 0; + case 3: + *s = mk_nongrouped_spec(mk_nongrouped_problem(2, 16, 32, 64, 128, 3, 3, 1, 1, "bf16")); + s->chiplet_swizzle = false; + return 0; + case 4: + *s = mk_nongrouped_spec(mk_nongrouped_problem(2, 16, 32, 64, 128, 3, 3, 1, 1, "bf16")); + s->iglp = ROCKE_DCONV_NONGROUPED_IGLP_NONE; + s->waves_per_eu = 3; + return 0; + case 5: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 32, 64, 64, 64, 3, 3, 1, 2, "bf16")); + return 0; + case 6: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 16, 32, 64, 64, 1, 1, 0, 1, "bf16")); + return 0; + case 7: + *s = mk_nongrouped_spec(mk_nongrouped_problem(2, 20, 40, 64, 64, 3, 3, 1, 1, "bf16")); + s->tile_w = 48; + s->tile_k = 32; + s->waves_m = 1; + s->atom = "16x16x32"; + return 0; + case 8: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 16, 32, 64, 64, 3, 3, 1, 1, "fp16")); + s->tile_w = 16; + s->tile_k = 32; + s->ck = 16; + s->waves_m = 1; + s->atom = "16x16x16"; + return 0; + case 9: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 16, 32, 64, 64, 3, 3, 1, 1, "fp16")); + s->atom = "32x32x8"; + s->ck = 16; + *arch = "gfx942"; + return 0; + case 10: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 16, 48, 32, 64, 3, 3, 1, 1, "bf16")); + s->tile_w = 48; + s->tile_k = 32; + s->ck = 16; + s->waves_m = 1; + s->waves_n = 4; + s->atom = "16x16x16"; + *arch = "gfx942"; + return 0; + case 11: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 16, 32, 64, 96, 5, 5, 2, 1, "bf16")); + s->swizzle_wgm = 1; + return 0; + case 12: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 20, 64, 64, 64, 3, 3, 1, 1, "bf16")); + s->tile_h = 16; + s->tile_w = 64; + s->waves_n = 4; + return 0; + case 13: + *s = mk_nongrouped_spec(mk_nongrouped_problem(4, 64, 64, 640, 640, 3, 3, 1, 1, "bf16")); + s->tile_h = 16; + s->tile_w = 64; + s->tile_k = 128; + s->ck = 16; + s->waves_m = 2; + s->waves_n = 4; + return 0; + case 14: + *s = mk_nongrouped_spec(mk_nongrouped_problem(1, 32, 64, 64, 64, 3, 3, 1, 2, "fp16")); + s->ck = 16; + s->double_buffer = true; + s->iglp = 1; + s->chiplet_swizzle = false; + return 0; + default: + return -1; + } +} + /* Fill the config for index `idx`. Returns 0 on success, -1 if unknown. * On success sets *kind, the matching spec struct, and *arch. */ static int make_cfg(int idx, @@ -48,6 +180,7 @@ static int make_cfg(int idx, rocke_direct_conv_dgrad_spec_t* sdgrad, rocke_direct_depthwise_dgrad_spec_t* sdw_dgrad, rocke_direct_conv_wgrad_spec_t* swg, + rocke_direct_conv_nongrouped_spec_t* snongrouped, const char** arch) { rocke_direct_conv_problem_t p = rocke_direct_conv_problem_default(); @@ -565,6 +698,12 @@ static int make_cfg(int idx, *arch = "gfx950"; return 0; default: + if(idx >= NONGROUPED_CFG_BASE + && make_nongrouped_cfg(idx - NONGROUPED_CFG_BASE, snongrouped, arch) == 0) + { + *kind = KIND_NONGROUPED; + return 0; + } return -1; } } @@ -589,8 +728,21 @@ int main(int argc, char** argv) rocke_direct_conv_dgrad_spec_t sdgrad; rocke_direct_depthwise_dgrad_spec_t sdw_dgrad; rocke_direct_conv_wgrad_spec_t swg; + rocke_direct_conv_nongrouped_spec_t snongrouped; const char* arch = "gfx950"; - if(make_cfg(idx, &kind, &s16, &s4, &s8, &s32, &sdw, &ssp, &sdgrad, &sdw_dgrad, &swg, &arch) + if(make_cfg(idx, + &kind, + &s16, + &s4, + &s8, + &s32, + &sdw, + &ssp, + &sdgrad, + &sdw_dgrad, + &swg, + &snongrouped, + &arch) != 0) { fprintf(stderr, "unknown config index %d\n", idx); @@ -615,6 +767,8 @@ int main(int argc, char** argv) kernel = rocke_build_direct_depthwise_dgrad_new(&b, &sdw_dgrad, arch); else if(kind == KIND_WGRAD) kernel = rocke_build_direct_conv_wgrad_new(&b, &swg, arch); + else if(kind == KIND_NONGROUPED) + kernel = rocke_build_direct_conv_nongrouped_new(&b, &snongrouped, arch); else kernel = rocke_build_direct_depthwise_new(&b, &sdw, arch); if(kernel == NULL) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py index 8e0d3e7446c2..6b2d57093bee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py @@ -6,10 +6,11 @@ # direct grouped convolution parity harness. Selects one of N sampled spec # configs by argv[1], builds the DirectConv16cSpec / DirectConv4cSpec / # DirectConv8cSpec / DirectConv32cSpec / DirectDepthwiseSpec / -# DirectConvDgradSpec / DirectDepthwiseDgradSpec / DirectConvWgradSpec, -# builds the kernel via the matching build_direct_conv_* function -# (arch=) and prints _native_lower(arch=) to stdout so it -# can be byte-compared with the C emitter conv_direct_grouped_emit.c. +# DirectConvDgradSpec / DirectDepthwiseDgradSpec / DirectConvWgradSpec (and, +# from index 32, the non-grouped DirectNongroupedConvSpec), builds the kernel +# via the matching build_direct_conv_* function (arch=) and prints +# _native_lower(arch=) to stdout so it can be byte-compared with +# the C emitter conv_direct_grouped_emit.c. import sys from kernels.common.conv_direct_grouped import ( @@ -33,6 +34,10 @@ build_direct_conv_dgrad, build_direct_depthwise_dgrad, ) +from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, + build_direct_conv_nongrouped, +) try: from rocke.core.lower_llvm import _lower_kernel_to_llvm_python as _native_lower @@ -42,6 +47,173 @@ from rocke.core.verify import verify +# --------------------------------------------------------------------------- +# Non-grouped (groups == 1) DirectNongroupedConvSpec configs, indices 33..47. +# +# Each pins one branch of build_direct_conv_nongrouped: the four MFMA atoms, +# fp16/bf16, stride 1/2, 1x1 / 3x3 / 5x5 filters, partial W/H/K tiles, staging +# passes that do / do not divide the block evenly (the scratch-tail select), +# single vs double buffered LDS, chiplet swizzle on/off, iglp and waves_per_eu. +# --------------------------------------------------------------------------- +_NONGROUPED_CFG_BASE = 33 + + +def _nongrouped_p(N, H, W, C, K, *, KH=3, KW=3, PAD=1, stride=1, dtype="bf16"): + return DirectConvProblem( + N=N, + H=H, + W=W, + groups=1, + cpg=C, + kpg=K, + KH=KH, + KW=KW, + PAD=PAD, + stride=stride, + dtype=dtype, + ) + + +# The geometry every non-grouped config shares unless it overrides it: the +# shape the sweep's best configs take (t8x32x64, ck32, 2x2 waves, 32x32x16, +# iglp0). +_NONGROUPED_BASE = dict( + tile_h=8, tile_w=32, tile_k=64, ck=32, waves_m=2, waves_n=2, iglp=0 +) + + +def _nongrouped_s(problem, **kw): + return DirectNongroupedConvSpec(problem=problem, **{**_NONGROUPED_BASE, **kw}) + + +def _nongrouped_spec(idx: int): + """Return (spec, arch) for non-grouped config ``idx`` (emitted as + ``_NONGROUPED_CFG_BASE + idx``), or None past the last one.""" + if idx == 0: + # bf16 baseline: even staging passes, chiplet swizzle, iglp 0. + return _nongrouped_s(_nongrouped_p(2, 16, 32, 64, 128)), "gfx950" + if idx == 1: + # fp16 operand / store path. + return _nongrouped_s(_nongrouped_p(2, 16, 32, 64, 128, dtype="fp16")), "gfx950" + if idx == 2: + # Ping-pong LDS: parity-selected buffers, one barrier per chunk. + return ( + _nongrouped_s(_nongrouped_p(2, 16, 32, 64, 128), double_buffer=True), + "gfx950", + ) + if idx == 3: + # No chiplet swizzle: plain mod/div grid decode. + return ( + _nongrouped_s(_nongrouped_p(2, 16, 32, 64, 128), chiplet_swizzle=False), + "gfx950", + ) + if idx == 4: + # waves_per_eu attribute, no iglp_opt. + return ( + _nongrouped_s(_nongrouped_p(2, 16, 32, 64, 128), iglp=None, waves_per_eu=3), + "gfx950", + ) + if idx == 5: + # Stride 2: input-row sharing across taps, strided staging window. + return _nongrouped_s(_nongrouped_p(1, 32, 64, 64, 64, stride=2)), "gfx950" + if idx == 6: + # 1x1 pointwise filter, PAD 0. + return ( + _nongrouped_s(_nongrouped_p(1, 16, 32, 64, 64, KH=1, KW=1, PAD=0)), + "gfx950", + ) + if idx == 7: + # 16x16x32 atom on a width that is not a multiple of 32 (partial W). + return ( + _nongrouped_s( + _nongrouped_p(2, 20, 40, 64, 64), + tile_w=48, + tile_k=32, + waves_m=1, + atom="16x16x32", + ), + "gfx950", + ) + if idx == 8: + # 16x16x16 atom (4-half fragments), ck 16, single-tile width. + return ( + _nongrouped_s( + _nongrouped_p(1, 16, 32, 64, 64, dtype="fp16"), + tile_w=16, + tile_k=32, + ck=16, + waves_m=1, + atom="16x16x16", + ), + "gfx950", + ) + if idx == 9: + # 32x32x8 atom on gfx942 (no 32x32x16 there). + return ( + _nongrouped_s( + _nongrouped_p(1, 16, 32, 64, 64, dtype="fp16"), atom="32x32x8", ck=16 + ), + "gfx942", + ) + if idx == 10: + # gfx942 bf16 16x16x16 with uneven staging passes. + return ( + _nongrouped_s( + _nongrouped_p(1, 16, 48, 32, 64), + tile_w=48, + tile_k=32, + ck=16, + waves_m=1, + waves_n=4, + atom="16x16x16", + ), + "gfx942", + ) + if idx == 11: + # 5x5 filter, PAD 2, partial K tile, swizzle_wgm 1. + return ( + _nongrouped_s( + _nongrouped_p(1, 16, 32, 64, 96, KH=5, KW=5, PAD=2), swizzle_wgm=1 + ), + "gfx950", + ) + if idx == 12: + # Partial H tile + multi column-block tile (tile_w = 2 atoms). + return ( + _nongrouped_s( + _nongrouped_p(1, 20, 64, 64, 64), tile_h=16, tile_w=64, waves_n=4 + ), + "gfx950", + ) + if idx == 13: + # A production-sized config: N4 C640 K640 64x64, t16x64x128 w2x4. + return ( + _nongrouped_s( + _nongrouped_p(4, 64, 64, 640, 640), + tile_h=16, + tile_w=64, + tile_k=128, + ck=16, + waves_m=2, + waves_n=4, + ), + "gfx950", + ) + if idx == 14: + # Double buffer + stride 2 + iglp 1 + fp16 + no swizzle together. + return ( + _nongrouped_s( + _nongrouped_p(1, 32, 64, 64, 64, stride=2, dtype="fp16"), + ck=16, + double_buffer=True, + iglp=1, + chiplet_swizzle=False, + ), + "gfx950", + ) + return None + + def _spec(idx: int): """Return (kind, spec, arch) for config index `idx`.""" if idx == 0: @@ -460,6 +632,10 @@ def _spec(idx: int): DirectDepthwiseSpec(problem=p, block_w=8, block_waves=2), "gfx950", ) + if idx >= _NONGROUPED_CFG_BASE: + sel = _nongrouped_spec(idx - _NONGROUPED_CFG_BASE) + if sel is not None: + return ("nongrouped", *sel) raise SystemExit(f"unknown config index {idx}") @@ -486,6 +662,8 @@ def main() -> int: kernel = build_direct_conv_dgrad(spec, arch=arch) elif kind == "dw_dgrad": kernel = build_direct_depthwise_dgrad(spec, arch=arch) + elif kind == "nongrouped": + kernel = build_direct_conv_nongrouped(spec, arch=arch) else: kernel = build_direct_depthwise(spec, arch=arch) if mode == "ll": 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 9d091f09b23d..6abffc9c6ad0 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 @@ -294,6 +294,7 @@ "kernels.common.conv_direct_grouped.DirectCoalescedWeightsDgradSpec": ("problem",), "kernels.common.conv_direct_grouped.DirectMfmaDgradSpec": ("problem",), "kernels.common.conv_direct_grouped.DirectConvWgradSpec": ("problem",), + "kernels.common.conv_direct_nongrouped.DirectNongroupedConvSpec": ("problem",), "kernels.common.deep_fused_conv_pool.FusedConvPoolProblem": ("conv",), "kernels.common.deep_fused_conv_pool.DeepFusedConvPoolSpec": ("problem",), "kernels.common.img2col.Img2ColSpec": ("problem",), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py index bf0ef2ec5ac9..5d6e7d95bca4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py @@ -537,6 +537,49 @@ def test_direct_conv_abi_and_shape_invariance(label, channels, direction, build) assert _names(signature) == expected, f"{label}: launch signature" +@pytest.mark.parametrize( + "stride,knobs", + [ + (1, dict()), + (2, dict(ck=16, double_buffer=True, chiplet_swizzle=False, iglp=None)), + (1, dict(tile_w=48, tile_k=32, waves_m=1, atom="16x16x32")), + ], + ids=["base", "s2_db_noswizzle", "atom16"], +) +def test_direct_nongrouped_abi_and_shape_invariance(stride, knobs): + """The non-grouped kernel bakes neither the extents nor the channel counts. + + Unlike the grouped variants, ``C`` and ``K`` are kernargs too, so they vary + here alongside the batch and the image -- including image sizes that leave + partial tiles and channel counts that leave a partial channel tile. + """ + from kernels.common.conv_direct_grouped import DirectConvProblem + from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, + build_direct_conv_nongrouped, + ) + + base = dict(tile_h=8, tile_w=32, tile_k=64, ck=32, waves_m=2, waves_n=2, iglp=0) + fingerprints = set() + kernel = None + for N, H, W, C, K in [ + (2, 16, 32, 64, 128), + (3, 29, 37, 128, 96), + (1, 64, 64, 640, 640), + ]: + problem = DirectConvProblem( + N=N, H=H, W=W, groups=1, cpg=C, kpg=K, stride=stride, dtype="bf16" + ) + spec = DirectNongroupedConvSpec(problem=problem, **{**base, **knobs}) + kernel = build_direct_conv_nongrouped(spec, arch=_ARCH) + fingerprints.add(_op_signature(kernel)) + + assert len(fingerprints) == 1, "non-grouped IR depends on the problem shape" + expected = [n for n, _ in conv_direct_arg_names(direction="fwd")] + assert [p.name for p in kernel.params] == expected + assert _names(conv_direct_args_signature("bf16")) == expected + + @pytest.mark.parametrize("fold_k32", [False, True]) def test_direct_dgrad_weight_transforms_are_shape_invariant(fold_k32): """The MFMA dgrad's weight transforms are cached once per filter/channels. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py index 9a53f3c2873a..d6a67aa33c32 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py @@ -2,8 +2,8 @@ # SPDX-License-Identifier: MIT """Correctness tests for direct grouped convolution across cpg variants. -Covers all four grouped variants (cpg = 4, 8, 16, 32) and the depthwise -variant (cpg = 1). Each test builds a kernel, compiles it, launches it on +Covers all four grouped variants (cpg = 4, 8, 16, 32), the depthwise +variant (cpg = 1), and the non-grouped (groups == 1) ``DirectNongroupedConvSpec``. Each test builds a kernel, compiles it, launches it on GPU, and compares the output against a float32 reference produced by torch.nn.functional.conv2d. @@ -22,7 +22,7 @@ import ctypes import importlib.util import unittest -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import List, Tuple from rocke.runtime.hip_module import get_device_arch @@ -982,6 +982,461 @@ def test_bf16_dgrad(self): self._run_dgrad(s) +# --------------------------------------------------------------------------- +# Non-grouped (groups == 1) direct conv +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _NgCase: + """One non-grouped correctness case: shape plus the geometry to build.""" + + id: str + N: int + H: int + W: int + C: int + K: int + KH: int = 3 + KW: int = 3 + PAD: int = 1 + stride: int = 1 + dtype: str = "bf16" + tile_h: int = 8 + tile_w: int = 32 + tile_k: int = 64 + ck: int = 32 + waves_m: int = 2 + waves_n: int = 2 + atom: str = "32x32x16" + # Schedule / pipeline knobs (DirectNongroupedConvSpec defaults). + double_buffer: bool = False + iglp: "int | None" = None + waves_per_eu: "int | None" = None + chiplet_swizzle: bool = True + swizzle_wgm: int = 8 + + +# Each case pins down one branch of the addressing: exact vs partial output +# tiles, K not covering tile_k, the 16-wide atom used for widths that are not a +# multiple of 32, stride 2, a 1x1 filter, and the fp16 operand path. The +# second group pins the pipeline/schedule knobs: ping-pong LDS, iglp, the +# occupancy hint, the grid decode with and without the chiplet swizzle, and a +# grid large enough (> num_xcds * chiplet_chunk = 512 WGs) that the chiplet +# remap is not the identity, with a partial tail block past the last full one. +_NG_CASES: List[_NgCase] = [ + _NgCase("ng_exact", N=1, H=16, W=32, C=64, K=64), + _NgCase("ng_partial_w", N=2, H=16, W=40, C=64, K=64), + _NgCase("ng_partial_h", N=1, H=20, W=32, C=64, K=64, tile_h=16, waves_n=4), + _NgCase("ng_partial_k", N=1, H=16, W=32, C=64, K=96), + _NgCase("ng_multi_colblock", N=1, H=16, W=64, C=64, K=64, tile_w=64), + _NgCase("ng_stride2", N=1, H=32, W=64, C=64, K=64, stride=2), + _NgCase("ng_pointwise", N=1, H=16, W=32, C=64, K=64, KH=1, KW=1, PAD=0), + _NgCase("ng_fp16", N=1, H=16, W=32, C=64, K=64, dtype="fp16"), + _NgCase( + "ng_atom16", + N=2, + H=20, + W=40, + C=64, + K=64, + tile_w=48, + tile_k=32, + waves_m=1, + atom="16x16x32", + ), + _NgCase( + "ng_atom16x16", + N=1, + H=16, + W=32, + C=64, + K=64, + tile_w=16, + tile_k=32, + ck=16, + waves_m=1, + atom="16x16x16", + ), + _NgCase("ng_32x32x8", N=1, H=16, W=32, C=64, K=64, ck=16, atom="32x32x8"), + _NgCase("ng_5x5", N=1, H=16, W=32, C=64, K=96, KH=5, KW=5, PAD=2), + _NgCase("ng_double_buffer", N=2, H=16, W=40, C=128, K=64, double_buffer=True), + _NgCase( + "ng_double_buffer_s2_fp16", + N=1, + H=32, + W=64, + C=64, + K=64, + stride=2, + ck=16, + dtype="fp16", + double_buffer=True, + iglp=1, + ), + _NgCase("ng_iglp0_we3", N=1, H=16, W=32, C=128, K=64, iglp=0, waves_per_eu=3), + _NgCase("ng_no_swizzle", N=2, H=16, W=64, C=64, K=128, chiplet_swizzle=False), + _NgCase("ng_wgm1", N=2, H=16, W=64, C=64, K=192, swizzle_wgm=1), + _NgCase( + "ng_chiplet_remap", + N=9, + H=64, + W=64, + C=32, + K=256, + iglp=0, + ), +] + + +# The cases are written for gfx950. gfx942 has neither wide-K atom for +# fp16/bf16 and a 64 KiB LDS, so there each case falls back to the narrower-K +# atom of the same tile size and then shrinks ck (and, failing that, tile_h) +# until the staged tiles fit. The addressing branch a case pins down depends on +# its shape, stride and filter -- not on ck -- so it still gets a numeric check +# instead of a skip. On gfx950 every case is used as written. +_NG_ATOM_FALLBACK = {"32x32x16": "32x32x8", "16x16x32": "16x16x16"} + + +def _fit_nongrouped_to_arch(spec, arch: str): + """``spec`` with its atom / ck / tile_h narrowed until ``arch`` can run it.""" + from kernels.common.conv_direct_nongrouped import _ATOMS, _X_LOAD_VEC + from rocke.core.arch import ArchTarget + + target = ArchTarget.from_gfx(arch) + ab = "bf16" if spec.problem.dtype == "bf16" else "f16" + t, k, _ = _ATOMS[spec.atom] + if not target.mma.has_shape(a_dtype=ab, b_dtype=ab, c_dtype="fp32", m=t, n=t, k=k): + spec = replace(spec, atom=_NG_ATOM_FALLBACK.get(spec.atom, spec.atom)) + min_ck = max(spec.atom_k, _X_LOAD_VEC) + while not target.fits_lds(spec.lds_bytes) and spec.ck // 2 >= min_ck: + spec = replace(spec, ck=spec.ck // 2) + while ( + not target.fits_lds(spec.lds_bytes) + and spec.tile_h > spec.waves_n + and (spec.tile_h // 2) % spec.waves_n == 0 + ): + spec = replace(spec, tile_h=spec.tile_h // 2) + return spec + + +def _run_nongrouped_one(arch: str, case: _NgCase) -> Tuple[bool, str]: + """Build, compile, launch, and verify one non-grouped direct-conv kernel.""" + from kernels.common.conv_direct_grouped import DirectConvProblem + from kernels.common.conv_direct_nongrouped import DirectNongroupedConvSpec + + p = DirectConvProblem( + N=case.N, + H=case.H, + W=case.W, + groups=1, + cpg=case.C, + kpg=case.K, + KH=case.KH, + KW=case.KW, + PAD=case.PAD, + stride=case.stride, + dtype=case.dtype, + ) + spec = DirectNongroupedConvSpec( + problem=p, + tile_h=case.tile_h, + tile_w=case.tile_w, + tile_k=case.tile_k, + ck=case.ck, + waves_m=case.waves_m, + waves_n=case.waves_n, + atom=case.atom, + double_buffer=case.double_buffer, + iglp=case.iglp, + waves_per_eu=case.waves_per_eu, + chiplet_swizzle=case.chiplet_swizzle, + swizzle_wgm=case.swizzle_wgm, + ) + return _run_nongrouped_spec(arch, _fit_nongrouped_to_arch(spec, arch), case.id) + + +def _run_nongrouped_spec( + arch: str, spec, case_id: str, artifact=None +) -> Tuple[bool, str]: + """Compile, launch and verify one ``DirectNongroupedConvSpec`` against conv2d. + + ``artifact`` reuses a binary compiled for another spec with the same + capabilities and tile; the shape of ``spec.problem`` then only reaches the + kernel through its kernargs. + """ + import torch + + from rocke import compile_kernel + from kernels.common.conv_abi import conv_direct_args_signature + from kernels.common.conv_args import ConvArgs + from kernels.common.conv_direct_nongrouped import ( + build_direct_conv_nongrouped, + is_valid_nongrouped_spec, + ) + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import HipError, Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + p = spec.problem + ok, why = is_valid_nongrouped_spec(spec, arch=arch) + if not ok: + return False, f"skip {why}" + + if artifact is None: + try: + artifact = compile_kernel( + build_direct_conv_nongrouped(spec, arch=arch), arch=arch + ) + except Exception as e: # noqa: BLE001 + return False, f"build/compile failed: {e}" + + td = torch.bfloat16 if p.dtype == "bf16" else torch.float16 + torch.manual_seed(0) + A_t = torch.empty(p.N, p.H, p.W, p.total_c, dtype=td).uniform_(-1.0, 1.0) + B_t = torch.empty(p.total_k, p.KH, p.KW, p.cpg, dtype=td).uniform_(-1.0, 1.0) + D_t = torch.empty(p.N, p.Ho, p.Wo, p.total_k, dtype=td) + ref = _conv_ref_grouped(A_t, B_t, p) + + rt = Runtime() + A_dev = rt.alloc(A_t.nbytes) + B_dev = rt.alloc(B_t.nbytes) + D_dev = rt.alloc(D_t.nbytes) + rt.memcpy_h2d(A_dev, _u8(A_t), A_t.nbytes) + rt.memcpy_h2d(B_dev, _u8(B_t), B_t.nbytes) + rt.memset(D_dev, 0, D_t.nbytes) + + try: + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=conv_direct_args_signature(p.dtype), + ) + except HipError as e: + rt.free(A_dev) + rt.free(B_dev) + rt.free(D_dev) + return False, f"kernel load failed: {e}" + + launcher( + ConvArgs.from_problem(p).to_launch_values( + int(A_dev), + int(B_dev), + int(D_dev), + A_t.nbytes, + B_t.nbytes, + D_t.nbytes, + ), + config=LaunchConfig( + grid=spec.grid(), block=(spec.threads_per_block, 1, 1), fence=True + ), + ) + + D_cpu = torch.empty_like(D_t) + rt.memcpy_d2h(_u8(D_cpu), D_dev, D_t.nbytes) + rt.free(A_dev) + rt.free(B_dev) + rt.free(D_dev) + synchronize_and_release(0) + + diff = (D_cpu.float() - ref.float().cpu()).abs() + rel_err = float(diff.max() / ref.abs().max().clamp(min=1.0)) + tol = _TOL_BF16 if p.dtype == "bf16" else _TOL + if rel_err >= tol: + return False, f"rel_err={rel_err:.3e} > tol={tol:.1e}" + print(f" PASS {case_id} {arch} rel_err={rel_err:.2e}", flush=True) + return True, "" + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestDirectConvNongroupedCorrectness(unittest.TestCase): + """Correctness sweep for the non-grouped (groups == 1) direct-conv family.""" + + def test_nongrouped(self): + for case in _NG_CASES: + with self.subTest(case=case.id): + passed, reason = _run_nongrouped_one(GPU_ARCH, case) + if reason.startswith("skip"): + self.skipTest(reason) + self.assertTrue(passed, f"FAIL {case.id} on {GPU_ARCH}: {reason}") + + def test_sweep_candidates(self): + """An evenly spaced sample of what ``nongrouped_specs`` hands the + benchmark. The winner of a sweep is whichever config is fastest, so + every candidate the sweep can pick has to be correct -- not just the + hand-picked geometries above. Wo=48 exercises both the exact-width + 16-wide atom and the padded 32-wide one.""" + from kernels.common.conv_direct_grouped import DirectConvProblem + from kernels.common.conv_direct_nongrouped import nongrouped_specs + + p = DirectConvProblem(N=2, H=16, W=48, groups=1, cpg=64, kpg=128, dtype="bf16") + specs = nongrouped_specs(p, arch=GPU_ARCH) + self.assertGreater(len(specs), 0) + for spec in specs[:: max(1, len(specs) // 8)]: + name = spec.kernel_name() + with self.subTest(spec=name): + passed, reason = _run_nongrouped_spec(GPU_ARCH, spec, name) + self.assertTrue(passed, f"FAIL {name} on {GPU_ARCH}: {reason}") + + def test_one_binary_serves_every_shape(self): + """AOT: a kernel compiled once runs any shape with its capabilities. + + The binary is built against a probe problem, then launched on shapes + that differ in batch, both spatial extents (exact and partial tiles, + so the runtime tile counts and bounds checks change) and both channel + counts (so the runtime channel-loop trip count and the K mask change). + """ + from rocke import compile_kernel + from kernels.common.conv_direct_grouped import DirectConvProblem + from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, + build_direct_conv_nongrouped, + ) + + def _p(N, H, W, C, K): + return DirectConvProblem( + N=N, H=H, W=W, groups=1, cpg=C, kpg=K, dtype="bf16" + ) + + spec = _fit_nongrouped_to_arch( + DirectNongroupedConvSpec( + problem=_p(1, 16, 32, 64, 64), + tile_h=8, + tile_w=32, + tile_k=64, + ck=32, + waves_m=2, + waves_n=2, + iglp=0, + ), + GPU_ARCH, + ) + artifact = compile_kernel( + build_direct_conv_nongrouped(spec, arch=GPU_ARCH), arch=GPU_ARCH + ) + for shape in [ + (1, 16, 32, 64, 64), + (3, 20, 40, 128, 96), + (2, 33, 17, 256, 192), + (9, 64, 64, 32, 256), + ]: + case = replace(spec, problem=_p(*shape)) + with self.subTest(shape=shape): + passed, reason = _run_nongrouped_spec( + GPU_ARCH, case, f"aot_{shape}", artifact=artifact + ) + self.assertTrue(passed, f"FAIL {shape} on {GPU_ARCH}: {reason}") + + +class TestDirectConvNongroupedValidation(unittest.TestCase): + """Non-grouped validation that does not require a GPU.""" + + @staticmethod + def _problem(**kw): + from kernels.common.conv_direct_grouped import DirectConvProblem + + base = dict(N=1, H=32, W=32, groups=1, cpg=64, kpg=64, dtype="bf16") + base.update(kw) + return DirectConvProblem(**base) + + def test_grouped_problem_rejected(self): + from kernels.common.conv_direct_nongrouped import DirectNongroupedConvSpec + + p = self._problem(groups=4, cpg=16, kpg=16) + with self.assertRaises(ValueError): + DirectNongroupedConvSpec(problem=p).validate() + + def test_oversized_accumulator_rejected(self): + """A tile needing more than 256 accumulator registers must not build. + + Left unguarded this is not just slow: a 2048-register tile hangs the + backend scheduler rather than failing, which stalls a whole sweep. + """ + from kernels.common.conv_direct_nongrouped import DirectNongroupedConvSpec + + spec = DirectNongroupedConvSpec( + problem=self._problem(), + tile_h=16, + tile_w=64, + tile_k=128, + ck=16, + waves_m=1, + waves_n=1, + ) + self.assertGreater(spec.acc_vgprs, 256) + with self.assertRaises(ValueError): + spec.validate() + + def test_flattened_grid_overflow_rejected(self): + """The 1-D grid (and the kernel's i32 workgroup count) must fit i32.""" + from kernels.common.conv_direct_nongrouped import DirectNongroupedConvSpec + + # 1 x 4 spatial tiles x 16 channel tiles per image, 2**25 images. + spec = DirectNongroupedConvSpec( + problem=self._problem(N=1 << 25, cpg=64, kpg=512), + tile_h=8, + tile_w=32, + tile_k=32, + ck=32, + waves_m=1, + waves_n=2, + ) + self.assertGreater(spec.grid()[0], (1 << 31) - 1) + with self.assertRaisesRegex(ValueError, "flattened grid"): + spec.validate() + + def test_is_valid_never_raises(self): + """``is_valid_nongrouped_spec`` reports, it does not raise -- including + when the arch-table queries past ``validate()`` fail.""" + from unittest import mock + + from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, + is_valid_nongrouped_spec, + ) + from rocke.core.arch import ArchTarget + + spec = DirectNongroupedConvSpec( + problem=self._problem(), tile_h=8, tile_k=64, ck=32, waves_n=2 + ) + self.assertTrue(is_valid_nongrouped_spec(spec)[0]) + + real = ArchTarget.from_gfx("gfx950") + broken = mock.MagicMock(wraps=real) + broken.mma.has_shape.side_effect = KeyError("no mma table") + with mock.patch.object(ArchTarget, "from_gfx", return_value=broken): + ok, why = is_valid_nongrouped_spec(spec) + self.assertFalse(ok) + self.assertIn("no mma table", why) + + ok, why = is_valid_nongrouped_spec(spec, arch="gfx_nonexistent") + self.assertFalse(ok) + + def test_tile_w_candidates_cover_exactly(self): + """Candidate widths must tile Wo without wasted columns when possible.""" + from kernels.common.conv_direct_nongrouped import tile_w_candidates + + for wo, at in ((96, 32), (160, 32), (112, 16), (208, 16)): + for tw in tile_w_candidates(wo, at): + self.assertEqual(wo % tw, 0, f"Wo={wo} atom={at} tw={tw}") + # 40 is not a multiple of 32, so the 32-wide atom falls back to its + # least-wasteful width rather than returning nothing. + self.assertEqual(tile_w_candidates(40, 32), [32]) + + def test_specs_are_unique_and_valid(self): + from kernels.common.conv_direct_nongrouped import ( + is_valid_nongrouped_spec, + nongrouped_specs, + ) + + specs = nongrouped_specs(self._problem(H=64, W=64, cpg=256, kpg=256)) + self.assertGreater(len(specs), 0) + names = [s.kernel_name() for s in specs] + self.assertEqual(len(names), len(set(names))) + for s in specs: + self.assertTrue(is_valid_nongrouped_spec(s)[0]) + + class TestDirectConvValidation(unittest.TestCase): """Validation-only tests that do not require a GPU.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py index fa1a0482fac6..bfbce7baf131 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py @@ -9,6 +9,8 @@ from __future__ import annotations +import json + import pytest from benchmarks.common import direct_kernel_sweep as dks @@ -91,7 +93,47 @@ def test_group_count_is_runtime(cache, groups): """Any group count is served; only tiles that do not divide it drop out.""" plans, rejected = dks.direct_plans(cache, _problem(groups=groups), "fwd", _ARCH) assert plans - assert all("block_groups" in why for _, why in rejected) + assert all( + "block_groups" in why + for ident, why in rejected + if ident.algorithm != "direct_nongrouped" + ) + + +def _nongrouped(plans): + return [p for p in plans if p.identity.algorithm == "direct_nongrouped"] + + +@pytest.mark.parametrize( + "C,K,H,W", [(64, 64, 56, 56), (640, 640, 64, 64), (96, 352, 28, 48)] +) +def test_nongrouped_channels_are_runtime(cache, C, K, H, W): + """One non-grouped binary per tile serves any C/K it divides.""" + problem = _problem(groups=1, cpg=C, kpg=K, H=H, W=W) + plans = _nongrouped(dks.direct_plans(cache, problem, "fwd", _ARCH)[0]) + assert plans + assert all(p.identity.cpg == 0 and p.identity.kpg == 0 for p in plans) + # Only the widths a sweep for this Wo would try are offered. + from kernels.common.conv_direct_nongrouped import tile_w_candidates + + for plan in plans: + spec = dks.make_spec( + "direct_nongrouped", problem, json.loads(plan.identity.knobs) + ) + assert spec.tile_w in tile_w_candidates(problem.Wo, spec.atom_tile) + (step,) = plan.steps + assert step.grid == spec.grid() + + +def test_nongrouped_only_serves_one_group(cache): + plans, rejected = dks.direct_plans(cache, _problem(groups=16), "fwd", _ARCH) + assert not _nongrouped(plans) + assert all(i.algorithm != "direct_nongrouped" for i, _ in rejected) + # Grouped forward has no stride-2 kernel; the non-grouped one does. + plans, _ = dks.direct_plans( + cache, _problem(groups=1, cpg=64, kpg=64, stride=2), "fwd", _ARCH + ) + assert plans and _variants(plans) == {"direct_nongrouped"} def test_capabilities_outside_the_list_are_not_served(cache): @@ -118,6 +160,7 @@ def test_dgrad_stride1_offers_mfma_pipeline_and_scalar(cache): "direction,problem", [ ("fwd", _problem()), + ("fwd", _problem(groups=1, cpg=128, kpg=96)), ("fwd", _problem(groups=64, cpg=1, kpg=1, KH=7, KW=7, PAD=3, stride=2)), ("dgrad", _problem(groups=64, cpg=1, kpg=1, KH=5, KW=5, PAD=0, stride=2)), ("dgrad", _problem(cpg=8, kpg=8, stride=2)), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index 720da7886a36..23202392aac8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -39,6 +39,7 @@ extern "C" { #include "rocke/instance_batched_gemm.h" #include "rocke/instance_block_scale_gemm.h" #include "rocke/instance_conv_direct_grouped.h" +#include "rocke/instance_conv_direct_nongrouped.h" #include "rocke/instance_conv_implicit_gemm.h" #include "rocke/instance_conv_implicit_gemm_wgrad.h" #include "rocke/instance_conv_wgrad_workspace_reduce.h" @@ -1714,6 +1715,95 @@ std::vector conv_direct_grouped_verify(const py::dict& d, const std return out; } +/* ======================== conv_direct_nongrouped ========================== */ + +/* Non-grouped (groups == 1) direct conv. The dict mirrors DirectNongroupedConvSpec + * with the problem nested; `iglp` / `waves_per_eu` may be None (absent knob). */ +rocke_direct_conv_nongrouped_spec_t dnongrouped_build_spec(const py::dict& d, + std::deque& store) +{ + auto keep = [&](const std::string& s) -> const char* { + store.push_back(s); + return store.back().c_str(); + }; + rocke_direct_conv_nongrouped_spec_t s = rocke_direct_conv_nongrouped_spec_default(); + if(d.contains("problem") && py::isinstance(d["problem"])) + { + py::dict pd = d["problem"].cast(); + std::string v; + fill_direct_conv_problem(&s.problem, pd); + if(dict_str(pd, "dtype", v)) + s.problem.dtype = keep(v); + } + { + std::string v; + if(dict_str(d, "name", v)) + s.name = keep(v); + if(dict_str(d, "atom", v)) + s.atom = keep(v); + } + s.tile_h = dict_int(d, "tile_h", s.tile_h); + s.tile_w = dict_int(d, "tile_w", s.tile_w); + s.tile_k = dict_int(d, "tile_k", s.tile_k); + s.ck = dict_int(d, "ck", s.ck); + s.waves_m = dict_int(d, "waves_m", s.waves_m); + s.waves_n = dict_int(d, "waves_n", s.waves_n); + s.wave_size = dict_int(d, "wave_size", s.wave_size); + s.lds_pad = dict_int(d, "lds_pad", s.lds_pad); + s.chiplet_swizzle = dict_bool(d, "chiplet_swizzle", s.chiplet_swizzle); + s.swizzle_wgm = dict_int(d, "swizzle_wgm", s.swizzle_wgm); + s.chiplet_chunk = dict_int(d, "chiplet_chunk", s.chiplet_chunk); + s.num_xcds = dict_int(d, "num_xcds", s.num_xcds); + s.double_buffer = dict_bool(d, "double_buffer", s.double_buffer); + /* The C spec encodes None as a sentinel (-1 / 0). An explicit value in the + * sentinel range would silently read as None here, so reject it with the + * message Python's validate() raises for the same spec. */ + if(d.contains("iglp") && !d["iglp"].is_none()) + { + s.iglp = d["iglp"].cast(); + if(s.iglp < 0) + throw std::runtime_error("iglp must be None or >= 0 (got " + std::to_string(s.iglp) + + ")"); + } + if(d.contains("waves_per_eu") && !d["waves_per_eu"].is_none()) + { + s.waves_per_eu = d["waves_per_eu"].cast(); + if(s.waves_per_eu < 1) + throw std::runtime_error("waves_per_eu must be None or >= 1 (got " + + std::to_string(s.waves_per_eu) + ")"); + } + return s; +} + +std::string conv_direct_nongrouped_lower_llvm(const py::dict& d, const std::string& arch) +{ + std::deque store; + rocke_direct_conv_nongrouped_spec_t s = dnongrouped_build_spec(d, store); + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + rocke_status_t st = rocke_direct_conv_nongrouped_lower_to_llvm( + &s, arch_or_default(arch), ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof err); + return take_lowered(st, ll, err, "rocke_engine.conv_direct_nongrouped_lower_llvm"); +} + +std::string conv_direct_nongrouped_serialize_ir(const py::dict& d, const std::string& arch) +{ + ROCKE_FAMILY_SERIALIZE_BODY( + "rocke_engine.conv_direct_nongrouped_serialize_ir", + rocke_direct_conv_nongrouped_spec_t, + dnongrouped_build_spec, + rocke_build_direct_conv_nongrouped_new(&b, &s, arch_or_default(arch))); +} + +std::vector conv_direct_nongrouped_verify(const py::dict& d, const std::string& arch) +{ + ROCKE_FAMILY_VERIFY_BODY("rocke_engine.conv_direct_nongrouped_verify", + rocke_direct_conv_nongrouped_spec_t, + dnongrouped_build_spec, + rocke_build_direct_conv_nongrouped_new(&b, &s, arch_or_default(arch))); +} + /* ======================= deep_fused_conv_pool ======================= */ /* This family is constructed through a factory in both reference emitters; the @@ -3659,6 +3749,10 @@ PYBIND11_MODULE(rocke_engine, m) &conv_direct_grouped_lower_llvm, &conv_direct_grouped_serialize_ir, &conv_direct_grouped_verify); + reg3("conv_direct_nongrouped", + &conv_direct_nongrouped_lower_llvm, + &conv_direct_nongrouped_serialize_ir, + &conv_direct_nongrouped_verify); register_img2col(m); /* separate TU; see note at top of file */ reg3("deep_fused_conv_pool", &deep_fused_conv_pool_lower_llvm, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp index 27365b859cfb..4907d0a1b2ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp @@ -490,6 +490,61 @@ const char* const* rocke_known_arches(int* count) return rocke_ati_known_arches; } +/* Reproduce str(KeyError(_build_target message)) for an unknown gfx target: + * + * Python _build_target: raise KeyError( + * f"unknown gfx target {gfx!r}; known: {sorted(specs)}. " + * f"Add a row to {_DATA_FILE.name}.") + * + * str(KeyError(msg)) == repr(msg); the single quotes make Python DOUBLE-quote + * the whole message. sorted(specs) renders as ['gfx...', 'gfx...']. */ +void rocke_set_unknown_arch_reason(char* out, size_t out_cap, const char* gfx) +{ + int count = 0; + const char* const* arches; + int i; + size_t pos = 0; + int wrote; + + if(out == NULL || out_cap == 0) + { + return; + } + + arches = rocke_known_arches(&count); + + wrote = snprintf(out + pos, out_cap - pos, "\"unknown gfx target '%s'; known: [", gfx); + if(wrote < 0) + { + out[0] = '\0'; + return; + } + pos += (size_t)wrote; + if(pos >= out_cap) + { + out[out_cap - 1] = '\0'; + return; + } + + for(i = 0; i < count; ++i) + { + wrote = snprintf(out + pos, out_cap - pos, "%s'%s'", (i == 0) ? "" : ", ", arches[i]); + if(wrote < 0) + { + out[out_cap - 1] = '\0'; + return; + } + pos += (size_t)wrote; + if(pos >= out_cap) + { + out[out_cap - 1] = '\0'; + return; + } + } + + snprintf(out + pos, out_cap - pos, "]. Add a row to arch_specs.json.\""); +} + /* Return the final gfx component without discarding profile or feature suffixes. */ static const char* rocke_target_id_start(const char* isa) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h index 12c52d012c2b..60251a2166d3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h @@ -341,6 +341,12 @@ int rocke_arch_max_threads_per_block(const rocke_arch_target_t* t); * (excluding the terminating NULL). */ const char* const* rocke_known_arches(int* count); +/* Write str(KeyError) of ArchTarget.from_gfx's miss for `gfx` into out: + * "unknown gfx target ''; known: ['gfx...', ...]. Add a row to arch_specs.json." + * wrapped in double quotes, exactly as Python's is_valid_spec(...) reports + * `return False, str(e)`. Truncates to out_cap; NULL/zero out is a no-op. */ +void rocke_set_unknown_arch_reason(char* out, size_t out_cap, const char* gfx); + /* Target identity helpers mirror rocke.core.arch.target. They do not validate * compiler support or query a GPU. Inputs and output buffers must not overlap. * Each writes at most out_cap - 1 bytes plus a NUL and returns out; a small diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_nongrouped.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_nongrouped.h new file mode 100644 index 000000000000..12e075157687 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_nongrouped.h @@ -0,0 +1,161 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT + * + * rocke/instance_conv_direct_nongrouped.h -- C++ engine port of the non-grouped + * (groups == 1) NHWC direct convolution builder in + * library/kernels/common/conv_direct_nongrouped.py. + * + * The kernel stages the input tile *with its halo* in LDS once per channel + * chunk and lets all KH*KW filter taps read shifted sub-tiles out of that one + * copy; weights are staged pre-swizzled into MFMA fragment order. The channel + * loop is a runtime scf.for whose loop-carried state is the accumulator tiles + * plus the global-load registers of the next chunk (software prefetch). + * + * Python (conv_direct_nongrouped.py) C (this header) + * ------------------------------------ ------------------------------------- + * @dataclass(frozen) DirectNongroupedConvSpec rocke_direct_conv_nongrouped_spec_t + * derived @property geometry rocke_direct_conv_nongrouped_{lds_bytes, + * acc_vgprs,threads_per_block,grid} + * .kernel_name() / .validate() rocke_direct_conv_nongrouped_kernel_name / + * rocke_direct_conv_nongrouped_validate + * is_valid_nongrouped_spec(spec, arch) rocke_direct_conv_nongrouped_is_valid_spec + * build_direct_conv_nongrouped(spec, arch) rocke_build_direct_conv_nongrouped(_new) + * (+ convenience: build -> lower .ll) rocke_direct_conv_nongrouped_lower_to_llvm + * + * The shared problem struct (rocke_direct_conv_problem_t, Python + * DirectConvProblem) comes from rocke/instance_conv_direct_grouped.h; this + * family requires problem.groups == 1, so cpg is the full C and kpg the full K. + * + * AOT: the kernel declares the direct-conv kernarg block of + * rocke_conv_direct_arg_names("fwd") (instance_conv_abi.h). N, H, W, Ho, Wo, + * C (p_total_c), K (p_total_k) and the NHWC / NHWK strides are runtime; the + * filter, stride, PAD, dtype and tile geometry are baked. The spec's problem + * only has to pass validate() -- the emitted IR does not depend on its extents + * or channel counts. validate() on the problem actually launched is what keeps + * every tensor below the masked-access offset, so run it per launch shape. + * + * Optional Python knobs (`iglp: int | None`, `waves_per_eu: int | None`) map to + * the sentinels ROCKE_DCONV_NONGROUPED_IGLP_NONE / ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE. + * Python rejects the sentinel values themselves (iglp < 0, waves_per_eu < 1) and + * this port rejects anything below them, so a spec means the same in both engines. + */ +#ifndef ROCKE_INSTANCE_CONV_DIRECT_NONGROUPED_H +#define ROCKE_INSTANCE_CONV_DIRECT_NONGROUPED_H + +#include +#include + +#include "rocke/instance_conv_direct_grouped.h" /* rocke_direct_conv_problem_t */ +#include "rocke/ir.h" +#include "rocke/lower_llvm.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* `iglp=None`: leave the backend scheduler alone (no iglp_opt emitted). */ +#define ROCKE_DCONV_NONGROUPED_IGLP_NONE (-1) +/* `waves_per_eu=None`: no amdgpu-waves-per-eu attribute. */ +#define ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE 0 + +/* ===================================================================== * + * DirectNongroupedConvSpec + * + * @dataclass(frozen=True) + * class DirectNongroupedConvSpec: + * problem: DirectConvProblem # groups must be 1 + * name: str = "direct_conv_nongrouped" + * tile_h: int = 16 # output rows per workgroup + * tile_w: int = 32 # output cols per workgroup + * tile_k: int = 128 # output channels per workgroup + * ck: int = 16 # input channels per LDS chunk + * waves_m: int = 2 # waves splitting tile_k + * waves_n: int = 4 # waves splitting tile_h + * atom: str = "32x32x16" # 32x32x16 | 32x32x8 | 16x16x32 | 16x16x16 + * wave_size: int = 64 + * lds_pad: int = 8 + * chiplet_swizzle: bool = True + * swizzle_wgm: int = 8 + * chiplet_chunk: int = 64 + * num_xcds: int = 8 + * double_buffer: bool = False + * iglp: int | None = None + * waves_per_eu: int | None = None + * ===================================================================== */ +typedef struct rocke_direct_conv_nongrouped_spec +{ + rocke_direct_conv_problem_t problem; + const char* name; /* default "direct_conv_nongrouped" */ + int tile_h; /* default 16 */ + int tile_w; /* default 32 */ + int tile_k; /* default 128 */ + int ck; /* default 16 */ + int waves_m; /* default 2 */ + int waves_n; /* default 4 */ + const char* atom; /* default "32x32x16" */ + int wave_size; /* default 64 */ + int lds_pad; /* default 8 */ + bool chiplet_swizzle; /* default true */ + int swizzle_wgm; /* default 8 */ + int chiplet_chunk; /* default 64 */ + int num_xcds; /* default 8 */ + bool double_buffer; /* default false */ + int iglp; /* default ROCKE_DCONV_NONGROUPED_IGLP_NONE */ + int waves_per_eu; /* default ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE */ +} rocke_direct_conv_nongrouped_spec_t; + +/* Spec with the Python dataclass defaults; problem == problem_default() with + * groups = 1. Caller fills N,H,W,cpg,kpg (and KH/KW/PAD/stride/dtype). */ +rocke_direct_conv_nongrouped_spec_t rocke_direct_conv_nongrouped_spec_default(void); + +/* Derived geometry (Python @property). Return 0 for an unknown atom. */ +int rocke_direct_conv_nongrouped_threads_per_block(const rocke_direct_conv_nongrouped_spec_t* spec); +long rocke_direct_conv_nongrouped_lds_bytes(const rocke_direct_conv_nongrouped_spec_t* spec); +int rocke_direct_conv_nongrouped_acc_vgprs(const rocke_direct_conv_nongrouped_spec_t* spec); +/* spec.grid() -> (n_w_tiles * n_h_tiles * N * n_k_tiles, 1, 1). */ +void rocke_direct_conv_nongrouped_grid(const rocke_direct_conv_nongrouped_spec_t* spec, + int grid[3]); + +/* kernel_name(): kernel_name_join(name, problem.short(), t{h}x{w}x{k}, ck{ck}, + * w{wm}x{wn}, a{atom}, g{wgm}|gnone, [db], [iglp{n}], [we{n}], [bf16]). */ +rocke_status_t rocke_direct_conv_nongrouped_kernel_name( + const rocke_direct_conv_nongrouped_spec_t* spec, char* out, size_t out_cap); + +/* validate(): the hard assertions of DirectNongroupedConvSpec.validate. On a violated + * invariant returns ROCKE_ERR_VALUE and writes the Python ValueError message + * into reason (if non-NULL); else ROCKE_OK and "ok". */ +rocke_status_t rocke_direct_conv_nongrouped_validate( + const rocke_direct_conv_nongrouped_spec_t* spec, char* reason, size_t reason_cap); + +/* is_valid_nongrouped_spec(spec, arch) -> (ok, reason). `arch` NULL => "gfx950". + * Adds to validate(): stride in {1, 2}, the MFMA atom present on arch, wave + * size / threads-per-block / LDS capacity within the arch limits. */ +bool rocke_direct_conv_nongrouped_is_valid_spec(const rocke_direct_conv_nongrouped_spec_t* spec, + const char* arch, + char* reason, + size_t reason_cap); + +/* build_direct_conv_nongrouped(spec, arch) into an already-initialised builder. */ +rocke_kernel_def_t* rocke_build_direct_conv_nongrouped( + rocke_ir_builder_t* b, const rocke_direct_conv_nongrouped_spec_t* spec, const char* arch); + +/* Initialise `b` with the spec's kernel name, then build (the Python + * IRBuilder(spec.kernel_name()) + build_direct_conv_nongrouped pair). */ +rocke_kernel_def_t* rocke_build_direct_conv_nongrouped_new( + rocke_ir_builder_t* b, const rocke_direct_conv_nongrouped_spec_t* spec, const char* arch); + +/* Convenience: build + lower to LLVM IR text. *out_ll is malloc'd (caller + * frees); on failure the builder/lowerer message is written to err. */ +rocke_status_t + rocke_direct_conv_nongrouped_lower_to_llvm(const rocke_direct_conv_nongrouped_spec_t* spec, + const char* arch, + rocke_llvm_flavor_t flavor, + char** out_ll, + char* err, + size_t err_cap); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* ROCKE_INSTANCE_CONV_DIRECT_NONGROUPED_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp index f1ea84c1dec8..8a73188036d7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp @@ -38,63 +38,6 @@ #include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx, has_shape */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join, sig entry */ -/* Reproduce str(KeyError(_build_target message)) for an unknown gfx target: - * - * Python _build_target: raise KeyError( - * f"unknown gfx target {gfx!r}; known: {sorted(specs)}. " - * f"Add a row to {_DATA_FILE.name}.") - * is_valid_spec: except KeyError as e: return False, str(e) - * - * str(KeyError(msg)) == repr(msg); the single quotes make Python DOUBLE-quote - * the whole message. sorted(specs) renders as ['gfx...', 'gfx...']. - * rocke_known_arches() == tuple(sorted(_load_specs())). Mirrors fmha_arch.cpp. */ -static void rocke_dconv__set_unknown_arch_reason(char* out, size_t out_cap, const char* gfx) -{ - int count = 0; - const char* const* arches; - int i; - size_t pos = 0; - int wrote; - - if(out == NULL || out_cap == 0) - { - return; - } - - arches = rocke_known_arches(&count); - - wrote = snprintf(out + pos, out_cap - pos, "\"unknown gfx target '%s'; known: [", gfx); - if(wrote < 0) - { - out[0] = '\0'; - return; - } - pos += (size_t)wrote; - if(pos >= out_cap) - { - out[out_cap - 1] = '\0'; - return; - } - - for(i = 0; i < count; ++i) - { - wrote = snprintf(out + pos, out_cap - pos, "%s'%s'", (i == 0) ? "" : ", ", arches[i]); - if(wrote < 0) - { - out[out_cap - 1] = '\0'; - return; - } - pos += (size_t)wrote; - if(pos >= out_cap) - { - out[out_cap - 1] = '\0'; - return; - } - } - - snprintf(out + pos, out_cap - pos, "]. Add a row to arch_specs.json.\""); -} - /* ===================================================================== * * DirectConvProblem * ===================================================================== */ @@ -361,7 +304,7 @@ bool rocke_direct_conv_16c_is_valid_spec(const rocke_direct_conv_16c_spec_t* spe if(target == NULL) { /* Full Python str(KeyError) text, reproduced verbatim. */ - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } @@ -592,7 +535,7 @@ bool rocke_direct_conv_4c_is_valid_spec(const rocke_direct_conv_4c_spec_t* spec, if(target == NULL) { /* Full Python str(KeyError) text, reproduced verbatim. */ - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } @@ -797,7 +740,7 @@ bool rocke_direct_conv_8c_is_valid_spec(const rocke_direct_conv_8c_spec_t* spec, target = rocke_archtarget_from_gfx(arch); if(target == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; @@ -1025,7 +968,7 @@ bool rocke_direct_conv_32c_is_valid_spec(const rocke_direct_conv_32c_spec_t* spe target = rocke_archtarget_from_gfx(arch); if(target == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; @@ -1227,7 +1170,7 @@ bool rocke_direct_depthwise_is_valid_spec(const rocke_direct_depthwise_spec_t* s } if(rocke_archtarget_from_gfx(arch) == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; @@ -1349,7 +1292,7 @@ bool rocke_direct_depthwise_spatial_is_valid_spec(const rocke_direct_depthwise_s } if(rocke_archtarget_from_gfx(arch) == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; @@ -1567,7 +1510,7 @@ bool rocke_direct_conv_dgrad_is_valid_spec(const rocke_direct_conv_dgrad_spec_t* } if(rocke_archtarget_from_gfx(arch) == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; @@ -1736,7 +1679,7 @@ bool rocke_direct_depthwise_dgrad_is_valid_spec(const rocke_direct_depthwise_dgr } if(rocke_archtarget_from_gfx(arch) == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; @@ -2004,7 +1947,7 @@ bool rocke_direct_conv_wgrad_is_valid_spec(const rocke_direct_conv_wgrad_spec_t* target = rocke_archtarget_from_gfx(arch); if(target == NULL) { - rocke_dconv__set_unknown_arch_reason(reason, reason_cap, arch); + rocke_set_unknown_arch_reason(reason, reason_cap, arch); return false; } p = &spec->problem; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_nongrouped.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_nongrouped.cpp new file mode 100644 index 000000000000..0fec1d6366ee --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_nongrouped.cpp @@ -0,0 +1,1415 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * conv_direct_nongrouped.cpp -- C++ engine port of build_direct_conv_nongrouped + * (library/kernels/common/conv_direct_nongrouped.py): the non-grouped (groups == 1) + * NHWC direct convolution with an LDS-staged, halo-inclusive activation tile. + * + * BYTE-IDENTITY CONTRACT. Every IR op (including each arith.constant) is + * emitted in exactly the order the Python builder emits it. Python evaluates + * call arguments left to right; C++ leaves argument evaluation order + * unspecified, so a nested call such as + * b.add(b.mul(x, b.const_i32(S)), b.const_i32(PAD)) + * is written here as one statement per op, in Python order. Never fold two + * emitting calls into one argument list in this file. + * + * Layout of this TU: + * - spec defaults / derived geometry / kernel_name / validate / is_valid + * (host-side, no IR) + * - the builder, split into the same phases as the Python function body + * (prologue, grid decode, staging metadata, read bases, channel loop, + * epilogue) sharing one file-local context struct + * - the _new / lower_to_llvm convenience entries + */ +#include "rocke/instance_conv_direct_nongrouped.h" + +#include +#include +#include +#include + +#include +#include + +#include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ +#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx, has_shape */ +#include "rocke/helper_rocke.helpers.grid.h" /* rocke_chiplet_aware_super_tile_dynamic */ +#include "rocke/helper_rocke.helpers.io.h" /* rocke_b_io_ir_type */ +#include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/instance_conv_direct_grouped_internal.h" /* rocke_dconv_emit_params */ +#include "rocke/lower_llvm.h" + +/* Global load width for the activation staging path, in halves (dwordx4). */ +#define ROCKE_DCONV_NONGROUPED_X_LOAD_VEC 8 + +/* Canonical i32 byte offset for every masked-out buffer access (staging loads + * and epilogue stores); see the Python `_OOB_BASE` comment. Loop-invariant + * predication is folded into the base offset, and a buffer access past + * num_records is dropped (a load returns zero). */ +#define ROCKE_DCONV_NONGROUPED_OOB_BASE 0x7F000000 + +/* ===================================================================== * + * Atom table (Python _ATOMS: name -> (tile, K, frag)) + * ===================================================================== */ +typedef rocke_value_t* (*rocke_dconv_nongrouped_mfma_fn)(rocke_ir_builder_t*, + rocke_value_t*, + rocke_value_t*, + rocke_value_t*); + +typedef struct rocke_dconv_nongrouped_atom +{ + const char* name; + int tile; /* MFMA M == N */ + int k; /* MFMA K */ + int frag; /* halves per lane in an A/B fragment */ + rocke_dconv_nongrouped_mfma_fn f16; + rocke_dconv_nongrouped_mfma_fn bf16; +} rocke_dconv_nongrouped_atom_t; + +static const rocke_dconv_nongrouped_atom_t k_nongrouped_atoms[] = { + {"32x32x16", 32, 16, 8, rocke_b_mfma_f32_32x32x16_f16, rocke_b_mfma_f32_32x32x16_bf16}, + {"32x32x8", 32, 8, 4, rocke_b_mfma_f32_32x32x8_f16, rocke_b_mfma_f32_32x32x8_bf16}, + {"16x16x32", 16, 32, 8, rocke_b_mfma_f32_16x16x32_f16, rocke_b_mfma_f32_16x16x32_bf16}, + {"16x16x16", 16, 16, 4, rocke_b_mfma_f32_16x16x16_f16, rocke_b_mfma_f32_16x16x16_bf16}, +}; + +static const rocke_dconv_nongrouped_atom_t* nongrouped_atom(const char* name) +{ + size_t i; + if(name == NULL) + { + return NULL; + } + for(i = 0; i < sizeof(k_nongrouped_atoms) / sizeof(k_nongrouped_atoms[0]); ++i) + { + if(strcmp(k_nongrouped_atoms[i].name, name) == 0) + { + return &k_nongrouped_atoms[i]; + } + } + return NULL; +} + +static const char* nongrouped_dtype(const rocke_direct_conv_nongrouped_spec_t* spec) +{ + return spec->problem.dtype ? spec->problem.dtype : "fp16"; +} + +static bool nongrouped_is_bf16(const rocke_direct_conv_nongrouped_spec_t* spec) +{ + return strcmp(nongrouped_dtype(spec), "bf16") == 0; +} + +static int nongrouped_ho(const rocke_direct_conv_problem_t* p) +{ + return (p->H + 2 * p->PAD - p->KH) / p->stride + 1; +} + +static int nongrouped_wo(const rocke_direct_conv_problem_t* p) +{ + return (p->W + 2 * p->PAD - p->KW) / p->stride + 1; +} + +static void nongrouped_set_reason(char* reason, size_t reason_cap, const char* msg) +{ + if(reason && reason_cap > 0) + { + snprintf(reason, reason_cap, "%s", msg); + } +} + +/* ===================================================================== * + * Spec defaults + derived geometry + * ===================================================================== */ + +rocke_direct_conv_nongrouped_spec_t rocke_direct_conv_nongrouped_spec_default(void) +{ + rocke_direct_conv_nongrouped_spec_t spec; + memset(&spec, 0, sizeof(spec)); + spec.problem = rocke_direct_conv_problem_default(); + spec.problem.groups = 1; + spec.name = "direct_conv_nongrouped"; + spec.tile_h = 16; + spec.tile_w = 32; + spec.tile_k = 128; + spec.ck = 16; + spec.waves_m = 2; + spec.waves_n = 4; + spec.atom = "32x32x16"; + spec.wave_size = 64; + spec.lds_pad = 8; + spec.chiplet_swizzle = true; + spec.swizzle_wgm = 8; + spec.chiplet_chunk = 64; + spec.num_xcds = 8; + spec.double_buffer = false; + spec.iglp = ROCKE_DCONV_NONGROUPED_IGLP_NONE; + spec.waves_per_eu = ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE; + return spec; +} + +int rocke_direct_conv_nongrouped_threads_per_block(const rocke_direct_conv_nongrouped_spec_t* spec) +{ + return spec->waves_m * spec->waves_n * spec->wave_size; +} + +long rocke_direct_conv_nongrouped_lds_bytes(const rocke_direct_conv_nongrouped_spec_t* spec) +{ + const rocke_dconv_nongrouped_atom_t* at = nongrouped_atom(spec->atom); + const rocke_direct_conv_problem_t* p = &spec->problem; + long lds_in_h, lds_in_w, x_halves, w_halves, one; + if(at == NULL) + { + return 0; + } + lds_in_h = (long)(spec->tile_h - 1) * p->stride + p->KH; + lds_in_w = (long)(spec->tile_w - 1) * p->stride + p->KW; + x_halves = lds_in_h * lds_in_w * (spec->ck + spec->lds_pad); + w_halves = (long)spec->tile_k * p->KH * p->KW * spec->ck; + /* + one scratch fragment per array (see the builder). */ + one = x_halves + ROCKE_DCONV_NONGROUPED_X_LOAD_VEC + w_halves + at->frag; + return 2 * one * (spec->double_buffer ? 2 : 1); +} + +int rocke_direct_conv_nongrouped_acc_vgprs(const rocke_direct_conv_nongrouped_spec_t* spec) +{ + const rocke_dconv_nongrouped_atom_t* at = nongrouped_atom(spec->atom); + int acc_per_lane, m_tiles_per_wave, n_tiles_per_wave; + if(at == NULL || spec->waves_m <= 0 || spec->waves_n <= 0 || spec->wave_size <= 0) + { + return 0; + } + acc_per_lane = at->tile * at->tile / spec->wave_size; + m_tiles_per_wave = (spec->tile_k / at->tile) / spec->waves_m; + n_tiles_per_wave = (spec->tile_h / spec->waves_n) * (spec->tile_w / at->tile); + return m_tiles_per_wave * n_tiles_per_wave * acc_per_lane; +} + +static void nongrouped_tile_counts(const rocke_direct_conv_nongrouped_spec_t* spec, + int* n_wt, + int* n_ht, + int* n_kt) +{ + const rocke_direct_conv_problem_t* p = &spec->problem; + *n_wt = (nongrouped_wo(p) + spec->tile_w - 1) / spec->tile_w; + *n_ht = (nongrouped_ho(p) + spec->tile_h - 1) / spec->tile_h; + *n_kt = (p->kpg + spec->tile_k - 1) / spec->tile_k; +} + +void rocke_direct_conv_nongrouped_grid(const rocke_direct_conv_nongrouped_spec_t* spec, int grid[3]) +{ + int n_wt, n_ht, n_kt; + nongrouped_tile_counts(spec, &n_wt, &n_ht, &n_kt); + grid[0] = n_wt * n_ht * spec->problem.N * n_kt; + grid[1] = 1; + grid[2] = 1; +} + +/* ===================================================================== * + * kernel_name() + * ===================================================================== */ +rocke_status_t rocke_direct_conv_nongrouped_kernel_name( + const rocke_direct_conv_nongrouped_spec_t* spec, char* out, size_t out_cap) +{ + char prob_short[128]; + char t_buf[64]; + char ck_buf[24]; + char w_buf[32]; + char a_buf[48]; + char g_buf[24]; + char iglp_buf[24]; + char we_buf[24]; + const char* parts[10]; + + if(spec == NULL || out == NULL || out_cap == 0) + { + return ROCKE_ERR_VALUE; + } + if(rocke_direct_conv_problem_short(&spec->problem, prob_short, sizeof(prob_short)) != ROCKE_OK) + { + return ROCKE_ERR_VALUE; + } + snprintf(t_buf, sizeof(t_buf), "t%dx%dx%d", spec->tile_h, spec->tile_w, spec->tile_k); + snprintf(ck_buf, sizeof(ck_buf), "ck%d", spec->ck); + snprintf(w_buf, sizeof(w_buf), "w%dx%d", spec->waves_m, spec->waves_n); + snprintf(a_buf, sizeof(a_buf), "a%s", spec->atom ? spec->atom : ""); + if(spec->chiplet_swizzle) + { + snprintf(g_buf, sizeof(g_buf), "g%d", spec->swizzle_wgm); + } + else + { + snprintf(g_buf, sizeof(g_buf), "gnone"); + } + iglp_buf[0] = '\0'; + if(spec->iglp != ROCKE_DCONV_NONGROUPED_IGLP_NONE) + { + snprintf(iglp_buf, sizeof(iglp_buf), "iglp%d", spec->iglp); + } + we_buf[0] = '\0'; + if(spec->waves_per_eu != ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE) + { + snprintf(we_buf, sizeof(we_buf), "we%d", spec->waves_per_eu); + } + parts[0] = prob_short; + parts[1] = t_buf; + parts[2] = ck_buf; + parts[3] = w_buf; + parts[4] = a_buf; + parts[5] = g_buf; + parts[6] = spec->double_buffer ? "db" : ""; + parts[7] = iglp_buf; + parts[8] = we_buf; + parts[9] = nongrouped_is_bf16(spec) ? "bf16" : ""; + return rocke_kernel_name_join(spec->name, parts, 10, NULL, NULL, 0, out, out_cap, NULL); +} + +/* ===================================================================== * + * validate() (Python DirectNongroupedConvSpec.validate, same check order) + * ===================================================================== */ +rocke_status_t rocke_direct_conv_nongrouped_validate( + const rocke_direct_conv_nongrouped_spec_t* spec, char* reason, size_t reason_cap) +{ + const rocke_direct_conv_problem_t* p; + const rocke_dconv_nongrouped_atom_t* at; + char msg[ROCKE_ERR_MSG_CAP]; + int t; + + if(spec == NULL) + { + nongrouped_set_reason(reason, reason_cap, "null spec"); + return ROCKE_ERR_VALUE; + } + p = &spec->problem; + +#define NONGROUPED_REJECT(...) \ + do \ + { \ + snprintf(msg, sizeof(msg), __VA_ARGS__); \ + nongrouped_set_reason(reason, reason_cap, msg); \ + return ROCKE_ERR_VALUE; \ + } while(0) + + if(strcmp(nongrouped_dtype(spec), "fp16") != 0 && strcmp(nongrouped_dtype(spec), "bf16") != 0) + { + NONGROUPED_REJECT("DirectNongroupedConvSpec: unsupported dtype '%s'", + nongrouped_dtype(spec)); + } + if(p->groups != 1) + { + NONGROUPED_REJECT("DirectNongroupedConvSpec is the groups==1 family (got groups=%d); " + "use conv_direct_grouped for grouped shapes", + p->groups); + } + at = nongrouped_atom(spec->atom); + if(at == NULL) + { + NONGROUPED_REJECT("unknown atom '%s'; expected one of ['32x32x16', '32x32x8', '16x16x32', " + "'16x16x16']", + spec->atom ? spec->atom : "None"); + } + /* Checked before any derived geometry: those divide by these. A positive C + * is also what lets the kernel assume its channel loop runs at least once. */ + if(spec->waves_m <= 0 || spec->waves_n <= 0 || spec->wave_size <= 0 || spec->ck <= 0 + || spec->tile_h <= 0 || spec->tile_w <= 0 || spec->tile_k <= 0 || p->stride <= 0 + || p->cpg <= 0 || p->kpg <= 0) + { + NONGROUPED_REJECT("DirectNongroupedConvSpec: tile, wave, stride and channel parameters " + "must be positive"); + } + t = at->tile; + if(spec->tile_w % t != 0) + { + NONGROUPED_REJECT("tile_w must be a multiple of %d (got %d)", t, spec->tile_w); + } + if(spec->tile_k % t != 0) + { + NONGROUPED_REJECT("tile_k must be a multiple of %d (got %d)", t, spec->tile_k); + } + if(spec->tile_h % spec->waves_n != 0) + { + NONGROUPED_REJECT("tile_h %d not divisible by waves_n %d", spec->tile_h, spec->waves_n); + } + if((spec->tile_k / t) % spec->waves_m != 0) + { + NONGROUPED_REJECT( + "tile_k/%d = %d not divisible by waves_m %d", t, spec->tile_k / t, spec->waves_m); + } + if(spec->ck % at->k != 0) + { + NONGROUPED_REJECT("ck %d not divisible by atom K %d", spec->ck, at->k); + } + if(spec->ck % ROCKE_DCONV_NONGROUPED_X_LOAD_VEC != 0) + { + NONGROUPED_REJECT("ck %d not divisible by %d", spec->ck, ROCKE_DCONV_NONGROUPED_X_LOAD_VEC); + } + if(p->cpg % spec->ck != 0) + { + NONGROUPED_REJECT("C %d not divisible by ck %d", p->cpg, spec->ck); + } + if(p->kpg % t != 0) + { + NONGROUPED_REJECT("K %d must be a multiple of %d (got %d)", p->kpg, t, p->kpg); + } + if(p->cpg % ROCKE_DCONV_NONGROUPED_X_LOAD_VEC != 0) + { + NONGROUPED_REJECT( + "C %d must be a multiple of %d", p->cpg, ROCKE_DCONV_NONGROUPED_X_LOAD_VEC); + } + if(rocke_direct_conv_nongrouped_threads_per_block(spec) > 1024) + { + NONGROUPED_REJECT("threads_per_block %d > 1024", + rocke_direct_conv_nongrouped_threads_per_block(spec)); + } + if(spec->lds_pad < 0 || spec->lds_pad % 8 != 0) + { + /* A negative pad shrinks the pixel stride below the ck halves each + * pixel stages, so neighbouring pixels overlap in LDS. The vec8 LDS + * store/load is emitted with align 16, so the pixel stride + * (ck + lds_pad halves) must be a multiple of 8 halves. */ + NONGROUPED_REJECT( + "lds_pad must be a non-negative multiple of 8 to keep ds_read_b128 aligned (got %d)", + spec->lds_pad); + } + if(spec->swizzle_wgm < 1) + { + NONGROUPED_REJECT("swizzle_wgm must be >= 1 (got %d)", spec->swizzle_wgm); + } + /* -1 / 0 are this port's "None" sentinels; anything below them has no + * Python spelling, and Python rejects the sentinel values themselves. */ + if(spec->iglp < ROCKE_DCONV_NONGROUPED_IGLP_NONE) + { + NONGROUPED_REJECT("iglp must be None or >= 0 (got %d)", spec->iglp); + } + if(spec->waves_per_eu < ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE) + { + NONGROUPED_REJECT("waves_per_eu must be None or >= 1 (got %d)", spec->waves_per_eu); + } + { + /* The kernel decodes its tile from a flat 1-D block id and computes the + * workgroup count in i32; past gridDim.x's 2**31 - 1 limit both the + * launch and that product would overflow. (The D-size bound below + * implies this one; it is checked on its own so the reason names the + * real limit.) */ + int n_wt, n_ht, n_kt; + long long grid_x; + nongrouped_tile_counts(spec, &n_wt, &n_ht, &n_kt); + grid_x = (long long)n_wt * n_ht * p->N * n_kt; + if(grid_x > INT32_MAX) + { + NONGROUPED_REJECT("flattened grid of %lld workgroups exceeds %d: grow the tile or " + "split the batch", + grid_x, + INT32_MAX); + } + } + { + /* Masked staging loads and epilogue stores go to OOB_BASE and rely on + * it being past num_records; a larger tensor would hand the loads real + * data and take the stores. */ + long long a_bytes = 2LL * p->N * p->H * p->W * p->cpg; + long long b_bytes = 2LL * p->kpg * p->KH * p->KW * p->cpg; + long long d_bytes = 2LL * p->N * nongrouped_ho(p) * nongrouped_wo(p) * p->kpg; + long long biggest = a_bytes > b_bytes ? a_bytes : b_bytes; + biggest = d_bytes > biggest ? d_bytes : biggest; + if(biggest > ROCKE_DCONV_NONGROUPED_OOB_BASE) + { + NONGROUPED_REJECT("tensors must fit below the %#x-byte masked-access offset (A %lld B, " + "B %lld B, D %lld B)", + ROCKE_DCONV_NONGROUPED_OOB_BASE, + a_bytes, + b_bytes, + d_bytes); + } + /* The last prefetch runs one chunk (two when double-buffered) past C. */ + if((long long)ROCKE_DCONV_NONGROUPED_OOB_BASE + 2LL * (p->cpg + spec->ck) > INT32_MAX) + { + NONGROUPED_REJECT("C %d too large: the masked-load offset plus the channel offset " + "would overflow i32", + p->cpg); + } + } + if(rocke_direct_conv_nongrouped_acc_vgprs(spec) > 256) + { + /* Past the 256-entry accumulator file the backend scheduler does not + * just slow down: a 2048-accumulator tile hung the compiler. */ + NONGROUPED_REJECT("accumulator tile needs %d registers per lane (max 256): shrink " + "tile_k/tile_h/tile_w or add waves", + rocke_direct_conv_nongrouped_acc_vgprs(spec)); + } +#undef NONGROUPED_REJECT + + nongrouped_set_reason(reason, reason_cap, "ok"); + return ROCKE_OK; +} + +/* ===================================================================== * + * is_valid_nongrouped_spec(spec, arch) + * ===================================================================== */ +bool rocke_direct_conv_nongrouped_is_valid_spec(const rocke_direct_conv_nongrouped_spec_t* spec, + const char* arch, + char* reason, + size_t reason_cap) +{ + const rocke_archtarget_t* target; + const rocke_direct_conv_problem_t* p; + const rocke_dconv_nongrouped_atom_t* at; + const char* ab; + char msg[ROCKE_ERR_MSG_CAP]; + + if(spec == NULL) + { + nongrouped_set_reason(reason, reason_cap, "null spec"); + return false; + } + if(arch == NULL) + { + arch = "gfx950"; + } + target = rocke_archtarget_from_gfx(arch); + if(target == NULL) + { + /* Same text as Python's str(KeyError) from ArchTarget.from_gfx. */ + rocke_set_unknown_arch_reason(reason, reason_cap, arch); + return false; + } + if(rocke_direct_conv_nongrouped_validate(spec, reason, reason_cap) != ROCKE_OK) + { + return false; + } + p = &spec->problem; + if(p->stride != 1 && p->stride != 2) + { + snprintf(msg, sizeof(msg), "stride %d is not supported (expected 1 or 2)", p->stride); + nongrouped_set_reason(reason, reason_cap, msg); + return false; + } + at = nongrouped_atom(spec->atom); + ab = nongrouped_is_bf16(spec) ? "bf16" : "f16"; + if(!rocke_mma_catalog_has_shape( + rocke_archtarget_mma(target), "mma", ab, ab, "fp32", at->tile, at->tile, at->k)) + { + snprintf(msg, sizeof(msg), "missing mfma_f32_%s_%s on %s", at->name, ab, arch); + nongrouped_set_reason(reason, reason_cap, msg); + return false; + } + if(spec->wave_size != target->wave_size) + { + snprintf(msg, + sizeof(msg), + "wave_size %d != %s wave %d", + spec->wave_size, + arch, + target->wave_size); + nongrouped_set_reason(reason, reason_cap, msg); + return false; + } + if(rocke_direct_conv_nongrouped_threads_per_block(spec) + > rocke_archtarget_max_threads_per_block(target)) + { + snprintf(msg, + sizeof(msg), + "threads_per_block %d exceeds arch limit", + rocke_direct_conv_nongrouped_threads_per_block(spec)); + nongrouped_set_reason(reason, reason_cap, msg); + return false; + } + if(!rocke_archtarget_fits_lds(target, rocke_direct_conv_nongrouped_lds_bytes(spec))) + { + snprintf(msg, + sizeof(msg), + "LDS %ld B exceeds %d B", + rocke_direct_conv_nongrouped_lds_bytes(spec), + target->lds_capacity_bytes); + nongrouped_set_reason(reason, reason_cap, msg); + return false; + } + nongrouped_set_reason(reason, reason_cap, "ok"); + return true; +} + +/* ===================================================================== * + * Builder context (the locals the Python closures share) + * ===================================================================== */ +namespace +{ + +struct nongrouped_stage_meta +{ + rocke_value_t* base; /* global byte offset base (chunk offset added per iter) */ + rocke_value_t* lds_idx; /* LDS half index of this thread's slot */ +}; + +struct nongrouped_ctx +{ + rocke_ir_builder_t* b; + const rocke_direct_conv_nongrouped_spec_t* spec; + rocke_direct_conv_problem_t p; + const rocke_type_t* io_type; + bool bf16; + rocke_dconv_nongrouped_mfma_fn mfma; + + /* Build-time geometry (Python all-caps locals). */ + int KH, KW, S, PAD; + int TH, TW, TK, CK, THREADS, WAVE, FRAG, AK, AT, ACC, QUADS; + int NCB, M_TILES, KATOMS, ROWS_W, MT_W, NT_W, CSTRIDE, LDS_IN_H, LDS_IN_W; + int X_CV, X_VECS, X_PASSES, W_SLOTS, W_PASSES; + int x_stage, w_stage, x_dump, w_dump; + bool DB; + + /* The AOT kernarg block. groups == 1, so p_total_c / p_total_k (C / K + * below) are the full channel counts. */ + rocke_dconv_params_t params; + rocke_value_t *C, *K; + + /* SSA values. */ + rocke_value_t *c0, *c1, *c_half, *oob_base; + rocke_value_t *a_rsrc, *b_rsrc, *d_rsrc; + rocke_value_t *X_smem, *W_smem; + rocke_value_t *tid, *lane, *wave_id, *wave_m, *wave_n, *lane_lo, *lane_hi; + rocke_value_t *k_tile, *cell, *n_img, *out_h0, *out_w0, *k_base, *in_h0, *in_w0; + rocke_value_t *x_read_base, *w_read_base; + + std::vector x_meta; + std::vector w_meta; +}; + +rocke_value_t* nongrouped_c(nongrouped_ctx& c, int64_t v) +{ + return rocke_b_const_i32(c.b, v); +} + +rocke_value_t* + nongrouped_buf_load(nongrouped_ctx& c, rocke_value_t* rsrc, rocke_value_t* voff, int dwords) +{ + if(c.bf16) + { + return rocke_b_buffer_load_vN_bf16(c.b, rsrc, voff, c.c0, dwords); + } + return rocke_b_buffer_load_vN_f16(c.b, rsrc, voff, c.c0, dwords); +} + +/* ---- prologue: attrs, params, constants, rsrcs, LDS, thread decode ---- */ +bool nongrouped_prologue(nongrouped_ctx& c, const char* arch) +{ + rocke_ir_builder_t* b = c.b; + const rocke_direct_conv_nongrouped_spec_t* spec = c.spec; + const rocke_dconv_nongrouped_atom_t* at; + char reason[ROCKE_ERR_MSG_CAP]; + + if(!rocke_direct_conv_nongrouped_is_valid_spec(spec, arch, reason, sizeof reason)) + { + if(b->status == ROCKE_OK) + { + b->status = ROCKE_ERR_VALUE; + std::string msg + = std::string("invalid DirectNongroupedConvSpec for ") + arch + ": " + reason; + snprintf(b->err, sizeof b->err, "%.*s", (int)(sizeof b->err - 1), msg.c_str()); + } + return false; + } + + c.p = spec->problem; + c.io_type = rocke_b_io_ir_type(b, nongrouped_dtype(spec)); + if(c.io_type == NULL) + { + return false; + } + c.bf16 = nongrouped_is_bf16(spec); + at = nongrouped_atom(spec->atom); + c.mfma = c.bf16 ? at->bf16 : at->f16; + + c.KH = c.p.KH; + c.KW = c.p.KW; + c.S = c.p.stride; + c.PAD = c.p.PAD; + + c.TH = spec->tile_h; + c.TW = spec->tile_w; + c.TK = spec->tile_k; + c.CK = spec->ck; + c.THREADS = rocke_direct_conv_nongrouped_threads_per_block(spec); + c.WAVE = spec->wave_size; + c.FRAG = at->frag; + c.AK = at->k; + c.AT = at->tile; + c.ACC = c.AT * c.AT / c.WAVE; + c.QUADS = c.ACC / 4; + c.NCB = c.TW / c.AT; + c.M_TILES = c.TK / c.AT; + c.KATOMS = c.CK / c.AK; + c.ROWS_W = c.TH / spec->waves_n; + c.MT_W = c.M_TILES / spec->waves_m; + c.NT_W = c.ROWS_W * c.NCB; + c.CSTRIDE = c.CK + spec->lds_pad; + c.LDS_IN_H = (c.TH - 1) * c.S + c.KH; + c.LDS_IN_W = (c.TW - 1) * c.S + c.KW; + + c.X_CV = c.CK / ROCKE_DCONV_NONGROUPED_X_LOAD_VEC; + c.X_VECS = c.LDS_IN_H * c.LDS_IN_W * c.X_CV; + c.X_PASSES = (c.X_VECS + c.THREADS - 1) / c.THREADS; + c.W_SLOTS = c.KH * c.KW * c.M_TILES * c.KATOMS * c.WAVE; + c.W_PASSES = (c.W_SLOTS + c.THREADS - 1) / c.THREADS; + + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", c.THREADS); + if(spec->waves_per_eu != ROCKE_DCONV_NONGROUPED_WAVES_PER_EU_NONE) + { + rocke_attr_set_int(b, &b->kernel->attrs, "waves_per_eu", spec->waves_per_eu); + } + + rocke_dconv_emit_params(b, &c.params, "fwd", c.io_type); + if(!rocke_ir_builder_ok(b)) + { + return false; + } + c.C = c.params.p_total_c; + c.K = c.params.p_total_k; + + c.c0 = nongrouped_c(c, 0); + c.c1 = nongrouped_c(c, 1); + c.c_half = nongrouped_c(c, 2); + /* Shared by the masked staging loads and the masked epilogue stores. */ + c.oob_base = nongrouped_c(c, ROCKE_DCONV_NONGROUPED_OOB_BASE); + + c.a_rsrc = rocke_b_buffer_rsrc(b, c.params.A, c.params.A_bytes); + c.b_rsrc = rocke_b_buffer_rsrc(b, c.params.Bp, c.params.B_bytes); + c.d_rsrc = rocke_b_buffer_rsrc(b, c.params.D, c.params.D_bytes); + + /* The staging loops are sized in whole thread-passes, so the last pass can + * own slots past the end of the tile; those write into a scratch tail. */ + { + long x_halves = (long)c.LDS_IN_H * c.LDS_IN_W * c.CSTRIDE; + long w_halves = (long)c.TK * c.KH * c.KW * c.CK; + int nbuf; + int shape[2]; + c.DB = spec->double_buffer; + c.x_stage = (int)(x_halves + ROCKE_DCONV_NONGROUPED_X_LOAD_VEC); + c.w_stage = (int)(w_halves + c.FRAG); + c.x_dump = (int)x_halves; + c.w_dump = (int)w_halves; + nbuf = c.DB ? 2 : 1; + shape[0] = 1; + shape[1] = c.x_stage * nbuf; + c.X_smem = rocke_b_smem_alloc(b, c.io_type, shape, 2, "lds_x"); + shape[1] = c.w_stage * nbuf; + c.W_smem = rocke_b_smem_alloc(b, c.io_type, shape, 2, "lds_w"); + } + + /* Thread / wave decomposition. */ + { + rocke_value_t* k; + c.tid = rocke_b_thread_id_x(b); + k = nongrouped_c(c, c.WAVE); + c.lane = rocke_b_mod(b, c.tid, k); + k = nongrouped_c(c, c.WAVE); + c.wave_id = rocke_b_div(b, c.tid, k); + k = nongrouped_c(c, spec->waves_n); + c.wave_m = rocke_b_div(b, c.wave_id, k); + k = nongrouped_c(c, spec->waves_n); + c.wave_n = rocke_b_mod(b, c.wave_id, k); + k = nongrouped_c(c, c.AT); + c.lane_lo = rocke_b_mod(b, c.lane, k); + k = nongrouped_c(c, c.AT); + c.lane_hi = rocke_b_div(b, c.lane, k); + } + return rocke_ir_builder_ok(b); +} + +/* ---- grid decode: flat 1-D grid of (spatial cell x channel tile) ---- + * The tile sizes are build-time; the tile counts follow the runtime extents. */ +void nongrouped_grid_decode(nongrouped_ctx& c) +{ + rocke_ir_builder_t* b = c.b; + const rocke_direct_conv_nongrouped_spec_t* spec = c.spec; + rocke_value_t *wgid, *k, *t, *n_wt, *n_ht, *n_kt, *n_hw, *n_cells, *st, *w_tile, *h_tile; + + k = nongrouped_c(c, c.TW - 1); + t = rocke_b_add(b, c.params.p_Wo, k); + k = nongrouped_c(c, c.TW); + n_wt = rocke_b_div(b, t, k); + k = nongrouped_c(c, c.TH - 1); + t = rocke_b_add(b, c.params.p_Ho, k); + k = nongrouped_c(c, c.TH); + n_ht = rocke_b_div(b, t, k); + k = nongrouped_c(c, c.TK - 1); + t = rocke_b_add(b, c.K, k); + k = nongrouped_c(c, c.TK); + n_kt = rocke_b_div(b, t, k); + n_hw = rocke_b_mul(b, n_wt, n_ht); + n_cells = rocke_b_mul(b, n_hw, c.params.p_N); + + wgid = rocke_b_block_id_x(b); + if(spec->chiplet_swizzle) + { + /* A swizzle_wgm above the runtime channel-tile count needs no clamp: + * the swizzle sizes its last group as min(wgm, tiles left). */ + rocke_super_tile_swizzle_result_t sw = rocke_chiplet_aware_super_tile_dynamic( + b, wgid, n_kt, n_cells, spec->swizzle_wgm, spec->num_xcds, spec->chiplet_chunk); + c.k_tile = sw.row; + c.cell = sw.col; + } + else + { + c.k_tile = rocke_b_mod(b, wgid, n_kt); + c.cell = rocke_b_div(b, wgid, n_kt); + } + + c.n_img = rocke_b_div(b, c.cell, n_hw); + st = rocke_b_mod(b, c.cell, n_hw); + w_tile = rocke_b_mod(b, st, n_wt); + h_tile = rocke_b_div(b, st, n_wt); + + /* Output-tile origins. */ + k = nongrouped_c(c, c.TH); + c.out_h0 = rocke_b_mul(b, h_tile, k); + k = nongrouped_c(c, c.TW); + c.out_w0 = rocke_b_mul(b, w_tile, k); + k = nongrouped_c(c, c.TK); + c.k_base = rocke_b_mul(b, c.k_tile, k); + + /* Input origin of the staged (halo-inclusive) activation tile. */ + k = nongrouped_c(c, c.S); + t = rocke_b_mul(b, c.out_h0, k); + k = nongrouped_c(c, c.PAD); + c.in_h0 = rocke_b_sub(b, t, k); + k = nongrouped_c(c, c.S); + t = rocke_b_mul(b, c.out_w0, k); + k = nongrouped_c(c, c.PAD); + c.in_w0 = rocke_b_sub(b, t, k); +} + +/* ---- _x_pass_meta(): per-pass (base, lds_idx) for the activation tile ---- */ +void nongrouped_x_pass_meta(nongrouped_ctx& c) +{ + rocke_ir_builder_t* b = c.b; + rocke_value_t *k, *img_base; + int j; + + img_base = rocke_b_mul(b, c.n_img, c.params.p_A_stride_n); + for(j = 0; j < c.X_PASSES; ++j) + { + rocke_value_t *v, *cv, *pos, *ih_l, *iw_l, *ih, *iw; + rocke_value_t *ge, *lt, *ok_h, *ok_w, *valid, *in_tile = NULL; + rocke_value_t *t0, *t1, *t2, *elems, *lds_idx, *base; + + k = nongrouped_c(c, (int64_t)j * c.THREADS); + v = rocke_b_add(b, c.tid, k); + k = nongrouped_c(c, c.X_CV); + cv = rocke_b_mod(b, v, k); + k = nongrouped_c(c, c.X_CV); + pos = rocke_b_div(b, v, k); + k = nongrouped_c(c, c.LDS_IN_W); + ih_l = rocke_b_div(b, pos, k); + k = nongrouped_c(c, c.LDS_IN_W); + iw_l = rocke_b_mod(b, pos, k); + ih = rocke_b_add(b, c.in_h0, ih_l); + iw = rocke_b_add(b, c.in_w0, iw_l); + + ge = rocke_b_cmp_ge(b, ih, c.c0); + lt = rocke_b_cmp_lt(b, ih, c.params.p_Hi); + ok_h = rocke_b_land(b, ge, lt); + ge = rocke_b_cmp_ge(b, iw, c.c0); + lt = rocke_b_cmp_lt(b, iw, c.params.p_Wi); + ok_w = rocke_b_land(b, ge, lt); + valid = rocke_b_land(b, ok_h, ok_w); + if(c.X_VECS % c.THREADS != 0) + { + k = nongrouped_c(c, c.X_VECS); + in_tile = rocke_b_cmp_lt(b, v, k); + valid = rocke_b_land(b, valid, in_tile); + } + + /* base = n*stride_n + ih*stride_hi + iw*stride_wi + cv*VEC; the chunk + * offset is added per iteration. */ + t0 = rocke_b_mul(b, ih, c.params.p_A_stride_hi); + t1 = rocke_b_mul(b, iw, c.params.p_A_stride_wi); + t0 = rocke_b_add(b, t0, t1); + t1 = rocke_b_add(b, img_base, t0); + k = nongrouped_c(c, ROCKE_DCONV_NONGROUPED_X_LOAD_VEC); + t2 = rocke_b_mul(b, cv, k); + elems = rocke_b_add(b, t1, t2); + + k = nongrouped_c(c, c.CSTRIDE); + t0 = rocke_b_mul(b, pos, k); + k = nongrouped_c(c, ROCKE_DCONV_NONGROUPED_X_LOAD_VEC); + t1 = rocke_b_mul(b, cv, k); + lds_idx = rocke_b_add(b, t0, t1); + if(in_tile != NULL) + { + k = nongrouped_c(c, c.x_dump); + lds_idx = rocke_b_select(b, in_tile, lds_idx, k); + } + t0 = rocke_b_mul(b, elems, c.c_half); + base = rocke_b_select(b, valid, t0, c.oob_base); + c.x_meta.push_back({base, lds_idx}); + } +} + +/* ---- _w_pass_meta(): per-pass (base, lds_idx) for the fragment-order W ---- */ +void nongrouped_w_pass_meta(nongrouped_ctx& c) +{ + rocke_ir_builder_t* b = c.b; + int j; + + for(j = 0; j < c.W_PASSES; ++j) + { + rocke_value_t *k, *g, *g_lane, *rest, *katom, *rest2, *m_tile, *tap, *r, *s; + rocke_value_t *t0, *t1, *k_out, *c_in_chunk, *valid, *in_tile = NULL; + rocke_value_t *krs_h, *krs_hr, *krs_w, *krs, *elems, *lds_idx, *base; + + k = nongrouped_c(c, (int64_t)j * c.THREADS); + g = rocke_b_add(b, c.tid, k); + k = nongrouped_c(c, c.WAVE); + g_lane = rocke_b_mod(b, g, k); + k = nongrouped_c(c, c.WAVE); + rest = rocke_b_div(b, g, k); + k = nongrouped_c(c, c.KATOMS); + katom = rocke_b_mod(b, rest, k); + k = nongrouped_c(c, c.KATOMS); + rest2 = rocke_b_div(b, rest, k); + k = nongrouped_c(c, c.M_TILES); + m_tile = rocke_b_mod(b, rest2, k); + k = nongrouped_c(c, c.M_TILES); + tap = rocke_b_div(b, rest2, k); + k = nongrouped_c(c, c.KW); + r = rocke_b_div(b, tap, k); + k = nongrouped_c(c, c.KW); + s = rocke_b_mod(b, tap, k); + + k = nongrouped_c(c, c.AT); + t0 = rocke_b_mul(b, m_tile, k); + k = nongrouped_c(c, c.AT); + t1 = rocke_b_mod(b, g_lane, k); + t0 = rocke_b_add(b, t0, t1); + k_out = rocke_b_add(b, c.k_base, t0); + + k = nongrouped_c(c, c.AK); + t0 = rocke_b_mul(b, katom, k); + k = nongrouped_c(c, c.AT); + t1 = rocke_b_div(b, g_lane, k); + k = nongrouped_c(c, c.FRAG); + t1 = rocke_b_mul(b, t1, k); + c_in_chunk = rocke_b_add(b, t0, t1); + + valid = rocke_b_cmp_lt(b, k_out, c.K); + if(c.W_SLOTS % c.THREADS != 0) + { + k = nongrouped_c(c, c.W_SLOTS); + in_tile = rocke_b_cmp_lt(b, g, k); + valid = rocke_b_land(b, valid, in_tile); + } + + /* base = ((k_out*KH + r)*KW + s)*C + c_in_chunk (KRSC) */ + k = nongrouped_c(c, c.KH); + krs_h = rocke_b_mul(b, k_out, k); + krs_hr = rocke_b_add(b, krs_h, r); + k = nongrouped_c(c, c.KW); + krs_w = rocke_b_mul(b, krs_hr, k); + krs = rocke_b_add(b, krs_w, s); + t0 = rocke_b_mul(b, krs, c.C); + elems = rocke_b_add(b, t0, c_in_chunk); + k = nongrouped_c(c, c.FRAG); + lds_idx = rocke_b_mul(b, g, k); + if(in_tile != NULL) + { + k = nongrouped_c(c, c.w_dump); + lds_idx = rocke_b_select(b, in_tile, lds_idx, k); + } + t0 = rocke_b_mul(b, elems, c.c_half); + base = rocke_b_select(b, valid, t0, c.oob_base); + c.w_meta.push_back({base, lds_idx}); + } +} + +/* ---- issue_stage_loads(c_off): global loads for one channel chunk ---- */ +void nongrouped_issue_stage_loads(nongrouped_ctx& c, + rocke_value_t* c_off, + std::vector& xs, + std::vector& ws) +{ + xs.clear(); + ws.clear(); + for(const nongrouped_stage_meta& m : c.x_meta) + { + rocke_value_t* voff = rocke_b_add(c.b, m.base, c_off); + xs.push_back(nongrouped_buf_load(c, c.a_rsrc, voff, ROCKE_DCONV_NONGROUPED_X_LOAD_VEC / 2)); + } + for(const nongrouped_stage_meta& m : c.w_meta) + { + rocke_value_t* voff = rocke_b_add(c.b, m.base, c_off); + ws.push_back(nongrouped_buf_load(c, c.b_rsrc, voff, c.FRAG / 2)); + } +} + +/* ---- commit_stage(xs, ws, x_buf, w_buf): publish one chunk into LDS ---- */ +void nongrouped_commit_stage(nongrouped_ctx& c, + const std::vector& xs, + const std::vector& ws, + rocke_value_t* x_buf, + rocke_value_t* w_buf) +{ + size_t i; + for(i = 0; i < c.x_meta.size(); ++i) + { + rocke_value_t* idx = c.x_meta[i].lds_idx; + rocke_value_t* ind[2]; + if(x_buf != NULL) + { + idx = rocke_b_add(c.b, idx, x_buf); + } + ind[0] = c.c0; + ind[1] = idx; + rocke_b_smem_store_vN(c.b, c.X_smem, ind, 2, xs[i], ROCKE_DCONV_NONGROUPED_X_LOAD_VEC); + } + for(i = 0; i < c.w_meta.size(); ++i) + { + rocke_value_t* idx = c.w_meta[i].lds_idx; + rocke_value_t* ind[2]; + if(w_buf != NULL) + { + idx = rocke_b_add(c.b, idx, w_buf); + } + ind[0] = c.c0; + ind[1] = idx; + rocke_b_smem_store_vN(c.b, c.W_smem, ind, 2, ws[i], c.FRAG); + } +} + +/* ---- per-lane LDS read bases ---- */ +void nongrouped_read_bases(nongrouped_ctx& c) +{ + rocke_ir_builder_t* b = c.b; + rocke_value_t *k, *t0, *t1, *x_wave_term, *x_lane_term; + + /* X: idx = wave row term + lane term + compile-time term. */ + k = nongrouped_c(c, (int64_t)c.ROWS_W * c.S); + t0 = rocke_b_mul(b, c.wave_n, k); + k = nongrouped_c(c, (int64_t)c.LDS_IN_W * c.CSTRIDE); + x_wave_term = rocke_b_mul(b, t0, k); + k = nongrouped_c(c, (int64_t)c.S * c.CSTRIDE); + t0 = rocke_b_mul(b, c.lane_lo, k); + k = nongrouped_c(c, c.FRAG); + t1 = rocke_b_mul(b, c.lane_hi, k); + x_lane_term = rocke_b_add(b, t0, t1); + c.x_read_base = rocke_b_add(b, x_wave_term, x_lane_term); + + /* W: idx = wave_m block + lane*FRAG + compile-time term. */ + k = nongrouped_c(c, (int64_t)c.MT_W * c.KATOMS * c.WAVE * c.FRAG); + t0 = rocke_b_mul(b, c.wave_m, k); + k = nongrouped_c(c, c.FRAG); + t1 = rocke_b_mul(b, c.lane, k); + c.w_read_base = rocke_b_add(b, t0, t1); +} + +rocke_value_t* + nongrouped_read_a_frag(nongrouped_ctx& c, rocke_value_t* base, int tap, int mt, int katom) +{ + int64_t off = ((int64_t)(tap * c.M_TILES + mt) * c.KATOMS + katom) * c.WAVE * c.FRAG; + rocke_value_t* k = nongrouped_c(c, off); + rocke_value_t* ind[2]; + ind[0] = c.c0; + ind[1] = rocke_b_add(c.b, base, k); + return rocke_b_smem_load_vN(c.b, c.W_smem, ind, 2, c.io_type, c.FRAG); +} + +/* One activation fragment, addressed by *input* row inside the wave window, so + * the KH taps of a column share ``(ROWS_W - 1) * stride + KH`` fragments. */ +rocke_value_t* nongrouped_read_b_frag( + nongrouped_ctx& c, rocke_value_t* base, int in_row, int cb, int s, int katom) +{ + int64_t pos = (int64_t)in_row * c.LDS_IN_W + (int64_t)cb * c.AT * c.S + s; + int64_t off = pos * c.CSTRIDE + (int64_t)katom * c.AK; + rocke_value_t* k = nongrouped_c(c, off); + rocke_value_t* ind[2]; + ind[0] = c.c0; + ind[1] = rocke_b_add(c.b, base, k); + return rocke_b_smem_load_vN(c.b, c.X_smem, ind, 2, c.io_type, c.FRAG); +} + +/* ---- emit_mfmas(): one channel chunk of MFMAs against the staged tiles ---- */ +void nongrouped_emit_mfmas(nongrouped_ctx& c, + std::vector& accs, + rocke_value_t* x_base, + rocke_value_t* w_base) +{ + const int n_in_rows = (c.ROWS_W - 1) * c.S + c.KH; + std::vector b_frags((size_t)n_in_rows * c.NCB); + std::vector a_frags((size_t)c.MT_W); + int katom, s_c, r_c, ir, cb, mt, row; + + /* Filter-column outermost: every activation fragment loaded for one + * (s, k_atom) is consumed by all KH tap rows and all MT_W channel tiles. */ + for(katom = 0; katom < c.KATOMS; ++katom) + { + for(s_c = 0; s_c < c.KW; ++s_c) + { + for(ir = 0; ir < n_in_rows; ++ir) + { + for(cb = 0; cb < c.NCB; ++cb) + { + b_frags[(size_t)ir * c.NCB + cb] + = nongrouped_read_b_frag(c, x_base, ir, cb, s_c, katom); + } + } + for(r_c = 0; r_c < c.KH; ++r_c) + { + int tap = r_c * c.KW + s_c; + for(mt = 0; mt < c.MT_W; ++mt) + { + a_frags[(size_t)mt] = nongrouped_read_a_frag(c, w_base, tap, mt, katom); + } + for(row = 0; row < c.ROWS_W; ++row) + { + int bf_row = row * c.S + r_c; + for(cb = 0; cb < c.NCB; ++cb) + { + int nt = row * c.NCB + cb; + for(mt = 0; mt < c.MT_W; ++mt) + { + size_t idx = (size_t)mt * c.NT_W + nt; + accs[idx] = c.mfma(c.b, + a_frags[(size_t)mt], + b_frags[(size_t)bf_row * c.NCB + cb], + accs[idx]); + } + } + } + } + } + } +} + +/* ---- the runtime channel loop; returns the accumulator results ---- */ +std::vector nongrouped_channel_loop(nongrouped_ctx& c) +{ + rocke_ir_builder_t* b = c.b; + const int n_acc = c.MT_W * c.NT_W; + rocke_value_t* zero_acc; + rocke_value_t* c_ck_bytes; + rocke_value_t* hi; + std::vector xs0, ws0; + std::vector iter_args; + std::vector names; + std::vector out; + rocke_for_t loop; + int i; + + zero_acc = rocke_b_zero_vec_f32(b, c.ACC); + + c_ck_bytes = nongrouped_c(c, (int64_t)c.CK * 2); + nongrouped_issue_stage_loads(c, c.c0, xs0, ws0); + if(c.DB) + { + /* Prologue: publish chunk 0 into buffer 0 and have chunk 1 in flight. */ + nongrouped_commit_stage(c, xs0, ws0, NULL, NULL); + nongrouped_issue_stage_loads(c, c_ck_bytes, xs0, ws0); + } + + for(i = 0; i < n_acc; ++i) + { + names.push_back("acc" + std::to_string(i)); + } + for(i = 0; i < (int)xs0.size(); ++i) + { + names.push_back("xs" + std::to_string(i)); + } + for(i = 0; i < (int)ws0.size(); ++i) + { + names.push_back("ws" + std::to_string(i)); + } + for(i = 0; i < n_acc; ++i) + { + iter_args.push_back({names[(size_t)i].c_str(), zero_acc}); + } + for(i = 0; i < (int)xs0.size(); ++i) + { + iter_args.push_back({names[(size_t)(n_acc + i)].c_str(), xs0[(size_t)i]}); + } + for(i = 0; i < (int)ws0.size(); ++i) + { + iter_args.push_back({names[(size_t)(n_acc + (int)xs0.size() + i)].c_str(), ws0[(size_t)i]}); + } + + /* Runtime trip count, clamped to one chunk. validate() guarantees C is a + * positive whole number of chunks, so the clamp never changes it; it lets + * LLVM prove the loop runs, without which it keeps a zero-trip guard and + * schedules the prefetch loads late. */ + hi = nongrouped_c(c, c.CK); + hi = rocke_b_div(b, c.C, hi); + hi = rocke_b_smax(b, hi, c.c1); + loop = rocke_b_scf_for_iter( + b, c.c0, hi, c.c1, iter_args.data(), (int)iter_args.size(), "c_iter", false, false); + if(loop.op == NULL || loop.iter_vars == NULL) + { + return out; + } + + rocke_b_region_enter(b, loop.body); + { + std::vector accs(loop.iter_vars, loop.iter_vars + n_acc); + std::vector xs(loop.iter_vars + n_acc, loop.iter_vars + n_acc + c.X_PASSES); + std::vector ws(loop.iter_vars + n_acc + c.X_PASSES, + loop.iter_vars + loop.num_iter_vars); + std::vector xs_n, ws_n, yields; + rocke_value_t *k, *t, *c_off; + + if(c.spec->iglp != ROCKE_DCONV_NONGROUPED_IGLP_NONE) + { + rocke_b_iglp_opt(b, c.spec->iglp); + } + + if(c.DB) + { + /* Parity of the induction variable picks the live buffer; every + * fragment offset stays a compile-time constant off these bases. */ + rocke_value_t *par, *cur_x, *cur_w, *nxt_x, *nxt_w, *xb, *wb; + k = nongrouped_c(c, 2); + par = rocke_b_mod(b, loop.iv, k); + k = nongrouped_c(c, c.x_stage); + cur_x = rocke_b_mul(b, par, k); + k = nongrouped_c(c, c.w_stage); + cur_w = rocke_b_mul(b, par, k); + k = nongrouped_c(c, c.x_stage); + nxt_x = rocke_b_sub(b, k, cur_x); + k = nongrouped_c(c, c.w_stage); + nxt_w = rocke_b_sub(b, k, cur_w); + + rocke_b_sync(b); + /* Stage chunk i+1 into the idle buffer while chunk i feeds MFMAs. */ + nongrouped_commit_stage(c, xs, ws, nxt_x, nxt_w); + k = nongrouped_c(c, 2); + t = rocke_b_add(b, loop.iv, k); + c_off = rocke_b_mul(b, t, c_ck_bytes); + nongrouped_issue_stage_loads(c, c_off, xs_n, ws_n); + xb = rocke_b_add(b, c.x_read_base, cur_x); + wb = rocke_b_add(b, c.w_read_base, cur_w); + nongrouped_emit_mfmas(c, accs, xb, wb); + } + else + { + /* Publish the chunk prefetched during the previous iteration. */ + rocke_b_sync(b); + nongrouped_commit_stage(c, xs, ws, NULL, NULL); + rocke_b_sync(b); + t = rocke_b_add(b, loop.iv, c.c1); + c_off = rocke_b_mul(b, t, c_ck_bytes); + nongrouped_issue_stage_loads(c, c_off, xs_n, ws_n); + nongrouped_emit_mfmas(c, accs, c.x_read_base, c.w_read_base); + } + + yields.insert(yields.end(), accs.begin(), accs.end()); + yields.insert(yields.end(), xs_n.begin(), xs_n.end()); + yields.insert(yields.end(), ws_n.begin(), ws_n.end()); + rocke_b_scf_yield(b, yields.data(), (int)yields.size()); + } + rocke_b_region_leave(b); + + for(i = 0; i < n_acc && i < loop.op->num_results; ++i) + { + out.push_back(loop.op->results[i]); + } + return out; +} + +/* ---- epilogue: square-atom accumulator -> packed dwordx2 quad stores ---- */ +void nongrouped_epilogue(nongrouped_ctx& c, const std::vector& accs_out) +{ + rocke_ir_builder_t* b = c.b; + rocke_value_t *k, *t0, *t1, *k_lane_base, *row_base, *img_out_base; + const rocke_dconv_params_t* prm = &c.params; + int mt, row, cb, q, jj; + + /* Slot i -> row = (i//4)*(AT//4) + lane_hi*4 + (i%4), col = lane_lo. + * Slots 4q..4q+3 are four consecutive output channels. */ + k = nongrouped_c(c, (int64_t)c.MT_W * c.AT); + t0 = rocke_b_mul(b, c.wave_m, k); + t0 = rocke_b_add(b, c.k_base, t0); + k = nongrouped_c(c, 4); + t1 = rocke_b_mul(b, c.lane_hi, k); + k_lane_base = rocke_b_add(b, t0, t1); + k = nongrouped_c(c, c.ROWS_W); + t0 = rocke_b_mul(b, c.wave_n, k); + row_base = rocke_b_add(b, c.out_h0, t0); + img_out_base = rocke_b_mul(b, c.n_img, prm->p_D_stride_n); + + for(mt = 0; mt < c.MT_W; ++mt) + { + rocke_value_t* k_mt; + k = nongrouped_c(c, (int64_t)mt * c.AT); + k_mt = rocke_b_add(b, k_lane_base, k); + for(row = 0; row < c.ROWS_W; ++row) + { + rocke_value_t *out_h, *h_ok, *row_off; + k = nongrouped_c(c, row); + out_h = rocke_b_add(b, row_base, k); + h_ok = rocke_b_cmp_lt(b, out_h, prm->p_Ho); + t0 = rocke_b_mul(b, out_h, prm->p_D_stride_ho); + row_off = rocke_b_add(b, img_out_base, t0); + for(cb = 0; cb < c.NCB; ++cb) + { + int nt = row * c.NCB + cb; + rocke_value_t* acc = accs_out[(size_t)mt * c.NT_W + nt]; + rocke_value_t *out_w, *hw_ok, *base_off; + + k = nongrouped_c(c, (int64_t)cb * c.AT); + t0 = rocke_b_add(b, c.out_w0, k); + out_w = rocke_b_add(b, t0, c.lane_lo); + t0 = rocke_b_cmp_lt(b, out_w, prm->p_Wo); + hw_ok = rocke_b_land(b, h_ok, t0); + t0 = rocke_b_mul(b, out_w, prm->p_D_stride_wo); + base_off = rocke_b_add(b, row_off, t0); + for(q = 0; q < c.QUADS; ++q) + { + rocke_value_t *k_out, *valid, *d_off, *safe, *quad, *narrow; + rocke_value_t* comps[4]; + + k = nongrouped_c(c, (int64_t)q * (c.AT / 4)); + k_out = rocke_b_add(b, k_mt, k); + t0 = rocke_b_cmp_lt(b, k_out, c.K); + valid = rocke_b_land(b, hw_ok, t0); + t0 = rocke_b_add(b, base_off, k_out); + d_off = rocke_b_mul(b, t0, c.c_half); + /* Same sentinel as the staging loads: validate() keeps D at + * or below OOB_BASE, so the store lies past num_records. */ + safe = rocke_b_select(b, valid, d_off, c.oob_base); + for(jj = 0; jj < 4; ++jj) + { + comps[jj] = rocke_b_vec_extract(b, acc, 4 * q + jj); + } + quad = rocke_b_vec_pack(b, comps, 4, rocke_f32()); + if(c.bf16) + { + narrow = rocke_b_vec_trunc_f32_to_bf16(b, quad); + rocke_b_buffer_store_vN_bf16(b, c.d_rsrc, safe, c.c0, narrow, 2); + } + else + { + narrow = rocke_b_vec_trunc_f32_to_f16(b, quad); + rocke_b_buffer_store_vN_f16(b, c.d_rsrc, safe, c.c0, narrow, 2); + } + } + } + } + } +} + +void nongrouped_set_err(char* err, size_t err_cap, const char* msg) +{ + if(err != NULL && err_cap > 0) + { + snprintf(err, err_cap, "%s", msg ? msg : ""); + } +} + +} // namespace + +/* ===================================================================== * + * build_direct_conv_nongrouped(spec, arch) + * ===================================================================== */ +rocke_kernel_def_t* rocke_build_direct_conv_nongrouped( + rocke_ir_builder_t* b, const rocke_direct_conv_nongrouped_spec_t* spec, const char* arch) +{ + nongrouped_ctx c = {}; + std::vector accs_out; + + if(b == NULL || spec == NULL) + { + return NULL; + } + if(arch == NULL) + { + arch = "gfx950"; + } + c.b = b; + c.spec = spec; + + if(!nongrouped_prologue(c, arch)) + { + return NULL; + } + nongrouped_grid_decode(c); + nongrouped_x_pass_meta(c); + nongrouped_w_pass_meta(c); + nongrouped_read_bases(c); + accs_out = nongrouped_channel_loop(c); + if((int)accs_out.size() != c.MT_W * c.NT_W || !rocke_ir_builder_ok(b)) + { + return NULL; + } + nongrouped_epilogue(c, accs_out); + return rocke_ir_builder_ok(b) ? b->kernel : NULL; +} + +rocke_kernel_def_t* rocke_build_direct_conv_nongrouped_new( + rocke_ir_builder_t* b, const rocke_direct_conv_nongrouped_spec_t* spec, const char* arch) +{ + if(b != NULL) + { + /* Zero first, so every early return below leaves a builder the caller + * can safely query and rocke_ir_builder_free(). */ + memset(b, 0, sizeof(*b)); + } + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + char name[512]; + if(b == NULL) + { + return NULL; + } + if(spec == NULL) + { + rocke_i_set_err_msg(b, ROCKE_ERR_VALUE, "build_direct_conv_nongrouped: null spec"); + return NULL; + } + if(rocke_direct_conv_nongrouped_kernel_name(spec, name, sizeof(name)) != ROCKE_OK) + { + rocke_i_set_err_msg( + b, ROCKE_ERR_VALUE, "build_direct_conv_nongrouped: kernel name does not fit"); + return NULL; + } + if(rocke_ir_builder_init(b, name) != ROCKE_OK) + { + return NULL; + } + return rocke_build_direct_conv_nongrouped(b, spec, arch); + }); +} + +rocke_status_t + rocke_direct_conv_nongrouped_lower_to_llvm(const rocke_direct_conv_nongrouped_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 = {}; + rocke_kernel_def_t* kernel; + rocke_status_t st; + + if(out_ll != NULL) + { + *out_ll = NULL; + } + if(spec == NULL || out_ll == NULL) + { + nongrouped_set_err(err, err_cap, "lower_to_llvm: null spec/out"); + return ROCKE_ERR_VALUE; + } + if(arch == NULL) + { + arch = "gfx950"; + } + kernel = rocke_build_direct_conv_nongrouped_new(&b, spec, arch); + if(kernel == NULL) + { + const char* m = rocke_ir_builder_error(&b); + st = rocke_ir_builder_status(&b); + nongrouped_set_err( + err, err_cap, (m != NULL && m[0] != '\0') ? m : "build_direct_conv_nongrouped failed"); + rocke_ir_builder_free(&b); + return (st == ROCKE_OK) ? ROCKE_ERR_VALUE : st; + } + st = rocke_lower_kernel_to_llvm_ex(kernel, flavor, arch, out_ll, err, err_cap); + rocke_ir_builder_free(&b); + return st; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md index 1d702137fbb9..1d2ad0dc9b13 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md @@ -87,6 +87,7 @@ and exact gfx targets. | `build_implicit_gemm_conv` | [L/common/conv_implicit_gemm.py](../../../library/kernels/common/conv_implicit_gemm.py) | convolution as GEMM | catalog-selected matrix MMA | owning validator | | `build_direct_conv_16c` | [L/common/conv_direct_grouped.py](../../../library/kernels/common/conv_direct_grouped.py) | grouped direct convolution | `mfma_f32_16x16x16_f16`, or `mfma_f32_16x16x32_f16` with K folding | owning validator | | `build_direct_conv_4c` | [L/common/conv_direct_grouped.py](../../../library/kernels/common/conv_direct_grouped.py) | small-channel grouped convolution | direct `mfma_f32_4x4x4_f16` intrinsic outside `MmaCatalog` | validated exact gfx target | +| `build_direct_conv_nongrouped` | [L/common/conv_direct_nongrouped.py](../../../library/kernels/common/conv_direct_nongrouped.py) | non-grouped (`groups == 1`) direct convolution, LDS halo reuse | square `mfma_f32_{32x32x16,32x32x8,16x16x32,16x16x16}` | owning validator | | `build_deep_fused_conv_pool` | [L/common/deep_fused_conv_pool.py](../../../library/kernels/common/deep_fused_conv_pool.py), [L/gfx1151/deep_fused_conv_pool.py](../../../library/kernels/gfx1151/deep_fused_conv_pool.py), [L/gfx1201/deep_fused_conv_pool.py](../../../library/kernels/gfx1201/deep_fused_conv_pool.py), [L/gfx950/deep_fused_conv_pool.py](../../../library/kernels/gfx950/deep_fused_conv_pool.py) | convolution plus pooling | target-specific matrix MMA plus VALU/window reduction | exact implementation selected by target | | `build_moe_gate_up_silu_gemm`, `build_moe_interleaved_gate_up_silu_gemm`, `build_moe_down_reduce_gemm`, `build_moe_down_silu_reduce_gemm` | [P/common/moe_gemm_fused.py](../../python/rocke/instances/common/moe_gemm_fused.py) | expert GEMM plus activation or reduction | catalog-selected matrix MMA plus fused VALU/reduction | owning validator | | `build_moe_fused_mega_gemm`, `build_moe_fused_mega_gemm_fp8`, `build_moe_fused_mega_wmma` | [P/common/moe_fused_mega.py](../../python/rocke/instances/common/moe_fused_mega.py), [P/common/moe_fused_mega_fp8.py](../../python/rocke/instances/common/moe_fused_mega_fp8.py), [P/gfx1250/fused_moe_mega_wmma.py](../../python/rocke/instances/gfx1250/fused_moe_mega_wmma.py) | persistent expert GEMM pipeline | target-selected matrix MMA plus fused routing/activation | owning validator or exact-target module | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index 31e93105c706..9af25c6d5e07 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 @@ -4,6 +4,7 @@ This page covers: - `library/kernels/common/conv_implicit_gemm.py` - `library/kernels/common/conv_direct_grouped.py` +- `library/kernels/common/conv_direct_nongrouped.py` - `library/kernels/common/img2col.py` - `instances/common/pooling.py` @@ -21,6 +22,9 @@ Implicit GEMM: Direct grouped: Specialized streaming kernels for grouped small-channel cases (16c, 4c) + +Direct non-grouped (groups == 1): + LDS-staged, halo-inclusive input tile shared by all KH*KW taps ``` ## Implicit-GEMM Convolution @@ -643,6 +647,80 @@ Streams the input rows through the same runtime `scf_for_iter` as Parity gate: configs 10 (stride=1, groups=3) and 11 (stride=2, groups=3) in `tests/instances/parity/conv_direct_grouped_emit.{c,py}`. +## Direct Non-grouped Convolution + +Source: `library/kernels/common/conv_direct_nongrouped.py`. Full write-up (every +optimization, rejected levers, gates): +`library/kernels/common/README_conv_direct_nongrouped.md`. + +The `groups == 1` counterpart of the grouped direct kernels, which draw all +their parallelism from `groups` and degenerate to one wave per output row when +there is a single group. It shares `DirectConvProblem` and the AOT direct-conv +kernarg block `conv_direct_arg_names(direction="fwd")` (launch values from +`ConvArgs.from_problem(problem)`). Batch, extents, `C` (`p_total_c`), `K` +(`p_total_k`) and the strides are runtime; the filter, stride, `PAD`, dtype and +tile are baked. Unlike the grouped kernels it does not bake the channel counts. + +```python +@dataclass(frozen=True) +class DirectNongroupedConvSpec: + problem: DirectConvProblem # groups must be 1 + tile_h: int = 16; tile_w: int = 32; tile_k: int = 128; ck: int = 16 + waves_m: int = 2; waves_n: int = 4 + atom: str = "32x32x16" # 32x32x16 | 32x32x8 | 16x16x32 | 16x16x16 + lds_pad: int = 8 + chiplet_swizzle: bool = True; swizzle_wgm: int = 8 + double_buffer: bool = False + iglp: int | None = None; waves_per_eu: int | None = None +``` + +Algorithm: + +```text +for each channel chunk (runtime scf.for, accumulators + next-chunk loads carried): + barrier; publish prefetched X (halo-inclusive, [pos][c + lds_pad]) and + W (pre-swizzled to MFMA fragment order) into LDS; barrier + issue global loads for chunk i+1 (software prefetch) + for k_atom, s: read X fragments by *input* row ((ROWS_W-1)*stride + KH of them) + for r: read W fragments for tap (r, s) + mfma into acc[m_tile][row][col_block] from X row row*stride + r +epilogue: each 4-channel accumulator quad -> one dwordx2 buffer store +``` + +Structural properties: + +- activation fragments shared across the `KH` taps of a filter column; +- staging predication folded into the base offset before the loop (out-of-range + buffer loads return zero), so the loop body has no compare/select; +- weights stored in fragment order, so the LDS store and read are both linear + in `lane`; activation channel stride padded by `lds_pad` for conflict-free + `ds_read_b128`; +- two barriers per chunk (one with `double_buffer`), `iglp_opt` leading the loop + body when set; +- each LDS array carries a scratch tail for staging slots past the tile end; +- `validate()` caps accumulators at 256 registers per lane (larger tiles hang + the backend scheduler rather than failing); +- the runtime channel-loop trip count is clamped to `max(C/ck, 1)` (a no-op + after `validate()`): without the clamp LLVM keeps a zero-trip guard and + schedules the prefetch loads late in the body. + +Grid: `spec.grid() = (n_w_tiles * n_h_tiles * N * n_k_tiles, 1, 1)`; the kernel +derives the tile counts from its kernargs and decodes the block id through +`chiplet_aware_super_tile_dynamic` with the channel tile fast-varying. +Candidates: `nongrouped_specs(problem, arch)`; +`tile_w_candidates(Wo, atom_tile)` prefers widths that divide `Wo` exactly, and +the 16-wide atoms exist for widths that are not a multiple of 32. The AOT cache +(`benchmarks/common/direct_kernel_sweep.py`, variant `direct_nongrouped`) builds +the shape-independent `nongrouped_knobs(arch, dtype)` grid for 3x3 stride 1/2 +and offers a problem only the widths `tile_w_candidates` would pick for it. + +C++ engine twin: `rocke_build_direct_conv_nongrouped` in +`cpp/instances/common/conv_direct_nongrouped.cpp`; byte-identity gated by +`library/tests/parity/conv_direct_grouped_emit.{c,py}` (configs 33+) +(`tools/check_byte_identity.py --only conv_direct_grouped`), Python IR golden +`conv_direct_nongrouped/*` in `tests/instances/rocke_ir_parity_harness.py`, binding +`rocke_engine.conv_direct_nongrouped_*` / `rocke.core.backend.lower_conv_direct_nongrouped`. + ## Img2Col Source: `library/kernels/common/img2col.py`. @@ -721,5 +799,7 @@ Grid: `ceil_div(total_output_elements, block_size)`. - Async loader writes lane-contiguous LDS but consumer assumes padded / swizzled physical layout. - K-packed MFMA fold packs S/C channels in the wrong order (close but not bit-correct). - Direct conv circular accumulator slot is not reset after store. +- Staging loop sized in whole thread-passes without a scratch tail: the last pass writes past the LDS array into the neighbouring allocation (first output row correct, later rows wrong). +- C++ port folds two emitting builder calls into one argument list: C++ argument evaluation order is unspecified, so `arith.constant`s reorder and byte-identity breaks. - Output descriptor uses NHWC instead of NHWK stride order. - Benchmark compares implicit-GEMM graph mode to direct per-launch mode without labeling launch overhead. 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 b76b36743426..deb72217ec72 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 @@ -42,6 +42,7 @@ implementation. |-----------------------------------|-------------------------------------------------------------------|------------------------------| | `conv_implicit_gemm.py` | `ConvProblem`, `ImplicitGemmConvSpec` | `instances/convolution.md` | | `conv_direct_grouped.py` | `DirectConvProblem`, `DirectConv16cSpec`, `DirectConv4cSpec` | `instances/convolution.md` | +| `conv_direct_nongrouped.py` | `DirectNongroupedConvSpec` (`groups == 1`) | `instances/convolution.md` | | `img2col.py` | `Img2ColSpec` | `instances/convolution.md` | | `pooling.py` | `PoolingProblem`, `Pooling2DSpec`, `PoolOp` | `instances/convolution.md` | @@ -180,6 +181,7 @@ From `helpers/README.md`: | grouped_gemm | - | planned | - | - | - | yes | - | (wraps gemm_universal) | | conv_implicit_gemm | yes (buffer) | full (unmerge+embed+pad) | - | - | - | yes | - | `AsyncTileLoader`, `CoalescedTileLoader`, `CShuffleEpilogue`, `MfmaAtom`, `WarpGrid`, `LdsLayout`, `SchedulePolicy` | | conv_direct_grouped | - | input/output/weight + H/W pad | - | - | - | yes | - | `MfmaAtom` (4x4x4 / 16x16x{16,32}) | +| conv_direct_nongrouped | - | - | - | - | - | yes | - | `chiplet_aware_super_tile_dynamic` (`helpers.grid`), `emit_direct_params` (AOT kernargs) | | img2col | - | reuses A descriptor | - | yes | - | yes | - | - | | pooling | - | full (input + pad) | - | yes | - | yes | - | - | | attention_unified | - | Q + output + paged-KV | - | - | - | yes | - | `OnlineSoftmaxState`, `PagedKvDescriptor` | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md index c05d84bc32ec..95a697d0a025 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md @@ -236,6 +236,7 @@ GroupedGemmLauncher, grouped_gemm_problems ConvProblem, ImplicitGemmConvSpec, build_implicit_gemm_conv DirectConvProblem, DirectConv16cSpec, build_direct_conv_16c DirectConv4cSpec, build_direct_conv_4c +DirectNongroupedConvSpec, build_direct_conv_nongrouped, is_valid_nongrouped_spec, nongrouped_specs, nongrouped_knobs Img2ColSpec, build_img2col PoolingProblem, Pooling2DSpec, PoolOp, build_pooling2d diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md index 0dbf98a6d40a..0a0705ddda02 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md @@ -113,6 +113,7 @@ prefix for brevity). | `instances/grouped_gemm.py` | `GroupedGemmProblem`, `GroupedGemmSpec`, `build_grouped_gemm`, `grouped_gemm_signature`, `GroupedGemmLauncher`, `grouped_gemm_problems`. | | `instances/conv_implicit_gemm.py` | `ConvProblem`, `ImplicitGemmConvSpec`, `make_a_descriptor`, `make_b_descriptor`, `make_d_descriptor`, `build_implicit_gemm_conv`. | | `instances/conv_direct_grouped.py` | `DirectConvProblem`, `DirectConv16cSpec`, `DirectConv4cSpec`, `build_direct_conv_16c`, `build_direct_conv_4c`. | +| `instances/conv_direct_nongrouped.py` | `DirectNongroupedConvSpec`, `is_valid_nongrouped_spec`, `build_direct_conv_nongrouped`, `nongrouped_specs`, `nongrouped_knobs`, `tile_w_candidates` (`groups == 1`; AOT, `C`/`K` runtime). | | `instances/img2col.py` | `Img2ColSpec`, `build_img2col`, `img2col_grid`, `img2col_signature`. (CK Tile 04.) | | `instances/pooling.py` | `PoolingProblem`, `Pooling2DSpec`, `PoolOp`, `build_pooling2d`, `pooling2d_grid`, `pooling2d_signature`. (CK Tile 36.) | | `instances/permute_nd.py` | `PermuteSpec`, `build_permute`, `permute_grid`, `permute_signature`. Rank-up-to-8 n-D permute. (CK Tile 06.) | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py index 8a29f78ffe77..970f8fdc7c94 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/backend.py @@ -922,6 +922,44 @@ def conv_direct_grouped_spec_to_dict(spec: Any, kind: str) -> Dict[str, Any]: return d +def conv_direct_nongrouped_spec_to_dict(spec: Any) -> Dict[str, Any]: + """:class:`DirectNongroupedConvSpec` -> flat dict (problem nested, dtype included). + ``iglp`` / ``waves_per_eu`` are forwarded as-is; ``None`` means "knob off".""" + p = spec.problem + return dict( + problem=dict( + N=p.N, + H=p.H, + W=p.W, + groups=p.groups, + cpg=p.cpg, + kpg=p.kpg, + KH=p.KH, + KW=p.KW, + PAD=p.PAD, + stride=p.stride, + dtype=p.dtype, + ), + name=spec.name, + tile_h=spec.tile_h, + tile_w=spec.tile_w, + tile_k=spec.tile_k, + ck=spec.ck, + waves_m=spec.waves_m, + waves_n=spec.waves_n, + atom=spec.atom, + wave_size=spec.wave_size, + lds_pad=spec.lds_pad, + chiplet_swizzle=spec.chiplet_swizzle, + swizzle_wgm=spec.swizzle_wgm, + chiplet_chunk=spec.chiplet_chunk, + num_xcds=spec.num_xcds, + double_buffer=spec.double_buffer, + iglp=spec.iglp, + waves_per_eu=spec.waves_per_eu, + ) + + def img2col_spec_to_dict(spec: Any) -> Dict[str, Any]: """:class:`Img2ColSpec` -> flat dict (problem nested).""" return dict( @@ -1497,6 +1535,43 @@ def py_fn(wi: bool) -> Tuple[str, str]: ) +def lower_conv_direct_nongrouped( + spec: Any, + *, + arch: str = "gfx950", + backend: Optional[str] = None, + want_ir: bool = False, +) -> "GemmLowerResult": + """Lower a :class:`DirectNongroupedConvSpec` (the ``groups == 1`` direct conv).""" + + def py_fn(wi: bool) -> Tuple[str, str]: + from kernels.common.conv_direct_nongrouped import build_direct_conv_nongrouped + from .lower_llvm import lower_kernel_to_llvm + + k = build_direct_conv_nongrouped(spec, arch=arch) + ll = lower_kernel_to_llvm(k, arch=arch) + ir = "" + if wi: + from .ir_serialize import serialize + + ir = serialize(k) + return ll, ir + + eng = _import_engine() + sd = conv_direct_nongrouped_spec_to_dict(spec) + return _lower_family( + "conv_direct_nongrouped", + spec, + arch, + backend, + want_ir, + py_fn, + lambda: eng.conv_direct_nongrouped_lower_llvm(sd, arch=arch), + lambda: eng.conv_direct_nongrouped_serialize_ir(sd, arch=arch), + _name_of(spec), + ) + + def lower_img2col( spec: Any, *, 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 ce43548f0038..b51fcbffb802 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 @@ -73,6 +73,7 @@ See the [packed FP6 input contract](../examples/gfx1250/gemm/FP6.md). | `conv_implicit_gemm_auto` | ✅ | ✅ | ❌ | MFMA-specialized autotuned path (raw `MfmaAtom`, K=32 kpack); not ported to WMMA | | `direct_conv_16c` | ❌ | ✅ | ❌ | `fold_k32` needs 16x16x32 atom (CDNA4) | | `direct_conv_4c` | ✅ | ✅ | ❌ | 4x4x4 MFMA atom not in WMMA catalog | +| `conv_direct_nongrouped` | ✅ | ✅ | ❌ | `groups=1`; gfx942 uses 32x32x8 / 16x16x16 atoms, gfx950 also 32x32x16 / 16x16x32; MFMA-only. Built from `library/kernels`; not yet wired into dispatch | --- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 5de9b23296ff..7798e50c3dfd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -698,6 +698,46 @@ "sha256": "f9b9c29edc1eee3570352a0dcc0ad87143ab1b145787b317e129369dc3c63db6", "status": "ok" }, + "conv_direct_nongrouped/gfx942/fp16_a32x32x8": { + "arch": "gfx942", + "bytes": 72050, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_fp16_a32x32x8_N1H16W32_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x8_g8_iglp0", + "sha256": "a531ad29885faff229610c22696421bfb13d43ed8d9a7592072ba45717a13681", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_a16x16x32_partial_w": { + "arch": "gfx950", + "bytes": 148077, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_a16x16x32_partial_w_N2H20W40_g1_c64k64_t8x48x32_ck32_w1x2_a16x16x32_g8_iglp0_bf16", + "sha256": "89dce2f5387e9bbdaf177b0f3fe089d1fa4a162746b96fffa42697e717aaeca6", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_noswizzle_we3": { + "arch": "gfx950", + "bytes": 79711, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_noswizzle_we3_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_gnone_we3_bf16", + "sha256": "76620d100a1b6493e123c316d2328297c31396156693186b595ab4a476192d61", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_t8x32x64_iglp0": { + "arch": "gfx950", + "bytes": 80608, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_t8x32x64_iglp0_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_g8_iglp0_bf16", + "sha256": "997aa9df0c1af2027df65366b175edec5bda52292912abf430a7f3092842ee60", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/fp16_db_s2": { + "arch": "gfx950", + "bytes": 74610, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_fp16_db_s2_N1H32W64_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x16_g8_db_iglp0", + "sha256": "c80479a6bbdb5c5e81decdc1ac5b7bfc773b017e02fd518b9cd4fa0b1d75d50a", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53578, @@ -1228,8 +1268,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 32, - "gfx950": 91 + "gfx942": 33, + "gfx950": 95 }, "expected_failures": 4, "families": { @@ -1241,6 +1281,7 @@ "conv_dgrad": 11, "conv_direct": 11, "conv_direct_dgrad": 5, + "conv_direct_nongrouped": 5, "conv_wgrad": 13, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -1250,7 +1291,7 @@ "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 150 + "ok_cases": 155 } }, "llvm22": { @@ -1951,6 +1992,46 @@ "sha256": "81cabd3943ad21bb38e37dbfb6569b3cea329f34e79fab5db4510b6cc8285a62", "status": "ok" }, + "conv_direct_nongrouped/gfx942/fp16_a32x32x8": { + "arch": "gfx942", + "bytes": 72189, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_fp16_a32x32x8_N1H16W32_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x8_g8_iglp0", + "sha256": "b1fc232fc6aa64821773df743a45131289633b9d9b7c3efdb489dc45d9609381", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_a16x16x32_partial_w": { + "arch": "gfx950", + "bytes": 148216, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_a16x16x32_partial_w_N2H20W40_g1_c64k64_t8x48x32_ck32_w1x2_a16x16x32_g8_iglp0_bf16", + "sha256": "8930d983effad4c839c5aeeef7bf64f706480f802d9476913a708f0f2bc73e59", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_noswizzle_we3": { + "arch": "gfx950", + "bytes": 79850, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_noswizzle_we3_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_gnone_we3_bf16", + "sha256": "cbb7ab6e0c51b516648ac1677a3752d5901631b45a1a70cb1ed318ba152c9dec", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_t8x32x64_iglp0": { + "arch": "gfx950", + "bytes": 80747, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_t8x32x64_iglp0_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_g8_iglp0_bf16", + "sha256": "2e72744b6c912ebd5326ca60a6c13b81d6ed19af0b30e70ced2249f5fc500a4b", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/fp16_db_s2": { + "arch": "gfx950", + "bytes": 74749, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_fp16_db_s2_N1H32W64_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x16_g8_db_iglp0", + "sha256": "65aecb5e9192cae49254347700f6aefdd11f47c805312d9fd9b4604a67f7755f", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53717, @@ -2481,8 +2562,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 32, - "gfx950": 91 + "gfx942": 33, + "gfx950": 95 }, "expected_failures": 4, "families": { @@ -2494,6 +2575,7 @@ "conv_dgrad": 11, "conv_direct": 11, "conv_direct_dgrad": 5, + "conv_direct_nongrouped": 5, "conv_wgrad": 13, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -2503,7 +2585,7 @@ "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 150 + "ok_cases": 155 } }, "llvm23": { @@ -3204,6 +3286,46 @@ "sha256": "d4b417486ade76f7aa923868d11ac7fa2ffd4f2e1850a5e9652af8da99c4843a", "status": "ok" }, + "conv_direct_nongrouped/gfx942/fp16_a32x32x8": { + "arch": "gfx942", + "bytes": 72193, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_fp16_a32x32x8_N1H16W32_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x8_g8_iglp0", + "sha256": "ec5e2d3974fbe6d94c8d194ce40a4d8acf98f114457447fad3e1bac2c8687f4c", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_a16x16x32_partial_w": { + "arch": "gfx950", + "bytes": 148220, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_a16x16x32_partial_w_N2H20W40_g1_c64k64_t8x48x32_ck32_w1x2_a16x16x32_g8_iglp0_bf16", + "sha256": "7e24f74e8ef1671f9d762cb62ca0a9631c0d125ebb4052ea0abc1ccbff268c81", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_noswizzle_we3": { + "arch": "gfx950", + "bytes": 79854, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_noswizzle_we3_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_gnone_we3_bf16", + "sha256": "3f574bb890191b9d5656120ea806af692b9d94d542608a87425211a91c290db2", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/bf16_t8x32x64_iglp0": { + "arch": "gfx950", + "bytes": 80751, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_bf16_t8x32x64_iglp0_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_g8_iglp0_bf16", + "sha256": "b4bb36d117c7af0b63e75b690fc192c9a12098665b5df20d96b050629531117e", + "status": "ok" + }, + "conv_direct_nongrouped/gfx950/fp16_db_s2": { + "arch": "gfx950", + "bytes": 74753, + "family": "conv_direct_nongrouped", + "kernel_name": "irhash_direct_nongrouped_fp16_db_s2_N1H32W64_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x16_g8_db_iglp0", + "sha256": "ee53aba1abe0cbbb3fea79653bc97f1cad65772492fec460beec20df2158b3c0", + "status": "ok" + }, "conv_wgrad/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 53721, @@ -3749,8 +3871,8 @@ "gfx1201": 7, "gfx1250": 8, "gfx90a": 5, - "gfx942": 32, - "gfx950": 91 + "gfx942": 33, + "gfx950": 95 }, "expected_failures": 0, "families": { @@ -3762,6 +3884,7 @@ "conv_dgrad": 11, "conv_direct": 11, "conv_direct_dgrad": 5, + "conv_direct_nongrouped": 5, "conv_wgrad": 13, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, @@ -3771,7 +3894,7 @@ "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 154 + "ok_cases": 159 } } }, 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 e5115172cd03..8c8ae1b58da9 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 @@ -1083,6 +1083,52 @@ def _build(): return _build +def build_direct_nongrouped( + name, + arch, + N, + H, + W, + C, + K, + KH=3, + KW=3, + PAD=1, + stride=1, + dtype="bf16", + **spec_kw, +): + """Non-grouped (groups == 1) direct conv; ``spec_kw`` overrides + DirectNongroupedConvSpec fields. Configs mirror parity/conv_direct_grouped_emit.* + (indices 33+), where the C++ engine is gated byte-identical by + check_byte_identity.py.""" + + def _build(): + from kernels.common.conv_direct_grouped import DirectConvProblem + from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, + build_direct_conv_nongrouped, + ) + + p = DirectConvProblem( + N=N, + H=H, + W=W, + groups=1, + cpg=C, + kpg=K, + KH=KH, + KW=KW, + PAD=PAD, + stride=stride, + dtype=dtype, + ) + spec = DirectNongroupedConvSpec(problem=p, name=name, **spec_kw) + return build_direct_conv_nongrouped(spec, arch=arch) + + return _build + + def build_grouped_gemm_case(name, arch, m, n, k, e): def _build(): from rocke.instances.gfx950.grouped_gemm import ( @@ -3075,6 +3121,53 @@ def add(family, case_id, arch, build): ), ) + # --- conv_direct_nongrouped: non-grouped (groups == 1) direct conv --- + # LDS halo-reuse tile, tap-shared activation fragments, fragment-order + # weights, hoisted staging predication. One case per structural branch; the + # full branch matrix (and the C++ twin) lives in + # library/tests/parity/conv_direct_grouped_emit.* (indices 33+). + _nongrouped_base = dict( + tile_h=8, tile_w=32, tile_k=64, ck=32, waves_m=2, waves_n=2, iglp=0 + ) + for _case_id, _arch, _shape, _over in ( + ("bf16_t8x32x64_iglp0", "gfx950", dict(N=2, H=16, W=32, C=64, K=128), {}), + ( + "fp16_db_s2", + "gfx950", + dict(N=1, H=32, W=64, C=64, K=64, stride=2, dtype="fp16"), + dict(ck=16, double_buffer=True), + ), + ( + "bf16_a16x16x32_partial_w", + "gfx950", + dict(N=2, H=20, W=40, C=64, K=64), + dict(tile_w=48, tile_k=32, waves_m=1, atom="16x16x32"), + ), + ( + "bf16_noswizzle_we3", + "gfx950", + dict(N=2, H=16, W=32, C=64, K=128), + dict(chiplet_swizzle=False, iglp=None, waves_per_eu=3), + ), + ( + "fp16_a32x32x8", + "gfx942", + dict(N=1, H=16, W=32, C=64, K=64, dtype="fp16"), + dict(atom="32x32x8", ck=16), + ), + ): + add( + "conv_direct_nongrouped", + f"conv_direct_nongrouped/{_arch}/{_case_id}", + _arch, + build_direct_nongrouped( + f"irhash_direct_nongrouped_{_case_id}", + _arch, + **_shape, + **{**_nongrouped_base, **_over}, + ), + ) + # gfx942 GQA head-fold (D128 sliding-window bf16). Registered SEPARATELY from # the D256 case above because that one early-returns into the lean D256 kernel # and never reaches the D128 fold body -- and the fold is default-ON for its From 3f308afce43bb91c85fff09bcd9bcd9b26b334d9 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 6 Oct 2026 23:07:49 +0000 Subject: [PATCH 6/7] fixes --- .../common/benchmark_conv_compare.py | 5 +- .../common/benchmark_direct_conv.py | 12 +- .../common/benchmark_implicit_gemm_conv.py | 18 +- .../benchmarks/common/direct_kernel_sweep.py | 13 + .../library/benchmarks/common/early_stop.py | 9 +- .../library/benchmarks/common/kernel_cache.py | 20 +- .../library/benchmarks/common/kernel_sweep.py | 13 +- .../common/README_conv_direct_nongrouped.md | 10 + .../kernels/common/conv_direct_grouped.py | 22 + .../kernels/common/conv_direct_nongrouped.py | 73 ++++ .../tests/test_direct_conv_correctness.py | 409 +++++++++++------- .../library/tests/test_direct_kernel_cache.py | 63 +++ .../rocke/library/tests/test_early_stop.py | 28 ++ .../rocke/instance_conv_direct_grouped.h | 7 + .../instance_conv_direct_grouped_internal.h | 4 +- .../conv_direct_grouped_spec_and_validate.cpp | 36 ++ 16 files changed, 571 insertions(+), 171 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 e2380aa965b9..ca8a82705a8a 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 @@ -41,7 +41,8 @@ ``--early-stop-after`` flags are forwarded to both scripts. Early stopping carries over from one script to the other: direct conv runs -first and records each case's best time, and the implicit-GEMM sweep starts +first and records each case's best time (a kernel that fails ``--verify`` is +never recorded), and the implicit-GEMM sweep starts each case from that time instead of from nothing -- so from its first kernel it skips any that cannot come within ``--early-stop`` times the direct result. ``--no-seed-from-direct`` turns that off. ``--dtype`` / ``--direction`` default to fp16 / fwd for explicit shapes; @@ -308,7 +309,7 @@ def main() -> int: "--top", type=int, default=10, - help="top-N results to show per benchmark (default: 5)", + help="top-N results to show per benchmark (default: 10)", ) parser.add_argument( "--warmup", type=int, default=3, help="warmup iterations (default: 3)" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py index f56aa41d57a9..e0ddd338cf40 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_direct_conv.py @@ -672,6 +672,7 @@ def _run_depthwise_sweep( warmup=args.warmup, iters=args.iters, stream=stream, + passed=kernel_passed, ) if ms is None: _stop.report(artifact.kernel_name) @@ -882,6 +883,7 @@ def _run_sweep( warmup=args.warmup, iters=args.iters, stream=stream, + passed=kernel_passed, ) if ms is None: _stop.report(artifact.kernel_name) @@ -1095,6 +1097,7 @@ def _run_nongrouped_sweep( warmup=args.warmup, iters=args.iters, stream=stream, + passed=kernel_passed, ) if ms is None: _stop.report(artifact.kernel_name) @@ -1346,6 +1349,7 @@ def _run_wgrad_sweep( warmup=args.warmup, iters=args.iters, stream=0, + passed=kernel_passed, ) if ms is None: _stop.report(artifact.kernel_name) @@ -1924,6 +1928,7 @@ def run_mfma_dgrad(): warmup=args.warmup, iters=args.iters, stream=0, + passed=kernel_passed, ) else: cfg = LaunchConfig(grid=grid, block=block_dim) @@ -1932,6 +1937,7 @@ def run_mfma_dgrad(): warmup=args.warmup, iters=args.iters, stream=0, + passed=kernel_passed, ) if ms is None: _stop.report(label) @@ -2205,7 +2211,11 @@ def run(bound=bound): try: ms = _stop.measure( - run, warmup=args.warmup, iters=args.iters, stream=0 + run, + warmup=args.warmup, + iters=args.iters, + stream=0, + passed=passed, ) except (HipError, RuntimeError) as e: print(f" [skip] {plan.label}: {e}", flush=True) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index 73a151f7efb2..25c128b32edd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -2358,6 +2358,7 @@ def _make(*shape): warmup=args.warmup, iters=args.iters, stream=stream, + passed=kernel_passed, ) if ms is None: _stop.report(artifact.kernel_name) @@ -2850,8 +2851,9 @@ def _degrees(pipeline: str, async_dma: bool) -> tuple: cfg = LaunchConfig(grid=grid, block=block, stream=stream) + kernel_passed = None if args.verify or args.dump_fail: - stopped, _ = _verify_kernel( + stopped, kernel_passed = _verify_kernel( rt=rt, launcher=launcher, values=values, @@ -2893,6 +2895,7 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): warmup=args.warmup, iters=args.iters, stream=stream, + passed=kernel_passed, ) if ms is None: _stop.report(f"{artifact.kernel_name} spk{_launch_sk}") @@ -3110,8 +3113,9 @@ def _launch_two_stage( _s1(_v1, config=_c1) _s2(_v2, config=_c2) + kernel_passed = None if args.verify or args.dump_fail: - stopped, _ = _verify_kernel( + stopped, kernel_passed = _verify_kernel( rt=rt2, launch_fn=_launch_two_stage, out_dev=dW_dev2, @@ -3138,6 +3142,7 @@ def _launch_two_stage( warmup=args.warmup, iters=args.iters, stream=0, + passed=kernel_passed, ) if ms is None: _stop.report(f"{s1_art.kernel_name} spk{_launch_sk}") @@ -3447,8 +3452,9 @@ def _make(*shape): _zero_init = spec.needs_atomic + kernel_passed = None if args.verify or args.dump_fail: - stopped, _ = _verify_kernel( + stopped, kernel_passed = _verify_kernel( rt=rt, launcher=launcher, values=values, @@ -3482,7 +3488,11 @@ def _launch_atomic(): timed_fn = launcher.bind(values, config=cfg) ms = _stop.measure( - timed_fn, warmup=args.warmup, iters=args.iters, stream=stream + timed_fn, + warmup=args.warmup, + iters=args.iters, + stream=stream, + passed=kernel_passed, ) if ms is None: _stop.report(artifact.kernel_name) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py index edf89358c8a9..4b4c58dafcfb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py @@ -312,6 +312,17 @@ def _uses_coalesced_weights(knobs: dict) -> bool: return knobs["waves_k"] > 1 or knobs["runtime_k_loop"] +def _nongrouped_spills(spec, arch: str) -> bool: + """Would this non-grouped kernel spill registers? Such a binary pays scratch + traffic every channel chunk and is the slowest to compile, so it is not + cached. + Its register footprint is fixed by the baked capabilities and the tile, not + by the probe shape, so one answer holds for every problem it would serve.""" + from kernels.common.conv_direct_nongrouped import nongrouped_register_reason + + return nongrouped_register_reason(spec, arch) is not None + + def validate_spec(variant: str, spec, arch: str) -> Tuple[bool, str]: """``(ok, reason)`` for ``spec`` on ``arch``, from the kernel's own validator.""" from kernels.common import conv_direct_grouped as dc @@ -522,6 +533,8 @@ def _caps_jobs( spec = make_spec(variant, probe, knobs) if not validate_spec(variant, spec, arch)[0]: continue + if variant == "direct_nongrouped" and _nongrouped_spills(spec, arch): + continue cell = [ _job( _identity( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py index bbd6f85d1f3f..7a0fc2332162 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/early_stop.py @@ -6,7 +6,9 @@ to the timed loops of kernels that were never going to rank. The warmup launches run anyway, so they are timed too: a kernel whose warmup is already ``factor`` times slower than the best kernel measured so far is reported and -skipped instead of being timed. +skipped instead of being timed. A kernel that failed verification is still +timed but never becomes the best: a wrong-but-fast kernel would otherwise prune +every correct one, here and, through the record file, in a seeded sweep. The best time can also come from another run: ``--early-stop-record FILE`` keeps each case's best time in FILE (rewritten whenever one improves, so a run @@ -120,6 +122,8 @@ class EarlyStop: more than ``factor`` times the best result so far, returns ``None`` without running the timed loop. The first ``after`` kernels of a sweep are always measured in full, so the best result is established before it is used. + A kernel measured with ``passed=False`` (it failed verification) neither + counts towards ``after`` nor updates the best result or the record file. """ def __init__( @@ -160,6 +164,7 @@ def measure( warmup: int, iters: int, stream: int = 0, + passed: Optional[bool] = None, ) -> Optional[float]: from rocke.runtime import time_launches @@ -179,6 +184,8 @@ def measure( ms = time_launches(fn, warmup=0, iters=iters, stream=stream) else: ms = time_launches(fn, warmup=warmup, iters=iters, stream=stream) + if passed is False: + return ms self.n_measured += 1 if self.best_ms is None or ms < self.best_ms: self.best_ms = ms diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py index 3648d6146d03..47571961526f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py @@ -25,11 +25,12 @@ import functools import hashlib +import itertools import json import os from dataclasses import asdict, dataclass, field from pathlib import Path -from typing import Dict, List, Optional, Tuple +from typing import Dict, Iterable, List, Optional, Tuple # --------------------------------------------------------------------------- @@ -503,9 +504,20 @@ def put( self.link(identity, key, meta) return path - def index(self) -> Dict[str, dict]: - """``{identity sha1: metadata}`` for every entry with a binary.""" - return {ident.stable_hash(): meta for ident, _, meta in self._entries(None)} + def index(self, directions: Optional[Iterable[str]] = None) -> Dict[str, dict]: + """``{identity sha1: metadata}`` for every entry with a binary. + + ``directions`` limits the walk to those directions' entries; every + direction is parsed per entry, so on a large cache reading only the + ones needed is the difference between a few files and ~10^6. + """ + if directions is None: + entries: Iterable = self._entries(None) + else: + entries = itertools.chain.from_iterable( + self._entries(d) for d in sorted(set(directions)) + ) + return {ident.stable_hash(): meta for ident, _, meta in entries} def meta(self, identity: KernelIdentity) -> Optional[dict]: """The identity's metadata, or ``None`` when it has no entry.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py index 2bf43e5f15e8..0fa198e40b4c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py @@ -70,6 +70,7 @@ from __future__ import annotations +import gc import itertools import time from concurrent.futures import ( @@ -1291,7 +1292,10 @@ def compile_jobs( """ digest = current_emitter_digest() comgr_id = current_comgr_id() - index = cache.index() + # Only the directions these jobs live in: a shared cache can hold ~10^6 + # entries of other families, and every one parsed here is resident in the + # process the compile workers fork from. + index = cache.index({j.identity.direction for j in all_jobs}) def _up_to_date(job) -> bool: meta = index.get(job.identity.stable_hash()) @@ -1400,6 +1404,13 @@ def on_compiled(result) -> None: # emit ran (and submitting them all takes seconds by itself). With # the window a compile waits behind at most ~2*jobs emits. window = 2 * jobs + # Forked workers share this process's heap copy-on-write, but a + # garbage collection in a worker writes to every tracked object's + # header and so copies every page holding one: with a large cache + # index resident, each of the 64 workers grew a private copy of it + # (GBs apiece) and the run was OOM-killed. Frozen objects are left + # alone by the collector. + gc.freeze() emit_iter = iter(emit_payloads) # Work handed over from a pool that broke (see below). carry: List[Tuple[str, tuple]] = [] diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md index 6ebf086b3add..ed49e59fb397 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_direct_nongrouped.md @@ -288,6 +288,16 @@ the AOT kernel cache builds (variant `direct_nongrouped` in offers a problem the binaries whose width `tile_w_candidates(Wo, atom_tile)` would pick and whose spec validates on the real shape. +The AOT grid also drops tiles that would spill registers +(`nongrouped_register_reason`): accumulators that fill the 256-entry +accumulator file, or an estimated live set (`nongrouped_live_regs`: +accumulators, carried prefetch, one step's operand fragments) above 80% of the +registers a wave gets once the block's waves share the SIMDs. On the gfx950 +grid this cuts 924 kernels to 658 and, with `waves_per_eu` unset, leaves no +spilling binary; the spilling ones were also the slowest to compile (up to +~90 s each). It is a heuristic, not a validity rule: the per-shape +`nongrouped_specs` sweep still offers those tiles. + ## Tests and gates | Gate | What it proves | Command | diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py index 9f9c25baab0c..1f8407aaeeae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_grouped.py @@ -5787,6 +5787,22 @@ def validate(self) -> None: ) +# The C++ engine holds the KH x KW depthwise weights (regular and spatial +# forward) in fixed-size register tables sized by ``ROCKE_DCONV_DW_MAX_K{H,W}`` +# (``platform/cpp/include/rocke/instance_conv_direct_grouped.h``): both engines' +# validators reject above it, so a spec that validates builds in either. +_DW_MAX_KH = 32 +_DW_MAX_KW = 32 + + +def _depthwise_filter_reason(p: "DirectConvProblem") -> Optional[str]: + if p.KH < 1 or p.KH > _DW_MAX_KH: + return f"KH must be in 1..{_DW_MAX_KH} (got {p.KH})" + if p.KW < 1 or p.KW > _DW_MAX_KW: + return f"KW must be in 1..{_DW_MAX_KW} (got {p.KW})" + return None + + def is_valid_depthwise_spec( spec: "DirectDepthwiseSpec", arch: str = "gfx950" ) -> Tuple[bool, str]: @@ -5813,6 +5829,9 @@ def is_valid_depthwise_spec( return False, why if p.cpg != 1 or p.kpg != 1: return False, f"cpg and kpg must both be 1 (got cpg={p.cpg}, kpg={p.kpg})" + why = _depthwise_filter_reason(p) + if why is not None: + return False, why return True, "ok" @@ -6217,6 +6236,9 @@ def is_valid_depthwise_spatial_spec( return False, why if p.cpg != 1 or p.kpg != 1: return False, f"cpg and kpg must both be 1 (got cpg={p.cpg}, kpg={p.kpg})" + why = _depthwise_filter_reason(p) + if why is not None: + return False, why if p.groups > spec.wave_size: return False, f"groups {p.groups} > wave_size {spec.wave_size}" if spec.n_w_per_wave == 0: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py index 5655d52b253d..7b98bce0aea3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_direct_nongrouped.py @@ -442,6 +442,79 @@ def is_valid_nongrouped_spec( return True, "ok" +# --------------------------------------------------------------------------- +# Register budget +# --------------------------------------------------------------------------- + +# SIMDs per CU on CDNA: a block's waves are spread over them, so a block of +# more than four waves puts several on one SIMD and they split its registers. +_SIMDS_PER_CU = 4 +# Share of a wave's register budget the estimate of nongrouped_live_regs may +# fill. The estimate counts only the big arrays; addresses, loop state and the +# scheduler's temporaries take the rest. Calibrated on the gfx950 AOT grid +# (both strides, waves_per_eu unset): past 0.8 configs start to spill, below it +# none do. +_REG_BUDGET_FILL = 0.8 + + +def nongrouped_live_regs(spec: DirectNongroupedConvSpec) -> int: + """Estimated registers per lane live across the channel loop. + + The accumulators, the prefetched staging loads carried to the next + iteration, and the operand fragments of one ``(k_atom, s)`` step: the + activation fragments of every input row of the wave window (all ``KH`` tap + rows reuse them) and one weight fragment per channel tile. + """ + p = spec.problem + threads = spec.threads_per_block + x_vecs = spec.lds_in_h * spec.lds_in_w * (spec.ck // _X_LOAD_VEC) + w_slots = p.KH * p.KW * spec.m_tiles_total * spec.k_atoms * spec.wave_size + x_passes = -(-x_vecs // threads) + w_passes = -(-w_slots // threads) + frag_regs = spec.frag // 2 + prefetch = x_passes * (_X_LOAD_VEC // 2) + w_passes * frag_regs + in_rows = (spec.rows_per_wave - 1) * p.stride + p.KH + b_frags = in_rows * spec.n_col_blocks * frag_regs + a_frags = spec.m_tiles_per_wave * frag_regs + return spec.acc_vgprs + prefetch + b_frags + a_frags + + +def nongrouped_register_reason( + spec: DirectNongroupedConvSpec, arch: str = "gfx950" +) -> "str | None": + """Why ``spec`` would spill registers on ``arch``, or None if it fits. + + Not a validity rule -- a spilling kernel is still correct -- but a spilled + MFMA loop is slow and takes the backend scheduler a minute or more to + compile, so a sweep should not spend a build on it. Two limits: + + * the accumulators must leave room in the accumulator register file: a tile + that fills all of it spills even with plenty of the unified budget left; + * :func:`nongrouped_live_regs` must fit ``_REG_BUDGET_FILL`` of the + registers one wave gets when the block's waves share the SIMDs. + + ``waves_per_eu`` is not modelled: the backend treats it as a hint and + mostly overrides it, so it neither predicts nor prevents a spill. + """ + from rocke.core.arch import ArchTarget + + limits = ArchTarget.from_gfx(arch).limits + if spec.acc_vgprs >= limits.agprs: + return ( + f"{spec.acc_vgprs} accumulator registers fill the {limits.agprs}-entry " + f"accumulator file" + ) + waves = spec.threads_per_block // spec.wave_size + budget = limits.vgprs // -(-waves // _SIMDS_PER_CU) + live = nongrouped_live_regs(spec) + if live > _REG_BUDGET_FILL * budget: + return ( + f"~{live} live registers per lane exceed {_REG_BUDGET_FILL:.0%} of the " + f"{budget} a wave gets with {waves} waves per block" + ) + return None + + # --------------------------------------------------------------------------- # Candidate generation # --------------------------------------------------------------------------- diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py index d6a67aa33c32..b9408648b79d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_conv_correctness.py @@ -70,7 +70,6 @@ class _Shape: For fprop ``cpg`` must equal ``kpg`` (both symmetric). For dgrad they may differ; ``kpg=0`` (the default) means kpg == cpg (symmetric). - ``stride`` must be 1 for depthwise fprop. """ id: str @@ -114,6 +113,48 @@ class _Shape: # groups=8 satisfies DirectConv16cSpec/DirectConvSpec default block_groups=8 _Shape("4c_N2H14W14_g8_s2", N=2, H=14, W=14, groups=8, cpg=4, stride=2), _Shape("16c_N2H14W14_g8_s2", N=2, H=14, W=14, groups=8, cpg=16, stride=2), + # Depthwise 7x7 "same": the FMA loop runs (H + KH - 1) // KH row groups and + # the load-free flush loop the rest up to (H + PAD + KH - 1) // KH, so the + # 3x3 case above (5 and 5) never enters the flush. H=14 gives 2 and 3, the + # non-multiple H=13 2 and 3, the small H=5 1 and 2 -- each shifts the + # accumulator slots the flush drains. + _Shape("dw_k7_N2H14W14_g64", N=2, H=14, W=14, groups=64, cpg=1, KH=7, KW=7, PAD=3), + _Shape( + "dw_k7s2_N2H14W14_g64", + N=2, + H=14, + W=14, + groups=64, + cpg=1, + KH=7, + KW=7, + PAD=3, + stride=2, + ), + _Shape( + "dw_k7s2_N2H13W11_g64", + N=2, + H=13, + W=11, + groups=64, + cpg=1, + KH=7, + KW=7, + PAD=3, + stride=2, + ), + _Shape( + "dw_k7s2_N1H5W9_g64", + N=1, + H=5, + W=9, + groups=64, + cpg=1, + KH=7, + KW=7, + PAD=3, + stride=2, + ), ] @@ -124,6 +165,22 @@ class _Shape: _Shape("sp_dw_N2H14W14_g3", N=2, H=14, W=14, groups=3, cpg=1), _Shape("sp_dw_N2H14W14_g64", N=2, H=14, W=14, groups=64, cpg=1), _Shape("sp_dw_N2H14W14_g3_s2", N=2, H=14, W=14, groups=3, cpg=1, stride=2), + # 7x7: the kernel prefetches PF = min(KH, ceil(16 / KW)) = 3 rows, fewer + # than KH, so the row window rotates within a filter's rows -- the 3x3 + # cases above (PF = KH) never do. + _Shape("sp_dw_k7_N2H14W14_g9", N=2, H=14, W=14, groups=9, cpg=1, KH=7, KW=7, PAD=3), + _Shape( + "sp_dw_k7s2_N2H13W11_g3", + N=2, + H=13, + W=11, + groups=3, + cpg=1, + KH=7, + KW=7, + PAD=3, + stride=2, + ), ] @@ -275,12 +332,17 @@ def _run_grouped_one(arch: str, shape: _Shape, dtype: str = "fp16") -> Tuple[boo return True, "" -def _run_depthwise_one(arch: str, shape: _Shape) -> Tuple[bool, str]: - """Build, compile, launch, and verify one depthwise direct-conv kernel. +def _run_depthwise_one( + arch: str, shape: _Shape, dtype: str = "fp16", spatial: bool = False +) -> Tuple[bool, str]: + """Build, compile, launch, and verify one depthwise forward kernel. - Uses ``DirectDepthwiseSpec`` (cpg = kpg = 1). Stride must be 1. + ``DirectDepthwiseSpec`` (one channel per lane) by default, + ``DirectDepthwiseSpatialSpec`` (lanes split over channels and output + columns, groups <= wave_size) with ``spatial``. cpg = kpg = 1. - Returns ``(passed, reason)``. + Returns ``(passed, reason)``. ``reason`` starts with ``"skip "`` when a + spatial shape is architecturally unsupported. """ import torch @@ -290,8 +352,11 @@ def _run_depthwise_one(arch: str, shape: _Shape) -> Tuple[bool, str]: from kernels.common.conv_direct_grouped import ( direct_launch_geometry, DirectConvProblem, + DirectDepthwiseSpatialSpec, DirectDepthwiseSpec, build_direct_depthwise, + build_direct_depthwise_spatial, + is_valid_depthwise_spatial_spec, is_valid_depthwise_spec, ) from rocke.runtime import synchronize_and_release @@ -311,142 +376,26 @@ def _run_depthwise_one(arch: str, shape: _Shape) -> Tuple[bool, str]: KW=shape.KW, PAD=shape.PAD, stride=shape.stride, + dtype=dtype, ) - spec = DirectDepthwiseSpec( - problem=p, - name=f"test_direct_dw_{shape.id}", - ) - - ok, reason = is_valid_depthwise_spec(spec, arch=arch) - if not ok: - return False, f"invalid spec (shapes should be pre-validated): {reason}" - - try: - kernel = build_direct_depthwise(spec, arch=arch) - except ValueError as e: - return False, f"build failed (shapes should be pre-validated): {e}" - - try: - artifact = compile_kernel(kernel, arch=arch) - except Exception as e: - return False, f"compile failed: {e}" - - torch.manual_seed(0) - total_c = shape.groups - total_k = shape.groups - A_t = torch.empty(p.N, p.H, p.W, total_c, dtype=torch.float16).uniform_(-1.0, 1.0) - B_t = torch.empty(total_k, p.KH, p.KW, 1, dtype=torch.float16).uniform_(-1.0, 1.0) - D_t = torch.empty(p.N, p.H, p.W, total_k, dtype=torch.float16) - - ref = _conv_ref_grouped(A_t, B_t, p) - - rt = Runtime() - A_dev = rt.alloc(A_t.nbytes) - B_dev = rt.alloc(B_t.nbytes) - D_dev = rt.alloc(D_t.nbytes) - rt.memcpy_h2d(A_dev, _u8(A_t), A_t.nbytes) - rt.memcpy_h2d(B_dev, _u8(B_t), B_t.nbytes) - rt.memset(D_dev, 0, D_t.nbytes) - - sig = conv_direct_args_signature("fp16") - _direct_args = ConvArgs.from_problem(p) - try: - launcher = KernelLauncher( - hsaco=artifact.hsaco, - kernel_name=artifact.kernel_name, - signature=sig, + if spatial: + spec = DirectDepthwiseSpatialSpec( + problem=p, name=f"test_direct_sp_dw_{shape.id}" ) - except HipError as e: - rt.free(A_dev) - rt.free(B_dev) - rt.free(D_dev) - return False, f"kernel load failed: {e}" - - grid, block = direct_launch_geometry(spec) - - values = _direct_args.to_launch_values( - int(A_dev), - int(B_dev), - int(D_dev), - A_t.nbytes, - B_t.nbytes, - D_t.nbytes, - ) - launcher(values, config=LaunchConfig(grid=grid, block=block, fence=True)) - - D_cpu = torch.empty_like(D_t) - rt.memcpy_d2h(_u8(D_cpu), D_dev, D_t.nbytes) - rt.free(A_dev) - rt.free(B_dev) - rt.free(D_dev) - synchronize_and_release(0) - - out_f32 = D_cpu.float() - ref_f32 = ref.float().cpu() - abs_diff = (out_f32 - ref_f32).abs() - ref_scale = ref_f32.abs().max().clamp(min=1.0) - rel_err = float(abs_diff.max() / ref_scale) - passed = rel_err < _TOL - if not passed: - return False, f"rel_err={rel_err:.3e} > tol={_TOL:.1e}" - print( - f" PASS {shape.id} {arch} rel_err={rel_err:.2e}", - flush=True, - ) - return True, "" - - -def _run_depthwise_spatial_one(arch: str, shape: _Shape) -> Tuple[bool, str]: - """Build, compile, launch, and verify one depthwise-spatial kernel. - - Uses ``DirectDepthwiseSpatialSpec`` (cpg = kpg = 1, groups <= wave_size). - - Returns ``(passed, reason)``. ``reason`` starts with ``"skip "`` when - the combination is architecturally unsupported. - """ - import torch - - from rocke import compile_kernel - from kernels.common.conv_args import ConvArgs - from kernels.common.conv_abi import conv_direct_args_signature - from kernels.common.conv_direct_grouped import ( - direct_launch_geometry, - DirectConvProblem, - DirectDepthwiseSpatialSpec, - build_direct_depthwise_spatial, - is_valid_depthwise_spatial_spec, - ) - from rocke.runtime import synchronize_and_release - from rocke.runtime.hip_module import HipError, Runtime - from rocke.runtime.launcher import KernelLauncher, LaunchConfig - - assert shape.cpg == 1, "spatial depthwise path requires cpg=1" - - p = DirectConvProblem( - N=shape.N, - H=shape.H, - W=shape.W, - groups=shape.groups, - cpg=1, - kpg=1, - KH=shape.KH, - KW=shape.KW, - PAD=shape.PAD, - stride=shape.stride, - ) - - spec = DirectDepthwiseSpatialSpec( - problem=p, - name=f"test_direct_sp_dw_{shape.id}", - ) - - ok, reason = is_valid_depthwise_spatial_spec(spec, arch=arch) - if not ok: - return False, f"skip {reason}" + ok, reason = is_valid_depthwise_spatial_spec(spec, arch=arch) + if not ok: + return False, f"skip {reason}" + build = build_direct_depthwise_spatial + else: + spec = DirectDepthwiseSpec(problem=p, name=f"test_direct_dw_{shape.id}") + ok, reason = is_valid_depthwise_spec(spec, arch=arch) + if not ok: + return False, f"invalid spec (shapes should be pre-validated): {reason}" + build = build_direct_depthwise try: - kernel = build_direct_depthwise_spatial(spec, arch=arch) + kernel = build(spec, arch=arch) except ValueError as e: return False, f"build failed: {e}" @@ -458,9 +407,10 @@ def _run_depthwise_spatial_one(arch: str, shape: _Shape) -> Tuple[bool, str]: torch.manual_seed(0) total_c = shape.groups total_k = shape.groups - A_t = torch.empty(p.N, p.H, p.W, total_c, dtype=torch.float16).uniform_(-1.0, 1.0) - B_t = torch.empty(total_k, p.KH, p.KW, 1, dtype=torch.float16).uniform_(-1.0, 1.0) - D_t = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=torch.float16) + _td = torch.bfloat16 if dtype == "bf16" else torch.float16 + A_t = torch.empty(p.N, p.H, p.W, total_c, dtype=_td).uniform_(-1.0, 1.0) + B_t = torch.empty(total_k, p.KH, p.KW, 1, dtype=_td).uniform_(-1.0, 1.0) + D_t = torch.empty(p.N, p.Ho, p.Wo, total_k, dtype=_td) ref = _conv_ref_grouped(A_t, B_t, p) @@ -472,7 +422,7 @@ def _run_depthwise_spatial_one(arch: str, shape: _Shape) -> Tuple[bool, str]: rt.memcpy_h2d(B_dev, _u8(B_t), B_t.nbytes) rt.memset(D_dev, 0, D_t.nbytes) - sig = conv_direct_args_signature("fp16") + sig = conv_direct_args_signature(dtype) _direct_args = ConvArgs.from_problem(p) try: launcher = KernelLauncher( @@ -510,11 +460,11 @@ def _run_depthwise_spatial_one(arch: str, shape: _Shape) -> Tuple[bool, str]: abs_diff = (out_f32 - ref_f32).abs() ref_scale = ref_f32.abs().max().clamp(min=1.0) rel_err = float(abs_diff.max() / ref_scale) - passed = rel_err < _TOL - if not passed: - return False, f"rel_err={rel_err:.3e} > tol={_TOL:.1e}" + tol = _TOL_BF16 if dtype == "bf16" else _TOL + if not rel_err < tol: + return False, f"rel_err={rel_err:.3e} > tol={tol:.1e}" print( - f" PASS {shape.id} {arch} rel_err={rel_err:.2e}", + f" PASS {shape.id} {arch} {dtype} rel_err={rel_err:.2e}", flush=True, ) return True, "" @@ -578,7 +528,7 @@ def test_depthwise(self): self._run_depthwise(s) def _run_depthwise_spatial(self, shape: _Shape) -> None: - passed, reason = _run_depthwise_spatial_one(GPU_ARCH, shape) + passed, reason = _run_depthwise_one(GPU_ARCH, shape, spatial=True) if reason.startswith("skip"): self.skipTest(reason) self.assertTrue( @@ -640,6 +590,44 @@ def test_depthwise_spatial(self): ), ] +# Ho-streaming depthwise dgrad (DirectDepthwiseDgradStreamSpec). Each loop +# iteration consumes KH dY rows with the next KH prefetched, and a dX row is +# flushed from its circular slot once its last (ho, r) pair is in, so the +# cases cover: stride 1 and 2, a filter taller than the remaining dY rows +# (7x7 on Ho=7), a height that is not a multiple of KH, non-"same" padding, +# and a multi-wave block. +_DW_DGRAD_STREAM_SHAPES: List[_Shape] = [ + _Shape("dw_dgrad_stream_s1_N2H14W14_g64", N=2, H=14, W=14, groups=64, cpg=1), + _Shape( + "dw_dgrad_stream_s2_N2H14W14_g64", N=2, H=14, W=14, groups=64, cpg=1, stride=2 + ), + _Shape( + "dw_dgrad_stream_k7s2_N2H13W11_g64", + N=2, + H=13, + W=11, + groups=64, + cpg=1, + KH=7, + KW=7, + PAD=3, + stride=2, + ), + _Shape( + "dw_dgrad_stream_k5p4_s2_N2H14W14_g64", + N=2, + H=14, + W=14, + groups=64, + cpg=1, + KH=5, + KW=5, + PAD=4, + stride=2, + ), + _Shape("dw_dgrad_stream_k3p0_N2H9W9_g96", N=2, H=9, W=9, groups=96, cpg=1, PAD=0), +] + def _run_dgrad_one(arch: str, shape: _Shape, dtype: str = "fp16") -> Tuple[bool, str]: """Build, compile, launch, and verify the direct dgrad kernel. @@ -776,8 +764,20 @@ def _run_dgrad_one(arch: str, shape: _Shape, dtype: str = "fp16") -> Tuple[bool, return True, "" -def _run_dw_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: - """Build, compile, launch, and verify the direct depthwise dgrad kernel.""" +def _run_dw_dgrad_one( + arch: str, + shape: _Shape, + dtype: str = "fp16", + streaming: bool = False, + block_waves: int = 1, +) -> Tuple[bool, str]: + """Build, compile, launch, and verify a direct depthwise dgrad kernel. + + ``DirectDepthwiseDgradSpec`` by default; with ``streaming`` the ho-streaming + ``DirectDepthwiseDgradStreamSpec`` the AOT cache and the benchmark sweep + build, whose prefetched dY row window and circular dX accumulator slots + are only exercised by running it. + """ import torch from rocke import compile_kernel @@ -787,7 +787,9 @@ def _run_dw_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: direct_launch_geometry, DirectConvProblem, DirectDepthwiseDgradSpec, + DirectDepthwiseDgradStreamSpec, build_direct_depthwise_dgrad, + build_direct_depthwise_dgrad_streaming, ) from rocke.runtime import synchronize_and_release from rocke.runtime.hip_module import HipError, Runtime @@ -804,11 +806,23 @@ def _run_dw_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: KW=shape.KW, PAD=shape.PAD, stride=shape.stride, + dtype=dtype, ) - spec = DirectDepthwiseDgradSpec(problem=p, name=f"test_dw_dgrad_{shape.id}") + if streaming: + spec = DirectDepthwiseDgradStreamSpec( + problem=p, + name=f"test_dw_dgrad_stream_{shape.id}", + block_waves=block_waves, + ) + build = build_direct_depthwise_dgrad_streaming + else: + spec = DirectDepthwiseDgradSpec( + problem=p, name=f"test_dw_dgrad_{shape.id}", block_waves=block_waves + ) + build = build_direct_depthwise_dgrad try: - kernel = build_direct_depthwise_dgrad(spec, arch=arch) + kernel = build(spec, arch=arch) except ValueError as e: return False, f"build failed: {e}" @@ -819,10 +833,11 @@ def _run_dw_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: torch.manual_seed(42) total_c = shape.groups # cpg=kpg=1 + _td = torch.bfloat16 if dtype == "bf16" else torch.float16 - dY = torch.empty(p.N, p.Ho, p.Wo, total_c, dtype=torch.float16).uniform_(-0.5, 0.5) - W = torch.empty(total_c, p.KH, p.KW, 1, dtype=torch.float16).uniform_(-0.5, 0.5) - dX = torch.zeros(p.N, p.H, p.W, total_c, dtype=torch.float16) + dY = torch.empty(p.N, p.Ho, p.Wo, total_c, dtype=_td).uniform_(-0.5, 0.5) + W = torch.empty(total_c, p.KH, p.KW, 1, dtype=_td).uniform_(-0.5, 0.5) + dX = torch.zeros(p.N, p.H, p.W, total_c, dtype=_td) dY_nchw = dY.permute(0, 3, 1, 2).float() W_nchw = W.permute(0, 3, 1, 2).float() @@ -846,7 +861,7 @@ def _run_dw_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: rt.memcpy_h2d(W_dev, _u8(W), W.nbytes) rt.memset(dX_dev, 0, dX.nbytes) - sig = conv_direct_args_signature("fp16", direction="dgrad") + sig = conv_direct_args_signature(dtype, direction="dgrad") _direct_args = ConvArgs.from_problem(p, direction="dgrad") try: launcher = KernelLauncher( @@ -884,9 +899,10 @@ def _run_dw_dgrad_one(arch: str, shape: _Shape) -> Tuple[bool, str]: abs_diff = (out_f32 - ref_f32).abs() ref_scale = ref_f32.abs().max().clamp(min=1.0) rel_err = float(abs_diff.max() / ref_scale) - if not (rel_err < _TOL): - return False, f"rel_err={rel_err:.3e} > tol={_TOL:.1e}" - print(f" PASS {shape.id} {arch} rel_err={rel_err:.2e}", flush=True) + tol = _TOL_BF16 if dtype == "bf16" else _TOL + if not (rel_err < tol): + return False, f"rel_err={rel_err:.3e} > tol={tol:.1e}" + print(f" PASS {shape.id} {arch} {dtype} rel_err={rel_err:.2e}", flush=True) return True, "" @@ -916,6 +932,18 @@ def test_dw_dgrad(self): with self.subTest(shape=s.id): self._run_dw_dgrad(s) + def test_dw_dgrad_streaming(self): + for s in _DW_DGRAD_STREAM_SHAPES: + for waves in (1, 2): + with self.subTest(shape=s.id, block_waves=waves): + passed, reason = _run_dw_dgrad_one( + GPU_ARCH, s, streaming=True, block_waves=waves + ) + self.assertTrue( + passed, + f"FAIL dw_dgrad_stream {s.id} bw{waves}wv on {GPU_ARCH}: {reason}", + ) + # --------------------------------------------------------------------------- # bf16 correctness tests @@ -981,6 +1009,38 @@ def test_bf16_dgrad(self): with self.subTest(shape=s.id): self._run_dgrad(s) + def test_bf16_depthwise(self): + for s in _SHAPES: + if s.cpg == 1: + with self.subTest(shape=s.id): + passed, reason = _run_depthwise_one(GPU_ARCH, s, dtype="bf16") + self.assertTrue( + passed, f"FAIL bf16 depthwise {s.id} on {GPU_ARCH}: {reason}" + ) + + def test_bf16_depthwise_spatial(self): + for s in _SPATIAL_SHAPES: + with self.subTest(shape=s.id): + passed, reason = _run_depthwise_one( + GPU_ARCH, s, dtype="bf16", spatial=True + ) + if reason.startswith("skip"): + self.skipTest(reason) + self.assertTrue( + passed, + f"FAIL bf16 depthwise spatial {s.id} on {GPU_ARCH}: {reason}", + ) + + def test_bf16_dw_dgrad_streaming(self): + for s in _DW_DGRAD_STREAM_SHAPES: + with self.subTest(shape=s.id): + passed, reason = _run_dw_dgrad_one( + GPU_ARCH, s, dtype="bf16", streaming=True + ) + self.assertTrue( + passed, f"FAIL bf16 dw_dgrad_stream {s.id} on {GPU_ARCH}: {reason}" + ) + # --------------------------------------------------------------------------- # Non-grouped (groups == 1) direct conv @@ -1452,6 +1512,43 @@ def test_cpg4_bf16_rejected(self): with self.assertRaises(ValueError): spec.validate() + def test_depthwise_filter_limit(self): + """Both depthwise forward variants accept filters up to 32x32 and reject + larger ones, as the C++ validators do: its builders keep the weights in + tables of ``ROCKE_DCONV_DW_MAX_K{H,W}`` entries.""" + from kernels.common.conv_direct_grouped import ( + DirectConvProblem, + DirectDepthwiseSpatialSpec, + DirectDepthwiseSpec, + is_valid_depthwise_spatial_spec, + is_valid_depthwise_spec, + ) + + cases = ( + (DirectDepthwiseSpec, is_valid_depthwise_spec, 64), + (DirectDepthwiseSpatialSpec, is_valid_depthwise_spatial_spec, 3), + ) + for spec_cls, is_valid, groups in cases: + # 31 is the largest odd extent under the cap (forward needs odd). + for kh, kw in ((31, 31), (33, 33), (33, 3), (3, 33)): + p = DirectConvProblem( + N=1, + H=40, + W=40, + groups=groups, + cpg=1, + kpg=1, + KH=kh, + KW=kw, + PAD=(kh - 1) // 2, + ) + with self.subTest(spec=spec_cls.__name__, KH=kh, KW=kw): + ok, why = is_valid(spec_cls(problem=p), "gfx950") + fits = kh <= 32 and kw <= 32 + self.assertEqual(ok, fits, why) + if not fits: + self.assertIn("must be in 1..32", why) + # --------------------------------------------------------------------------- # Wgrad shapes diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py index bfbce7baf131..94d258972932 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_direct_kernel_cache.py @@ -100,6 +100,30 @@ def test_group_count_is_runtime(cache, groups): ) +@pytest.mark.parametrize("stride", [1, 2]) +@pytest.mark.parametrize("groups", [3, 9]) +def test_small_group_depthwise_offers_spatial(cache, groups, stride): + """A depthwise problem with fewer groups than a wave gets the spatial + kernel, launched for its own group count: the binary was compiled on the + groups=8 probe, but each wave covers ``wave_size // groups`` columns of + the real problem.""" + # At W=226 every grid differs from the one the groups=8 probe would get. + problem = _problem(H=8, W=226, groups=groups, cpg=1, kpg=1, stride=stride) + plans, _ = dks.direct_plans(cache, problem, "fwd", _ARCH) + spatial = [p for p in plans if p.identity.algorithm == "direct_depthwise_spatial"] + assert {json.loads(p.identity.knobs)["block_waves"] for p in spatial} == set( + dks.DW_BLOCK_WAVES + ) + for plan in spatial: + waves = json.loads(plan.identity.knobs)["block_waves"] + block_w = waves * (plan.identity.wave_size // groups) + (step,) = plan.steps + assert step.grid == (-(-problem.Wo // block_w), 1, problem.N) + assert step.block == (waves * plan.identity.wave_size, 1, 1) + probe_w = waves * (plan.identity.wave_size // 8) + assert step.grid[0] != -(-problem.Wo // probe_w) + + def _nongrouped(plans): return [p for p in plans if p.identity.algorithm == "direct_nongrouped"] @@ -125,6 +149,45 @@ def test_nongrouped_channels_are_runtime(cache, C, K, H, W): assert step.grid == spec.grid() +def test_nongrouped_grid_skips_spilling_tiles(jobs): + """Tiles that would spill are not cached, and the cut leaves both atoms and + both strides covered.""" + from kernels.common.conv_direct_nongrouped import ( + DirectNongroupedConvSpec, + nongrouped_register_reason, + ) + + built = [j for j in jobs if j.identity.algorithm == "direct_nongrouped"] + assert built + for job in built: + caps = dks.DirectCaps(**job.caps) + spec = dks.make_spec( + "direct_nongrouped", + dks.probe_problem(caps, dks.DIRECT_DTYPE, "direct_nongrouped"), + job.spec_kwargs, + ) + assert nongrouped_register_reason(spec, _ARCH) is None, spec.kernel_name() + assert {json.loads(j.identity.knobs)["atom"] for j in built} == { + "32x32x16", + "16x16x32", + } + assert {j.identity.stride_h for j in built} == {1, 2} + + # A tile filling the whole accumulator file is rejected however few waves + # share the SIMD: it spilled in every such build on gfx950. + full = DirectNongroupedConvSpec( + problem=_problem(groups=1, cpg=64, kpg=256), + tile_h=8, + tile_w=32, + tile_k=256, + ck=16, + waves_m=2, + waves_n=2, + ) + assert full.acc_vgprs == 256 + assert "accumulator file" in nongrouped_register_reason(full, _ARCH) + + def test_nongrouped_only_serves_one_group(cache): plans, rejected = dks.direct_plans(cache, _problem(groups=16), "fwd", _ARCH) assert not _nongrouped(plans) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py index 70669335a9d3..a3dc8160e3ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_early_stop.py @@ -118,3 +118,31 @@ def test_all_stopped_against_a_seed_says_so(monkeypatch): assert stop.measure(slow, warmup=3, iters=10) is None assert "all 1 early-stopped" in stop.summary() assert "1.000 ms" in stop.summary() + + +def test_failed_kernels_neither_prune_nor_record(monkeypatch, tmp_path): + """A kernel that failed --verify is timed but cannot become the best: a + wrong-but-fast one would otherwise prune every correct kernel, here and in + the sweep seeded from the record.""" + import json + from argparse import Namespace + + from kernels.common._conv_implicit_gemm_common import ConvProblem + + wrong, ok, slow = object(), object(), object() + _fake_timer(monkeypatch, {wrong: 0.1, ok: 2.0, slow: 9.0}) + problem = ConvProblem(N=2, Hi=16, Wi=16, C=32, K=32, Y=3, X=3, pH=1, pW=1) + record = tmp_path / "best.json" + + es = _reset_handoff(monkeypatch) + args = Namespace(early_stop=5.0, early_stop_after=1, early_stop_record=str(record)) + stop = es.EarlyStop.for_case(args, problem, "fp16", "fwd") + assert stop.measure(wrong, warmup=3, iters=10, passed=False) == 0.1 + assert stop.best_ms is None and stop.n_measured == 0 + assert not record.exists() + # The failed kernel did not count towards `after` either: this one is the + # first eligible kernel and is measured in full. + assert stop.measure(slow, warmup=3, iters=10, passed=True) == 9.0 + assert stop.measure(ok, warmup=3, iters=10) == 2.0 + assert stop.best_ms == 2.0 + assert list(json.loads(record.read_text()).values()) == [2.0] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h index bfe543fdb048..01908b873051 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped.h @@ -487,6 +487,13 @@ typedef struct rocke_direct_conv_wgrad_spec #define ROCKE_DCONV_WGRAD_MAX_KH 8 #define ROCKE_DCONV_WGRAD_MAX_KW 8 +/* Largest filter the depthwise forward variants (regular and spatial) accept. + * Their builders hold the KH x KW weights in fixed-size register tables sized + * by these, so both validators enforce the cap (Python: _DW_MAX_KH / + * _DW_MAX_KW) and a validated spec always builds. */ +#define ROCKE_DCONV_DW_MAX_KH 32 +#define ROCKE_DCONV_DW_MAX_KW 32 + rocke_direct_conv_wgrad_spec_t rocke_direct_conv_wgrad_spec_default(void); /* @property block_k -> waves_k * wave_tile_k. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h index 3ea225e89ede..fd0e1c7a96e0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h @@ -653,8 +653,8 @@ typedef struct rocke_dconv_32c_ctx * KH*KW scalar FMA products into a BLOCK_W * KH circular register array. * ===================================================================== */ #define ROCKE_DCONV_DW_MAX_BLOCK_W 64 -#define ROCKE_DCONV_DW_MAX_KH 32 -#define ROCKE_DCONV_DW_MAX_KW 32 +/* ROCKE_DCONV_DW_MAX_KH / _MAX_KW live in the public header: they bound the + * weight tables below AND are enforced by both validators. */ typedef struct rocke_dconv_dw_ctx { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp index 8a73188036d7..2ef52d045af0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp @@ -1195,6 +1195,24 @@ bool rocke_direct_depthwise_is_valid_spec(const rocke_direct_depthwise_spec_t* s } return false; } + if(p->KH < 1 || p->KH > ROCKE_DCONV_DW_MAX_KH) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "KH must be in 1..%d (got %d)", ROCKE_DCONV_DW_MAX_KH, p->KH); + } + return false; + } + if(p->KW < 1 || p->KW > ROCKE_DCONV_DW_MAX_KW) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "KW must be in 1..%d (got %d)", ROCKE_DCONV_DW_MAX_KW, p->KW); + } + return false; + } block_ch = rocke_direct_depthwise_block_ch(spec); (void)block_ch; if(reason && reason_cap > 0) @@ -1317,6 +1335,24 @@ bool rocke_direct_depthwise_spatial_is_valid_spec(const rocke_direct_depthwise_s } return false; } + if(p->KH < 1 || p->KH > ROCKE_DCONV_DW_MAX_KH) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "KH must be in 1..%d (got %d)", ROCKE_DCONV_DW_MAX_KH, p->KH); + } + return false; + } + if(p->KW < 1 || p->KW > ROCKE_DCONV_DW_MAX_KW) + { + if(reason && reason_cap > 0) + { + snprintf( + reason, reason_cap, "KW must be in 1..%d (got %d)", ROCKE_DCONV_DW_MAX_KW, p->KW); + } + return false; + } if(p->groups > spec->wave_size) { if(reason && reason_cap > 0) From d5930db5bfd8f5970a7cd8d0c200e5c0bfc1ac42 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 6 Oct 2026 23:10:24 +0000 Subject: [PATCH 7/7] ir --- .../rocke_representative_ir_sha256.json | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index f486dfda8f5f..cf8b49b82e4a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -3184,10 +3184,10 @@ }, "conv_direct/gfx942/dw_n4h16": { "arch": "gfx942", - "bytes": 180823, + "bytes": 180883, "family": "conv_direct", "kernel_name": "irhash_directdw_942_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "a930c98b2e0e1d4eb6dc4eb059af54ddd58fdd9450731be6d993a76a668d8ea5", + "sha256": "ad41b7cd3109c200f5dac3878b96039e4217c6d4f7a2df95a02e04f08ed0cf09", "status": "ok" }, "conv_direct/gfx950/16c_n32h8_bg4_fold": { @@ -3240,10 +3240,10 @@ }, "conv_direct/gfx950/dw_n4h16": { "arch": "gfx950", - "bytes": 180823, + "bytes": 180883, "family": "conv_direct", "kernel_name": "irhash_directdw_950_a_N4H16W16_g64_c1k1_bw16_bw1wv", - "sha256": "b0aff29da91dbf33235cd2ad25f3e27df7892f72f864ecfa699c297033c379f0", + "sha256": "c3e27f9925174e0f81fb97f2f29ca9a0aced69757a138bcf49672bc9a0c09e59", "status": "ok" }, "conv_direct_dgrad/gfx942/dgrad_n1h8_cpg16_bg8": { @@ -3288,42 +3288,42 @@ }, "conv_direct_nongrouped/gfx942/fp16_a32x32x8": { "arch": "gfx942", - "bytes": 72193, + "bytes": 72253, "family": "conv_direct_nongrouped", "kernel_name": "irhash_direct_nongrouped_fp16_a32x32x8_N1H16W32_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x8_g8_iglp0", - "sha256": "ec5e2d3974fbe6d94c8d194ce40a4d8acf98f114457447fad3e1bac2c8687f4c", + "sha256": "8c9f1c90726eeb3d44e351ca3a09babebb9769ec097913c1f5dcddf847cd49ea", "status": "ok" }, "conv_direct_nongrouped/gfx950/bf16_a16x16x32_partial_w": { "arch": "gfx950", - "bytes": 148220, + "bytes": 148280, "family": "conv_direct_nongrouped", "kernel_name": "irhash_direct_nongrouped_bf16_a16x16x32_partial_w_N2H20W40_g1_c64k64_t8x48x32_ck32_w1x2_a16x16x32_g8_iglp0_bf16", - "sha256": "7e24f74e8ef1671f9d762cb62ca0a9631c0d125ebb4052ea0abc1ccbff268c81", + "sha256": "b6844959b92d991d7319361d143653d80a8d8bdbd886b395845660aff678ffe7", "status": "ok" }, "conv_direct_nongrouped/gfx950/bf16_noswizzle_we3": { "arch": "gfx950", - "bytes": 79854, + "bytes": 79914, "family": "conv_direct_nongrouped", "kernel_name": "irhash_direct_nongrouped_bf16_noswizzle_we3_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_gnone_we3_bf16", - "sha256": "3f574bb890191b9d5656120ea806af692b9d94d542608a87425211a91c290db2", + "sha256": "4d637b2a3ea4f0da1df58950d34e960bfa05212e76468d4129a864b1abdcf994", "status": "ok" }, "conv_direct_nongrouped/gfx950/bf16_t8x32x64_iglp0": { "arch": "gfx950", - "bytes": 80751, + "bytes": 80811, "family": "conv_direct_nongrouped", "kernel_name": "irhash_direct_nongrouped_bf16_t8x32x64_iglp0_N2H16W32_g1_c64k128_t8x32x64_ck32_w2x2_a32x32x16_g8_iglp0_bf16", - "sha256": "b4bb36d117c7af0b63e75b690fc192c9a12098665b5df20d96b050629531117e", + "sha256": "8fcc44fc6ce064760c476a2bc679f6ad2c108611e56e9c32be6fcc74b645f919", "status": "ok" }, "conv_direct_nongrouped/gfx950/fp16_db_s2": { "arch": "gfx950", - "bytes": 74753, + "bytes": 74813, "family": "conv_direct_nongrouped", "kernel_name": "irhash_direct_nongrouped_fp16_db_s2_N1H32W64_g1_c64k64_t8x32x64_ck16_w2x2_a32x32x16_g8_db_iglp0", - "sha256": "ee53aba1abe0cbbb3fea79653bc97f1cad65772492fec460beec20df2158b3c0", + "sha256": "21f519428620240701ef490ea0c66989d993fa27a28f234b3b86069477530a8e", "status": "ok" }, "conv_wgrad/gfx1151/n1h8c16k32r3": {