From 2f92cd556e326f9e93f8274776e46e8bdbaf064c Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Sun, 25 Oct 2026 12:43:29 +0000 Subject: [PATCH 1/9] bench --- .../benchmarks/common/batched_timing.py | 370 ++++++++++++++++ .../common/benchmark_conv_compare.py | 120 +++++- .../common/benchmark_direct_conv.py | 391 +++++++++++++---- .../common/benchmark_implicit_gemm_conv.py | 389 ++++++++++++++++- .../benchmarks/common/direct_kernel_sweep.py | 12 +- .../library/benchmarks/common/early_stop.py | 19 + .../library/benchmarks/common/kernel_cache.py | 132 +++++- .../library/benchmarks/common/module_cache.py | 181 ++++++++ .../library/kernels/conv/kernel_identity.py | 141 ++++++ .../library/tests/test_batched_timing.py | 264 ++++++++++++ .../library/tests/test_capability_keys.py | 402 ++++++++++++++++++ .../rocke/library/tests/test_module_cache.py | 160 +++++++ .../platform/python/rocke/runtime/launcher.py | 40 +- 13 files changed, 2519 insertions(+), 102 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py new file mode 100644 index 000000000000..876b0b8720a1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py @@ -0,0 +1,370 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Rank a sweep's kernels in two rounds, with one device drain per chunk. + +:class:`~benchmarks.common.early_stop.EarlyStop` times every kernel the same +way -- warmup launches, then a timed loop -- and each of those is its own +synchronous round trip: :func:`rocke.runtime.time_launches` syncs the device, +records a pair of events, waits on the second and drains the stream. On a +sweep whose kernels run for tens of microseconds those drains, not the +kernels, are the clock. This module times a whole chunk of kernels between two +drains instead, and spends the full warmup/iters protocol only on the few +candidates that could still rank. + +The protocol +------------ + +**Round 1, the probe.** Every candidate runs ``probe_warmup`` warmup launches +and ``probe_iters`` timed ones, all queued on one stream with no host sync in +between; a pair of HIP events brackets each candidate's timed launches and the +chunk is drained once at the end. Only the best ``finalists`` candidates are +carried forward, held with their bound launches so round 2 reloads nothing. + +**Round 2, the measurement.** The finalists are timed exactly as the +per-kernel path does, with the caller's full ``warmup`` and ``iters``. Every +reported number comes from that round: a probe time ranks a candidate, it +never describes one. + +Why the probe still warms up +---------------------------- + +Two cold effects sit on a first launch, and they need different answers. + +*Per kernel*: the first launch of a code object faults its instructions in and +allocates its scratch. That cost lands on whichever launch goes first, so each +candidate gets ``probe_warmup`` launches of its own before its events are +recorded. The probe is cheap because it drops the **syncs**, not the warmup. + +*Per device*: clocks ramp over the first few milliseconds of sustained work. +No per-candidate warmup can cover that one -- it penalises whoever is measured +first rather than whoever is slow, and in a sweep that is simply whoever the +cache listed first. :func:`prime_device` runs one candidate enough times to +bring the clocks up before the first measured chunk, and the chunk drains are +too short for them to fall back. + +What is left is sampling noise on a one- or two-launch measurement, which is +why the cut is deliberately loose: the probe only has to get a candidate into +the top ``finalists``, not to rank it there. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, List, Optional, Sequence + +# Candidates timed between two device drains. The chunk bounds three things at +# once: the HIP events live at any moment (two per candidate), how long the +# shared output buffer stays in flight, and the blast radius of a kernel that +# faults -- a chunk that raises is re-timed one candidate at a time so the +# failure is charged to the one that caused it. +DEFAULT_CHUNK = 128 + +# Launches issued before a case's first measured chunk; see "Per device" above. +DEFAULT_PRIMER = 50 + +# Candidates carried from the probe into the full-fidelity round. They are held +# with their bound launches, so this also bounds what stays bound at once. +DEFAULT_FINALISTS = 64 + +# The probe's own protocol. One warmup launch covers the per-kernel cold cost; +# two timed launches cost little more and halve the sampling noise. +DEFAULT_PROBE_WARMUP = 1 +DEFAULT_PROBE_ITERS = 2 + + +@dataclass +class Ranked: + """What the ranking measured for one candidate. + + ``probe_ms`` is set for every candidate that ran at all; ``ms`` only for a + finalist, and it is the number to report. ``failed`` marks a candidate + whose launch raised -- its ``probe_ms`` is meaningless and it never ranks. + ``passed`` is what the caller's finalist hook returned (``None`` when it + did not check). + """ + + index: int + probe_ms: Optional[float] = None + ms: Optional[float] = None + passed: Optional[bool] = None + failed: bool = False + + @property + def is_finalist(self) -> bool: + return self.ms is not None + + +def prime_device( + run: Callable[[], None], *, launches: int = DEFAULT_PRIMER, stream: int = 0 +) -> None: + """Bring the device's clocks up before the first measured chunk. + + ``run`` is any candidate's bound launch -- the ramp is a property of the + device, not of the kernel, so which one does the priming does not matter. + """ + from rocke.runtime import synchronize_and_release + from rocke.runtime.launcher import no_fence + + if launches <= 0: + return + with no_fence(): + for _ in range(int(launches)): + run() + synchronize_and_release(stream) + + +def _time_chunk( + runs: Sequence[Callable[[], None]], + *, + warmup: int, + iters: int, + stream: int, +) -> List[float]: + """Time every callable in ``runs`` between one pair of device drains. + + The launches are queued in order on one stream, so same-stream FIFO makes + each candidate's event pair bracket exactly its own timed launches. Only + the last event is waited on. + + Raises whatever a launch raises, after draining: the queued launches are + fire-and-forget, so the caller must not free their buffers until the + device is idle. + """ + from rocke.runtime import resolve_stream + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import no_fence + + rt = Runtime() + resolved = resolve_stream(stream) + events = [(rt.event(), rt.event()) for _ in runs] + try: + with no_fence(): + for run, (start, end) in zip(runs, events): + for _ in range(warmup): + run() + start.record(stream=resolved) + for _ in range(iters): + run() + end.record(stream=resolved) + events[-1][1].synchronize() + out = [start.elapsed_to(end) / iters for start, end in events] + except Exception as exc: # noqa: BLE001 - re-raised after the drain + _drain_after_failure(rt, exc) + _destroy(events) + raise + _destroy(events) + rt.wait_stream(resolved) + return out + + +def _destroy(events) -> None: + for start, end in events: + start.destroy() + end.destroy() + + +def _drain_after_failure(rt, exc: BaseException) -> None: + """Drain the device after a chunk raised, before ``exc`` propagates. + + Mirrors :func:`rocke.runtime.launcher._drain_after_failure`: a drain that + fails is not raised, because the exception worth seeing is the launch + error, not the drain that the same faulted kernel broke. + """ + try: + rt.sync() + except Exception as drain_exc: # noqa: BLE001 - reported on ``exc`` instead + add_note = getattr(exc, "add_note", None) + if add_note is not None: + add_note(f"draining the device after this error also failed: {drain_exc}") + + +def _time_one_by_one( + runs: Sequence[Callable[[], None]], + *, + warmup: int, + iters: int, + stream: int, +) -> List[Optional[float]]: + """Fall back to one drain per candidate, to attribute a chunk's failure.""" + from rocke.runtime import synchronize_and_release, time_launches + from rocke.runtime.hip_module import HipError + + out: List[Optional[float]] = [] + for run in runs: + try: + out.append(time_launches(run, warmup=warmup, iters=iters, stream=stream)) + except (HipError, RuntimeError): + # The launch raised; time_launches drained before it did, but a + # failed drain is not what it reports, so drain again -- a drain + # that fails here ends the sweep rather than letting the caller + # free buffers a queued kernel still uses. + synchronize_and_release(stream) + out.append(None) + return out + + +def time_batch( + runs: Sequence[Callable[[], None]], + *, + warmup: int, + iters: int, + stream: int = 0, + chunk: int = DEFAULT_CHUNK, +) -> List[Optional[float]]: + """Per-candidate times from ``runs``, one device drain per ``chunk``. + + ``None`` for a candidate whose launch raised. A chunk that raises is + re-timed one candidate at a time: the chunk's event pairs are all past the + failure point, so none of its times can be trusted or attributed. + """ + from rocke.runtime.hip_module import HipError + + out: List[Optional[float]] = [] + for start in range(0, len(runs), max(1, chunk)): + part = runs[start : start + max(1, chunk)] + try: + out.extend(_time_chunk(part, warmup=warmup, iters=iters, stream=stream)) + except (HipError, RuntimeError): + out.extend( + _time_one_by_one(part, warmup=warmup, iters=iters, stream=stream) + ) + return out + + +def rank( + n: int, + make_run: Callable[[int], Optional[Callable[[], None]]], + *, + warmup: int, + iters: int, + probe_warmup: int = DEFAULT_PROBE_WARMUP, + probe_iters: int = DEFAULT_PROBE_ITERS, + finalists: int = DEFAULT_FINALISTS, + chunk: int = DEFAULT_CHUNK, + primer: int = DEFAULT_PRIMER, + stream: int = 0, + on_probe: Optional[Callable[[Ranked], None]] = None, + on_finalist: Optional[Callable[[int, Callable[[], None]], Optional[bool]]] = None, + on_drop: Optional[Callable[[int], None]] = None, + held_bytes: Optional[Callable[[], int]] = None, + chunk_bytes: int = 0, +) -> List[Ranked]: + """Probe all ``n`` candidates, then time the best ``finalists`` in full. + + ``make_run(i)`` binds candidate ``i``'s launch, or returns ``None`` when it + cannot be loaded at all; it is called once per candidate, during the probe, + and the finalists' callables are kept so nothing is bound twice. + + ``on_probe`` is called with each candidate's :class:`Ranked` as its probe + result lands, so a long case can report progress. ``on_finalist(i, run)`` + runs just before a finalist is timed -- the hook the caller verifies in -- + and a ``False`` return times the candidate anyway (so its row still carries + a number) but keeps it from being reported as the best. + + ``on_drop(i)`` is called for every candidate whose launch this function + will not call again: one that failed, and one the probe ranked out. A + caller whose bind allocates device scratch frees it there -- the probe + holds a candidate's launch for as long as it could still rank, which is + longer than a serial sweep holds it and is why the scratch has to be + accounted for rather than freed before the next bind. + + ``held_bytes`` and ``chunk_bytes`` bound that scratch. When both are set, + the chunk being bound is closed as soon as ``held_bytes()`` -- the device + memory the caller is holding for this ranking, retained finalists included + -- passes ``chunk_bytes``. A candidate whose own scratch is over the budget + gets a chunk to itself, which is what a serial sweep would have done with + it anyway. ``chunk`` still caps the count. + + The finalists' scratch is still held when this returns -- they were timed + with it -- so the caller frees that along with the rest of the case. + + Returns one :class:`Ranked` per candidate, in candidate order. + """ + results = [Ranked(index=i) for i in range(n)] + # The finalists' bound launches, kept alive from the probe. A launch holds + # its packed kernargs and its kernel handle, so keeping it is what makes + # round 2 free; the dict is capped at ``finalists`` so a case with tens of + # thousands of candidates still only ever holds that many. + kept: dict = {} + primed = False + + def drop(i: int) -> None: + kept.pop(i, None) + if on_drop is not None: + on_drop(i) + + def trim() -> None: + """Keep only the best ``finalists`` of everything probed so far.""" + if len(kept) <= finalists: + return + ordered = sorted(kept, key=lambda i: results[i].probe_ms) + for i in ordered[finalists:]: + drop(i) + + nxt = 0 + while nxt < n: + runs = [] + live = [] + # Bind until the chunk is full by count, or until the caller is holding + # more device scratch than it wants to; a chunk always takes at least + # one candidate so an oversized one cannot stall the walk. + while nxt < n and len(runs) < max(1, chunk): + i = nxt + nxt += 1 + run = make_run(i) + if run is None: + results[i].failed = True + drop(i) + if on_probe is not None: + on_probe(results[i]) + continue + runs.append(run) + live.append(i) + if chunk_bytes and held_bytes is not None and held_bytes() > chunk_bytes: + break + if not runs: + continue + if not primed: + prime_device(runs[0], launches=primer, stream=stream) + primed = True + times = time_batch( + runs, warmup=probe_warmup, iters=probe_iters, stream=stream, chunk=chunk + ) + for i, run, ms in zip(live, runs, times): + if ms is None: + results[i].failed = True + else: + results[i].probe_ms = ms + kept[i] = run + if on_probe is not None: + on_probe(results[i]) + # Trim to the best ``finalists`` as we go, so a candidate that cannot + # rank gives its launch -- and the caller its scratch -- back at the + # end of the chunk that measured it rather than at the end of the case. + trim() + for i in live: + if results[i].failed: + drop(i) + + # Round 2. The finalists are verified first and then timed as one batch: + # a candidate's timed launches start right after the previous one's, so + # the clocks that the priming brought up never fall back between them. + order = sorted(kept, key=lambda i: results[i].probe_ms) + if on_finalist is not None: + for i in order: + # A kernel that fails here is still timed, so its row carries a + # number; the caller reads ``passed`` and keeps it from ranking. + results[i].passed = on_finalist(i, kept[i]) + times = time_batch( + [kept[i] for i in order], + warmup=warmup, + iters=iters, + stream=stream, + chunk=chunk, + ) + for i, ms in zip(order, times): + if ms is None: + results[i].failed = True + else: + results[i].ms = ms + return results 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 8fbbfae2abb9..911e74fe4734 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 @@ -38,7 +38,7 @@ The ``--top`` / ``--warmup`` / ``--iters`` / ``--jobs`` / ``--verify`` / ``--dtype`` / ``--direction`` / ``--run-from-cache`` / ``--early-stop`` / -``--early-stop-after`` flags are forwarded to both scripts. +``--early-stop-after`` / ``--fast-rank`` 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 (a kernel that fails ``--verify`` is @@ -48,6 +48,13 @@ ``--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. +``--fast-rank`` (cache runs only) replaces that mechanism rather than tuning +it: each script probes its cached kernels in batches and measures only the +best few, so the seed no longer prunes anything on the implicit side. The +record file is still written -- it is the comparison's own record of the best +direct-conv time -- and passing ``--early-stop``/``--early-stop-after`` +alongside ``--fast-rank`` has no effect. + Notes ----- - Without a cache, direct conv requires ``cpg == kpg`` (C/groups == K/groups) @@ -214,6 +221,23 @@ def _build_mode_args(args) -> list[str]: a += ["--early-stop", str(args.early_stop)] if args.early_stop_after is not None: a += ["--early-stop-after", str(args.early_stop_after)] + # Both scripts take the fast-ranking flags with the same names and the + # same defaults (they share benchmarks.common.batched_timing), so a knob + # is forwarded only when it was given and the scripts' own default stands + # otherwise -- one place for the default, not three. + if args.fast_rank: + a += ["--fast-rank"] + for flag, value in ( + ("--fast-rank-finalists", args.fast_rank_finalists), + ("--fast-rank-chunk", args.fast_rank_chunk), + ("--fast-rank-warmup", args.fast_rank_warmup), + ("--fast-rank-iters", args.fast_rank_iters), + ("--fast-rank-primer", args.fast_rank_primer), + ): + if value is not None: + a += [flag, str(value)] + if args.module_cache_bytes is not None: + a += ["--module-cache-bytes", str(args.module_cache_bytes)] return a @@ -242,6 +266,11 @@ def _build_implicit_only_args(args) -> list[str]: a += ["--streamk-reductions", args.streamk_reductions] if args.streamk_ctas != -1: a += ["--streamk-ctas", str(args.streamk_ctas)] + # Only the implicit-GEMM sweep allocates device scratch per launch config + # (two-stage wgrad workspaces, dgrad sub-GEMM tables), so only it has a + # budget to cap the probe's chunks with. + if args.fast_rank and args.fast_rank_scratch is not None: + a += ["--fast-rank-scratch", str(args.fast_rank_scratch)] return a @@ -396,6 +425,85 @@ def main() -> int: help="skip the implicit-GEMM benchmark (direct-conv only)", ) + rank_grp = parser.add_argument_group( + "Fast ranking", + "Forwarded to both scripts; needs --run-from-cache. Each times its " + "cached kernels in batches between two device drains and spends the " + "full --warmup/--iters only on the best few. A knob left unset keeps " + "the benchmark scripts' own default.", + ) + rank_grp.add_argument( + "--fast-rank", + action="store_true", + dest="fast_rank", + help="rank each script's cached kernels with a batched probe round " + "instead of timing every one of them in full. Replaces --early-stop " + "(and the implicit sweep's --split-k-prune) for the run, and makes " + "--verify check the finalists rather than every kernel.", + ) + rank_grp.add_argument( + "--fast-rank-finalists", + type=int, + default=None, + dest="fast_rank_finalists", + metavar="N", + help="kernels carried from the probe into the full-fidelity round", + ) + rank_grp.add_argument( + "--fast-rank-chunk", + type=int, + default=None, + dest="fast_rank_chunk", + metavar="N", + help="kernels timed between two device drains", + ) + rank_grp.add_argument( + "--fast-rank-scratch", + type=int, + default=None, + dest="fast_rank_scratch", + metavar="MiB", + help="device scratch the probe may hold before it closes a chunk. " + "Forwarded to implicit-GEMM only -- direct conv shares one set of " + "buffers across its kernels and allocates no per-kernel scratch.", + ) + rank_grp.add_argument( + "--fast-rank-warmup", + type=int, + default=None, + dest="fast_rank_warmup", + metavar="N", + help="probe warmup launches per kernel", + ) + rank_grp.add_argument( + "--fast-rank-iters", + type=int, + default=None, + dest="fast_rank_iters", + metavar="N", + help="probe timed launches per kernel", + ) + rank_grp.add_argument( + "--fast-rank-primer", + type=int, + default=None, + dest="fast_rank_primer", + metavar="N", + help="launches issued before a case's first measured chunk, to bring " + "the device's clocks up", + ) + + parser.add_argument( + "--module-cache-bytes", + type=int, + default=None, + dest="module_cache_bytes", + metavar="N", + help="AOT: HSACO bytes each script keeps loaded across cases, so a " + "shape-generic binary is loaded once per run rather than once per " + "case (unset keeps the scripts' own default; 0 loads per use)", + ) + implicit_grp = parser.add_argument_group( "Implicit-GEMM options", "Flags forwarded only to benchmark_implicit_gemm_conv.py.", @@ -501,6 +609,16 @@ def main() -> int: args = parser.parse_args() + if args.fast_rank and args.run_from_cache is None: + # Both scripts reject it anyway; catching it here beats launching two + # subprocesses to watch them both exit 2. + print( + "error: --fast-rank ranks the kernels of an AOT cache; " + "pass --run-from-cache DIR", + file=sys.stderr, + ) + return 2 + using_miopen = args.miopen_cmd is not None or args.miopen_file is not None implicit_only = _build_implicit_only_args(args) 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 f523542ce7db..ce7ba6f89ecb 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 @@ -40,7 +40,9 @@ os.environ.setdefault("ROCKE_CPP_QUIET_FALLBACK", "1") +from benchmarks.common import batched_timing from benchmarks.common.early_stop import EarlyStop, add_early_stop_arg +from benchmarks.common.module_cache import add_module_cache_arg from builders.common.conv_reference import conv_reference as _conv_reference from builders.common.conv_reference import dgrad_reference as _dgrad_reference_shared from builders.common.conv_reference import wgrad_reference as _wgrad_reference_shared @@ -2158,6 +2160,212 @@ def _no_plan_reason(rejected: list, dtype: str, variants: "frozenset[str]") -> s ) +def _measure_plans( + *, + args, + plans, + modules, + rt, + ptrs, + sizes, + out_name, + Out_t, + ref_out, + flop, + bytes_xfer, + stop, + u8, +): + """Time every plan of one case. Returns ``(results, rc)``. + + ``results`` holds one ``(TFLOPS, GB/s, ms, passed, label)`` row per plan + that ran; ``rc`` is 1 when a plan failed verification. A plan is bound + here rather than by the caller because every plan of a case reads and + writes the same buffers: ``ptrs`` / ``sizes`` are the case's, shared. + + ``--fast-rank`` swaps the per-plan protocol (warmup plus timed loop, each + its own device drain) for :func:`_fast_rank_plans`. Everything else about + the case -- the rows, the ranked table, the record file -- is unchanged. + + ``modules`` is the run's loaded-binary cache: a plan's binaries are + shape-generic, so a step's module is loaded on the first case that uses + it and every later case binds against the same one. + """ + import torch + from benchmarks.common.direct_kernel_sweep import launch_values + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import HipError + from rocke.runtime.launcher import LaunchConfig + + results: list = [] + rc = 0 + n_plans = len(plans) + width = len(str(n_plans)) + + def bind(plan): + """The plan's launch, or ``None`` if a step cannot be loaded.""" + try: + launches = [ + ( + modules.launcher( + hsaco_path=st.hsaco_path, + kernel_name=st.kernel_name, + signature=st.signature, + ), + launch_values(st, ptrs, sizes), + LaunchConfig(grid=st.grid, block=st.block, stream=0), + ) + for st in plan.steps + ] + except HipError as e: + print(f" [skip] {plan.label}: {e}", flush=True) + return None + + # 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() + + return run + + def verify(plan, run) -> bool: + """Peak-normalised rel-err of one plan against the case's reference.""" + nonlocal rc + + rt.memset(ptrs[out_name], 0, Out_t.nbytes) + run() + synchronize_and_release(0) + out_cpu = torch.empty_like(Out_t) + rt.memcpy_d2h(u8(out_cpu), ptrs[out_name], Out_t.nbytes) + ref_f32 = ref_out.float().cpu() + rel_err = float( + (out_cpu.float() - ref_f32).abs().max() + / ref_f32.abs().max().clamp(min=1.0) + ) + passed = rel_err < 5e-2 + if not passed: + rc = 1 + print(f" verify {plan.label}: FAIL(rel_err={rel_err:.2e})", flush=True) + return passed + + def row(ms, passed, label, idx) -> None: + """Append one measured plan's row and print its status line.""" + results.append((flop / ms / 1e9, bytes_xfer / ms / 1e6, ms, passed, label)) + # One line per measured launch, as the implicit-GEMM sweep prints, so a + # long run shows its progress and partial results rather than only the + # ranked table at the end. + verdict = "" if passed is None else (" PASS" if passed else " FAIL") + print( + f" [{idx:>{width}}/{n_plans}] " + f"{results[-1][0]:7.1f} TFLOPS {ms:8.3f} ms{verdict} {label}", + flush=True, + ) + + if args.fast_rank: + _fast_rank_plans( + args=args, + plans=plans, + bind=bind, + verify=verify if ref_out is not None else None, + stop=stop, + row=row, + width=width, + ) + return results, rc + + for plan_idx, plan in enumerate(plans, 1): + run = bind(plan) + if run is None: + continue + + passed = None if ref_out is None else verify(plan, run) + + try: + ms = stop.measure( + run, + warmup=args.warmup, + iters=args.iters, + stream=0, + passed=passed, + ) + except (HipError, RuntimeError) as e: + print(f" [skip] {plan.label}: {e}", flush=True) + # time_launches drained before raising, but it does not raise a + # failed drain (the launch error is the one it reports). Drain + # again so a failed drain ends the sweep here, before the shared + # buffers are freed with launches possibly still using them. + synchronize_and_release(0) + continue + if ms is None: + stop.report(plan.label) + continue + synchronize_and_release(0) + row(ms, passed, plan.label, plan_idx) + + return results, rc + + +def _fast_rank_plans(*, args, plans, bind, verify, stop, row, width) -> None: + """Rank one case's plans with :mod:`benchmarks.common.batched_timing`. + + The probe times every plan with one device drain per chunk instead of + three per plan, and keeps only the best ``--fast-rank-finalists``; those + are timed again with the full ``--warmup``/``--iters``, so every reported + number comes from the same protocol the per-plan path uses. A plan the + probe ranked out is reported with its probe time and never measured again. + + Two consequences worth knowing, both of them the point of the flag: + ``--verify`` checks the finalists rather than every plan (a plan that is + 30x off the best is not worth a reference comparison), and the ranked + table holds at most ``--fast-rank-finalists`` rows. + """ + from benchmarks.common import batched_timing + + n = len(plans) + + def _on_probe(r) -> None: + if r.failed: + return + print( + f" [{r.index + 1:>{width}}/{n}] probe {r.probe_ms:8.3f} ms " + f"{plans[r.index].label}", + flush=True, + ) + + def _on_finalist(i, run): + return None if verify is None else verify(plans[i], run) + + ranked = batched_timing.rank( + n, + lambda i: bind(plans[i]), + warmup=args.warmup, + iters=args.iters, + probe_warmup=args.fast_rank_warmup, + probe_iters=args.fast_rank_iters, + finalists=args.fast_rank_finalists, + chunk=args.fast_rank_chunk, + primer=args.fast_rank_primer, + on_probe=_on_probe, + on_finalist=_on_finalist, + ) + n_probed = sum(1 for r in ranked if r.probe_ms is not None) + n_final = sum(1 for r in ranked if r.is_finalist) + print( + f" probe: {n_probed}/{n} timed, {n_final} finalist(s) measured in full", + flush=True, + ) + for r in ranked: + if r.ms is None: + continue + stop.note(r.ms, passed=r.passed) + row(r.ms, r.passed, plans[r.index].label, r.index + 1) + + def _cache_dispatch(args, arch: str, cases: list) -> int: """Handle the two cache modes. @@ -2235,14 +2443,16 @@ def _run_from_cache(args, arch: str, cases: list, cache, directions) -> int: import torch - from benchmarks.common.direct_kernel_sweep import direct_plans, launch_values - from rocke.runtime import synchronize_and_release - from rocke.runtime.hip_module import HipError, Runtime - from rocke.runtime.launcher import KernelLauncher, LaunchConfig + from benchmarks.common.direct_kernel_sweep import direct_plans + from benchmarks.common.module_cache import ModuleCache + from rocke.runtime.hip_module import Runtime def _u8(t): return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + # The binaries are shape-generic, so their modules outlive a case. + modules = ModuleCache(args.module_cache_bytes) + from kernels.conv.direct_aot_jobs import DIRECT_DTYPES variants = _cache_variants(args) @@ -2327,81 +2537,23 @@ def _u8(t): flop = float(p.flops) bytes_xfer = float(sum(t.nbytes for t in inputs.values()) + Out_t.nbytes) - results = [] _stop = EarlyStop.for_case(args, p, dtype, direction) - for plan in plans: - try: - launches = [ - ( - KernelLauncher( - hsaco=st.hsaco_path.read_bytes(), - kernel_name=st.kernel_name, - signature=st.signature, - ), - launch_values(st, ptrs, sizes), - LaunchConfig(grid=st.grid, block=st.block, stream=0), - ) - for st in plan.steps - ] - except HipError as e: - print(f" [skip] {plan.label}: {e}", flush=True) - continue - - # 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: - rt.memset(ptrs[out_name], 0, Out_t.nbytes) - run() - synchronize_and_release(0) - out_cpu = torch.empty_like(Out_t) - rt.memcpy_d2h(_u8(out_cpu), ptrs[out_name], Out_t.nbytes) - ref_f32 = ref_out.float().cpu() - rel_err = float( - (out_cpu.float() - ref_f32).abs().max() - / ref_f32.abs().max().clamp(min=1.0) - ) - passed = rel_err < 5e-2 - if not passed: - overall_rc = 1 - print( - f" verify {plan.label}: FAIL(rel_err={rel_err:.2e})", - flush=True, - ) - - try: - ms = _stop.measure( - run, - warmup=args.warmup, - iters=args.iters, - stream=0, - passed=passed, - ) - except (HipError, RuntimeError) as e: - print(f" [skip] {plan.label}: {e}", flush=True) - # time_launches drained before raising, but it does not - # raise a failed drain (the launch error is the one it - # reports). Drain again so a failed drain ends the sweep - # here, before the shared buffers below are freed with - # launches possibly still using them. - synchronize_and_release(0) - continue - if ms is None: - _stop.report(plan.label) - continue - synchronize_and_release(0) - results.append( - (flop / ms / 1e9, bytes_xfer / ms / 1e6, ms, passed, plan.label) - ) + results, case_rc = _measure_plans( + args=args, + plans=plans, + modules=modules, + rt=rt, + ptrs=ptrs, + sizes=sizes, + out_name=out_name, + Out_t=Out_t, + ref_out=ref_out, + flop=flop, + bytes_xfer=bytes_xfer, + stop=_stop, + u8=_u8, + ) + overall_rc = overall_rc or case_rc for ptr in ptrs.values(): rt.free(ptr) @@ -2425,6 +2577,8 @@ def run(bound=bound): # 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) + if len(cases) > 1: + print(f"\n{modules.summary()}", flush=True) return overall_rc @@ -2508,6 +2662,77 @@ def main() -> int: help="on the first verify FAIL, dump tensors to PATH/ and stop the sweep.", ) + rank_grp = parser.add_argument_group( + "Fast ranking (AOT --run-from-cache only)", + "Time every cached kernel in batches between two device drains, then " + "spend the full --warmup/--iters only on the best few. Replaces " + "--early-stop for the run; --early-stop-record still gets the case's " + "best time.", + ) + rank_grp.add_argument( + "--fast-rank", + action="store_true", + dest="fast_rank", + help="probe every cached kernel with a batched 1+N launch round and " + "measure only the finalists in full. --verify then checks the " + "finalists rather than every kernel.", + ) + rank_grp.add_argument( + "--fast-rank-finalists", + type=int, + default=batched_timing.DEFAULT_FINALISTS, + dest="fast_rank_finalists", + metavar="N", + help=f"kernels carried from the probe into the full-fidelity round " + f"(default: {batched_timing.DEFAULT_FINALISTS}). This also caps the " + f"ranked table and bounds what stays bound at once.", + ) + rank_grp.add_argument( + "--fast-rank-chunk", + type=int, + default=batched_timing.DEFAULT_CHUNK, + dest="fast_rank_chunk", + metavar="N", + help=f"kernels timed between two device drains " + f"(default: {batched_timing.DEFAULT_CHUNK})", + ) + rank_grp.add_argument( + "--fast-rank-warmup", + type=int, + default=batched_timing.DEFAULT_PROBE_WARMUP, + dest="fast_rank_warmup", + metavar="N", + help=f"probe warmup launches per kernel, covering its cold code object " + f"and scratch (default: {batched_timing.DEFAULT_PROBE_WARMUP})", + ) + rank_grp.add_argument( + "--fast-rank-iters", + type=int, + default=batched_timing.DEFAULT_PROBE_ITERS, + dest="fast_rank_iters", + metavar="N", + help=f"probe timed launches per kernel " + f"(default: {batched_timing.DEFAULT_PROBE_ITERS})", + ) + rank_grp.add_argument( + "--fast-rank-primer", + type=int, + default=batched_timing.DEFAULT_PRIMER, + dest="fast_rank_primer", + metavar="N", + help=f"launches issued before the first measured chunk of a case, to " + f"bring the device's clocks up so the kernels probed first are not " + f"measured at idle clocks (default: {batched_timing.DEFAULT_PRIMER})", + ) + + add_module_cache_arg( + parser.add_argument_group( + "Loaded binaries (AOT --run-from-cache only)", + "A cached binary is shape-generic, so its HIP module is loaded " + "once per run instead of once per case.", + ) + ) + miopen_grp = parser.add_argument_group( "MIOpen input", "Load the conv problem from a MIOpenDriver command instead of explicit shape flags. " @@ -2623,6 +2848,16 @@ def main() -> int: ) args = parser.parse_args() + if args.fast_rank and not args.run_from_cache: + # The JIT sweep builds and loads its kernels one at a time, so + # batching their launches would save nothing; saying so beats + # accepting the flag and ignoring it. + print( + "error: --fast-rank ranks the kernels of an AOT cache; " + "pass --run-from-cache DIR", + file=sys.stderr, + ) + return 2 # --dtype unset: one fp16 run for shape flags, every dtype for --compile-all. from kernels.conv.direct_aot_jobs import DIRECT_DTYPES 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 a65c93921dd9..798e72ae4aa8 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 @@ -76,10 +76,12 @@ # validators on other targets. # --------------------------------------------------------------------------- +from benchmarks.common import batched_timing # noqa: E402 from benchmarks.common.early_stop import ( # noqa: E402 EarlyStop, add_early_stop_arg, ) +from benchmarks.common.module_cache import add_module_cache_arg # noqa: E402 from kernels.conv.aot_jobs import _AOT_ALIAS_PIPELINES # noqa: E402 from kernels.conv.aot_jobs import ( # noqa: E402 CACHE_EPILOGUES as _EPILOGUES, @@ -1224,6 +1226,88 @@ def main() -> int: "the sweep. Implies --verify." ), ) + rank_grp = parser.add_argument_group( + "Fast ranking (AOT --run-from-cache only)", + "Time every cached kernel in batches between two device drains, then " + "spend the full --warmup/--iters only on the best few. Replaces " + "--early-stop and --split-k-prune for the run; --early-stop-record " + "still gets the case's best time. Stream-K kernels keep the per-kernel " + "path: their workspace contract is per launch config.", + ) + rank_grp.add_argument( + "--fast-rank", + action="store_true", + dest="fast_rank", + help="probe every cached kernel with a batched 1+N launch round and " + "measure only the finalists in full. --verify then checks the " + "finalists rather than every kernel.", + ) + rank_grp.add_argument( + "--fast-rank-finalists", + type=int, + default=batched_timing.DEFAULT_FINALISTS, + dest="fast_rank_finalists", + metavar="N", + help=f"kernels carried from the probe into the full-fidelity round " + f"(default: {batched_timing.DEFAULT_FINALISTS}). This also caps the " + f"ranked table.", + ) + rank_grp.add_argument( + "--fast-rank-chunk", + type=int, + default=batched_timing.DEFAULT_CHUNK, + dest="fast_rank_chunk", + metavar="N", + help=f"kernels timed between two device drains " + f"(default: {batched_timing.DEFAULT_CHUNK})", + ) + rank_grp.add_argument( + "--fast-rank-scratch", + type=int, + default=2048, + dest="fast_rank_scratch", + metavar="MiB", + help="device scratch the probe may hold before it closes a chunk " + "(default: 2048). Only two-stage wgrad allocates enough for this to " + "bite; fwd allocates none at all.", + ) + rank_grp.add_argument( + "--fast-rank-warmup", + type=int, + default=batched_timing.DEFAULT_PROBE_WARMUP, + dest="fast_rank_warmup", + metavar="N", + help=f"probe warmup launches per kernel, covering its cold code object " + f"and scratch (default: {batched_timing.DEFAULT_PROBE_WARMUP})", + ) + rank_grp.add_argument( + "--fast-rank-iters", + type=int, + default=batched_timing.DEFAULT_PROBE_ITERS, + dest="fast_rank_iters", + metavar="N", + help=f"probe timed launches per kernel " + f"(default: {batched_timing.DEFAULT_PROBE_ITERS})", + ) + rank_grp.add_argument( + "--fast-rank-primer", + type=int, + default=batched_timing.DEFAULT_PRIMER, + dest="fast_rank_primer", + metavar="N", + help=f"launches issued before the first measured chunk of a case, to " + f"bring the device's clocks up so the kernels probed first are not " + f"measured at idle clocks (default: {batched_timing.DEFAULT_PRIMER})", + ) + + add_module_cache_arg( + parser.add_argument_group( + "Loaded binaries (AOT --run-from-cache only)", + "A cached binary is shape-generic, so its HIP module is loaded " + "once per run instead of once per case.", + ) + ) + parser.add_argument( "--debug-init", nargs="?", @@ -1538,6 +1622,16 @@ def main() -> int: ) args = parser.parse_args() + if args.fast_rank and not args.run_from_cache: + # The JIT sweep builds and loads its kernels one at a time, so + # batching their launches would save nothing; saying so beats + # accepting the flag and ignoring it. + print( + "error: --fast-rank ranks the kernels of an AOT cache; " + "pass --run-from-cache DIR", + file=sys.stderr, + ) + return 2 # --dtype unset: one fp16 run for shape flags, every dtype for --compile-all. args.compile_dtypes = (args.dtype,) if args.dtype is not None else _COMPILE_DTYPES if args.dtype is None: @@ -4535,12 +4629,17 @@ def _run_from_cache(args, arch, target, cases, cache, directions) -> int: import torch + from benchmarks.common.module_cache import ModuleCache from rocke.runtime import synchronize_and_release from rocke.runtime.hip_module import HipError, Runtime from rocke.runtime.launcher import KernelLauncher, LaunchConfig _TORCH_DT = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + # A cached binary is shape-generic: its module is loaded on the first case + # that offers it and reused by every later one. + modules = ModuleCache(args.module_cache_bytes) + def _u8(t): return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) @@ -4734,11 +4833,59 @@ def _u8(t): _stop = EarlyStop.for_case(args, problem, case_dtype, direction) n_failed = 0 - n_cand = len(candidates) # Device scratch of the launch configs in flight (two-stage # workspaces, dgrad sub-GEMM tables): one per (kernel, degree), # so it is freed as the sweep goes rather than per case. live_extras: list = [] + + # --fast-rank takes every kernel whose scratch it can account for. + # A stream-K binary allocates and zeroes per launch config through + # its own contract (_cached_streamk_run returns the free), so it + # stays on the serial path below whichever way the flag is set. + if args.fast_rank: + streamk = [ + c + for c in candidates + if direction == "wgrad" and c[0].streamk != "off" + ] + fast = [ + c + for c in candidates + if not (direction == "wgrad" and c[0].streamk != "off") + ] + if fast: + failed, stopped = _fast_rank_candidates( + args=args, + candidates=fast, + modules=modules, + direction=direction, + problem=problem, + run_problem=run_problem, + batch_chunks=batch_chunks, + case_dtype=case_dtype, + arch=arch, + rt=rt, + ptrs=(A_dev, B_dev, D_dev), + sizes=(A_t.nbytes, B_t.nbytes, D_t.nbytes), + stage2_kernels=stage2_kernels, + ref_out=ref_out, + D_t=D_t, + flop=flop, + bytes_xfer=bytes_xfer, + results=results, + stop=_stop, + u8=_u8, + ) + n_failed += failed + if failed: + overall_rc = 1 + if stopped: + for buf in (A_dev, B_dev, D_dev): + rt.free(buf) + return 1 + candidates = streamk + + n_cand = len(candidates) for cand_idx, (ident, hsaco_path, meta) in enumerate(candidates, 1): kernel_name = meta.get("kernel_name") if not kernel_name: @@ -4749,8 +4896,8 @@ def _u8(t): ) continue try: - launcher = KernelLauncher( - hsaco=hsaco_path.read_bytes(), + launcher = modules.launcher( + hsaco_path=hsaco_path, kernel_name=kernel_name, signature=_signature_for_identity(ident, case_dtype), ) @@ -4961,6 +5108,8 @@ def _u8(t): tflops, _, ms, label = results[0] print(f"\nBest: {tflops:.1f} TFLOPS — {label} {ms:.3f} ms", flush=True) + if len(cases) > 1: + print(f"\n{modules.summary()}", flush=True) return overall_rc @@ -4980,6 +5129,236 @@ def _launch_problem(direction, problem): return fold_width(problem) +def _extras_bytes(extras) -> int: + """Device bytes :func:`_extras_and_grid` allocated for one launch config.""" + return int(extras.get("ws_bytes", 0)) + int(extras.get("sub_gemm_bytes", 0)) + + +def _fast_rank_candidates( + *, + args, + candidates, + modules, + direction, + problem, + run_problem, + batch_chunks, + case_dtype, + arch, + rt, + ptrs, + sizes, + stage2_kernels, + ref_out, + D_t, + flop, + bytes_xfer, + results, + stop, + u8, +): + """Rank one case's cached kernels with :mod:`benchmarks.common.batched_timing`. + + The probe times every kernel with one device drain per chunk instead of + three per kernel, and only the best ``--fast-rank-finalists`` are timed + again with the full ``--warmup``/``--iters``; those are the numbers that + reach ``results``. + + What makes this harder here than in the direct-conv sweep is scratch. A + serial sweep frees a launch config's workspace before binding the next, so + it holds one; a probe holds a kernel's launch for as long as it could + still rank. Only two directions allocate any: dgrad's sub-GEMM table (a + few KB) and two-stage wgrad's f32 workspace (which is not small). So the + chunk is closed on a byte budget as well as a count, and a kernel the + probe rules out gives its scratch back at the end of the chunk that + measured it. + + Three differences from the serial path, all of them the point of the flag: + ``--verify`` checks the finalists rather than every kernel, + ``--split-k-prune`` does not apply (a kernel's degrees are ranked like any + other candidates, and a probe launch costs about what the prune saved), + and stream-K kernels are not handled here at all -- their workspace and + zeroing contract is per launch config, so the caller leaves them on the + serial path. + + Returns ``(n_failed, stopped)``; ``stopped`` is a ``--dump-fail`` dump, + which ends the sweep as it does on the serial path. + """ + from benchmarks.common import batched_timing + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import HipError + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + + # (identity, path, meta, split_k). The split-K sweep is flattened into the + # candidate list so a degree ranks like any other kernel. + items = [] + for ident, hsaco_path, meta in candidates: + if not meta.get("kernel_name"): + print( + f" [skip] {ident.label()}: cache entry has no kernel_name " + f"(rebuild the cache)", + flush=True, + ) + continue + if direction == "wgrad" and ident.split_k > 1: + degrees = _runtime_split_ks( + problem, ident, args.split_k, args.split_k_degrees + ) + if not degrees: + print( + f" [skip] {ident.label()}: no split-K degree fits this " + f"problem (see --split-k)", + flush=True, + ) + continue + else: + degrees = (None,) + for split_k in degrees: + items.append((ident, hsaco_path, meta, split_k)) + + n = len(items) + if not n: + return 0, False + width = len(str(n)) + owned: dict = {} # item index -> the extras its bind allocated + labels: dict = {} + held = {"bytes": 0} + state = {"stopped": False} + + def _release(mine) -> None: + for extras in mine: + held["bytes"] -= _extras_bytes(extras) + for key in ("ws_ptr", "sub_gemm_buf"): + if key in extras: + rt.free(extras[key]) + + def make_run(i): + ident, hsaco_path, meta, split_k = items[i] + # A kernel swept at several split-K degrees is one binary, and so is + # one offered to the next case: the module cache holds it for both. + try: + launcher = modules.launcher( + hsaco_path=hsaco_path, + kernel_name=meta["kernel_name"], + signature=_signature_for_identity(ident, case_dtype), + ) + except HipError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + return None + mine: list = [] + bound = _cached_bind_batch( + batch_chunks, + direction, + run_problem, + ident, + meta, + rt, + launcher, + case_dtype, + split_k, + ptrs, + sizes, + arch, + stage2_kernels, + KernelLauncher, + LaunchConfig, + mine, + ) + if bound is None: + _release(mine) + return None + owned[i] = mine + held["bytes"] += sum(_extras_bytes(e) for e in mine) + run, labels[i] = bound + return run + + def on_drop(i) -> None: + _release(owned.pop(i, ())) + + def on_probe(r) -> None: + if r.failed: + return + print( + f" [{r.index + 1:>{width}}/{n}] probe {r.probe_ms:8.3f} ms " + f"{labels[r.index]}", + flush=True, + ) + + def on_finalist(i, run): + if ref_out is None: + return None + try: + stopped, passed = _verify_kernel( + rt=rt, + launch_fn=run, + out_dev=ptrs[2], + out_t=D_t, + zero_init_out=True, + ref_out=ref_out, + kernel_name=labels[i], + dump_fail=args.dump_fail, + u8=u8, + arch=arch, + compute_dtype=case_dtype, + ) + except (HipError, RuntimeError) as e: + print(f" verify {labels[i]}: FAIL({e})", flush=True) + # A launch that raised can leave queued work in flight on this + # config's scratch; drain before anything frees it. + synchronize_and_release() + return False + state["stopped"] = state["stopped"] or stopped + return passed + + ranked = batched_timing.rank( + n, + make_run, + warmup=args.warmup, + iters=args.iters, + probe_warmup=args.fast_rank_warmup, + probe_iters=args.fast_rank_iters, + finalists=args.fast_rank_finalists, + chunk=args.fast_rank_chunk, + primer=args.fast_rank_primer, + on_probe=on_probe, + on_finalist=on_finalist, + on_drop=on_drop, + held_bytes=lambda: held["bytes"], + chunk_bytes=args.fast_rank_scratch * 1024 * 1024, + ) + + n_probed = sum(1 for r in ranked if r.probe_ms is not None) + n_final = sum(1 for r in ranked if r.is_finalist) + print( + f" probe: {n_probed}/{n} timed, {n_final} finalist(s) measured in full", + flush=True, + ) + n_failed = 0 + for r in ranked: + if r.ms is None: + continue + if r.passed is False: + # As on the serial path: a kernel that failed verification is + # reported and never ranked. + n_failed += 1 + continue + stop.note(r.ms, passed=r.passed) + results.append( + (flop / (r.ms * 1e9), bytes_xfer / (r.ms * 1e6), r.ms, labels[r.index]) + ) + print( + f" [{r.index + 1:>{width}}/{n}] {results[-1][0]:7.1f} TFLOPS " + f"{r.ms:8.3f} ms {labels[r.index]}", + flush=True, + ) + + # The finalists were timed holding theirs; nothing outside this call owns + # any of it. + for i in list(owned): + _release(owned.pop(i)) + return n_failed, state["stopped"] + + def _free_extras(rt, live_extras: list) -> None: """Free the device buffers :func:`_extras_and_grid` allocated, and forget them.""" for extras in live_extras: @@ -5665,6 +6044,10 @@ def _extras_and_grid(direction, problem, ident, rt, dtype, split_k=None): rt.memcpy_h2d(buf, flat.ctypes.data, flat.nbytes) extras["sub_gemm_buf"] = buf extras["num_sub_gemms"] = len(sub_gemms) + # Not a kernarg: what a batched ranking charges this config's scratch + # budget (see _extras_bytes). The table is small, but the accounting + # should name everything it holds. + extras["sub_gemm_bytes"] = int(flat.nbytes) blocks = sum( -(-sg.gemm_m // ident.tile_m) * -(-sg.gemm_n // ident.tile_n) for sg in sub_gemms 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 b2d06d7817f4..65b11376d303 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 @@ -513,7 +513,7 @@ def plan_for( reduce_id = _split_reduce_identity( identity.arch, identity.wave_size, spec.split_c, identity.dtype_a ) - entry = cache.get(reduce_id) + entry = cache.entry_meta(reduce_id) if entry is None: return None, f"missing cached {_SPLIT_REDUCE} helper (rebuild the cache)" reduce_spec = make_spec(_SPLIT_REDUCE, problem, {"split_c": spec.split_c}) @@ -521,7 +521,7 @@ def plan_for( identity, hsaco_path, meta, spec, problem, "fwd", ("x", "w", "ws_split") ) fold = DirectStep( - kernel_name=entry[1]["kernel_name"], + kernel_name=entry["kernel_name"], hsaco_path=cache.hsaco_path(reduce_id), signature=conv.signature, grid=reduce_spec.grid(problem), @@ -556,7 +556,7 @@ def plan_for( fold_knobs, identity.dtype_a, ) - entry = cache.get(fold_id) + entry = cache.entry_meta(fold_id) if entry is None: return None, "missing cached wgrad fold (rebuild the cache)" fold_spec = make_spec(_WGRAD_FOLD, problem, fold_knobs) @@ -568,7 +568,7 @@ def plan_for( identity, hsaco_path, meta, spec, problem, "wgrad", ("dy", "x", "ws") ) fold = DirectStep( - kernel_name=entry[1]["kernel_name"], + kernel_name=entry["kernel_name"], hsaco_path=cache.hsaco_path(fold_id), signature=dw_wgrad_fold_signature(fold_spec), grid=grid, @@ -619,7 +619,7 @@ def plan_for( helper_knobs, identity.dtype_a, ) - entry = cache.get(helper_id) + entry = cache.entry_meta(helper_id) if entry is None: return None, f"missing cached {helper} helper (rebuild the cache)" helper_path = cache.hsaco_path(helper_id) @@ -627,7 +627,7 @@ def plan_for( grid, block = dc.direct_launch_geometry(helper_spec) steps.append( DirectStep( - kernel_name=entry[1]["kernel_name"], + kernel_name=entry["kernel_name"], hsaco_path=helper_path, signature=conv_direct_transform_signature( identity.dtype_a, runtime_channels 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 cd1829178847..9bfe61a3daf4 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 @@ -239,6 +239,25 @@ def measure( _record(self.record_key, ms) return ms + def note(self, ms: float, *, passed: Optional[bool] = None) -> None: + """Record a time this object did not measure itself. + + For a sweep that ranks its kernels another way (see + :mod:`benchmarks.common.batched_timing`) and so never calls + :meth:`measure`, but still has to leave ``--early-stop-record`` the + case's best time -- ``benchmark_conv_compare.py`` seeds the + implicit-GEMM sweep from that file. A kernel that failed verification + is ignored here for the same reason it is in :meth:`measure`: a + wrong-but-fast kernel would prune every correct one in the seeded run. + """ + if passed is False: + return + 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) + 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 "" 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 2deefbe9287e..c2277a1acea8 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 @@ -28,13 +28,16 @@ import itertools import json import os +from collections import OrderedDict from pathlib import Path from typing import Dict, Iterable, List, Optional, Tuple from kernels.conv.kernel_identity import ( # noqa: F401 (re-exported) KernelIdentity, current_llvm_flavor, + identity_capability_key, identity_supports_problem, + problem_capability_key, ) @@ -122,6 +125,13 @@ def comgr_input_key(comgr_input) -> str: # KernelCache # --------------------------------------------------------------------------- +# How many ``(identity, path, meta)`` rows :meth:`KernelCache.compatible` +# keeps memoised across every shape key. Bounding rows rather than keys is +# what makes the memo's footprint independent of how wide the cache is: ~2M +# tuple references is ~16 MB, a few dozen shapes' worth of a 20k-entry cache +# or thousands of a narrow one. +_COMPATIBLE_MEMO_ROWS = 2_000_000 + # Direction → subdirectory name. _DIRECTION_DIRS: Dict[str, str] = { "fwd": "conv_fwd", @@ -174,6 +184,24 @@ def __init__(self, root: Path, arch: str) -> None: # whole cache several times (describe, stale check, compatible), and # each pass parses every metadata file. Dropped on every write. self._loaded: Dict[str, list] = {} + # Entries of one subdirectory grouped by capability key, built on the + # first :meth:`compatible` for that direction (see _groups). + self._groups: Dict[str, list] = {} + # ``compatible`` answers, by (direction, problem capability key). + self._compatible: "OrderedDict[tuple, list]" = OrderedDict() + self._compatible_rows = 0 + # Parsed metadata by path, for the lookups that only want the entry + # point symbol (:meth:`entry_meta`) and would otherwise re-read and + # re-parse the same file once per candidate per shape. + self._meta: Dict[Path, Optional[dict]] = {} + + def _forget(self) -> None: + """Drop everything read from disk. Called on every write.""" + self._loaded.clear() + self._groups.clear() + self._compatible.clear() + self._compatible_rows = 0 + self._meta.clear() def _dir_for(self, identity: KernelIdentity) -> Path: subdir = _DIRECTION_DIRS.get(identity.direction, identity.direction) @@ -234,12 +262,12 @@ def link(self, identity: KernelIdentity, key: str, meta: dict) -> None: self._meta_path(identity), json.dumps(full_meta, indent=2, sort_keys=True).encode("utf-8"), ) - self._loaded.clear() + self._forget() def discard(self, identity: KernelIdentity) -> None: """Remove ``identity``'s entry. Its binary stays: others may share it.""" self._meta_path(identity).unlink(missing_ok=True) - self._loaded.clear() + self._forget() def put( self, @@ -275,14 +303,39 @@ def index(self, directions: Optional[Iterable[str]] = None) -> Dict[str, dict]: 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.""" + """The identity's metadata, or ``None`` when it has no entry. + + Memoised by path: a direct-conv plan looks up the same transform or + fold helper for every candidate of every shape, and parsing its one + file again each time is pure overhead. Dropped on every write. + """ meta_path = self._meta_path(identity) - if not meta_path.exists(): + if meta_path in self._meta: + return self._meta[meta_path] + parsed: Optional[dict] = None + if meta_path.exists(): + try: + parsed = json.loads(meta_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + parsed = None + self._meta[meta_path] = parsed + return parsed + + def entry_meta(self, identity: KernelIdentity) -> Optional[dict]: + """The identity's metadata when its binary is present, else ``None``. + + :meth:`get` with the bytes left on disk. A caller that only needs the + entry point symbol of a helper kernel -- which is every caller that + builds a launch plan, since the binary is read later, once, by the + launcher -- asks here instead and reads no HSACO at all. + """ + meta = self.meta(identity) + if meta is None: return None - try: - return json.loads(meta_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): + path = self._resolve(self._meta_path(identity), meta) + if not path.exists() or path.stat().st_size == 0: return None + return meta def _resolve(self, meta_path: Path, meta: dict) -> Path: if "blob" in meta: @@ -375,6 +428,40 @@ def supports_problem(identity: KernelIdentity, problem: object) -> Tuple[bool, s """ return identity_supports_problem(identity, problem) + def _groups_for(self, direction: Optional[str]) -> list: + """``[(identity, [entries...])]`` of one direction, by capability key. + + The identity heading a group is a representative: every entry in it + answers :func:`identity_supports_problem` alike, so the predicate runs + once per group rather than once per entry. A cache sweeps one + geometry over the whole tuning grid, so the groups are much wider than + one entry -- tile, pipeline, LDS layout and scheduling knobs all leave + the key alone. + """ + cache_key = direction or "" + groups = self._groups.get(cache_key) + if groups is None: + by_key: Dict[tuple, list] = {} + for pos, entry in enumerate(self._entries(direction)): + by_key.setdefault(identity_capability_key(entry[0]), []).append( + (pos, entry) + ) + groups = self._groups[cache_key] = [ + (members[0][1][0], members) for members in by_key.values() + ] + return groups + + def _remember(self, key: tuple, out: list) -> None: + """Memoise one ``compatible`` answer, oldest out when the budget is.""" + self._compatible[key] = out + self._compatible_rows += len(out) + while ( + self._compatible_rows > _COMPATIBLE_MEMO_ROWS + and len(self._compatible) > 1 + ): + _, dropped = self._compatible.popitem(last=False) + self._compatible_rows -= len(dropped) + def compatible( self, problem: object, *, direction: Optional[str] = None ) -> List[Tuple[KernelIdentity, Path, dict]]: @@ -382,10 +469,35 @@ def compatible( The metadata is returned because the caller needs ``kernel_name`` from it: the HSACO's entry-point symbol is not derivable from the identity. + + Answers are grouped and memoised, which a sweep over thousands of + shapes needs and a one-shot caller does not notice: see + :func:`~kernels.conv.kernel_identity.problem_capability_key`. The + memo is bounded by the number of rows it holds, not by the number of + shapes, and is dropped whenever the cache is written. The returned + list is the caller's own -- mutating it does not disturb the memo. """ - out: List[Tuple[KernelIdentity, Path, dict]] = [] - for identity, hsaco_path, meta in self._entries(direction): + groups = self._groups_for(direction) + is_direct = all(ident.is_direct for ident, _ in groups) + key = problem_capability_key(problem, is_direct=is_direct) + if key is not None: + key = (direction, key) + hit = self._compatible.get(key) + if hit is not None: + self._compatible.move_to_end(key) + return list(hit) + picked: List[Tuple[int, Tuple[KernelIdentity, Path, dict]]] = [] + for identity, members in groups: ok, _ = self.supports_problem(identity, problem) if ok: - out.append((identity, hsaco_path, meta)) + picked.extend(members) + # Back into the order :meth:`_entries` yields (identity sha1 within a + # direction), which is what every caller downstream -- ``--sample``'s + # fixed-seed subset, ``--early-stop-after``, the probe order -- sees. + # Grouping is an implementation detail and must not reshuffle it. + picked.sort(key=lambda item: item[0]) + out = [entry for _, entry in picked] + if key is not None: + self._remember(key, out) + return list(out) return out diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py new file mode 100644 index 000000000000..1b1d4b623abc --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py @@ -0,0 +1,181 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""One HIP module per binary per run, instead of one per binary per shape. + +An AOT cache holds shape-generic binaries: the same HSACO serves every shape +its capabilities cover. The sweeps nonetheless built a +:class:`~rocke.runtime.launcher.KernelLauncher` -- ``read_bytes`` plus +``hipModuleLoadData`` plus ``hipModuleGetFunction`` -- inside the loop over +candidates, inside the loop over cases, so a run of *C* cases over *K* +candidates loaded *C x K* modules of *K* distinct binaries. At a few hundred +microseconds each that is the single largest fixed cost per candidate in a +run that compiles nothing. + +:class:`ModuleCache` keys launchers by the binary and the entry point, so a +blob is read and loaded once per run and every later shape reuses the loaded +code object. + +Admission, not eviction +----------------------- +The budget is bytes of HSACO held, and when it is reached the cache stops +admitting rather than evicting. That is deliberate. Every case walks the same +candidate list in the same order, which is exactly the access pattern that +defeats LRU: with a capacity below the working set, each lookup evicts the +entry the next case will ask for and the hit rate is zero. Keeping a stable +prefix of the list instead gives a hit rate of ``budget / working set`` -- +and never does worse than not caching at all. A cache that does not evict +also never has to prove no launch is in flight on a module it unloads. + +The cost of a miss after the budget is reached is exactly the old behavior: +the binary is read and loaded for that candidate alone. :meth:`clear` drains +the device and unloads everything; a run that wants the memory back between +phases calls it, and process exit covers the rest. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Sequence, Tuple + +#: Default HSACO bytes to hold loaded. Conv binaries run from tens to a few +#: hundred kB, so this keeps a few thousand of them -- for most caches the +#: whole working set, and for the widest ones the part the probe reaches +#: first. +DEFAULT_BUDGET_BYTES = 512 * 1024 * 1024 + + +def _signature_key(signature: Sequence[Mapping[str, Any]]) -> Optional[tuple]: + """Hashable form of a launch signature, or ``None`` if it has no such form. + + The signature decides the launcher's packer, so two launchers over one + binary and entry point are still different launchers when it differs (a + conv kernel and the dgrad transform that shares its blob, say). + """ + try: + key = tuple(tuple(sorted(arg.items())) for arg in signature) + hash(key) + except TypeError: + return None + return key + + +class ModuleCache: + """Loaded :class:`KernelLauncher`\\s by ``(binary, entry point, signature)``. + + ``budget_bytes`` of 0 disables the cache: every :meth:`launcher` call + builds a fresh one, which is what the sweeps did before this existed. + """ + + def __init__(self, budget_bytes: int = DEFAULT_BUDGET_BYTES) -> None: + self._budget = max(0, int(budget_bytes)) + self._held: Dict[tuple, Any] = {} + self._bytes = 0 + self._hits = 0 + self._misses = 0 + self._full = False + + # -- query ------------------------------------------------------------ + + @property + def nbytes(self) -> int: + """HSACO bytes of the loaded modules.""" + return self._bytes + + @property + def hits(self) -> int: + return self._hits + + @property + def misses(self) -> int: + return self._misses + + def __len__(self) -> int: + return len(self._held) + + def summary(self) -> str: + """One line of accounting, for the end of a run.""" + total = self._hits + self._misses + pct = 100.0 * self._hits / total if total else 0.0 + line = ( + f"module cache: {len(self._held)} module(s), " + f"{self._bytes / 1024 / 1024:.1f} MiB, " + f"{self._hits}/{total} reused ({pct:.0f}%)" + ) + if self._full: + line += ( + f" -- budget of {self._budget / 1024 / 1024:.0f} MiB reached, " + f"the rest reload per use (--module-cache-bytes)" + ) + return line + + # -- use -------------------------------------------------------------- + + def launcher( + self, + *, + hsaco_path: Path, + kernel_name: str, + signature: Sequence[Mapping[str, Any]], + ) -> Any: + """The launcher for this binary and entry point, loading it if new. + + Raises whatever :class:`KernelLauncher` raises on a binary the device + rejects (``HipError``); callers already handle that per candidate. + """ + from rocke.runtime.launcher import KernelLauncher + + key: Optional[Tuple] = None + sig_key = _signature_key(signature) + if self._budget and sig_key is not None: + key = (os.fspath(hsaco_path), kernel_name, sig_key) + held = self._held.get(key) + if held is not None: + self._hits += 1 + return held + + self._misses += 1 + hsaco = Path(hsaco_path).read_bytes() + launcher = KernelLauncher( + hsaco=hsaco, kernel_name=kernel_name, signature=signature + ) + if key is not None: + if self._bytes + len(hsaco) <= self._budget: + self._held[key] = launcher + self._bytes += len(hsaco) + else: + self._full = True + return launcher + + def clear(self) -> None: + """Drain the device, then unload every module held. + + The drain is why this is not called implicitly: the cache cannot know + that the launches it handed out have finished, and unloading a module + under a kernel in flight faults. + """ + if not self._held: + return + from rocke.runtime import synchronize_and_release + + synchronize_and_release() + for launcher in self._held.values(): + launcher.unload() + self._held.clear() + self._bytes = 0 + self._full = False + + +def add_module_cache_arg(parser, *, default: int = DEFAULT_BUDGET_BYTES) -> None: + """Add ``--module-cache-bytes`` to an AOT sweep's parser.""" + parser.add_argument( + "--module-cache-bytes", + type=int, + default=default, + dest="module_cache_bytes", + metavar="N", + help=f"HSACO bytes to keep loaded across cases, so a shape-generic " + f"binary is loaded once per run rather than once per shape " + f"(default: {default // 1024 // 1024} MiB; 0 loads per use)", + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py index 1a767c062c39..f6682da710d4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/kernel_identity.py @@ -21,16 +21,21 @@ from __future__ import annotations +import dataclasses import hashlib import json import math from dataclasses import asdict, dataclass, field __all__ = [ + "DIRECT_PROBLEM_CAPABILITY_ATTRS", + "IDENTITY_CAPABILITY_FIELDS", "KernelIdentity", "current_llvm_flavor", "dgrad_sub_gemms", + "identity_capability_key", "identity_supports_problem", + "problem_capability_key", ] @@ -631,3 +636,139 @@ def identity_supports_problem( return False, why return True, "ok" + + +# --------------------------------------------------------------------------- +# Capability keys +# --------------------------------------------------------------------------- +# +# :func:`identity_supports_problem` is a pure function of the identity and the +# problem, but it reads only a part of each: most of an identity is tuning +# (tile shapes it does read, but also pipeline, LDS layout, scheduling) and +# most of a direct problem is extents the baked kernel takes as kernargs. A +# caller filtering one cache of ~10^4 entries per shape over ~10^3 shapes pays +# for that asymmetry ~10^7 times, so both sides get a *capability key*: the +# projection onto what the predicate actually reads. Two identities with the +# same key answer alike for every problem, and two problems with the same key +# are answered alike by every identity -- so the predicate can be evaluated +# once per (identity key, problem key) pair instead of once per entry and +# shape. +# +# Both keys are deliberately conservative: a field listed here that the +# predicate does not read only costs reuse, while one it reads and the key +# omits is a wrong answer. ``library/tests/test_capability_keys.py`` holds the +# two sides to that by perturbing the fields the keys leave out and asserting +# the predicate does not move. + +#: Fields of :class:`KernelIdentity` that :func:`identity_supports_problem` +#: reads (directly, or through ``_grid_fits`` / ``_group_merge_fits`` / +#: :func:`dgrad_sub_gemms`). ``is_direct`` is a property over ``algorithm`` +#: and stands in for it. +IDENTITY_CAPABILITY_FIELDS = ( + "llvm_flavor", + "direction", + "is_direct", + "is_3d", + "is_pointwise", + "grouped", + "patch", + "group_merge", + "tile_m", + "tile_n", + "tile_k", + "vector_size_a", + "vector_size_b", + "vector_size_c", + "dtype_d", + "split_k", + "two_stage", + "filter_h", + "filter_w", + "stride_h", + "stride_w", + "dilation_h", + "dilation_w", + "pad_h", + "pad_w", + "cpg", + "kpg", + "async_dma", + "async_chunk_a", + "async_chunk_b", + "max_sub_gemms", + "tap_outer_k", + "dy_halo", + "halo_max_wo", +) + +#: Problem attributes the predicate reads when the identity is a *direct* one. +#: A direct binary bakes its whole geometry and takes batch, image size and +#: group count as kernargs, so the predicate never looks at an extent: it +#: stops at the baked-geometry equalities. Every branch that reads one +#: (``_grid_fits``, the dgrad sub-GEMM table, the patch and dy-halo rules, the +#: wgrad atomic epilogue) is behind ``not identity.is_direct`` or behind a +#: field only an implicit-GEMM identity sets. That is what makes this key +#: coarse enough to be worth having: thousands of shapes collapse onto the few +#: hundred geometries the cache was built for. +DIRECT_PROBLEM_CAPABILITY_ATTRS = ( + "cpg", + "kpg", + "groups", + "is_3d", + "is_pointwise", + "Y", + "X", + "Z", + "KH", + "KW", + "stride", + "sH", + "sW", + "dilation", + "dH", + "dW", + "dil_h", + "dil_w", + "PAD", + "pH", + "pW", +) + + +def identity_capability_key(identity: KernelIdentity) -> tuple: + """``identity`` projected onto :data:`IDENTITY_CAPABILITY_FIELDS`. + + Identities sharing a key are indistinguishable to + :func:`identity_supports_problem` -- it is told the same thing about all + of them -- so a filter over a cache can answer once per key. + """ + return tuple(getattr(identity, name) for name in IDENTITY_CAPABILITY_FIELDS) + + +def problem_capability_key(problem: object, *, is_direct: bool) -> tuple | None: + """``problem`` projected onto what the predicate reads, or ``None``. + + ``is_direct`` says which projection applies: a direct identity reads only + :data:`DIRECT_PROBLEM_CAPABILITY_ATTRS`, an implicit-GEMM one reads + extents too (the launch grid, the dgrad sub-GEMM decomposition, the + dy-halo and patch rules), and rather than track which of those a future + rule adds, the implicit key is the whole problem. It is then exact rather + than coarse, which still collapses a shape that repeats across dtypes, + directions or runs. + + ``None`` means "not keyable" -- an implicit-GEMM problem that is not a + dataclass, or one whose fields are not hashable. Callers fall back to + evaluating the predicate. + """ + if is_direct: + return tuple(getattr(problem, a, None) for a in DIRECT_PROBLEM_CAPABILITY_ATTRS) + if not dataclasses.is_dataclass(problem) or isinstance(problem, type): + return None + key = tuple( + (f.name, getattr(problem, f.name, None)) for f in dataclasses.fields(problem) + ) + try: + hash(key) + except TypeError: + return None + return key diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py new file mode 100644 index 000000000000..bad7972e3530 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py @@ -0,0 +1,264 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +# +# Selection logic of benchmarks.common.batched_timing.rank, with the device +# stubbed out. What needs a GPU is the timing itself; what needs testing is +# what the ranking does with the times -- which candidates reach the +# full-fidelity round, which are dropped, and what happens to one that cannot +# be loaded or whose launch fails. Those are host decisions, so they are +# tested on the host: time_batch and prime_device are replaced by fakes that +# report a fixed time per candidate. + +from __future__ import annotations + +import pytest +from benchmarks.common import batched_timing + + +class _FakeDevice: + """Stands in for the device: candidate ``i`` runs for ``times[i]`` ms. + + ``None`` in ``times`` is a candidate whose launch raises -- ``time_batch`` + reports that as a ``None`` time, so the fake does the same. + """ + + def __init__(self, times, *, probe_error: float = 1.0): + self.times = list(times) + self.probe_error = probe_error + self.probed = 0 + self.measured = 0 + self.primed = 0 + self.batches = 0 + + def install(self, monkeypatch) -> None: + monkeypatch.setattr(batched_timing, "time_batch", self._time_batch) + monkeypatch.setattr(batched_timing, "prime_device", self._prime) + + def _prime(self, run, **_kwargs) -> None: + self.primed += 1 + + def _time_batch(self, runs, *, warmup, iters, stream=0, chunk=0): + self.batches += 1 + probing = iters == batched_timing.DEFAULT_PROBE_ITERS + out = [] + for run in runs: + i = run() + ms = self.times[i] + if ms is None: + out.append(None) + continue + if probing: + self.probed += 1 + # A probe is a one- or two-launch sample: close, not exact. + out.append(ms * self.probe_error) + else: + self.measured += 1 + out.append(ms) + return out + + def make_run(self, i): + return lambda i=i: i + + +def _rank(device, n, **kwargs): + kwargs.setdefault("warmup", 3) + kwargs.setdefault("iters", 10) + return batched_timing.rank(n, device.make_run, **kwargs) + + +def test_only_the_fastest_reach_the_full_round(monkeypatch): + times = [10.0, 3.0, 7.0, 1.0, 9.0, 2.0, 8.0, 4.0, 6.0, 5.0] + device = _FakeDevice(times) + device.install(monkeypatch) + + ranked = _rank(device, len(times), finalists=3, chunk=4) + + # Every candidate is probed once; only the three fastest are measured. + assert device.probed == len(times) + assert device.measured == 3 + assert sorted(r.index for r in ranked if r.is_finalist) == [1, 3, 5] + # A finalist reports the full-fidelity time, never the probe's. + assert [r.ms for r in ranked if r.index == 3] == [1.0] + + +def test_a_candidate_dropped_by_the_probe_keeps_its_probe_time(monkeypatch): + times = [1.0, 5.0, 6.0] + device = _FakeDevice(times, probe_error=1.1) + device.install(monkeypatch) + + ranked = _rank(device, len(times), finalists=1, chunk=8) + + dropped = [r for r in ranked if not r.is_finalist] + assert [r.index for r in dropped] == [1, 2] + # It is reportable -- that is what the sweep prints for it -- but it is + # the probe's number, so it is kept apart from ``ms``. + assert all(r.probe_ms is not None and r.ms is None for r in dropped) + assert dropped[0].probe_ms == pytest.approx(5.5) + + +def test_the_finalist_cut_spans_chunks(monkeypatch): + # The fastest candidate sits in the last chunk, so a cut applied per chunk + # rather than across the case would already have dropped it. + times = [9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0] + device = _FakeDevice(times) + device.install(monkeypatch) + + ranked = _rank(device, len(times), finalists=2, chunk=3) + + assert sorted(r.index for r in ranked if r.is_finalist) == [7, 8] + + +def test_a_candidate_that_cannot_be_loaded_is_marked_failed(monkeypatch): + times = [4.0, 1.0, 2.0] + device = _FakeDevice(times) + device.install(monkeypatch) + + def make_run(i): + return None if i == 1 else device.make_run(i) + + ranked = batched_timing.rank( + len(times), make_run, warmup=3, iters=10, finalists=8, chunk=8 + ) + + assert [r.index for r in ranked if r.failed] == [1] + assert ranked[1].probe_ms is None and ranked[1].ms is None + # The rest are unaffected: the load failure costs no launch at all. + assert device.probed == 2 + + +def test_a_launch_that_fails_never_ranks(monkeypatch): + # ``None`` from time_batch is a candidate whose launch raised. + times = [4.0, None, 2.0] + device = _FakeDevice(times) + device.install(monkeypatch) + + ranked = _rank(device, len(times), finalists=8, chunk=8) + + assert [r.index for r in ranked if r.failed] == [1] + assert sorted(r.index for r in ranked if r.is_finalist) == [0, 2] + + +def test_the_device_is_primed_once_per_case(monkeypatch): + device = _FakeDevice([3.0] * 12) + device.install(monkeypatch) + + _rank(device, 12, finalists=2, chunk=3) + + # Four probe chunks, one priming: the clocks are brought up before the + # first measured chunk and the chunk drains are too short to lose them. + assert device.primed == 1 + + +def test_the_finalist_hook_decides_what_may_rank(monkeypatch): + times = [3.0, 1.0, 2.0] + device = _FakeDevice(times) + device.install(monkeypatch) + checked = [] + + def on_finalist(i, run): + checked.append(i) + return i != 1 # the fastest one is wrong + + ranked = _rank(device, len(times), finalists=3, chunk=8, on_finalist=on_finalist) + + # Checked fastest-first, and every finalist is checked before it is timed. + assert checked == [1, 2, 0] + assert ranked[1].passed is False + # It is still timed, so its row carries a number for the caller to report. + assert ranked[1].ms == 1.0 + assert ranked[0].passed is True + + +def test_scratch_is_released_for_everything_that_cannot_rank(monkeypatch): + times = [9.0, 1.0, 8.0, None, 2.0, 7.0] + device = _FakeDevice(times) + device.install(monkeypatch) + dropped = [] + + ranked = _rank( + device, len(times), finalists=2, chunk=3, on_drop=dropped.append + ) + + # Everything but the two finalists: the slow ones and the one that failed. + assert sorted(dropped) == [0, 2, 3, 5] + assert sorted(r.index for r in ranked if r.is_finalist) == [1, 4] + # A finalist's scratch is still held -- it was timed with it. + assert not set(dropped) & {1, 4} + + +def test_a_full_scratch_budget_closes_the_chunk_early(monkeypatch): + device = _FakeDevice([1.0] * 8) + device.install(monkeypatch) + held = {"bytes": 0} + sizes = [] + + def make_run(i): + held["bytes"] += 100 + sizes.append(len(sizes)) + return device.make_run(i) + + batches = [] + real_time_batch = device._time_batch + + def counting(runs, **kwargs): + if kwargs["iters"] == batched_timing.DEFAULT_PROBE_ITERS: + batches.append(len(runs)) + return real_time_batch(runs, **kwargs) + + monkeypatch.setattr(batched_timing, "time_batch", counting) + + batched_timing.rank( + 8, + make_run, + warmup=3, + iters=10, + finalists=8, + chunk=100, + held_bytes=lambda: held["bytes"], + chunk_bytes=250, + # Nothing is ever released here, so the budget is passed for good + # after the third bind; every later chunk holds one candidate. + on_drop=lambda i: None, + ) + + # The chunk that crosses the budget keeps the candidate that crossed it. + assert batches == [3, 1, 1, 1, 1, 1] + + +def test_without_a_budget_the_count_alone_chunks(monkeypatch): + device = _FakeDevice([1.0] * 5) + device.install(monkeypatch) + batches = [] + real_time_batch = device._time_batch + + def counting(runs, **kwargs): + if kwargs["iters"] == batched_timing.DEFAULT_PROBE_ITERS: + batches.append(len(runs)) + return real_time_batch(runs, **kwargs) + + monkeypatch.setattr(batched_timing, "time_batch", counting) + + batched_timing.rank( + 5, device.make_run, warmup=3, iters=10, finalists=5, chunk=2 + ) + + assert batches == [2, 2, 1] + + +def test_probe_results_are_reported_as_they_land(monkeypatch): + device = _FakeDevice([5.0, 4.0, 3.0, 2.0, 1.0]) + device.install(monkeypatch) + # Read inside the callback: the Ranked handed to it is the live one, and a + # finalist's gets its ``ms`` filled in later. + seen = [] + + def on_probe(r): + seen.append((r.index, r.probe_ms, r.ms)) + + _rank(device, 5, finalists=1, chunk=2, on_probe=on_probe) + + # One callback per candidate, in candidate order, before any full-fidelity + # timing -- this is what a long case prints its progress from. + assert [i for i, _, _ in seen] == [0, 1, 2, 3, 4] + assert all(probe is not None for _, probe, _ in seen) + assert all(ms is None for _, _, ms in seen) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py new file mode 100644 index 000000000000..cdcce6865e41 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py @@ -0,0 +1,402 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""The capability keys say what ``identity_supports_problem`` reads. + +``KernelCache.compatible`` evaluates the predicate once per (identity key, +problem key) pair instead of once per entry and shape, which is only sound if +the keys cover everything the predicate reads. A field the keys leave out but +the predicate reads would make one cached answer stand in for two different +ones -- a binary offered to a shape it cannot run, reported as a launch +failure or, worse, as a wrong result. + +So the keys are checked the way they are used: perturb what a key leaves out +and assert the predicate does not move. Nothing is compiled and no GPU is +touched. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from benchmarks.common.kernel_cache import KernelCache +from kernels.common._conv_implicit_gemm_common import ConvProblem +from kernels.common.conv_direct_grouped import DirectConvProblem +from kernels.conv.kernel_identity import ( + DIRECT_PROBLEM_CAPABILITY_ATTRS, + IDENTITY_CAPABILITY_FIELDS, + KernelIdentity, + identity_capability_key, + identity_supports_problem, + problem_capability_key, +) + +_ARCH = "gfx942" + + +# --------------------------------------------------------------------------- +# Corpora +# --------------------------------------------------------------------------- + + +def _identity(**kw) -> KernelIdentity: + base = dict( + arch=_ARCH, + direction="fwd", + algorithm="implicit_gemm", + dtype_a="bf16", + dtype_b="bf16", + dtype_d="bf16", + tile_m=128, + tile_n=128, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=32, + warp_tile_n=32, + warp_tile_k=16, + pipeline="compv3", + epilogue="cshuffle", + wave_size=64, + vector_size_a=8, + vector_size_b=8, + vector_size_c=8, + ) + base.update(kw) + return KernelIdentity(**base) + + +def _identities() -> list: + """Identities spanning every branch of the predicate. + + One per capability class that reads a problem differently: the three + directions, the grouped and group-merged forms, and each of the + implicit-GEMM dgrad specialisations (patch, tilde table, tap-outer-K, + dY halo), plus the direct kernels, whose geometry is baked. + """ + return [ + _identity(), + _identity(grouped=True), + _identity(is_3d=True), + _identity(is_pointwise=True, filter_h=1, filter_w=1), + _identity(async_dma=True, async_chunk_a=8, async_chunk_b=8), + _identity(direction="wgrad", split_k=4), + _identity(direction="wgrad", split_k=4, two_stage=True), + _identity(direction="wgrad", group_merge=4, grouped=True), + _identity(direction="fwd", group_merge=4, grouped=True), + _identity(direction="dgrad", grouped=True), + _identity(direction="dgrad", patch=True), + _identity(direction="dgrad", max_sub_gemms=4), + _identity(direction="dgrad", tap_outer_k=True), + _identity(direction="dgrad", dy_halo=1, halo_max_wo=64, filter_h=3, filter_w=3), + _identity( + direction="direct_fwd", + algorithm="direct_nongrouped", + filter_h=3, + filter_w=3, + pad_h=1, + pad_w=1, + stride_h=1, + stride_w=1, + dilation_h=1, + dilation_w=1, + vector_size_a=0, + vector_size_b=0, + vector_size_c=0, + ), + _identity( + direction="direct_wgrad", + algorithm="direct_dw_wgrad", + filter_h=3, + filter_w=3, + pad_h=1, + pad_w=1, + stride_h=1, + stride_w=1, + dilation_h=1, + dilation_w=1, + cpg=1, + kpg=1, + vector_size_a=0, + vector_size_b=0, + vector_size_c=0, + ), + _identity( + direction="direct_dgrad", + algorithm="direct_dgrad_mfma_rtc", + filter_h=3, + filter_w=3, + pad_h=1, + pad_w=1, + stride_h=1, + stride_w=1, + dilation_h=1, + dilation_w=1, + vector_size_a=0, + vector_size_b=0, + vector_size_c=0, + ), + ] + + +def _conv_problems() -> list: + return [ + ConvProblem(N=2, Hi=28, Wi=28, C=64, K=64, Y=3, X=3, pH=1, pW=1), + ConvProblem(N=1, Hi=56, Wi=56, C=32, K=64, Y=1, X=1), + ConvProblem(N=4, Hi=14, Wi=14, C=128, K=128, Y=3, X=3, pH=1, pW=1, groups=128), + ConvProblem(N=2, Hi=32, Wi=32, C=64, K=64, Y=3, X=3, sH=2, sW=2, pH=1, pW=1), + ConvProblem(N=2, Hi=32, Wi=32, C=64, K=64, Y=3, X=3, dH=2, dW=2, pH=2, pW=2), + ConvProblem(N=2, Hi=224, Wi=224, C=3, K=768, Y=16, X=16, sH=16, sW=16), + ConvProblem( + N=1, + Hi=16, + Wi=16, + C=32, + K=32, + Y=3, + X=3, + pH=1, + pW=1, + Di=16, + Z=3, + sD=1, + pD=1, + dD=1, + ), + ] + + +def _direct_problems() -> list: + return [ + DirectConvProblem(N=2, H=28, W=28, groups=1, cpg=64, kpg=64), + DirectConvProblem(N=1, H=56, W=56, groups=16, cpg=16, kpg=16), + DirectConvProblem( + N=4, H=14, W=14, groups=128, cpg=1, kpg=1, KH=3, KW=3, PAD=1 + ), + DirectConvProblem( + N=2, H=32, W=32, groups=1, cpg=32, kpg=64, KH=1, KW=1, PAD=0, stride=2 + ), + DirectConvProblem( + N=2, H=32, W=32, groups=8, cpg=8, kpg=8, dil_h=2, dil_w=2, PAD=2 + ), + ] + + +def _answers(identity, problems=None) -> list: + """The predicate's verdicts (not its reasons) over a problem list. + + Each family is asked about its own problem type: the predicate reads a + ``DirectConvProblem`` for a direct identity and a ``ConvProblem`` for an + implicit-GEMM one, and the two do not share an attribute set. + """ + if problems is None: + problems = _direct_problems() if identity.is_direct else _conv_problems() + return [identity_supports_problem(identity, problem)[0] for problem in problems] + + +def _perturbed(value): + """A different value of the same kind, or ``None`` when there is none.""" + if isinstance(value, bool): + return not value + if isinstance(value, int): + return value + 1 + if isinstance(value, str): + return value + "_x" + if value is None: + return 3 + return None + + +# --------------------------------------------------------------------------- +# The identity side +# --------------------------------------------------------------------------- + +# ``is_direct`` is a property over ``algorithm``; the key carries the property, +# so the field behind it is not free to move. +_IDENTITY_KEY_FIELDS = set(IDENTITY_CAPABILITY_FIELDS) | {"algorithm"} + + +def test_identity_key_covers_every_field_the_predicate_reads(): + """Moving a field outside the key must not move the predicate. + + If this fails, the named field belongs in ``IDENTITY_CAPABILITY_FIELDS``: + the cache groups entries by that key and answers once per group, so two + identities differing only in it would be told the same thing. + """ + free = [ + f.name + for f in dataclasses.fields(KernelIdentity) + if f.name not in _IDENTITY_KEY_FIELDS + ] + assert free, "every field is in the key -- the grouping would be pointless" + for identity in _identities(): + expected = _answers(identity) + for name in free: + other = _perturbed(getattr(identity, name)) + if other is None: + continue + moved = dataclasses.replace(identity, **{name: other}) + assert _answers(moved) == expected, ( + f"{identity.direction}/{identity.algorithm}: the predicate " + f"reads {name!r}, which the identity capability key leaves out" + ) + + +def test_identity_key_separates_what_the_predicate_separates(): + """Two identities with one key is the claim; different keys for different + capabilities is what makes the key useful rather than merely sound.""" + keys = {identity_capability_key(i) for i in _identities()} + assert len(keys) == len(_identities()) + + # ... and tuning that the predicate never reads must not split a group. + base = _identity() + assert identity_capability_key(base) == identity_capability_key( + dataclasses.replace(base, pipeline="compv4", lds_k_pad=8, waves_per_eu=2) + ) + + +# --------------------------------------------------------------------------- +# The problem side +# --------------------------------------------------------------------------- + + +def test_direct_problem_key_covers_every_extent_the_predicate_reads(): + """A direct binary bakes its geometry and takes extents as kernargs, so + the predicate must not read one. That is what lets thousands of shapes + share a few hundred keys -- and it is a claim, so it is checked.""" + direct = [i for i in _identities() if i.is_direct] + assert direct + free = [ + f.name + for f in dataclasses.fields(DirectConvProblem) + if f.name not in DIRECT_PROBLEM_CAPABILITY_ATTRS + ] + assert set(free) == {"N", "H", "W", "dtype"}, free + for problem in _direct_problems(): + expected = [_answers(i, [problem]) for i in direct] + for name in free: + other = _perturbed(getattr(problem, name)) + moved = dataclasses.replace(problem, **{name: other}) + assert [_answers(i, [moved]) for i in direct] == expected, ( + f"the direct predicate reads {name!r}, which " + f"DIRECT_PROBLEM_CAPABILITY_ATTRS leaves out" + ) + + +def test_direct_problem_key_collapses_shapes_onto_geometries(): + p = DirectConvProblem(N=2, H=28, W=28, groups=16, cpg=16, kpg=16) + same = dataclasses.replace(p, N=64, H=7, W=7, dtype="bf16") + assert problem_capability_key(p, is_direct=True) == problem_capability_key( + same, is_direct=True + ) + other = dataclasses.replace(p, PAD=0) + assert problem_capability_key(p, is_direct=True) != problem_capability_key( + other, is_direct=True + ) + + +def test_implicit_problem_key_is_the_whole_shape(): + """An implicit-GEMM binary is shape-generic, so its admission does read + extents (the grid, the dgrad sub-GEMM table). The key is exact there.""" + p = _conv_problems()[0] + assert problem_capability_key(p, is_direct=False) == problem_capability_key( + dataclasses.replace(p), is_direct=False + ) + assert problem_capability_key(p, is_direct=False) != problem_capability_key( + dataclasses.replace(p, N=p.N + 1), is_direct=False + ) + + +def test_unkeyable_problem_is_reported_as_such(): + """A caller must be able to tell "no key" from "this key", so it can fall + back to evaluating the predicate rather than memoise everything as one.""" + + class NotADataclass: + cpg = kpg = 1 + + assert problem_capability_key(NotADataclass(), is_direct=False) is None + + +# --------------------------------------------------------------------------- +# What the cache does with them +# --------------------------------------------------------------------------- + + +@pytest.fixture +def cache(tmp_path): + """A cache holding every identity of the corpus with a placeholder blob.""" + c = KernelCache(tmp_path, _ARCH) + for i, identity in enumerate(_identities()): + c.put(identity, b"placeholder", {"kernel_name": f"k{i}"}) + # A second entry per capability class, differing only in tuning: the + # grouping must return both, in the order a plain scan would. + c.put( + dataclasses.replace(identity, pipeline="compv4"), + b"placeholder", + {"kernel_name": f"k{i}_v4"}, + ) + return c + + +def _scan(cache, problem, direction): + """``compatible`` as it read before grouping: one predicate per entry.""" + return [ + entry + for entry in cache._entries(direction) + if identity_supports_problem(entry[0], problem)[0] + ] + + +@pytest.mark.parametrize( + "direction", ["fwd", "wgrad", "dgrad", "direct_fwd", "direct_wgrad", "direct_dgrad"] +) +def test_compatible_matches_a_plain_scan(cache, direction): + """Grouping and memoising must not change which kernels are offered, nor + the order -- ``--sample``'s fixed-seed subset and the probe order are + taken off this list.""" + direct = direction.startswith("direct_") + problems = _direct_problems() if direct else _conv_problems() + offered = 0 + for problem in problems: + expected = _scan(cache, problem, direction) + offered += len(expected) + assert cache.compatible(problem, direction=direction) == expected + # Again, from the memo this time. + assert cache.compatible(problem, direction=direction) == expected + # An agreement on "nothing matches" everywhere would prove nothing. + assert offered + + +def test_compatible_returns_the_callers_own_list(cache): + problem = _direct_problems()[0] + first = cache.compatible(problem, direction="direct_fwd") + first.clear() + assert cache.compatible(problem, direction="direct_fwd") + + +def test_a_write_drops_the_memo(cache): + """An entry added after a shape was answered has to be offered to it.""" + # The depthwise shape: the corpus's wgrad kernel bakes cpg = kpg = 1. + problem = _direct_problems()[2] + before = cache.compatible(problem, direction="direct_wgrad") + assert before + added = dataclasses.replace( + [i for i in _identities() if i.direction == "direct_wgrad"][0], + tile_k=32, + ) + cache.put(added, b"placeholder", {"kernel_name": "added"}) + after = cache.compatible(problem, direction="direct_wgrad") + assert len(after) == len(before) + 1 + + +def test_entry_meta_reads_no_binary(cache, monkeypatch): + """Plan building asks for a helper's entry point, not its code.""" + identity = _identities()[0] + monkeypatch.setattr( + type(cache.blob_path("x")), + "read_bytes", + lambda self: pytest.fail("entry_meta read an HSACO"), + ) + assert cache.entry_meta(identity)["kernel_name"] == "k0" + assert cache.entry_meta(dataclasses.replace(identity, tile_k=13)) is None diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py new file mode 100644 index 000000000000..d4dd50059bc9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py @@ -0,0 +1,160 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Admission policy of benchmarks.common.module_cache, with HIP stubbed out. + +What needs a GPU is loading the code object; what needs testing is *when* the +cache loads one -- which lookups reuse a module, which key separates two, and +what the byte budget does once it is reached. Those are host decisions, so +``KernelLauncher`` is replaced by a fake that counts its constructions. +""" + +from __future__ import annotations + +import pytest + +from benchmarks.common.module_cache import ModuleCache + +_SIG = [{"name": "A", "type": "ptr"}, {"name": "n", "type": "i32"}] +_OTHER_SIG = [{"name": "A", "type": "ptr"}, {"name": "n", "type": "i32"}] + + +class _FakeLauncher: + """Counts loads, and records whether it was unloaded.""" + + loads = 0 + + def __init__(self, *, hsaco, kernel_name, signature): + type(self).loads += 1 + self.hsaco = hsaco + self.kernel_name = kernel_name + self.signature = signature + self.unloaded = False + + def unload(self): + self.unloaded = True + + +@pytest.fixture +def fake_hip(monkeypatch): + """``ModuleCache`` builds launchers through this instead of HIP.""" + import rocke.runtime + import rocke.runtime.launcher + + _FakeLauncher.loads = 0 + monkeypatch.setattr(rocke.runtime.launcher, "KernelLauncher", _FakeLauncher) + monkeypatch.setattr( + rocke.runtime, "synchronize_and_release", lambda *a, **k: None, raising=False + ) + return _FakeLauncher + + +@pytest.fixture +def blobs(tmp_path): + """Three distinct binaries of 1 kB each.""" + out = [] + for i in range(3): + path = tmp_path / f"k{i}.hsaco" + path.write_bytes(bytes([i]) * 1024) + out.append(path) + return out + + +def test_a_binary_is_loaded_once_however_many_shapes_ask(fake_hip, blobs): + cache = ModuleCache() + first = cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=_SIG) + for _ in range(10): + assert ( + cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=_SIG) + is first + ) + assert fake_hip.loads == 1 + assert (cache.hits, cache.misses) == (10, 1) + + +@pytest.mark.parametrize( + "second", + [ + pytest.param({"kernel_name": "other"}, id="entry-point"), + pytest.param({"signature": _OTHER_SIG}, id="signature"), + ], +) +def test_one_blob_two_entry_points_are_two_launchers(fake_hip, blobs, second): + """The launcher bakes the entry point and the packer, so neither may be + shared across a key that differs in them.""" + cache = ModuleCache() + call = dict(hsaco_path=blobs[0], kernel_name="k", signature=_SIG) + a = cache.launcher(**call) + b = cache.launcher(**{**call, **second}) + assert a is not b + assert fake_hip.loads == 2 + + +def test_distinct_binaries_are_distinct_entries(fake_hip, blobs): + cache = ModuleCache() + held = [ + cache.launcher(hsaco_path=p, kernel_name="k", signature=_SIG) for p in blobs + ] + assert len({id(x) for x in held}) == 3 + assert len(cache) == 3 + assert cache.nbytes == 3 * 1024 + + +def test_the_budget_stops_admitting_rather_than_evicting(fake_hip, blobs): + """Every case walks the same candidate list in the same order, so evicting + the entry the next case will ask for is the one policy guaranteed to give + no hits. A full cache keeps what it has.""" + cache = ModuleCache(2 * 1024) + first_pass = [ + cache.launcher(hsaco_path=p, kernel_name="k", signature=_SIG) for p in blobs + ] + assert len(cache) == 2 + second_pass = [ + cache.launcher(hsaco_path=p, kernel_name="k", signature=_SIG) for p in blobs + ] + # The two that fit are reused; the third is loaded again, as it would be + # with no cache at all. + assert second_pass[0] is first_pass[0] + assert second_pass[1] is first_pass[1] + assert second_pass[2] is not first_pass[2] + assert fake_hip.loads == 4 + assert "budget" in cache.summary() + + +def test_a_budget_of_zero_loads_per_use(fake_hip, blobs): + cache = ModuleCache(0) + a = cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=_SIG) + b = cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=_SIG) + assert a is not b + assert fake_hip.loads == 2 + assert len(cache) == 0 + + +def test_clear_unloads_every_module(fake_hip, blobs, monkeypatch): + drained = [] + import rocke.runtime + + monkeypatch.setattr( + rocke.runtime, + "synchronize_and_release", + lambda *a, **k: drained.append(True), + raising=False, + ) + cache = ModuleCache() + held = [ + cache.launcher(hsaco_path=p, kernel_name="k", signature=_SIG) for p in blobs + ] + cache.clear() + assert drained, "a module must not be unloaded with launches possibly in flight" + assert all(launcher.unloaded for launcher in held) + assert len(cache) == 0 and cache.nbytes == 0 + + +def test_an_unhashable_signature_is_not_a_crash(fake_hip, blobs): + """A signature the key cannot be built from loads per use instead of + raising: the cache is an optimisation, not a dependency.""" + cache = ModuleCache() + weird = [{"name": "A", "shape": [1, 2]}] + a = cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=weird) + b = cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=weird) + assert a is not b + assert len(cache) == 0 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 0fcc8a7bf8f8..f308fc76bd0d 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 @@ -467,6 +467,15 @@ def signature(self) -> Sequence[Mapping[str, Any]]: def cluster_dims(self) -> Optional[Tuple[int, int, int]]: return self._cluster_dims + def _live_fn(self): + """The kernel handle, or a clear error if :meth:`unload` took it.""" + if self._fn is None: + raise RuntimeError( + f"KernelLauncher({self._kernel_name!r}) was unloaded; " + f"build a new one to launch it again" + ) + return self._fn + def _resolve_cluster(self, config: LaunchConfig) -> Optional[Tuple[int, int, int]]: if config.cluster is None: return self._cluster_dims @@ -501,7 +510,7 @@ def __call__( # there's no bucket growth and no need to attach tensor # refs to a completion event. rt.launch_blocking( - self._fn, + self._live_fn(), config.grid, config.block, args, @@ -522,7 +531,7 @@ def __call__( # us and ~1 us per launch in tight benchmark loops, and # multiple-percent on tiny kernels like the conv bake-off. rt.launch( - self._fn, + self._live_fn(), config.grid, config.block, args, @@ -551,7 +560,7 @@ def bind( lives; drain the stream before dropping it. """ launch = _runtime().prepare_launch( - self._fn, + self._live_fn(), config.grid, config.block, self._packer(values), @@ -565,9 +574,32 @@ def run(_keep=keep) -> None: return run + def unload(self) -> None: + """Unload the HIP module this launcher owns. The launcher is dead after. + + A launcher normally lives as long as its problem shape, and the module + goes when the process does. A harness that loads thousands of distinct + binaries in one run (an AOT sweep over a kernel cache) needs to give + the code objects back rather than hold every one it ever launched, so + it owns the lifetime explicitly and calls this. + + **The device must be drained first.** Unloading a module whose kernels + are still in flight faults; this call cannot check that for the caller + (there is no stream to wait on here), so the caller synchronizes -- + :func:`synchronize_and_release` -- before it unloads. + + Idempotent: a second call does nothing. Calling the launcher after it + is unloaded raises. + """ + module, self._module = self._module, None + self._fn = None + if module is not None: + module.unload() + 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})" + dead = "" if self._module is not None else ", unloaded" + return f"KernelLauncher({self._kernel_name!r}{key_str}{dead})" class PipelineLauncher: From 2a28cb8c1e3c2593808962a7c64ed18a14cb4f34 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Sun, 25 Oct 2026 23:35:54 +0000 Subject: [PATCH 2/9] impr --- .../common/benchmark_implicit_gemm_conv.py | 172 +++++++++++------- .../library/tests/test_aot_launch_zeroing.py | 118 ++++++++++++ .../python/rocke/runtime/hip_module.py | 29 +++ 3 files changed, 253 insertions(+), 66 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_aot_launch_zeroing.py 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 798e72ae4aa8..2e811029bbc9 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 @@ -4838,25 +4838,14 @@ def _u8(t): # so it is freed as the sweep goes rather than per case. live_extras: list = [] - # --fast-rank takes every kernel whose scratch it can account for. - # A stream-K binary allocates and zeroes per launch config through - # its own contract (_cached_streamk_run returns the free), so it - # stays on the serial path below whichever way the flag is set. + # --fast-rank takes every kernel of the case, stream-K included: + # its workspace is accounted like any other launch config's, and + # its per-launch zeroing is queued rather than drained for. if args.fast_rank: - streamk = [ - c - for c in candidates - if direction == "wgrad" and c[0].streamk != "off" - ] - fast = [ - c - for c in candidates - if not (direction == "wgrad" and c[0].streamk != "off") - ] - if fast: + if candidates: failed, stopped = _fast_rank_candidates( args=args, - candidates=fast, + candidates=candidates, modules=modules, direction=direction, problem=problem, @@ -4883,7 +4872,7 @@ def _u8(t): for buf in (A_dev, B_dev, D_dev): rt.free(buf) return 1 - candidates = streamk + candidates = [] n_cand = len(candidates) for cand_idx, (ident, hsaco_path, meta) in enumerate(candidates, 1): @@ -4909,10 +4898,11 @@ def _u8(t): # kernarg) and has its own ABI, workspace and zeroing # contract per reduction, so it gets its own launch setup and # one launch; its workspace is released after the loop. - sk_bound = sk_free = None + sk_bound = None + sk_live: list = [] if direction == "wgrad" and ident.streamk != "off": try: - sk_run, sk_free, sk_label = _cached_streamk_run( + sk_run, sk_extras, sk_label = _cached_streamk_run( args, problem, ident, @@ -4928,6 +4918,7 @@ def _u8(t): except ValueError as e: print(f" [skip] {ident.label()}: {e}", flush=True) continue + sk_live.append(sk_extras) sk_bound = (sk_run, sk_label) # A split-K wgrad binary is one kernel for every degree > 1 @@ -5025,8 +5016,7 @@ def _u8(t): stopped, passed = False, False if stopped: _free_extras(rt, live_extras) - if sk_free is not None: - sk_free() + _free_extras(rt, sk_live) for buf in (A_dev, B_dev, D_dev): rt.free(buf) return 1 @@ -5076,8 +5066,7 @@ def _u8(t): flush=True, ) - if sk_free is not None: - sk_free() + _free_extras(rt, sk_live) synchronize_and_release() _free_extras(rt, live_extras) @@ -5173,13 +5162,19 @@ def _fast_rank_candidates( probe rules out gives its scratch back at the end of the chunk that measured it. - Three differences from the serial path, all of them the point of the flag: - ``--verify`` checks the finalists rather than every kernel, + Stream-K kernels rank here too. Their workspace is per launch config, as + the two-stage one is, so it is accounted and released the same way; what + kept them out was the zeroing their reduction needs before every launch, + which used to drain the device and so serialised the batch around them. + It is queued on the launch stream now (see :func:`_streamk_launch`), and + same-stream FIFO is what makes it safe for the ``atomic`` reduction to + clear the shared dW while the previous candidate's launches are still in + flight. + + Two differences from the serial path, both of them the point of the flag: + ``--verify`` checks the finalists rather than every kernel, and ``--split-k-prune`` does not apply (a kernel's degrees are ranked like any - other candidates, and a probe launch costs about what the prune saved), - and stream-K kernels are not handled here at all -- their workspace and - zeroing contract is per launch config, so the caller leaves them on the - serial path. + other candidates, and a probe launch costs about what the prune saved). Returns ``(n_failed, stopped)``; ``stopped`` is a ``--dump-fail`` dump, which ends the sweep as it does on the serial path. @@ -5200,7 +5195,11 @@ def _fast_rank_candidates( flush=True, ) continue - if direction == "wgrad" and ident.split_k > 1: + if direction == "wgrad" and ident.streamk != "off": + # A stream-K binary partitions K over its CTA pool itself; the + # split-K kernarg is not its knob, so it ranks once. + degrees = (None,) + elif direction == "wgrad" and ident.split_k > 1: degrees = _runtime_split_ks( problem, ident, args.split_k, args.split_k_degrees ) @@ -5246,24 +5245,48 @@ def make_run(i): print(f" [skip] {ident.label()}: {e}", flush=True) return None mine: list = [] - bound = _cached_bind_batch( - batch_chunks, - direction, - run_problem, - ident, - meta, - rt, - launcher, - case_dtype, - split_k, - ptrs, - sizes, - arch, - stage2_kernels, - KernelLauncher, - LaunchConfig, - mine, - ) + if direction == "wgrad" and ident.streamk != "off": + # Its own ABI, workspace and zeroing contract; one launch, on the + # case's problem rather than a batch slice, as the serial path + # runs it. + try: + sk_run, sk_extras, sk_label = _cached_streamk_run( + args, + problem, + ident, + rt, + arch, + launcher, + ptrs, + sizes, + stage2_kernels, + KernelLauncher, + LaunchConfig, + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + return None + mine.append(sk_extras) + bound = (sk_run, sk_label) + else: + bound = _cached_bind_batch( + batch_chunks, + direction, + run_problem, + ident, + meta, + rt, + launcher, + case_dtype, + split_k, + ptrs, + sizes, + arch, + stage2_kernels, + KernelLauncher, + LaunchConfig, + mine, + ) if bound is None: _release(mine) return None @@ -5574,12 +5597,22 @@ def _streamk_launch( ``spec`` is the stream-K :class:`WgradConvSpec` (its ``streamk_ctas`` is the CTA pool, -1 = the arch default); ``problem`` the launch shape; ``ptrs`` / ``sizes`` the dY, X, dW buffers. Allocates the reduction's - workspace and returns ``(run, free, pool)``: ``run()`` zeroes what the + workspace and returns ``(run, extras, pool)``: ``run()`` zeroes what the reduction accumulates into -- dW for ``atomic``, the f32 scratch for ``workspace``, only the flag table for ``linear``/``tree`` -- launches the - kernel and, for ``workspace``, the Stage 2 reduce into dW; ``free()`` - releases the workspace; ``pool`` is the resolved CTA pool. + kernel and, for ``workspace``, the Stage 2 reduce into dW; ``extras`` + carries the workspace in the shape :func:`_free_extras` frees and + :func:`_extras_bytes` accounts, as the other legs' scratch is; ``pool`` is + the resolved CTA pool. + + The zeroing is queued on the launch stream, not drained for: see + :meth:`rocke.runtime.hip_module.Runtime.memset_async`. Same-stream FIFO + orders it against this kernel's launches, and against those of any other + kernel queued in the same batch -- which is what lets the probe time + stream-K kernels back to back on one shared dW. """ + from rocke.runtime import resolve_stream + from kernels.common.conv_implicit_gemm_wgrad import ( wgrad_streamk_grid, wgrad_streamk_launch, @@ -5596,7 +5629,9 @@ def _streamk_launch( # An all-data-parallel partition (no stream-K CTAs) needs no workspace, # but the kernargs still carry the pointers, so a 4-byte stub stands in # for the empty allocation. - ws = rt.alloc(max(ws_bytes, 4)) + alloc_bytes = max(ws_bytes, 4) + ws = rt.alloc(alloc_bytes) + stream = resolve_stream(0) try: extra = {} if reduction == "workspace": @@ -5633,19 +5668,16 @@ def _streamk_launch( def run(): if reduction == "atomic": - rt.memset(ptrs[2], 0, sizes[2]) + rt.memset_async(ptrs[2], 0, sizes[2], stream) elif reduction == "workspace": - rt.memset(ws, 0, ws_bytes) + rt.memset_async(ws, 0, ws_bytes, stream) elif flags_bytes: - rt.memset(ws, 0, flags_bytes) + rt.memset_async(ws, 0, flags_bytes, stream) launcher(values, config=cfg) if s2 is not None: s2[0](s2[1], config=s2[2]) - def free(): - rt.free(ws) - - return run, free, pool + return run, {"ws_ptr": ws, "ws_bytes": alloc_bytes}, pool def _streamk_mode_label(mode: str, reduction: str) -> str: @@ -5668,7 +5700,7 @@ def _cached_streamk_run( KernelLauncher, LaunchConfig, ): - """``(run, free, label)`` for a cached stream-K wgrad binary. + """``(run, extras, label)`` for a cached stream-K wgrad binary. The launch helpers take a spec; this one describes the cached binary on the launch shape. Nothing is built from it -- the partition, workspace and @@ -5699,7 +5731,7 @@ def _cached_streamk_run( streamk_ctas=args.streamk_ctas, ws_replicas=_ws_replicas_of(ident), ) - run, free, pool = _streamk_launch( + run, extras, pool = _streamk_launch( rt=rt, launcher=launcher, spec=spec, @@ -5711,7 +5743,7 @@ def _cached_streamk_run( KernelLauncher=KernelLauncher, LaunchConfig=LaunchConfig, ) - return run, free, f"{ident.label()} @streamk_ctas={pool}" + return run, extras, f"{ident.label()} @streamk_ctas={pool}" def _split_k_degrees( @@ -5931,6 +5963,8 @@ def _cached_bind( :func:`_free_extras`). Returns None, after saying why, when the binary cannot launch ``problem``. """ + from rocke.runtime import resolve_stream + from benchmarks.common.kernel_sweep import _launch_values_for try: @@ -5952,6 +5986,9 @@ def _cached_bind( stream=0, ) run = launcher.bind(values, config=cfg) + # The stream ``bind`` resolved ``cfg.stream`` to, so a zeroing step below + # queues ahead of this kernel's launches rather than draining for them. + stream = resolve_stream(cfg.stream) if ( direction == "wgrad" and ident.split_k > 1 @@ -5961,11 +5998,13 @@ def _cached_bind( # The packed-atomic split-K epilogue accumulates into dW, so every # launch needs dW zeroed first -- part of the algorithm's cost, as the # JIT sweep times it and as the two-stage leg below times its scratch - # clear. + # clear. Queued on the launch stream: dW is shared by every candidate + # of the case, and FIFO is what keeps one candidate's clear behind the + # previous one's launches when the probe batches them. _launch = run - def run(_launch=_launch, dw_ptr=ptrs[2], dw_bytes=sizes[2]): - rt.memset(dw_ptr, 0, dw_bytes) + def run(_launch=_launch, dw_ptr=ptrs[2], dw_bytes=sizes[2], stream=stream): + rt.memset_async(dw_ptr, 0, dw_bytes, stream) _launch() elif direction == "wgrad" and ident.two_stage: @@ -5995,8 +6034,9 @@ def run( s2_launcher=s2_launcher, s2_values=s2_values, s2_cfg=s2_cfg, + stream=stream, ): - rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) + rt.memset_async(extras["ws_ptr"], 0, extras["ws_bytes"], stream) launcher(values, config=cfg) s2_launcher(s2_values, config=s2_cfg) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_aot_launch_zeroing.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_aot_launch_zeroing.py new file mode 100644 index 000000000000..116f558a8c04 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_aot_launch_zeroing.py @@ -0,0 +1,118 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""A per-launch clear in the AOT sweep must be queued, never drained for. + +Three wgrad legs have to zero something before every launch: the packed-atomic +split-K epilogue accumulates into dW, the two-stage path accumulates into its +f32 scratch, and a stream-K reduction clears dW, the scratch or its flag table +depending on which reduction it uses. That clear is part of the algorithm's +cost and is inside the timed region on purpose. + +``hipMemset`` is synchronous with respect to the whole device, so one inside a +timed ``run()`` drains everything queued before it. On the serial path that +costs little -- the next thing was a drain anyway. Under ``--fast-rank`` it +undoes the flag: :func:`benchmarks.common.batched_timing.rank` queues a chunk +of candidates between two drains precisely so they do not serialise, and one +blocking clear in the middle serialises them again. Queued on the launch +stream (``hipMemsetAsync``), the clear is also what makes a shared output safe +to reuse across a batch, since same-stream FIFO orders it behind the previous +candidate's launches. + +This is an easy thing to undo by writing the obvious ``rt.memset`` in a new +leg, and the symptom is a silent loss of speed rather than a failure, so the +rule is checked by reading the source rather than left to review. +""" + +from __future__ import annotations + +import ast +import inspect +import textwrap + +import pytest + +from benchmarks.common import benchmark_implicit_gemm_conv as sweep + +#: The AOT functions that build a launch closure for a cached binary. +_LAUNCH_BUILDERS = ("_cached_bind", "_streamk_launch") + + +def _function(name: str) -> ast.FunctionDef: + src = textwrap.dedent(inspect.getsource(getattr(sweep, name))) + return ast.parse(src).body[0] + + +def _blocking_memsets(node: ast.AST) -> list: + """``rt.memset(...)`` calls anywhere under ``node``.""" + out = [] + for sub in ast.walk(node): + if ( + isinstance(sub, ast.Call) + and isinstance(sub.func, ast.Attribute) + and sub.func.attr == "memset" + ): + out.append(sub) + return out + + +def _run_closures(fn: ast.FunctionDef) -> list: + """The nested ``run`` definitions -- the callable the timing loop calls.""" + return [ + node + for node in ast.walk(fn) + if isinstance(node, ast.FunctionDef) and node.name == "run" and node is not fn + ] + + +@pytest.mark.parametrize("name", _LAUNCH_BUILDERS) +def test_a_timed_launch_closure_never_drains_the_device(name): + fn = _function(name) + closures = _run_closures(fn) + assert closures, f"{name} builds no run() closure -- has it been renamed?" + for closure in closures: + blocking = _blocking_memsets(closure) + assert not blocking, ( + f"{name}: run() calls a blocking memset at line " + f"{blocking[0].lineno} of the function. A clear inside a timed " + f"launch has to be rt.memset_async(..., stream) on the launch " + f"stream, or it drains every candidate the probe queued with it." + ) + + +@pytest.mark.parametrize("name", _LAUNCH_BUILDERS) +def test_the_clear_goes_on_the_launch_stream(name): + """An async clear on the wrong stream is worse than a blocking one: it + would be unordered against the launches it has to precede.""" + fn = _function(name) + for closure in _run_closures(fn): + for call in ast.walk(closure): + if ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Attribute) + and call.func.attr == "memset_async" + ): + assert len(call.args) == 4, ( + f"{name}: memset_async without an explicit stream falls " + f"back to the null stream, which synchronises with every " + f"other stream on ROCm" + ) + assert isinstance(call.args[3], ast.Name) + + +def test_the_runtime_offers_a_queued_fill(): + from rocke.runtime.hip_module import Runtime + + assert hasattr(Runtime, "memset_async") + params = list(inspect.signature(Runtime.memset_async).parameters) + assert params == ["self", "ptr", "value", "nbytes", "stream"] + + +def test_setup_outside_a_launch_still_drains(): + """The rule is about the timed region, not about memset. Allocating a + two-stage workspace clears it once, before anything is queued, and that + one must stay blocking -- the launch that reads it is bound afterwards.""" + fn = _function("_extras_and_grid") + assert _blocking_memsets(fn), ( + "_extras_and_grid no longer clears the workspace it allocates; if the " + "clear moved into the launch, it has to be queued instead" + ) 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 5580c573e67d..aea9559d7776 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 @@ -164,6 +164,13 @@ def _b(name: str, *argtypes, restype=ctypes.c_int) -> _LazyFn: "hipMemcpy", ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int ) _hipMemset = _b("hipMemset", ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t) +_hipMemsetAsync = _b( + "hipMemsetAsync", + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_size_t, + ctypes.c_void_p, +) _hipDeviceSynchronize = _b("hipDeviceSynchronize") _hipStreamSynchronize = _b("hipStreamSynchronize", ctypes.c_void_p) _hipEventCreate = _b("hipEventCreate", ctypes.POINTER(_HipEventHandle)) @@ -521,6 +528,28 @@ def memcpy_d2h(self, dst_buf: ctypes.Array, src: int, nbytes: int) -> None: def memset(self, ptr: int, value: int, nbytes: int) -> None: _check(_hipMemset(ctypes.c_void_p(ptr), value, nbytes), "hipMemset") + def memset_async(self, ptr: int, value: int, nbytes: int, stream: int = 0) -> None: + """Queue a fill on ``stream`` instead of draining the device for it. + + :meth:`memset` is ``hipMemset``, which is synchronous with respect to + the host and the device: it waits for everything already queued. That + is the right primitive for setting a buffer up between measurements, + and the wrong one *inside* one -- a kernel whose contract is "zero + this, then launch" (a split-K epilogue accumulating into its output, + a stream-K flag table) would drain the device once per launch, which + serialises a batch that was queued precisely so it would not. + + Queued on one stream, the fill is ordered against that stream's + launches by FIFO alone, so it is also what makes a batch of such + kernels safe to queue back to back on shared buffers. + """ + _check( + _hipMemsetAsync( + ctypes.c_void_p(ptr), value, nbytes, ctypes.c_void_p(stream) + ), + "hipMemsetAsync", + ) + def _reap_completed(self, stream: int) -> None: """Drop bucket entries whose completion events have fired. From 7aa19997b9b2f61101b1f50dd32778ac00070502 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Mon, 26 Oct 2026 00:26:13 +0000 Subject: [PATCH 3/9] impr --- .../common/benchmark_conv_compare.py | 42 ++ .../common/benchmark_direct_conv.py | 72 ++- .../common/benchmark_implicit_gemm_conv.py | 71 ++- .../library/benchmarks/common/case_results.py | 342 ++++++++++ .../benchmarks/common/run_conv_shards.py | 442 +++++++++++++ .../library/tests/test_run_conv_shards.py | 589 ++++++++++++++++++ 6 files changed, 1555 insertions(+), 3 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py 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 911e74fe4734..c6bef927a4f9 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 @@ -72,6 +72,8 @@ import tempfile from pathlib import Path +from benchmarks.common import case_results + _SCRIPT_DIR = Path(__file__).parent # Regex patterns that match the "Best:" summary printed by both benchmarks. @@ -301,6 +303,25 @@ def _build_miopen_args(args) -> list[str]: _ALL_STOPPED_RE = re.compile(r"all \d+ early-stopped: .*") +def _write_results(path: str, rows_direct: str, rows_implicit: str) -> None: + """Join the two sweeps' per-case rows into ``path`` (and a ``.csv``). + + The two files are what the children wrote under ``--results``; a sweep that + was skipped (``--skip-direct`` / ``--skip-implicit``) or that reached no + case leaves none, and :func:`case_results.read_rows` reads that as no rows, + so the merged table is simply one-sided. + """ + rows = case_results.merge( + case_results.read_rows(rows_direct) + case_results.read_rows(rows_implicit) + ) + if not rows: + print(f"\n[info] no per-case results to write to {path}", file=sys.stderr) + return + csv_path = case_results.write(rows, path) + print(f"\n{case_results.summary(rows)}") + print(f"results: {path} (and {csv_path})", flush=True) + + def _print_summary( direct: tuple[float | None, str], implicit: tuple[float | None, str], @@ -493,6 +514,9 @@ def main() -> int: "the device's clocks up", ) + case_results.add_results_arg(parser, merged=True) + case_results.add_skip_measured_arg(parser) + parser.add_argument( "--module-cache-bytes", type=int, @@ -665,6 +689,21 @@ def main() -> int: if seed: direct_args += ["--early-stop-record", seed_file] + # Each sweep writes its own per-case rows; they are joined below on + # the case key both of them record. The files are the children's, not + # the user's: --results names the merged table. + rows_direct = str(Path(tmp) / "direct_rows.jsonl") + rows_implicit = str(Path(tmp) / "implicit_rows.jsonl") + if args.results: + direct_args += ["--results", rows_direct] + implicit_args += ["--results", rows_implicit] + if args.skip_measured: + # Each sweep skips what *it* already recorded: a run killed + # between the two has the case on one side only, and resuming has + # to finish the other rather than skip it. + direct_args += ["--skip-measured", args.skip_measured] + implicit_args += ["--skip-measured", args.skip_measured] + if not args.skip_direct: tflops, name, _ = _run_script( _SCRIPT_DIR / "benchmark_direct_conv.py", @@ -686,6 +725,9 @@ def main() -> int: if tflops is None and m: implicit_note = m.group(0) + if args.results: + _write_results(args.results, rows_direct, rows_implicit) + _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 ce7ba6f89ecb..29d3d1760551 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 @@ -41,7 +41,14 @@ os.environ.setdefault("ROCKE_CPP_QUIET_FALLBACK", "1") from benchmarks.common import batched_timing -from benchmarks.common.early_stop import EarlyStop, add_early_stop_arg +from benchmarks.common.case_results import ( + RUNNERS_UP, + ResultsWriter, + add_results_arg, + add_skip_measured_arg, + measured_keys, +) +from benchmarks.common.early_stop import EarlyStop, add_early_stop_arg, case_key from benchmarks.common.module_cache import add_module_cache_arg from builders.common.conv_reference import conv_reference as _conv_reference from builders.common.conv_reference import dgrad_reference as _dgrad_reference_shared @@ -2452,6 +2459,12 @@ def _u8(t): # The binaries are shape-generic, so their modules outlive a case. modules = ModuleCache(args.module_cache_bytes) + # One row per case, for the comparison driver and the shard runner to + # merge. Does nothing without --results. + results_out = ResultsWriter(args.results, "direct") + # --skip-measured: cases a previous run of this sweep already recorded. + done = measured_keys(args.skip_measured, "direct") + n_skipped = 0 from kernels.conv.direct_aot_jobs import DIRECT_DTYPES @@ -2466,6 +2479,11 @@ def _u8(t): continue torch_dt = torch.bfloat16 if dtype == "bf16" else torch.float16 for direction in directions or (case_direction,): + if case_key(p, dtype, direction) in done: + # Counted, not printed: resuming a few thousand cases would + # otherwise open with a few thousand lines of nothing. + n_skipped += 1 + continue plans, rejected = _filter_variants( *direct_plans(cache, p, direction, arch), variants ) @@ -2575,10 +2593,33 @@ def _u8(t): 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] + tflops, gbps, ms, passed, label = results[0] print(f"\nBest: {tflops:.1f} TFLOPS — {label} {ms:.3f} ms", flush=True) + results_out.row( + key=case_key(p, dtype, direction), + case=case_idx, + shape=p.short(), + dtype=dtype, + direction=direction, + tflops=tflops, + ms=ms, + gbps=gbps, + kernel=label, + candidates=len(plans), + passed=passed, + # The tail of the ranking just printed, names only. + runners_up=[r[-1] for r in results[1 : 1 + RUNNERS_UP]], + ) + results_out.close() + if n_skipped: + print( + f"\n{n_skipped} case(s) skipped: already in {args.skip_measured}", + flush=True, + ) if len(cases) > 1: print(f"\n{modules.summary()}", flush=True) + if results_out.n_rows: + print(f"{results_out.n_rows} row(s) written to {args.results}", flush=True) return overall_rc @@ -2725,6 +2766,17 @@ def main() -> int: f"measured at idle clocks (default: {batched_timing.DEFAULT_PRIMER})", ) + add_results_arg( + parser.add_argument_group( + "Per-case results (AOT --run-from-cache only)", + "A machine-readable row per case, which benchmark_conv_compare.py " + "merges with the implicit-GEMM sweep's and run_conv_shards.py " + "concatenates across cards.", + ), + merged=False, + ) + add_skip_measured_arg(parser) + add_module_cache_arg( parser.add_argument_group( "Loaded binaries (AOT --run-from-cache only)", @@ -2848,6 +2900,22 @@ def main() -> int: ) args = parser.parse_args() + if args.skip_measured and not args.run_from_cache: + print( + "error: --skip-measured resumes an AOT sweep; pass " + "--run-from-cache DIR", + file=sys.stderr, + ) + return 2 + if args.results and not args.run_from_cache: + # The JIT sweep has no cached kernels to name, and its per-case best is + # a kernel it just built -- a different artifact. Saying so beats + # accepting the flag and writing an empty file. + print( + "error: --results records an AOT sweep; pass --run-from-cache DIR", + file=sys.stderr, + ) + return 2 if args.fast_rank and not args.run_from_cache: # The JIT sweep builds and loads its kernels one at a time, so # batching their launches would save nothing; saying so beats 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 2e811029bbc9..b88906759043 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 @@ -77,9 +77,17 @@ # --------------------------------------------------------------------------- from benchmarks.common import batched_timing # noqa: E402 +from benchmarks.common.case_results import ( # noqa: E402 + RUNNERS_UP, + ResultsWriter, + add_results_arg, + add_skip_measured_arg, + measured_keys, +) from benchmarks.common.early_stop import ( # noqa: E402 EarlyStop, add_early_stop_arg, + case_key, ) from benchmarks.common.module_cache import add_module_cache_arg # noqa: E402 from kernels.conv.aot_jobs import _AOT_ALIAS_PIPELINES # noqa: E402 @@ -1300,6 +1308,17 @@ def main() -> int: f"measured at idle clocks (default: {batched_timing.DEFAULT_PRIMER})", ) + add_results_arg( + parser.add_argument_group( + "Per-case results (AOT --run-from-cache only)", + "A machine-readable row per case, which benchmark_conv_compare.py " + "merges with the direct-conv sweep's and run_conv_shards.py " + "concatenates across cards.", + ), + merged=False, + ) + add_skip_measured_arg(parser) + add_module_cache_arg( parser.add_argument_group( "Loaded binaries (AOT --run-from-cache only)", @@ -1622,6 +1641,23 @@ def main() -> int: ) args = parser.parse_args() + if args.skip_measured and not args.run_from_cache: + print( + "error: --skip-measured resumes an AOT sweep; pass " + "--run-from-cache DIR", + file=sys.stderr, + ) + return 2 + if args.results and not args.run_from_cache: + # The JIT sweep has no cached kernels to name, and its per-case best is + # a kernel it just built -- a different artifact, and --csv already + # covers it. Saying so beats accepting the flag and writing nothing. + print( + "error: --results records an AOT sweep; pass --run-from-cache DIR " + "(the JIT sweep writes --csv instead)", + file=sys.stderr, + ) + return 2 if args.fast_rank and not args.run_from_cache: # The JIT sweep builds and loads its kernels one at a time, so # batching their launches would save nothing; saying so beats @@ -4639,6 +4675,12 @@ def _run_from_cache(args, arch, target, cases, cache, directions) -> int: # A cached binary is shape-generic: its module is loaded on the first case # that offers it and reused by every later one. modules = ModuleCache(args.module_cache_bytes) + # One row per case, for the comparison driver and the shard runner to + # merge. Does nothing without --results. + results_out = ResultsWriter(args.results, "implicit") + # --skip-measured: cases a previous run of this sweep already recorded. + done = measured_keys(args.skip_measured, "implicit") + n_skipped = 0 def _u8(t): return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) @@ -4652,6 +4694,11 @@ def _u8(t): case_direction = case[2] if len(case) > 2 else args.direction for direction in directions or (case_direction,): + if case_key(problem, case_dtype, direction) in done: + # Counted, not printed: resuming a few thousand cases would + # otherwise open with a few thousand lines of nothing. + n_skipped += 1 + continue # The problem the kernels launch on (dgrad runs width-folded, as # the dispatcher launches it); the tensors, the reference and the # FLOP count stay the case's own -- the fold keeps every byte. @@ -5094,11 +5141,33 @@ def _u8(t): flush=True, ) # The line benchmark_conv_compare.py reads, as the JIT sweep prints. - tflops, _, ms, label = results[0] + tflops, gbps, ms, label = results[0] print(f"\nBest: {tflops:.1f} TFLOPS — {label} {ms:.3f} ms", flush=True) + results_out.row( + key=case_key(problem, case_dtype, direction), + case=case_idx, + shape=problem.short(), + dtype=case_dtype, + direction=direction, + tflops=tflops, + ms=ms, + gbps=gbps, + kernel=label, + candidates=n_compatible, + # The tail of the Top-N just printed, names only. + runners_up=[r[-1] for r in results[1 : 1 + RUNNERS_UP]], + ) + results_out.close() + if n_skipped: + print( + f"\n{n_skipped} case(s) skipped: already in {args.skip_measured}", + flush=True, + ) if len(cases) > 1: print(f"\n{modules.summary()}", flush=True) + if results_out.n_rows: + print(f"{results_out.n_rows} row(s) written to {args.results}", flush=True) return overall_rc diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py new file mode 100644 index 000000000000..812fa76151d1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py @@ -0,0 +1,342 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""The per-case result row of an AOT conv sweep, and the merge across engines. + +A ``--run-from-cache`` sweep leaves text: a block per case and a ``Best:`` +line. That is readable one case at a time and useless for a few thousand of +them spread over eight cards, which is what +:mod:`benchmarks.common.run_conv_shards` produces. (The ``--csv`` of the +implicit-GEMM sweep does not cover this: it is wired into the JIT sweep's +per-case ranking, not into ``_run_from_cache``.) + +So each sweep writes one JSON object per case as the case finishes -- +:class:`ResultsWriter`, one call site per script, next to the ``Best:`` print -- +and :func:`merge` joins the two engines' rows into one table. + +The join key is :func:`benchmarks.common.early_stop.case_key`: the normalised +geometry plus dtype and direction, the key ``--early-stop-seed`` already uses to +carry a direct-conv time into the implicit-GEMM sweep. The two sweeps describe a +shape with different types (``DirectConvProblem`` / ``ConvProblem``) whose +``short()`` strings do not even have the same fields, so that key is the only +thing the two sides can be matched on -- and reusing it means a row cannot be +keyed differently from the seed that primed it. + +Rows are appended and flushed per case: a shard that is killed after four hours +keeps the four hours. +""" + +from __future__ import annotations + +import csv +import json +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence + +#: The two sweeps, by the name their rows carry. +ENGINES = ("direct", "implicit") + +#: How many kernels below the winner each engine records, per case. +#: +#: Names only, no times: when the top few kernels are within the noise of the +#: measurement the winner is not stable between runs -- two runs of the same +#: case pick different instances at the same TFLOPS to three digits. What that +#: costs you is the *identity* of the kernel, so that is what the runner-up +#: columns carry; their times would only repeat the winner's. +RUNNERS_UP = 2 + + +def _engine_fields(engine: str) -> tuple: + """The merged table's columns for one engine, in order.""" + return ( + f"{engine}_tflops", + f"{engine}_ms", + f"{engine}_kernel", + *(f"{engine}_kernel_{rank}" for rank in range(2, RUNNERS_UP + 2)), + ) + + +#: Columns of the merged table, in order. +MERGED_FIELDS = ( + "key", + "shape", + "dtype", + "direction", + *_engine_fields("direct"), + *_engine_fields("implicit"), + "winner", + "speedup", +) + + +def add_results_arg(parser, *, merged: bool) -> None: + """Add ``--results FILE`` to a sweep's parser. + + ``merged`` distinguishes the comparison driver (which writes one row per + case with both engines on it) from a single sweep (one row per case, its + own engine). + """ + if merged: + what = ( + "carrying the best direct-conv and implicit-GEMM kernel side by " + f"side (plus the next {RUNNERS_UP} kernels of each by name, which " + "is what a near-tie costs you), and a .csv of the same rows next " + "to it. Written once both sweeps have finished, from the rows " + "each of them recorded." + ) + else: + what = ( + f"carrying the best kernel, its time, the next {RUNNERS_UP} " + "kernels by name, and the key benchmark_conv_compare.py joins the " + "two sweeps on. Rows are appended as each case finishes, so an " + "interrupted run keeps what it measured." + ) + parser.add_argument( + "--results", + default=None, + metavar="FILE", + help=f"AOT (--run-from-cache) only: write FILE (JSON Lines), one row " + f"per case {what}", + ) + + +def add_skip_measured_arg(parser) -> None: + """Add ``--skip-measured FILE`` to a sweep's parser.""" + parser.add_argument( + "--skip-measured", + default=None, + dest="skip_measured", + metavar="FILE", + help="AOT (--run-from-cache) only: skip every case this sweep already " + "recorded in FILE (rows written by --results, this engine's only). " + "Resumes a run that was interrupted without re-measuring what it " + "finished.", + ) + + +def measured_keys(path, engine: str) -> set: + """Case keys ``engine`` already has a row for in ``path``. + + One engine's rows only: a run killed between the direct sweep and the + implicit one has the case measured on one side and not the other, and + resuming has to finish the second rather than skip it. + """ + if not path: + return set() + return {row["key"] for row in read_rows(path) if row.get("engine") == engine} + + +class ResultsWriter: + """Appends one row per case to a JSON Lines file. + + A sweep builds one of these for the whole run (``None`` path: a writer that + does nothing, so the call site needs no branch) and calls :meth:`row` where + it prints the case's ``Best:`` line. A case that produced no result writes + no row -- its absence is the finding, and the log says why. + """ + + def __init__(self, path: Optional[str], engine: str) -> None: + if engine not in ENGINES: + raise ValueError(f"engine must be one of {ENGINES}, got {engine!r}") + self._engine = engine + self._path = Path(path) if path else None + self._fh = None + self._n = 0 + if self._path is not None: + self._path.parent.mkdir(parents=True, exist_ok=True) + self._fh = self._path.open("w", encoding="utf-8") + + @property + def n_rows(self) -> int: + return self._n + + def row( + self, + *, + key: str, + case: int, + shape: str, + dtype: str, + direction: str, + tflops: float, + ms: float, + gbps: float, + kernel: str, + candidates: int, + passed: Optional[bool] = None, + runners_up: Sequence[str] = (), + ) -> None: + """Record one case's best kernel. No-op without a path. + + ``runners_up`` are the names of the kernels ranked below it, best + first; at most :data:`RUNNERS_UP` are kept. The caller passes the tail + of the same ranking it prints, so the row says what the winner beat. + """ + if self._fh is None: + return + record: Dict[str, Any] = { + "key": key, + "engine": self._engine, + "case": int(case), + "shape": shape, + "dtype": dtype, + "direction": direction, + "tflops": round(float(tflops), 3), + "ms": round(float(ms), 6), + "gbps": round(float(gbps), 3), + "kernel": kernel, + "candidates": int(candidates), + } + if passed is not None: + record["passed"] = bool(passed) + if runners_up: + record["runners_up"] = [str(name) for name in runners_up][:RUNNERS_UP] + self._fh.write(json.dumps(record, sort_keys=True) + "\n") + # Per case, not per run: a shard killed mid-sweep keeps its rows. + self._fh.flush() + self._n += 1 + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + self._fh = None + + def __enter__(self) -> "ResultsWriter": + return self + + def __exit__(self, *_exc) -> None: + self.close() + + +def read_rows(path) -> List[dict]: + """Rows of a JSON Lines results file; ``[]`` when there is none. + + A line that does not parse is skipped rather than fatal: the file is + append-as-you-go, so a run killed mid-write can leave a partial last line. + """ + p = Path(path) + if not p.is_file(): + return [] + out: List[dict] = [] + for line in p.read_text(encoding="utf-8", errors="replace").splitlines(): + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(row, dict) and "key" in row: + out.append(row) + return out + + +def merge(rows: Iterable[dict]) -> List[dict]: + """Join per-engine rows into one row per case. + + Rows may come from both engines and from several shards. A key only one + engine reached still gets a row with the other side left ``None`` -- that + is the common and interesting case (a shape the direct kernels do not + cover), not an error. Within one engine and key, the faster row wins, so + merging a case measured twice is idempotent. + + A row's runner-up names travel with it: a slower row is dropped whole + rather than having its ranking blended into the faster one's. A ranking + only means something within the run that measured it -- two runs of one + case can order the same near-tied kernels differently, and a mix of the + two would be a ranking that never happened. + + Sorted by key, so two runs of the same case list produce comparable files. + """ + by_key: Dict[str, dict] = {} + for row in rows: + key = row.get("key") + engine = row.get("engine") + if not key or engine not in ENGINES: + continue + merged = by_key.setdefault( + key, + { + "key": key, + "shape": row.get("shape"), + "dtype": row.get("dtype"), + "direction": row.get("direction"), + }, + ) + ms = row.get("ms") + prev = merged.get(f"{engine}_ms") + if prev is not None and (ms is None or ms >= prev): + continue + merged[f"{engine}_tflops"] = row.get("tflops") + merged[f"{engine}_ms"] = ms + merged[f"{engine}_kernel"] = row.get("kernel") + runners_up = row.get("runners_up") or () + for rank in range(2, RUNNERS_UP + 2): + # Assigned even when absent: a faster row with a shorter ranking + # (or none) has to clear the columns the row it replaced left. + index = rank - 2 + merged[f"{engine}_kernel_{rank}"] = ( + runners_up[index] if index < len(runners_up) else None + ) + + out = [] + for key in sorted(by_key): + merged = by_key[key] + for field in MERGED_FIELDS: + merged.setdefault(field, None) + merged["winner"], merged["speedup"] = _verdict( + merged["direct_ms"], merged["implicit_ms"] + ) + out.append({field: merged[field] for field in MERGED_FIELDS}) + return out + + +def _verdict(direct_ms, implicit_ms): + """``(winner, speedup)``: which engine was faster and by how much. + + ``speedup`` is the loser's time over the winner's, so it reads as "the + winner is N times faster" whichever side won. Both ``None`` when only one + engine reached the case -- a one-sided row is a fact about coverage, not a + comparison. + """ + if direct_ms is None or implicit_ms is None: + return None, None + if direct_ms <= implicit_ms: + return "direct", round(implicit_ms / direct_ms, 3) if direct_ms else None + return "implicit", round(direct_ms / implicit_ms, 3) if implicit_ms else None + + +def write(rows: Sequence[dict], path) -> Path: + """Write merged ``rows`` as JSON Lines, plus a ``.csv`` beside it. + + Returns the CSV's path. Both are written whole (the merge has every row by + then), unlike :class:`ResultsWriter`, which appends as a sweep goes. + """ + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + with p.open("w", encoding="utf-8") as fh: + for row in rows: + fh.write(json.dumps(row, sort_keys=True) + "\n") + csv_path = p.with_suffix(p.suffix + ".csv") if p.suffix else p.with_suffix(".csv") + with csv_path.open("w", encoding="utf-8", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(MERGED_FIELDS)) + writer.writeheader() + writer.writerows(rows) + return csv_path + + +def summary(rows: Sequence[dict]) -> str: + """One line: how many cases, and how the two engines split them.""" + both = sum(1 for r in rows if r["direct_ms"] is not None and r["implicit_ms"]) + direct_only = sum( + 1 for r in rows if r["direct_ms"] is not None and r["implicit_ms"] is None + ) + implicit_only = sum( + 1 for r in rows if r["direct_ms"] is None and r["implicit_ms"] is not None + ) + wins = sum(1 for r in rows if r["winner"] == "direct") + return ( + f"{len(rows)} case(s): {both} measured by both " + f"(direct faster on {wins}), {direct_only} direct-only, " + f"{implicit_only} implicit-only" + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py new file mode 100644 index 000000000000..962250b3682f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py @@ -0,0 +1,442 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Run one MIOpenDriver case file across several GPUs, one process per card. + +A cache sweep is per-case independent work, so a file of a few thousand cases +is embarrassingly parallel across cards -- and in the AOT path there is nothing +to build first, so a shard starts launching as soon as it has read the cache. +This runner splits the file, pins one process per card with +``HIP_VISIBLE_DEVICES``, and merges the shards' per-case rows at the end. + +Usage (from ``library/``):: + + python benchmarks/common/run_conv_shards.py \\ + --miopen-file ../cases.txt --devices 8 --log-dir /scratch/sweep -- \\ + --run-from-cache /scratch/cache --arch gfx942 --fast-rank + +Everything after ``--`` is forwarded verbatim to the benchmark script +(``--script``, default ``compare``); the runner adds ``--miopen-file`` and +``--results`` per shard. + +More shards than cards, on purpose +---------------------------------- +Case cost spans orders of magnitude -- a 1x1 convolution on a 200-pixel row +against a 7x7 on 224x224 -- so one static shard per card leaves cards idle for +however long the unluckiest shard runs over. ``--shards`` defaults to four per +card and a card takes the next shard off the queue when it frees up, which +bounds the tail at roughly one shard rather than one quarter of the run. + +The cost is that each shard re-reads the kernel cache: ~18 s for a 447k-entry +cache, so ~72 s per card over a run. Lower ``--shards`` to one per card to pay +it once, at the price of the tail. + +Lines are dealt round-robin (line *i* to shard *i mod shards*) rather than in +blocks: a case file is usually grouped by model and resolution, so contiguous +blocks would put all the expensive shapes in one shard. + +What a shard leaves behind +-------------------------- +In ``--log-dir``: ``shard-NN.cases`` (its exact input, re-runnable by hand), +``shard-NN.log`` (its whole output), and ``shard-NN.jsonl`` (its per-case +rows). The runner then writes ``results.jsonl`` and ``results.csv`` over all +shards. A shard that dies keeps the rows it had written -- they are flushed +per case -- so a re-run of the remaining cases merges cleanly with them. + +Devices (``--devices``): a count ``N`` (devices ``0..N-1``), a list ``2,3``, +or ``auto`` (default): the entries of ``HIP_VISIBLE_DEVICES``, else the count +of ``ROCR_VISIBLE_DEVICES`` (``0..k-1``), else one device. This duplicates the +rule in ``library/tests/run_device_shards.py`` rather than importing it: +``tests/`` sits above ``benchmarks/`` in the layering +``library/tests/test_library_layering.py`` enforces, so the import would be a +back-edge. The runner itself imports neither torch nor the HIP runtime. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Dict, List, Mapping, Optional, Sequence, Tuple + +from benchmarks.common import case_results + +_SCRIPT_DIR = Path(__file__).parent + +SCRIPTS = { + "compare": "benchmark_conv_compare.py", + "direct": "benchmark_direct_conv.py", + "implicit": "benchmark_implicit_gemm_conv.py", +} + +#: Shards per card when ``--shards`` is not given. See the module docstring: +#: enough to keep the tail short, few enough that re-reading the cache per +#: shard stays a rounding error against hours of sweeping. +SHARDS_PER_DEVICE = 4 + +#: How often the scheduler looks for a finished shard. A shard runs for +#: minutes; a second of latency on starting the next one is not worth a busier +#: loop. +POLL_SECONDS = 1.0 + + +# --------------------------------------------------------------------------- devices +def resolve_devices(spec: str, env: Optional[Mapping[str, str]] = None) -> List[str]: + """``--devices`` to the device ids to pin shards to.""" + env = os.environ if env is None else env + spec = spec.strip() + if spec != "auto": + if "," in spec: + return [d.strip() for d in spec.split(",") if d.strip()] + n = int(spec) + if n < 1: + raise ValueError("--devices needs at least one device") + return [str(i) for i in range(n)] + hip = env.get("HIP_VISIBLE_DEVICES", "").strip() + if hip: + return [d.strip() for d in hip.split(",") if d.strip()] + rocr = env.get("ROCR_VISIBLE_DEVICES", "").strip() + if rocr: + return [str(i) for i in range(len([d for d in rocr.split(",") if d]))] + return ["0"] + + +# --------------------------------------------------------------------------- cases +def read_cases(path) -> List[str]: + """The case lines of a MIOpenDriver file: blanks and ``#`` comments out. + + The benchmark scripts skip those themselves, but dropping them here keeps + a shard's line count equal to its case count, so the split is even in + cases rather than in file lines. + """ + out = [] + for line in Path(path).read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line and not line.startswith("#"): + out.append(line) + return out + + +def split_cases(cases: Sequence[str], shards: int) -> List[List[str]]: + """Deal ``cases`` round-robin into ``shards`` lists (see the module doc). + + Empty shards are dropped: asking for more shards than cases is not an + error, it just gives fewer shards -- a process with no case to run would + read the cache for nothing. + """ + if shards < 1: + raise ValueError("--shards needs at least one shard") + dealt: List[List[str]] = [[] for _ in range(shards)] + for i, case in enumerate(cases): + dealt[i % shards].append(case) + return [shard for shard in dealt if shard] + + +# --------------------------------------------------------------------------- runs +def previous_rows(log_dir: Path) -> List[dict]: + """Every per-case row a previous run left in ``log_dir``. + + The shard files are overwritten by the new run, so this is read first. + ``measured.jsonl`` -- what a previous ``--resume`` itself carried -- is + read too, so resuming a resumed run does not lose the first attempt's + work. + """ + out: List[dict] = [] + for path in sorted(log_dir.glob("shard-*.jsonl")) + [log_dir / "measured.jsonl"]: + out.extend(case_results.read_rows(path)) + return out + + +def shard_paths(log_dir: Path, index: int, width: int) -> Dict[str, Path]: + stem = f"shard-{index:0{width}d}" + return { + "cases": log_dir / f"{stem}.cases", + "log": log_dir / f"{stem}.log", + "rows": log_dir / f"{stem}.jsonl", + } + + +def shard_command( + script: Path, forwarded: Sequence[str], paths: Mapping[str, Path] +) -> List[str]: + """The child's argv: the user's arguments plus this shard's two files.""" + return [ + sys.executable, + "-u", # a piped stdout is block-buffered otherwise; the log lags hours + str(script), + *forwarded, + "--miopen-file", + str(paths["cases"]), + "--results", + str(paths["rows"]), + ] + + +def _tail(path: Path, n: int = 30) -> str: + try: + return "\n".join( + path.read_text(encoding="utf-8", errors="replace").splitlines()[-n:] + ) + except OSError: + return "" + + +def run_shards( + plan: Sequence[Tuple[int, Mapping[str, Path]]], + *, + devices: Sequence[str], + command, + popen=subprocess.Popen, + poll_seconds: float = POLL_SECONDS, + on_done=None, +) -> List[Tuple[int, str, int, float]]: + """Run every shard, at most one per device at a time. + + A device that frees up takes the next shard off the queue, so a long shard + delays only the card running it. ``command(index, paths)`` builds the + child's argv; ``popen`` is the process factory (a stub, under test); + ``on_done(index, device, rc, seconds)`` is called as each shard ends. + + Returns ``(shard index, device, rc, seconds)`` per shard, in the order + they finished. + """ + queue = list(plan) + running: Dict[str, Tuple[int, Mapping[str, Path], object, object, float]] = {} + free = list(devices) + done: List[Tuple[int, str, int, float]] = [] + + while queue or running: + while queue and free: + index, paths = queue.pop(0) + device = free.pop(0) + env = dict(os.environ) + env["HIP_VISIBLE_DEVICES"] = device + fh = open(paths["log"], "w", encoding="utf-8") + proc = popen( + command(index, paths), + env=env, + stdout=fh, + stderr=subprocess.STDOUT, + ) + running[device] = (index, paths, proc, fh, time.monotonic()) + if not running: + break + time.sleep(poll_seconds) + for device, (index, paths, proc, fh, started) in list(running.items()): + rc = proc.poll() + if rc is None: + continue + fh.close() + del running[device] + free.append(device) + entry = (index, device, int(rc), time.monotonic() - started) + done.append(entry) + if on_done is not None: + on_done(*entry) + return done + + +# --------------------------------------------------------------------------- cli +def make_parser() -> argparse.ArgumentParser: + ap = argparse.ArgumentParser( + description=__doc__.splitlines()[0], + epilog="Everything after -- is forwarded to the benchmark script.", + ) + ap.add_argument( + "--miopen-file", + required=True, + metavar="FILE", + dest="miopen_file", + help="MIOpenDriver case file to split across the devices", + ) + ap.add_argument( + "--devices", + default="auto", + help="device count, comma-separated ids, or auto (default: auto)", + ) + ap.add_argument( + "--shards", + type=int, + default=None, + metavar="N", + help=f"case shards to deal out (default: {SHARDS_PER_DEVICE} per " + f"device, so a slow shard holds up one card rather than a quarter of " + f"the run; 1 per device re-reads the cache once instead of " + f"{SHARDS_PER_DEVICE} times)", + ) + ap.add_argument( + "--script", + choices=sorted(SCRIPTS), + default="compare", + help="benchmark to run in each shard (default: compare, both engines)", + ) + ap.add_argument( + "--log-dir", + default=None, + dest="log_dir", + metavar="DIR", + help="where the shards' case files, logs and rows go " + "(default: a new directory under the system temp dir)", + ) + ap.add_argument( + "--results", + default=None, + metavar="FILE", + help="merged per-case table (default: results.jsonl in --log-dir); " + "a .csv of the same rows is written beside it", + ) + ap.add_argument( + "--resume", + action="store_true", + help="re-run into an existing --log-dir, skipping the cases its rows " + "already cover. A sweep that died after four hours keeps the four " + "hours: rows are flushed per case. Each engine skips only what it " + "recorded itself, so a case measured by one sweep and not the other " + "finishes.", + ) + ap.add_argument("forwarded", nargs=argparse.REMAINDER) + return ap + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = make_parser().parse_args(argv) + forwarded = list(args.forwarded) + if forwarded[:1] == ["--"]: + forwarded = forwarded[1:] + for flag in ("--miopen-file", "--results", "--skip-measured"): + if flag in forwarded: + print( + f"error: {flag} is the runner's to set per shard; drop it from " + f"the forwarded arguments", + file=sys.stderr, + ) + return 2 + if "--resume" in forwarded: + # Easy to put on the wrong side of the separator, and the child's + # "unrecognized arguments" would send the reader looking in the wrong + # script. + print( + "error: --resume is the runner's own flag; it goes before --", + file=sys.stderr, + ) + return 2 + if not any(a.startswith("--run-from-cache") for a in forwarded): + # Sharding is for the AOT path: a JIT sweep would build the same + # kernels again on every card, and the per-case rows this runner + # merges are written only by a cache run. Catching it here beats + # starting a process per card to watch them all exit 2. + print( + "error: sharding runs an AOT sweep; forward --run-from-cache DIR " + "after --", + file=sys.stderr, + ) + return 2 + + devices = resolve_devices(args.devices) + shards = ( + args.shards + if args.shards is not None + else SHARDS_PER_DEVICE * len(devices) + ) + cases = read_cases(args.miopen_file) + if not cases: + print(f"error: {args.miopen_file}: no cases", file=sys.stderr) + return 2 + dealt = split_cases(cases, shards) + + import tempfile + + if args.resume and not args.log_dir: + # Resuming into a fresh temporary directory would find no rows and + # silently re-run everything. + print( + "error: --resume needs the --log-dir of the run it continues", + file=sys.stderr, + ) + return 2 + log_dir = Path(args.log_dir or tempfile.mkdtemp(prefix="rocke-conv-shards-")) + log_dir.mkdir(parents=True, exist_ok=True) + + # What the previous attempt measured, gathered before its shard files are + # overwritten. The sweeps skip by case key, which only they can compute + # (it comes off a parsed problem), so the runner hands them the rows and + # does no matching of its own. + carried: List[dict] = [] + if args.resume: + carried = previous_rows(log_dir) + if not carried: + print(f"[info] --resume: no rows in {log_dir}; running everything") + else: + measured = log_dir / "measured.jsonl" + measured.write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in carried), + encoding="utf-8", + ) + forwarded = list(forwarded) + ["--skip-measured", str(measured)] + print(f"[resume] {len(carried)} row(s) carried from {log_dir}") + width = len(str(len(dealt) - 1)) or 1 + plan = [] + for index, shard in enumerate(dealt): + paths = shard_paths(log_dir, index, width) + paths["cases"].write_text("\n".join(shard) + "\n", encoding="utf-8") + plan.append((index, paths)) + + script = _SCRIPT_DIR / SCRIPTS[args.script] + print( + f"{len(cases)} case(s) over {len(dealt)} shard(s) on " + f"{len(devices)} device(s): {args.script}", + flush=True, + ) + print(f"logs: {log_dir}", flush=True) + + def _report(index: int, device: str, rc: int, seconds: float) -> None: + paths = shard_paths(log_dir, index, width) + rows = len(case_results.read_rows(paths["rows"])) + print( + f"shard {index} device={device} rc={rc} " + f"{rows} row(s) {seconds / 60:.1f} min", + flush=True, + ) + + started = time.monotonic() + done = run_shards( + plan, + devices=devices, + command=lambda index, paths: shard_command(script, forwarded, paths), + poll_seconds=POLL_SECONDS, + on_done=_report, + ) + + rc_all = 0 + for index, _device, rc, _seconds in done: + if rc: + rc_all = rc_all or rc + print(f"\nshard {index} failed (rc={rc}); last lines of its log:") + print(_tail(shard_paths(log_dir, index, width)["log"])) + + rows = case_results.merge( + carried + + [ + row + for _index, paths in plan + for row in case_results.read_rows(paths["rows"]) + ] + ) + if rows: + out = Path(args.results or (log_dir / "results.jsonl")) + csv_path = case_results.write(rows, out) + print(f"\n{case_results.summary(rows)}") + print(f"results: {out} (and {csv_path})", flush=True) + else: + print("\n[info] no per-case rows were written by any shard") + rc_all = rc_all or 1 + print(f"wall: {(time.monotonic() - started) / 60:.1f} min", flush=True) + return rc_all + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py new file mode 100644 index 000000000000..af20456ae2e2 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py @@ -0,0 +1,589 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Host decisions of the per-GPU conv sweep runner, with the devices stubbed. + +What needs GPUs is the sweeping; what needs testing is how the work is split, +how a freed card picks up the next shard, and whether the two engines' rows +actually join -- all host-side, so ``Popen`` is replaced by a fake and no +benchmark is run. + +The join is the load-bearing assumption of the whole table: the two sweeps +describe a shape with different types, and only +:func:`benchmarks.common.early_stop.case_key` makes them meet. If that ever +stopped agreeing, the merge would silently report every case twice, once per +engine, instead of once with a winner. +""" + +from __future__ import annotations + +import json +import pathlib +import sys + +import pytest + +from benchmarks.common import case_results, run_conv_shards +from benchmarks.common.early_stop import case_key +from kernels.common._conv_implicit_gemm_common import ConvProblem +from kernels.common.conv_direct_grouped import DirectConvProblem + + +# --------------------------------------------------------------------------- split +def test_every_case_lands_in_exactly_one_shard(): + cases = [f"convbfp16 -n {i}" for i in range(23)] + dealt = run_conv_shards.split_cases(cases, 4) + assert len(dealt) == 4 + flat = [case for shard in dealt for case in shard] + assert sorted(flat) == sorted(cases) + assert len(flat) == len(set(flat)) + + +def test_shards_are_balanced_in_count(): + dealt = run_conv_shards.split_cases([str(i) for i in range(100)], 8) + sizes = {len(shard) for shard in dealt} + assert max(sizes) - min(sizes) <= 1 + + +def test_round_robin_spreads_a_grouped_file(): + """A case file is grouped by model and resolution, so a contiguous split + would hand one shard every expensive shape. Dealing round-robin is what + keeps the cost mixed, and that is the reason this split was chosen.""" + cheap, pricey = ["cheap"] * 50, ["pricey"] * 50 + dealt = run_conv_shards.split_cases(cheap + pricey, 4) + for shard in dealt: + assert shard.count("pricey") == pytest.approx(len(shard) / 2, abs=1) + + +def test_more_shards_than_cases_drops_the_empty_ones(): + """A shard with no case would read the whole kernel cache for nothing.""" + dealt = run_conv_shards.split_cases(["a", "b"], 8) + assert dealt == [["a"], ["b"]] + + +def test_comments_and_blanks_are_not_cases(tmp_path): + f = tmp_path / "cases.txt" + f.write_text("# header\n\nconvbfp16 -n 1\n\n# mid\nconvbfp16 -n 2\n") + assert run_conv_shards.read_cases(f) == ["convbfp16 -n 1", "convbfp16 -n 2"] + + +# --------------------------------------------------------------------------- devices +@pytest.mark.parametrize( + "spec,env,expected", + [ + ("4", {}, ["0", "1", "2", "3"]), + ("2,5", {}, ["2", "5"]), + ("auto", {"HIP_VISIBLE_DEVICES": "3,4"}, ["3", "4"]), + ("auto", {"ROCR_VISIBLE_DEVICES": "0,1,2"}, ["0", "1", "2"]), + ("auto", {}, ["0"]), + ], +) +def test_resolve_devices(spec, env, expected): + assert run_conv_shards.resolve_devices(spec, env) == expected + + +def test_resolve_devices_rejects_nothing_to_run_on(): + with pytest.raises(ValueError): + run_conv_shards.resolve_devices("0") + + +# --------------------------------------------------------------------------- queue +class _FakeProc: + """A process that finishes after ``ticks`` polls with ``rc``.""" + + def __init__(self, ticks: int, rc: int = 0): + self.ticks = ticks + self.rc = rc + + def poll(self): + if self.ticks > 0: + self.ticks -= 1 + return None + return self.rc + + +def _plan(tmp_path, n): + return [(i, run_conv_shards.shard_paths(tmp_path, i, 2)) for i in range(n)] + + +def test_a_freed_card_takes_the_next_shard(tmp_path): + """The point of dealing more shards than cards: a long shard holds up the + card running it, not a fixed quarter of the work.""" + starts = [] + durations = {0: 6, 1: 1, 2: 1, 3: 1, 4: 1} + + def popen(cmd, **_kw): + index = int(cmd[-1]) + starts.append(index) + return _FakeProc(durations[index]) + + done = run_shards_with(tmp_path, 5, popen, starts) + assert starts == [0, 1, 2, 3, 4], "shards are dealt in order" + card = {index: device for index, device, _rc, _s in done} + # Shard 0 holds its card for the whole run; every other shard goes to the + # other one, one after the next, instead of waiting for a fixed slot. + assert len(card) == 5 + assert {card[i] for i in (1, 2, 3, 4)} == {"0", "1"} - {card[0]} + finished = [index for index, _d, _rc, _s in done] + assert finished.index(0) >= 3, "the slow shard outlasts the quick ones" + + +def run_shards_with(tmp_path, n, popen, _starts): + return run_conv_shards.run_shards( + _plan(tmp_path, n), + devices=["0", "1"], + command=lambda index, paths: ["python", str(index)], + popen=popen, + poll_seconds=0, + ) + + +def test_every_shard_runs_on_one_card_at_a_time(tmp_path): + live = {"now": 0, "max": 0} + + def popen(cmd, **_kw): + live["now"] += 1 + live["max"] = max(live["max"], live["now"]) + + class P: + def poll(_self): + live["now"] -= 1 + return 0 + + return P() + + run_shards_with(tmp_path, 8, popen, []) + assert live["max"] <= 2 + + +def test_a_failing_shard_does_not_stop_the_others(tmp_path): + def popen(cmd, **_kw): + return _FakeProc(0, rc=1 if cmd[-1] == "2" else 0) + + done = run_shards_with(tmp_path, 4, popen, []) + assert len(done) == 4 + assert {index: rc for index, _d, rc, _s in done} == {0: 0, 1: 0, 2: 1, 3: 0} + + +def test_the_shard_command_carries_the_user_arguments(tmp_path): + paths = run_conv_shards.shard_paths(tmp_path, 3, 2) + cmd = run_conv_shards.shard_command( + tmp_path / "bench.py", ["--run-from-cache", "/c", "--fast-rank"], paths + ) + assert "--fast-rank" in cmd + assert cmd[cmd.index("--miopen-file") + 1] == str(paths["cases"]) + assert cmd[cmd.index("--results") + 1] == str(paths["rows"]) + + +@pytest.mark.parametrize("flag", ["--miopen-file", "--results"]) +def test_the_runners_own_flags_are_refused_in_the_forwarded_args( + tmp_path, flag, capsys +): + f = tmp_path / "cases.txt" + f.write_text("convbfp16 -n 1\n") + rc = run_conv_shards.main( + ["--miopen-file", str(f), "--devices", "1", "--", flag, "x"] + ) + assert rc == 2 + assert flag in capsys.readouterr().err + + +def test_a_jit_run_is_refused_before_any_process_starts(tmp_path, capsys): + """Without a cache every card would build the same kernels again, and no + rows would be written -- the runner says so instead of starting N + processes to watch them all exit 2.""" + f = tmp_path / "cases.txt" + f.write_text("convbfp16 -n 1\n") + rc = run_conv_shards.main(["--miopen-file", str(f), "--devices", "2"]) + assert rc == 2 + assert "--run-from-cache" in capsys.readouterr().err + assert not list(tmp_path.glob("shard-*")) + + +# --------------------------------------------------------------------------- merge +def _row(engine, key, ms, kernel, tflops=100.0, runners_up=None): + row = { + "key": key, + "engine": engine, + "case": 1, + "shape": "shape", + "dtype": "bf16", + "direction": "fwd", + "tflops": tflops, + "ms": ms, + "gbps": 1.0, + "kernel": kernel, + "candidates": 7, + } + if runners_up: + row["runners_up"] = list(runners_up) + return row + + +def test_the_two_engines_join_on_one_row(): + rows = case_results.merge( + [_row("direct", "k", 0.5, "d-kernel"), _row("implicit", "k", 1.0, "i-kernel")] + ) + assert len(rows) == 1 + assert rows[0]["direct_kernel"] == "d-kernel" + assert rows[0]["implicit_kernel"] == "i-kernel" + assert rows[0]["winner"] == "direct" + assert rows[0]["speedup"] == 2.0 + + +def test_a_case_only_one_engine_reached_still_gets_a_row(): + """A shape the direct kernels do not cover is the interesting case, not an + error -- the table has to say so rather than drop it.""" + rows = case_results.merge([_row("implicit", "k", 1.0, "i-kernel")]) + assert len(rows) == 1 + assert rows[0]["direct_ms"] is None + assert rows[0]["winner"] is None and rows[0]["speedup"] is None + + +def test_the_same_case_measured_twice_keeps_the_faster(): + """Shards do not overlap, but a re-run after a killed shard merges with + the rows it left, so the merge has to be idempotent.""" + rows = case_results.merge( + [_row("direct", "k", 0.9, "slow"), _row("direct", "k", 0.4, "fast")] + ) + assert rows[0]["direct_kernel"] == "fast" + assert case_results.merge( + [_row("direct", "k", 0.4, "fast"), _row("direct", "k", 0.9, "slow")] + ) == rows + + +def test_the_runner_up_kernels_get_a_column_each(): + """Names only -- what a near-tie costs you is which instance won, not a + second TFLOPS column that repeats the first.""" + rows = case_results.merge( + [_row("direct", "k", 0.5, "best", runners_up=["second", "third"])] + ) + assert rows[0]["direct_kernel"] == "best" + assert rows[0]["direct_kernel_2"] == "second" + assert rows[0]["direct_kernel_3"] == "third" + + +def test_a_shorter_ranking_leaves_the_rest_of_the_columns_empty(): + rows = case_results.merge([_row("implicit", "k", 1.0, "best", runners_up=["x"])]) + assert rows[0]["implicit_kernel_2"] == "x" + assert rows[0]["implicit_kernel_3"] is None + + +def test_the_faster_row_brings_its_whole_ranking(): + """A ranking means something only within the run that measured it: the + slower row goes whole, rather than leaving its runner-ups behind to be read + as the faster run's.""" + rows = case_results.merge( + [ + _row("direct", "k", 0.9, "slow", runners_up=["slow-2", "slow-3"]), + _row("direct", "k", 0.4, "fast", runners_up=["fast-2"]), + ] + ) + assert rows[0]["direct_kernel"] == "fast" + assert rows[0]["direct_kernel_2"] == "fast-2" + assert rows[0]["direct_kernel_3"] is None + + +def test_merge_is_ordered_by_key(): + rows = case_results.merge( + [_row("direct", k, 1.0, "x") for k in ("c", "a", "b")] + ) + assert [r["key"] for r in rows] == ["a", "b", "c"] + + +def test_a_partial_last_line_is_skipped_not_fatal(tmp_path): + """Rows are flushed per case, so a killed shard can leave half a line.""" + f = tmp_path / "rows.jsonl" + f.write_text(json.dumps(_row("direct", "k", 1.0, "x")) + '\n{"key": "trunc') + assert [r["key"] for r in case_results.read_rows(f)] == ["k"] + + +def test_write_leaves_a_csv_beside_the_jsonl(tmp_path): + rows = case_results.merge( + [_row("direct", "k", 0.5, "d"), _row("implicit", "k", 1.0, "i")] + ) + out = tmp_path / "results.jsonl" + csv_path = case_results.write(rows, out) + assert csv_path.is_file() and csv_path != out + header, body = csv_path.read_text().splitlines()[:2] + assert header.split(",") == list(case_results.MERGED_FIELDS) + assert "d" in body and "i" in body + + +def test_a_writer_without_a_path_writes_nothing(tmp_path): + writer = case_results.ResultsWriter(None, "direct") + writer.row( + key="k", + case=1, + shape="s", + dtype="bf16", + direction="fwd", + tflops=1.0, + ms=1.0, + gbps=1.0, + kernel="x", + candidates=1, + ) + writer.close() + assert writer.n_rows == 0 + assert not list(tmp_path.iterdir()) + + +# --------------------------------------------------------------------------- end to end +_STUB = '''\ +import argparse, json, sys +ap = argparse.ArgumentParser() +ap.add_argument("--miopen-file") +ap.add_argument("--results") +ap.add_argument("--run-from-cache") +args = ap.parse_args() +with open(args.results, "w") as out: + for line in open(args.miopen_file): + line = line.strip() + if not line: + continue + n = int(line.split()[-1]) + for engine, ms in (("direct", 1.0 + n), ("implicit", 2.0 + n)): + out.write(json.dumps({ + "key": f"fwd/bf16/case{n}", "engine": engine, "case": n, + "shape": f"s{n}", "dtype": "bf16", "direction": "fwd", + "tflops": 100.0, "ms": ms, "gbps": 1.0, + "kernel": f"{engine}-kernel-{n}", "candidates": 3, + }) + "\\n") +print("stub done") +''' + + +def test_a_whole_run_splits_schedules_and_merges(tmp_path, monkeypatch, capsys): + """main() end to end against a stub benchmark: the only thing missing is + the GPU, and the runner never touches one.""" + stub = tmp_path / "stub_bench.py" + stub.write_text(_STUB) + monkeypatch.setattr(run_conv_shards, "_SCRIPT_DIR", tmp_path) + monkeypatch.setattr(run_conv_shards, "SCRIPTS", {"compare": stub.name}) + monkeypatch.setattr(run_conv_shards, "POLL_SECONDS", 0) + + cases = tmp_path / "cases.txt" + cases.write_text("".join(f"convbfp16 -n {i}\n" for i in range(10))) + logs = tmp_path / "logs" + + rc = run_conv_shards.main( + [ + "--miopen-file", str(cases), + "--devices", "2", + "--shards", "4", + "--log-dir", str(logs), + "--", "--run-from-cache", "/nonexistent", + ] + ) + assert rc == 0, capsys.readouterr() + + # Each shard's input is kept, and together they are the original file. + shard_cases = sorted(logs.glob("shard-*.cases")) + assert len(shard_cases) == 4 + dealt = [ln for f in shard_cases for ln in f.read_text().split() if ln.isdigit()] + assert sorted(int(n) for n in dealt) == list(range(10)) + + rows = case_results.read_rows(logs / "results.jsonl") + assert len(rows) == 10 + assert (logs / "results.jsonl.csv").is_file() + # Direct was given the lower time for every case, so it wins all of them. + assert all(r["winner"] == "direct" for r in rows) + assert {r["key"] for r in rows} == {f"fwd/bf16/case{i}" for i in range(10)} + assert "10 case(s)" in capsys.readouterr().out + + +def test_a_run_whose_shards_measured_nothing_is_a_failure(tmp_path, monkeypatch): + """Eight cards and an empty table means the run did not do what was asked; + exiting 0 would hide that in a nightly job.""" + stub = tmp_path / "stub_bench.py" + stub.write_text("import sys; open(sys.argv[sys.argv.index('--results')+1],'w')") + monkeypatch.setattr(run_conv_shards, "_SCRIPT_DIR", tmp_path) + monkeypatch.setattr(run_conv_shards, "SCRIPTS", {"compare": stub.name}) + monkeypatch.setattr(run_conv_shards, "POLL_SECONDS", 0) + cases = tmp_path / "cases.txt" + cases.write_text("convbfp16 -n 1\n") + rc = run_conv_shards.main( + ["--miopen-file", str(cases), "--devices", "1", + "--log-dir", str(tmp_path / "logs"), "--", "--run-from-cache", "/c"] + ) + assert rc != 0 + + +# --------------------------------------------------------------------------- resume +def test_only_this_engines_rows_are_skipped(): + """A run killed between the direct sweep and the implicit one has the case + on one side only; resuming must finish the other, not skip it.""" + rows = [_row("direct", "k", 1.0, "d"), _row("implicit", "other", 1.0, "i")] + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + f = pathlib.Path(tmp) / "rows.jsonl" + f.write_text("".join(json.dumps(r) + "\n" for r in rows)) + assert case_results.measured_keys(f, "direct") == {"k"} + assert case_results.measured_keys(f, "implicit") == {"other"} + assert case_results.measured_keys(None, "direct") == set() + + +def test_resume_after_the_separator_is_caught_here(tmp_path, capsys): + """argparse.REMAINDER swallows anything after --, so a misplaced runner + flag would reach the child and be reported by the wrong script.""" + f = tmp_path / "cases.txt" + f.write_text("convbfp16 -n 1\n") + rc = run_conv_shards.main( + ["--miopen-file", str(f), "--devices", "1", "--", + "--run-from-cache", "/c", "--resume"] + ) + assert rc == 2 + assert "goes before --" in capsys.readouterr().err + + +def test_resume_carries_the_earlier_rows_and_skips_their_cases( + tmp_path, monkeypatch, capsys +): + """Second pass: the stub records which cases it was actually given, and + the merged table still holds the first pass's rows.""" + stub = tmp_path / "stub_bench.py" + stub.write_text(_RESUME_STUB) + monkeypatch.setattr(run_conv_shards, "_SCRIPT_DIR", tmp_path) + monkeypatch.setattr(run_conv_shards, "SCRIPTS", {"compare": stub.name}) + monkeypatch.setattr(run_conv_shards, "POLL_SECONDS", 0) + cases = tmp_path / "cases.txt" + cases.write_text("".join(f"convbfp16 -n {i}\n" for i in range(6))) + logs = tmp_path / "logs" + runner = [ + "--miopen-file", str(cases), "--devices", "1", "--shards", "2", + "--log-dir", str(logs), + ] + forwarded = ["--", "--run-from-cache", "/c"] + argv = runner + forwarded + + # First pass: the stub records cases 0-2 and then "dies" on the rest. + monkeypatch.setenv("STUB_LIMIT", "2") + run_conv_shards.main(argv) + first = {r["key"] for r in case_results.read_rows(logs / "results.jsonl")} + assert first and len(first) < 6 + + # Second pass with --resume: the earlier cases are handed to the sweep as + # already measured, and only the rest are run. + monkeypatch.delenv("STUB_LIMIT") + (logs / "ran.json").unlink() # the stub appends; this pass's cases only + capsys.readouterr() + rc = run_conv_shards.main(runner + ["--resume"] + forwarded) + out = capsys.readouterr().out + assert rc == 0, out + assert "row(s) carried" in out + + rows = case_results.read_rows(logs / "results.jsonl") + assert {r["key"] for r in rows} == {f"fwd/bf16/case{i}" for i in range(6)} + # Nothing measured twice, and the first pass's rows survived. + assert len(rows) == 6 + ran = json.loads((logs / "ran.json").read_text()) + assert not (set(ran) & first), f"re-measured {set(ran) & first}" + + +_RESUME_STUB = '''\ +import argparse, json, os, pathlib, sys +ap = argparse.ArgumentParser() +ap.add_argument("--miopen-file") +ap.add_argument("--results") +ap.add_argument("--skip-measured") +ap.add_argument("--run-from-cache") +args = ap.parse_args() +done = set() +if args.skip_measured: + for line in open(args.skip_measured): + done.add(json.loads(line)["key"]) +limit = int(os.environ.get("STUB_LIMIT", "999")) +ran = [] +with open(args.results, "w") as out: + for n_done, line in enumerate(open(args.miopen_file)): + line = line.strip() + if not line: + continue + n = int(line.split()[-1]) + key = f"fwd/bf16/case{n}" + if key in done: + continue + if n_done >= limit: + break + ran.append(key) + for engine in ("direct", "implicit"): + out.write(json.dumps({ + "key": key, "engine": engine, "case": n, "shape": f"s{n}", + "dtype": "bf16", "direction": "fwd", "tflops": 1.0, + "ms": 1.0 if engine == "direct" else 2.0, "gbps": 1.0, + "kernel": f"{engine}-{n}", "candidates": 1, + }) + "\\n") +log = pathlib.Path(args.results).parent / "ran.json" +prev = json.loads(log.read_text()) if log.is_file() else [] +log.write_text(json.dumps(prev + ran)) +''' + + +# --------------------------------------------------------------------------- compare +def test_compare_merges_one_engine_when_the_other_was_skipped(tmp_path, capsys): + """--skip-direct leaves no direct rows; the table is then one-sided rather + than missing, and the driver must not trip over the absent file.""" + from benchmarks.common import benchmark_conv_compare as compare + + implicit = tmp_path / "implicit.jsonl" + implicit.write_text(json.dumps(_row("implicit", "k", 1.0, "i-kernel")) + "\n") + out = tmp_path / "merged.jsonl" + + compare._write_results(str(out), str(tmp_path / "absent.jsonl"), str(implicit)) + + rows = case_results.read_rows(out) + assert len(rows) == 1 and rows[0]["implicit_kernel"] == "i-kernel" + assert "results:" in capsys.readouterr().out + + +def test_compare_says_so_when_there_is_nothing_to_write(tmp_path, capsys): + from benchmarks.common import benchmark_conv_compare as compare + + out = tmp_path / "merged.jsonl" + compare._write_results(str(out), str(tmp_path / "a"), str(tmp_path / "b")) + assert not out.exists() + assert "no per-case results" in capsys.readouterr().err + + +def test_the_comparison_driver_needs_no_torch(): + """conv_compare is a launcher: it has to run where torch is not installed + (the benchmark scripts it spawns import it, it does not).""" + import subprocess + + probe = ( + "import sys; sys.argv=['x'];" + "import benchmarks.common.benchmark_conv_compare;" + "print('torch' in sys.modules)" + ) + out = subprocess.run( + [sys.executable, "-c", probe], capture_output=True, text=True, check=True + ) + assert out.stdout.strip() == "False" + + +# --------------------------------------------------------------------------- the key +def test_the_two_problem_types_meet_on_one_key(): + """The assumption the merged table rests on: the same convolution, read by + the direct sweep and by the implicit-GEMM sweep, keys the same. + + Their ``short()`` strings do not even carry the same fields, so this is the + only thing that can join them -- and it is the key ``--early-stop-seed`` + already uses to carry a direct time into the implicit sweep. + """ + direct = DirectConvProblem( + N=16, H=28, W=28, groups=4, cpg=16, kpg=16, KH=3, KW=3, PAD=1, stride=1 + ) + implicit = ConvProblem( + N=16, Hi=28, Wi=28, C=64, K=64, Y=3, X=3, pH=1, pW=1, groups=4 + ) + assert direct.short() != implicit.short() + assert case_key(direct, "bf16", "fwd") == case_key(implicit, "bf16", "fwd") + + +def test_a_different_shape_keys_differently(): + a = DirectConvProblem(N=16, H=28, W=28, groups=4, cpg=16, kpg=16) + b = DirectConvProblem(N=16, H=28, W=28, groups=4, cpg=16, kpg=16, stride=2) + assert case_key(a, "bf16", "fwd") != case_key(b, "bf16", "fwd") + assert case_key(a, "bf16", "fwd") != case_key(a, "fp16", "fwd") + assert case_key(a, "bf16", "fwd") != case_key(a, "bf16", "dgrad") From f3f8279aea62a588046d223f6feb1e186effa692 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Mon, 26 Oct 2026 12:48:16 +0000 Subject: [PATCH 4/9] impr --- .../common/benchmark_conv_compare.py | 23 +- .../library/benchmarks/common/case_results.py | 171 ++++++++++++- .../benchmarks/common/run_conv_shards.py | 117 ++++++++- .../library/tests/test_run_conv_shards.py | 239 +++++++++++++++++- 4 files changed, 539 insertions(+), 11 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 c6bef927a4f9..88705eba2dec 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 @@ -303,16 +303,27 @@ def _build_miopen_args(args) -> list[str]: _ALL_STOPPED_RE = re.compile(r"all \d+ early-stopped: .*") -def _write_results(path: str, rows_direct: str, rows_implicit: str) -> None: +def _write_results( + path: str, + rows_direct: str, + rows_implicit: str, + *, + miopen_file: str | None = None, + miopen_cmd: str | None = None, +) -> None: """Join the two sweeps' per-case rows into ``path`` (and a ``.csv``). The two files are what the children wrote under ``--results``; a sweep that was skipped (``--skip-direct`` / ``--skip-implicit``) or that reached no case leaves none, and :func:`case_results.read_rows` reads that as no rows, so the merged table is simply one-sided. + + The children key their rows by geometry, not by the line they read, so the + MIOpenDriver command goes back on here, from the input this run was given. """ rows = case_results.merge( - case_results.read_rows(rows_direct) + case_results.read_rows(rows_implicit) + case_results.read_rows(rows_direct) + case_results.read_rows(rows_implicit), + case_results.miopen_cmds_by_key(miopen_file=miopen_file, miopen_cmd=miopen_cmd), ) if not rows: print(f"\n[info] no per-case results to write to {path}", file=sys.stderr) @@ -726,7 +737,13 @@ def main() -> int: implicit_note = m.group(0) if args.results: - _write_results(args.results, rows_direct, rows_implicit) + _write_results( + args.results, + rows_direct, + rows_implicit, + miopen_file=args.miopen_file, + miopen_cmd=args.miopen_cmd, + ) _print_summary(direct_result, implicit_result, implicit_note) return 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py index 812fa76151d1..7b76709917cd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py @@ -22,6 +22,13 @@ thing the two sides can be matched on -- and reusing it means a row cannot be keyed differently from the seed that primed it. +The merge writes two files: the JSON keeps every field, and the ``.csv`` beside +it is the view you read -- the MIOpenDriver command of the case, the instance +each engine picked, what they cost, and the winner. The command is not in the +rows (they are keyed by geometry, which is what joins the two engines), so the +merge puts it back from the run's own input, or writes one from the key when +the run was given a shape rather than a command. + Rows are appended and flushed per case: a shard that is killed after four hours keeps the four hours. """ @@ -56,18 +63,43 @@ def _engine_fields(engine: str) -> tuple: ) -#: Columns of the merged table, in order. +#: Fields of a merged row, in order. The JSON file carries all of them. MERGED_FIELDS = ( "key", + "miopen_cmd", "shape", "dtype", "direction", *_engine_fields("direct"), *_engine_fields("implicit"), "winner", + "winner_ms", + "winner_kernel", "speedup", ) +#: Columns of the ``.csv`` beside it, in order. +#: +#: A narrower view of the same rows, in the order a comparison is read: what +#: was run, which instance each engine picked, what they cost, who won. The +#: JSON keeps the rest (tflops, the key, the shape and its dtype/direction) +#: for anything that reads these files by field rather than by eye. +#: +#: The implicit side lists its whole ranking because that is where near-ties +#: live -- the three names are what re-picking the winner would cost you. The +#: direct side lists only its winner: the two tables are joined per case, and +#: a row wide enough to hold both rankings is one nobody reads. +CSV_FIELDS = ( + "miopen_cmd", + "direct_kernel", + "implicit_kernel", + *(f"implicit_kernel_{rank}" for rank in range(2, RUNNERS_UP + 2)), + "direct_ms", + "implicit_ms", + "winner_ms", + "winner_kernel", +) + def add_results_arg(parser, *, merged: bool) -> None: """Add ``--results FILE`` to a sweep's parser. @@ -81,7 +113,9 @@ def add_results_arg(parser, *, merged: bool) -> None: "carrying the best direct-conv and implicit-GEMM kernel side by " f"side (plus the next {RUNNERS_UP} kernels of each by name, which " "is what a near-tie costs you), and a .csv of the same rows next " - "to it. Written once both sweeps have finished, from the rows " + "to it -- the .csv narrowed to the reading columns: the " + "MIOpenDriver command, each engine's pick, their times, and the " + "winner. Written once both sweeps have finished, from the rows " "each of them recorded." ) else: @@ -231,7 +265,121 @@ def read_rows(path) -> List[dict]: return out -def merge(rows: Iterable[dict]) -> List[dict]: +#: MIOpenDriver keyword per dtype, and ``-F`` value per direction. +_MIOPEN_DRIVER = {"fp32": "conv", "fp16": "convfp16", "bf16": "convbfp16"} +_MIOPEN_FORW = {"fwd": 1, "dgrad": 2, "wgrad": 4} + +#: ``(MIOpenDriver flag, geometry field)`` of a 2-D command, in driver order. +_MIOPEN_2D_FLAGS = ( + ("-n", "N"), + ("-c", "C"), + ("-H", "Hi"), + ("-W", "Wi"), + ("-k", "K"), + ("-y", "Y"), + ("-x", "X"), + ("-p", "pH"), + ("-q", "pW"), + ("-u", "sH"), + ("-v", "sW"), + ("-l", "dH"), + ("-j", "dW"), + ("-g", "groups"), +) + +#: The depth flags a 3-D command adds, long-form. +_MIOPEN_3D_FLAGS = ( + ("--in_d", "Di"), + ("--fil_d", "Z"), + ("--conv_stride_d", "sD"), + ("--pad_d", "pD"), + ("--dilation_d", "dD"), +) + + +def miopen_cmd_from_key(key: str) -> str: + """A MIOpenDriver command for ``key``, synthesised from its geometry. + + A case measured from explicit shape flags (or a JSON case file) has no + command behind it, and a row identified only by its kernel names is one + you cannot re-run. The key carries every field + :func:`benchmarks.common.early_stop.case_key` put in it, so the command is + a rewrite of the key rather than new information: 2-D commands round-trip + back through ``parse_miopen_cmd``. + + 3-D keys get the depth extents as long-form flags -- the sweeps' MIOpen + parsers are 2-D only, so nothing here reads them back; they are there so + the row still says what was run. Returns ``key`` unchanged when it does not + parse (a file written by a future version, say). + """ + try: + direction, dtype, fields = key.split("/", 2) + geom = dict(field.split("=", 1) for field in fields.split(",")) + driver = _MIOPEN_DRIVER[dtype] + forw = _MIOPEN_FORW[direction] + flags = [f"{flag} {geom[name]}" for flag, name in _MIOPEN_2D_FLAGS] + if geom.get("is_3d") == "True": + flags += ["--spatial_dim 3"] + flags += [f"{flag} {geom[name]}" for flag, name in _MIOPEN_3D_FLAGS] + except (KeyError, ValueError): + return key + return f"MIOpenDriver {driver} {' '.join(flags)} -F {forw} -t 1 -in_layout=NHWC" + + +def miopen_cmds_by_key( + *, miopen_file=None, miopen_cmd: Optional[str] = None +) -> Dict[str, str]: + """Map each case key back to the MIOpenDriver line it was read from. + + The sweeps key their rows by geometry, not by input line, so the line is + not in them; re-parsing the input the run was given is what puts it back. + One line can carry several cases (``-F 0`` is three directions), and each + of them maps to that same line. + + Best effort by design: a line this does not parse is simply absent from the + map, and its row falls back to :func:`miopen_cmd_from_key`. Writing the + results of a finished sweep must not fail over the label on a row. + """ + try: + from benchmarks.common.benchmark_implicit_gemm_conv import ( + _FORW_TO_DIRECTIONS, + parse_miopen_cmd, + ) + from benchmarks.common.early_stop import case_key + except ImportError: + return {} + + lines: List[str] = [] + if miopen_file: + path = Path(miopen_file) + if path.is_file(): + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + elif miopen_cmd: + lines = [miopen_cmd] + + out: Dict[str, str] = {} + for line in lines: + line = line.strip() + if not line or line.startswith("#"): + continue + try: + problem, dtype, forw = parse_miopen_cmd(line) + for direction in _FORW_TO_DIRECTIONS.get(forw, ["fwd"]): + # setdefault: two lines describing one case are one case, and + # it is the first that the sweeps measured. + out.setdefault(case_key(problem, dtype, direction), line) + except (ValueError, KeyError, SystemExit): + # SystemExit among them: the MIOpen parsers hand the flags to an + # argparse of their own, which exits rather than raises on a + # malformed value -- and a line nobody can read is a line to pass + # over, not the end of writing the results. + continue + return out + + +def merge( + rows: Iterable[dict], miopen_cmds: Optional[Dict[str, str]] = None +) -> List[dict]: """Join per-engine rows into one row per case. Rows may come from both engines and from several shards. A key only one @@ -240,6 +388,10 @@ def merge(rows: Iterable[dict]) -> List[dict]: cover), not an error. Within one engine and key, the faster row wins, so merging a case measured twice is idempotent. + ``miopen_cmds`` is :func:`miopen_cmds_by_key` of the input the run was + given, when it was given one; a key it does not cover is labelled with + :func:`miopen_cmd_from_key`, so every row says what to re-run. + A row's runner-up names travel with it: a slower row is dropped whole rather than having its ranking blended into the faster one's. A ranking only means something within the run that measured it -- two runs of one @@ -287,6 +439,13 @@ def merge(rows: Iterable[dict]) -> List[dict]: merged["winner"], merged["speedup"] = _verdict( merged["direct_ms"], merged["implicit_ms"] ) + if merged["winner"] is not None: + merged["winner_ms"] = merged[f"{merged['winner']}_ms"] + merged["winner_kernel"] = merged[f"{merged['winner']}_kernel"] + # A row the other engine never reached keeps both empty, as ``winner`` + # does: one time is not a comparison, and the engine that did run it is + # already on the row under its own name. + merged["miopen_cmd"] = (miopen_cmds or {}).get(key) or miopen_cmd_from_key(key) out.append({field: merged[field] for field in MERGED_FIELDS}) return out @@ -309,6 +468,10 @@ def _verdict(direct_ms, implicit_ms): def write(rows: Sequence[dict], path) -> Path: """Write merged ``rows`` as JSON Lines, plus a ``.csv`` beside it. + The JSON carries every field of :data:`MERGED_FIELDS`; the CSV is the + reading view, :data:`CSV_FIELDS` of the same rows. Nothing is lost by the + narrower one -- the two files are written from the same list. + Returns the CSV's path. Both are written whole (the merge has every row by then), unlike :class:`ResultsWriter`, which appends as a sweep goes. """ @@ -319,7 +482,7 @@ def write(rows: Sequence[dict], path) -> Path: fh.write(json.dumps(row, sort_keys=True) + "\n") csv_path = p.with_suffix(p.suffix + ".csv") if p.suffix else p.with_suffix(".csv") with csv_path.open("w", encoding="utf-8", newline="") as fh: - writer = csv.DictWriter(fh, fieldnames=list(MERGED_FIELDS)) + writer = csv.DictWriter(fh, fieldnames=list(CSV_FIELDS), extrasaction="ignore") writer.writeheader() writer.writerows(rows) return csv_path diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py index 962250b3682f..5716e31c5e74 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py @@ -43,6 +43,17 @@ shards. A shard that dies keeps the rows it had written -- they are flushed per case -- so a re-run of the remaining cases merges cleanly with them. +Progress +-------- +A shard reports nothing until it exits, and one shard is minutes to hours, so +the runner counts the case keys in the shards' row files every +``PROGRESS_SECONDS`` and prints a line across the whole run:: + + [progress] 142/4218 cases (3%) | 7/22 shards | 8/8 busy | 12 min elapsed | ETA 2.1 h + +For one shard's detail, follow its ``shard-NN.log``: the sweeps print a banner +per case into it. + Devices (``--devices``): a count ``N`` (devices ``0..N-1``), a list ``2,3``, or ``auto`` (default): the entries of ``HIP_VISIBLE_DEVICES``, else the count of ``ROCR_VISIBLE_DEVICES`` (``0..k-1``), else one device. This duplicates the @@ -61,7 +72,7 @@ import sys import time from pathlib import Path -from typing import Dict, List, Mapping, Optional, Sequence, Tuple +from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple from benchmarks.common import case_results @@ -83,6 +94,11 @@ #: loop. POLL_SECONDS = 1.0 +#: How often the aggregate progress line is printed. Without it the console +#: says nothing between one shard finishing and the next -- a quarter of an +#: hour of silence on a long sweep, which reads the same as a hang. +PROGRESS_SECONDS = 30.0 + # --------------------------------------------------------------------------- devices def resolve_devices(spec: str, env: Optional[Mapping[str, str]] = None) -> List[str]: @@ -136,6 +152,66 @@ def split_cases(cases: Sequence[str], shards: int) -> List[List[str]]: return [shard for shard in dealt if shard] +# --------------------------------------------------------------------------- progress +def measured_keys(paths: Iterable[Path]) -> Set[str]: + """The case keys recorded across *paths*, as a set. + + By key rather than by row: the comparison writes one row per engine, so a + case both sweeps measured is one case done, not two. The rows are flushed + per case, so reading the shard files mid-run is how the progress line sees + inside a shard without the children reporting anything back. + + This is a floor on the cases attempted, not a count of them: a case no + cached kernel supports finishes without writing a row. A sweep with many + of those shows a percentage that lags the real position and so an ETA on + the long side. + """ + keys: Set[str] = set() + for path in paths: + for row in case_results.read_rows(path): + keys.add(row["key"]) + return keys + + +def _duration(seconds: float) -> str: + """Minutes up to an hour and a half, hours past it.""" + if seconds < 90 * 60: + return f"{seconds / 60:.0f} min" + return f"{seconds / 3600:.1f} h" + + +def progress_line( + *, + measured: int, + total: int, + baseline: int, + shards_done: int, + shards: int, + busy: int, + devices: int, + elapsed: float, +) -> str: + """One line of aggregate progress over every shard. + + ``baseline`` is what a ``--resume`` carried in. Those cases cost this run + no time, so the rate the ETA extrapolates from counts only the cases this + run measured itself -- otherwise resuming a nearly-finished sweep would + report minutes left before it had timed a single case. + """ + pct = 100.0 * measured / total if total else 0.0 + parts = [ + f"{measured}/{total} cases ({pct:.0f}%)", + f"{shards_done}/{shards} shards", + f"{busy}/{devices} busy", + f"{_duration(elapsed)} elapsed", + ] + fresh = measured - baseline + remaining = total - measured + if fresh > 0 and remaining > 0: + parts.append(f"ETA {_duration(elapsed / fresh * remaining)}") + return "[progress] " + " | ".join(parts) + + # --------------------------------------------------------------------------- runs def previous_rows(log_dir: Path) -> List[dict]: """Every per-case row a previous run left in ``log_dir``. @@ -193,13 +269,17 @@ def run_shards( popen=subprocess.Popen, poll_seconds: float = POLL_SECONDS, on_done=None, + on_progress=None, + progress_seconds: float = PROGRESS_SECONDS, ) -> List[Tuple[int, str, int, float]]: """Run every shard, at most one per device at a time. A device that frees up takes the next shard off the queue, so a long shard delays only the card running it. ``command(index, paths)`` builds the child's argv; ``popen`` is the process factory (a stub, under test); - ``on_done(index, device, rc, seconds)`` is called as each shard ends. + ``on_done(index, device, rc, seconds)`` is called as each shard ends and + ``on_progress(shards_done, busy)`` every ``progress_seconds`` while any + shard is still running. Returns ``(shard index, device, rc, seconds)`` per shard, in the order they finished. @@ -208,6 +288,7 @@ def run_shards( running: Dict[str, Tuple[int, Mapping[str, Path], object, object, float]] = {} free = list(devices) done: List[Tuple[int, str, int, float]] = [] + last_progress = time.monotonic() while queue or running: while queue and free: @@ -237,6 +318,10 @@ def run_shards( done.append(entry) if on_done is not None: on_done(*entry) + now = time.monotonic() + if on_progress is not None and now - last_progress >= progress_seconds: + last_progress = now + on_progress(len(done), len(running)) return done @@ -403,12 +488,35 @@ def _report(index: int, device: str, rc: int, seconds: float) -> None: ) started = time.monotonic() + + # The shards' own row files, re-read on each progress tick. A shard reports + # nothing until it exits, so its rows are the only live view into it. + row_files = [paths["rows"] for _index, paths in plan] + baseline = {row["key"] for row in carried} + + def _progress(shards_done: int, busy: int) -> None: + print( + progress_line( + measured=len(baseline | measured_keys(row_files)), + total=len(cases), + baseline=len(baseline), + shards_done=shards_done, + shards=len(dealt), + busy=busy, + devices=len(devices), + elapsed=time.monotonic() - started, + ), + flush=True, + ) + done = run_shards( plan, devices=devices, command=lambda index, paths: shard_command(script, forwarded, paths), poll_seconds=POLL_SECONDS, on_done=_report, + on_progress=_progress, + progress_seconds=PROGRESS_SECONDS, ) rc_all = 0 @@ -424,7 +532,10 @@ def _report(index: int, device: str, rc: int, seconds: float) -> None: row for _index, paths in plan for row in case_results.read_rows(paths["rows"]) - ] + ], + # The shards' rows are keyed by geometry; the case file is what puts + # each row's MIOpenDriver line back on it. + case_results.miopen_cmds_by_key(miopen_file=args.miopen_file), ) if rows: out = Path(args.results or (log_dir / "results.jsonl")) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py index af20456ae2e2..5e1e53f785a4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py @@ -305,10 +305,102 @@ def test_write_leaves_a_csv_beside_the_jsonl(tmp_path): csv_path = case_results.write(rows, out) assert csv_path.is_file() and csv_path != out header, body = csv_path.read_text().splitlines()[:2] - assert header.split(",") == list(case_results.MERGED_FIELDS) + assert header.split(",") == list(case_results.CSV_FIELDS) assert "d" in body and "i" in body +def test_the_csv_is_narrower_than_the_json_it_is_written_from(tmp_path): + """The CSV is the reading view; the JSON keeps the fields it leaves out, so + nothing is lost by the narrowing.""" + rows = case_results.merge( + [_row("direct", "k", 0.5, "d"), _row("implicit", "k", 1.0, "i")] + ) + out = tmp_path / "results.jsonl" + csv_path = case_results.write(rows, out) + assert set(case_results.CSV_FIELDS) < set(case_results.MERGED_FIELDS) + assert "tflops" not in csv_path.read_text().splitlines()[0] + assert json.loads(out.read_text().splitlines()[0])["direct_tflops"] == 100.0 + + +def test_the_winner_columns_repeat_the_winning_engines_row(tmp_path): + rows = case_results.merge( + [_row("direct", "k", 0.5, "d"), _row("implicit", "k", 1.0, "i")] + ) + assert rows[0]["winner_ms"] == 0.5 and rows[0]["winner_kernel"] == "d" + + +def test_a_one_sided_case_names_no_winner(): + """One time is not a comparison: the engine that did reach the case is + already on the row under its own name.""" + rows = case_results.merge([_row("implicit", "k", 1.0, "i")]) + assert rows[0]["winner_ms"] is None and rows[0]["winner_kernel"] is None + + +# ------------------------------------------------------------------ miopen cmd +_MIOPEN_LINE = ( + "convbfp16 -F 1 -n 4 -c 64 -k 128 -H 320 -W 320 -y 3 -x 3 -p 1 -q 1 " + "-u 2 -v 2 -l 1 -j 1 -g 1 -m conv -I NHWC -O NHWC -f NHWC" +) + + +def _key_of(line, direction="fwd"): + from benchmarks.common.benchmark_implicit_gemm_conv import parse_miopen_cmd + + problem, dtype, _forw = parse_miopen_cmd(line) + return case_key(problem, dtype, direction) + + +def test_the_case_file_puts_its_own_line_on_the_row(tmp_path): + """The sweeps key their rows by geometry, so the line the user wrote is not + in them -- the merge reads it back off the input the run was given.""" + f = tmp_path / "cases.txt" + f.write_text(f"# a comment\n\n{_MIOPEN_LINE}\n") + key = _key_of(_MIOPEN_LINE) + rows = case_results.merge( + [_row("direct", key, 0.5, "d")], + case_results.miopen_cmds_by_key(miopen_file=f), + ) + assert rows[0]["miopen_cmd"] == _MIOPEN_LINE + + +def test_a_case_with_no_command_behind_it_gets_one_written_from_its_key(): + """A run given shape flags rather than a command still has to leave a row + you can re-run.""" + from benchmarks.common.benchmark_implicit_gemm_conv import parse_miopen_cmd + + key = _key_of(_MIOPEN_LINE) + rows = case_results.merge([_row("direct", key, 0.5, "d")]) + cmd = rows[0]["miopen_cmd"] + assert cmd != _MIOPEN_LINE and "MIOpenDriver convbfp16" in cmd + # It describes the same case: round-tripping it lands back on the key. + problem, dtype, forw = parse_miopen_cmd(cmd) + assert forw == 1 and case_key(problem, dtype, "fwd") == key + + +def test_a_key_the_case_file_does_not_cover_falls_back_to_a_written_command(tmp_path): + f = tmp_path / "cases.txt" + f.write_text(_MIOPEN_LINE + "\n") + rows = case_results.merge( + [_row("direct", _key_of(_MIOPEN_LINE, "wgrad"), 0.5, "d")], + case_results.miopen_cmds_by_key(miopen_file=f), + ) + assert rows[0]["miopen_cmd"].endswith("-F 4 -t 1 -in_layout=NHWC") + + +def test_a_case_file_line_that_does_not_parse_is_skipped_not_fatal(tmp_path): + """Writing the results of a finished sweep must not fail over the label on + a row.""" + f = tmp_path / "cases.txt" + f.write_text(f"not a MIOpenDriver command\n{_MIOPEN_LINE}\n") + assert list(case_results.miopen_cmds_by_key(miopen_file=f)) == [ + _key_of(_MIOPEN_LINE) + ] + + +def test_a_key_that_does_not_parse_is_left_as_it_is(): + assert case_results.miopen_cmd_from_key("nonsense") == "nonsense" + + def test_a_writer_without_a_path_writes_nothing(tmp_path): writer = case_results.ResultsWriter(None, "direct") writer.row( @@ -520,6 +612,109 @@ def test_resume_carries_the_earlier_rows_and_skips_their_cases( ''' +# --------------------------------------------------------------------------- progress +def _rows_file(path, keys, engines=("direct", "implicit")): + path.write_text( + "".join( + json.dumps(_row(engine, key, 1.0, f"{engine}-{key}")) + "\n" + for key in keys + for engine in engines + ) + ) + return path + + +def test_a_case_both_engines_measured_counts_once(tmp_path): + """The row files hold one row per engine; progress is in cases.""" + a = _rows_file(tmp_path / "shard-0.jsonl", ["k1", "k2"]) + b = _rows_file(tmp_path / "shard-1.jsonl", ["k3"], engines=("implicit",)) + + assert run_conv_shards.measured_keys([a, b]) == {"k1", "k2", "k3"} + + +def test_measured_keys_ignores_a_shard_that_has_written_nothing(tmp_path): + """Every shard's row file is read on each tick, including the ones whose + card has not reached its first case yet.""" + a = _rows_file(tmp_path / "shard-0.jsonl", ["k1"]) + + assert run_conv_shards.measured_keys([a, tmp_path / "absent.jsonl"]) == {"k1"} + + +def test_the_progress_line_reports_the_run(tmp_path): + line = run_conv_shards.progress_line( + measured=142, + total=4218, + baseline=0, + shards_done=7, + shards=22, + busy=8, + devices=8, + elapsed=12 * 60, + ) + assert "142/4218 cases (3%)" in line + assert "7/22 shards" in line and "8/8 busy" in line + # 142 cases in 12 min leaves 4076 -- hours, so the ETA says hours. + assert "ETA 5.7 h" in line + + +def test_a_resumed_run_does_not_extrapolate_from_carried_cases(tmp_path): + """4000 cases carried in by --resume took this run no time. Dividing the + elapsed minute by them would promise the remaining 218 in seconds.""" + line = run_conv_shards.progress_line( + measured=4010, + total=4218, + baseline=4000, + shards_done=0, + shards=22, + busy=8, + devices=8, + elapsed=60, + ) + # 10 cases in a minute, 208 to go: ~21 min, not the ~3 s the carried rows + # would imply. + assert "ETA 21 min" in line + + +def test_the_progress_line_omits_an_eta_it_cannot_compute(tmp_path): + """Nothing measured yet: no rate, so no guess.""" + line = run_conv_shards.progress_line( + measured=0, + total=4218, + baseline=0, + shards_done=0, + shards=22, + busy=8, + devices=8, + elapsed=30, + ) + assert "0/4218 cases (0%)" in line + assert "ETA" not in line + + +def test_progress_is_reported_while_the_shards_run(tmp_path): + """The point of the line: the console says something between one shard + finishing and the next, which on a long sweep is a quarter of an hour.""" + ticks = [] + + def popen(cmd, **_kw): + return _FakeProc(3) + + run_conv_shards.run_shards( + _plan(tmp_path, 4), + devices=["0", "1"], + command=lambda index, paths: ["python", str(index)], + popen=popen, + poll_seconds=0, + on_progress=lambda shards_done, busy: ticks.append((shards_done, busy)), + progress_seconds=0, + ) + + assert ticks, "no progress was reported" + # Reported while cards were still working, not only at the end. + assert any(busy > 0 for _done, busy in ticks) + assert ticks[-1][0] == 4, "the last tick sees every shard finished" + + # --------------------------------------------------------------------------- compare def test_compare_merges_one_engine_when_the_other_was_skipped(tmp_path, capsys): """--skip-direct leaves no direct rows; the table is then one-sided rather @@ -546,6 +741,48 @@ def test_compare_says_so_when_there_is_nothing_to_write(tmp_path, capsys): assert "no per-case results" in capsys.readouterr().err +def _compare_argv(cache) -> list: + """The smallest argv that gets ``main`` to the sweep it will stub out.""" + return ["x", "--run-from-cache", str(cache), "--skip-direct"] + + +def test_compare_fails_when_a_sweep_dies(tmp_path, monkeypatch, capsys): + """A sweep that crashed on every case prints no ``Best:`` line -- but then + neither does a shape none of the cached kernels support. The exit status is + the only thing that separates them, and ``run_conv_shards.py`` reports it + per shard.""" + from benchmarks.common import benchmark_conv_compare as compare + + monkeypatch.setattr(compare, "_run_script", lambda *a, **k: (None, "", "boom", 1)) + monkeypatch.setattr(sys, "argv", _compare_argv(tmp_path)) + + assert compare.main() == 1 + assert "exited 1" in capsys.readouterr().err + + +def test_compare_turns_a_killed_sweep_into_a_usable_status(tmp_path, monkeypatch): + """A child killed by a signal reports a negative status. Returned verbatim + it would be truncated to an arbitrary byte -- zero among them, which is the + one value that must not come out of a failed run.""" + from benchmarks.common import benchmark_conv_compare as compare + + monkeypatch.setattr(compare, "_run_script", lambda *a, **k: (None, "", "", -9)) + monkeypatch.setattr(sys, "argv", _compare_argv(tmp_path)) + + assert compare.main() == 1 + + +def test_compare_succeeds_when_the_sweep_does(tmp_path, monkeypatch): + """The other half of the status: a healthy sweep still exits zero, so the + shard runner does not start reporting every shard as failed.""" + from benchmarks.common import benchmark_conv_compare as compare + + monkeypatch.setattr(compare, "_run_script", lambda *a, **k: (1.0, "kern", "", 0)) + monkeypatch.setattr(sys, "argv", _compare_argv(tmp_path)) + + assert compare.main() == 0 + + def test_the_comparison_driver_needs_no_torch(): """conv_compare is a launcher: it has to run where torch is not installed (the benchmark scripts it spawns import it, it does not).""" From aa5d590a71da5be910370f15a7ef0261db4ea2ef Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Mon, 26 Oct 2026 18:54:06 -0500 Subject: [PATCH 5/9] feat(rocke): time conv sweeps the way hipDNN does by default The conv sweeps timed N launches back to back between one event pair and averaged. hipDNN's dnn-benchmark (and MIOpen, where it records events around a launch) times each call on its own, so it pays the fixed device cost of an event pair per call that a back-to-back loop amortises away. On small kernels that made the sweep's numbers look faster than the same kernels measure through hipDNN, which skewed every rocKE-vs-MIOpen comparison built on them. - rocke.runtime.StallGate / time_launches_staged: the ctypes twin of hipdnn_data_sdk's StallGate (hipStreamWaitValue32 on signal memory, released by the host, with a watchdog) and of dnn-benchmark's staged loop: one stalled start/stop event pair per call, host submission outside it. - --timing {staged,block} on both conv sweeps and benchmark_conv_compare.py, default staged: every reported time is the median of stalled single-call samples. Early stopping and the --fast-rank probe still rank with cheap block launches; only the reported number changes protocol. - Per-case rows record their timing (rows without the field are block); --skip-measured / --resume skip only cases measured the same way, and a merge takes one protocol's rows instead of letting the cheaper protocol win the faster-row rule. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../benchmarks/common/batched_timing.py | 119 +++++++- .../common/benchmark_conv_compare.py | 17 +- .../common/benchmark_direct_conv.py | 20 +- .../common/benchmark_implicit_gemm_conv.py | 15 +- .../library/benchmarks/common/case_results.py | 62 +++- .../library/benchmarks/common/early_stop.py | 35 ++- .../benchmarks/common/run_conv_shards.py | 38 ++- .../library/tests/test_batched_timing.py | 64 +++- .../rocke/library/tests/test_early_stop.py | 38 +++ .../library/tests/test_run_conv_shards.py | 155 ++++++++-- .../platform/python/rocke/runtime/__init__.py | 8 + .../python/rocke/runtime/hip_module.py | 14 + .../python/rocke/runtime/stall_gate.py | 286 ++++++++++++++++++ .../platform/tests/runtime/test_stall_gate.py | 176 +++++++++++ 14 files changed, 976 insertions(+), 71 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/stall_gate.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_stall_gate.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py index 876b0b8720a1..6450620eb39b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py @@ -45,6 +45,29 @@ What is left is sampling noise on a one- or two-launch measurement, which is why the cut is deliberately loose: the probe only has to get a candidate into the top ``finalists``, not to rank it there. + +Timing the finalists like hipDNN +-------------------------------- + +Round 2 has two protocols, chosen by ``timing``. ``block`` is +:func:`rocke.runtime.time_launches`: ``iters`` launches back to back between +one event pair, averaged. ``staged`` (the default) is what hipDNN's +``dnn-benchmark`` does by default, and what a number compared against hipDNN +or MIOpen has to be: +every call is its own sample, a start/stop event pair around one call queued +behind a host-released stall (:func:`rocke.runtime.time_launches_staged`), and +the reported time is the median. The span holds no host submission, but it +does hold the fixed device cost of one event pair around a single call: a +back-to-back loop amortises it away, while hipDNN's timing (and MIOpen's, +where it records events around a launch) pays it on every call. On small +kernels that fixed cost is a large share of the time, so the two protocols +must not be compared with each other. + +Only round 2 changes. The probe stays a batched block measurement either way: +it only orders the candidates, and the event cost is the same constant on all +of them, so it does not change which reach the finalists. Staged samples drain +the device after every call, so the clocks the priming brought up are allowed +to settle between samples -- deliberately, as dnn-benchmark does. """ from __future__ import annotations @@ -71,6 +94,30 @@ DEFAULT_PROBE_WARMUP = 1 DEFAULT_PROBE_ITERS = 2 +#: How a reported time is measured; see "Timing the finalists like hipDNN" +#: above. The sweeps' ``--timing`` flag picks one for every reported number, +#: ranked or not (:meth:`benchmarks.common.early_stop.EarlyStop.measure`). +TIMINGS = ("block", "staged") +DEFAULT_TIMING = "staged" +#: The protocol of :func:`rocke.runtime.time_launches`, and of every result +#: recorded before the choice existed. +BLOCK_TIMING = "block" + + +def add_timing_arg(parser) -> None: + """Add ``--timing`` to a sweep's parser.""" + parser.add_argument( + "--timing", + choices=TIMINGS, + default=DEFAULT_TIMING, + help="how every reported time is measured. staged: each launch its " + "own stalled event pair, median of --iters -- what hipDNN's " + "dnn-benchmark measures, so the numbers compare with hipDNN and " + "MIOpen. block: --iters launches back to back between one event " + "pair, averaged. Early stopping and the --fast-rank probe still rank " + f"with block launches (default: {DEFAULT_TIMING})", + ) + @dataclass class Ranked: @@ -231,6 +278,48 @@ def time_batch( return out +def time_batch_staged( + runs: Sequence[Callable[[], None]], + *, + warmup: int, + iters: int, + stream: int = 0, +) -> List[Optional[float]]: + """Per-candidate median of ``iters`` stalled single-call samples. + + Every sample is :func:`rocke.runtime.time_launches_staged`'s: a start/stop + event pair around one call, queued behind a stall the host releases once + the call is enqueued. ``warmup`` samples of each candidate are discarded. + The process's shared gate serves every candidate. + + ``None`` for a candidate whose launch raised, and for one whose call + blocked the host on the stalled stream (the watchdog had to release it, so + the sample holds host time); the second is reported, as it means the bound + launch synchronizes when it should not. + """ + from statistics import median + + from rocke.runtime.hip_module import HipError + from rocke.runtime.stall_gate import StallTimeout, time_launches_staged + + out: List[Optional[float]] = [] + for i, run in enumerate(runs): + try: + samples = time_launches_staged( + run, warmup=warmup, iters=iters, stream=stream + ) + except StallTimeout as e: + print(f" [staged] finalist {i}: {e}; not ranked", flush=True) + out.append(None) + continue + except (HipError, RuntimeError): + # time_launches_staged drained before raising. + out.append(None) + continue + out.append(median(samples) if samples else None) + return out + + def rank( n: int, make_run: Callable[[int], Optional[Callable[[], None]]], @@ -248,6 +337,7 @@ def rank( on_drop: Optional[Callable[[int], None]] = None, held_bytes: Optional[Callable[[], int]] = None, chunk_bytes: int = 0, + timing: str = DEFAULT_TIMING, ) -> List[Ranked]: """Probe all ``n`` candidates, then time the best ``finalists`` in full. @@ -278,8 +368,15 @@ def rank( The finalists' scratch is still held when this returns -- they were timed with it -- so the caller frees that along with the rest of the case. + ``timing`` picks round 2's protocol (:data:`TIMINGS`): ``block`` times the + finalists with :func:`time_batch`, ``staged`` with + :func:`time_batch_staged`, so ``ms`` is then a median of single-call + samples. The probe is the same either way. + Returns one :class:`Ranked` per candidate, in candidate order. """ + if timing not in TIMINGS: + raise ValueError(f"timing must be one of {TIMINGS}, got {timing!r}") results = [Ranked(index=i) for i in range(n)] # The finalists' bound launches, kept alive from the probe. A launch holds # its packed kernargs and its kernel handle, so keeping it is what makes @@ -346,22 +443,24 @@ def trim() -> None: if results[i].failed: drop(i) - # Round 2. The finalists are verified first and then timed as one batch: - # a candidate's timed launches start right after the previous one's, so - # the clocks that the priming brought up never fall back between them. + # Round 2. The finalists are verified first and then timed as one batch. + # In block timing a candidate's timed launches start right after the + # previous one's, so the clocks that the priming brought up never fall + # back between them; staged timing drains after every sample, as hipDNN's + # does (see the module docstring). order = sorted(kept, key=lambda i: results[i].probe_ms) if on_finalist is not None: for i in order: # A kernel that fails here is still timed, so its row carries a # number; the caller reads ``passed`` and keeps it from ranking. results[i].passed = on_finalist(i, kept[i]) - times = time_batch( - [kept[i] for i in order], - warmup=warmup, - iters=iters, - stream=stream, - chunk=chunk, - ) + finals = [kept[i] for i in order] + if timing == "staged": + times = time_batch_staged(finals, warmup=warmup, iters=iters, stream=stream) + else: + times = time_batch( + finals, warmup=warmup, iters=iters, stream=stream, chunk=chunk + ) for i, ms in zip(order, times): if ms is None: results[i].failed = True 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 88705eba2dec..551f22b2dc64 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 @@ -55,6 +55,14 @@ direct-conv time -- and passing ``--early-stop``/``--early-stop-after`` alongside ``--fast-rank`` has no effect. +Both scripts report times measured the way hipDNN's ``dnn-benchmark`` times a +call (``--timing staged``, the default: one stalled event pair per launch, the +median; see :mod:`benchmarks.common.batched_timing`), so the numbers compare +with hipDNN and MIOpen. ``--timing block`` restores the back-to-back average. +Early stopping and the ``--fast-rank`` probe rank with block launches either +way. ``--timing`` is forwarded to both scripts, and the merged results record +which protocol measured them. + Notes ----- - Without a cache, direct conv requires ``cpg == kpg`` (C/groups == K/groups) @@ -72,7 +80,7 @@ import tempfile from pathlib import Path -from benchmarks.common import case_results +from benchmarks.common import batched_timing, case_results _SCRIPT_DIR = Path(__file__).parent @@ -223,6 +231,7 @@ def _build_mode_args(args) -> list[str]: a += ["--early-stop", str(args.early_stop)] if args.early_stop_after is not None: a += ["--early-stop-after", str(args.early_stop_after)] + a += ["--timing", args.timing] # Both scripts take the fast-ranking flags with the same names and the # same defaults (they share benchmarks.common.batched_timing), so a knob # is forwarded only when it was given and the scripts' own default stands @@ -310,6 +319,7 @@ def _write_results( *, miopen_file: str | None = None, miopen_cmd: str | None = None, + timing: str | None = None, ) -> None: """Join the two sweeps' per-case rows into ``path`` (and a ``.csv``). @@ -320,10 +330,13 @@ def _write_results( The children key their rows by geometry, not by the line they read, so the MIOpenDriver command goes back on here, from the input this run was given. + ``timing`` is how both children timed their kernels; the merge takes only + rows measured that way (``None``: the rows must agree among themselves). """ rows = case_results.merge( case_results.read_rows(rows_direct) + case_results.read_rows(rows_implicit), case_results.miopen_cmds_by_key(miopen_file=miopen_file, miopen_cmd=miopen_cmd), + timing=timing, ) if not rows: print(f"\n[info] no per-case results to write to {path}", file=sys.stderr) @@ -416,6 +429,7 @@ def main() -> int: metavar="N", help="forwarded to both scripts (default: theirs)", ) + batched_timing.add_timing_arg(parser) parser.add_argument( "--no-seed-from-direct", action="store_true", @@ -743,6 +757,7 @@ def main() -> int: rows_implicit, miopen_file=args.miopen_file, miopen_cmd=args.miopen_cmd, + timing=args.timing, ) _print_summary(direct_result, implicit_result, implicit_note) 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 8c19af397e9c..af878d829939 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 @@ -1159,7 +1159,7 @@ def _run_nongrouped_sweep( ) results: List[NonGroupedResult] = [] - _stop = EarlyStop(args.early_stop, args.early_stop_after) + _stop = EarlyStop(args.early_stop, args.early_stop_after, timing=args.timing) n_run = 0 for spec, kernel in pending: artifact = artifact_map.get(kernel.name) @@ -2259,8 +2259,7 @@ def verify(plan, run) -> bool: rt.memcpy_d2h(u8(out_cpu), ptrs[out_name], Out_t.nbytes) ref_f32 = ref_out.float().cpu() rel_err = float( - (out_cpu.float() - ref_f32).abs().max() - / ref_f32.abs().max().clamp(min=1.0) + (out_cpu.float() - ref_f32).abs().max() / ref_f32.abs().max().clamp(min=1.0) ) passed = rel_err < 5e-2 if not passed: @@ -2365,13 +2364,15 @@ def _on_finalist(i, run): finalists=args.fast_rank_finalists, chunk=args.fast_rank_chunk, primer=args.fast_rank_primer, + timing=args.timing, on_probe=_on_probe, on_finalist=_on_finalist, ) n_probed = sum(1 for r in ranked if r.probe_ms is not None) n_final = sum(1 for r in ranked if r.is_finalist) print( - f" probe: {n_probed}/{n} timed, {n_final} finalist(s) measured in full", + f" probe: {n_probed}/{n} timed, {n_final} finalist(s) measured in full " + f"({args.timing} timing)", flush=True, ) for r in ranked: @@ -2469,9 +2470,10 @@ def _u8(t): modules = ModuleCache(args.module_cache_bytes) # One row per case, for the comparison driver and the shard runner to # merge. Does nothing without --results. - results_out = ResultsWriter(args.results, "direct") - # --skip-measured: cases a previous run of this sweep already recorded. - done = measured_keys(args.skip_measured, "direct") + results_out = ResultsWriter(args.results, "direct", args.timing) + # --skip-measured: cases a previous run of this sweep already recorded, + # timed the way this run times them. + done = measured_keys(args.skip_measured, "direct", args.timing) n_skipped = 0 from kernels.conv.direct_aot_jobs import DIRECT_DTYPES @@ -2675,6 +2677,7 @@ def main() -> int: "--warmup", type=int, default=3, help="warmup iterations (default: 3)" ) add_early_stop_arg(parser) + batched_timing.add_timing_arg(parser) parser.add_argument( "--iters", type=int, default=10, help="timed iterations (default: 10)" ) @@ -2910,8 +2913,7 @@ def main() -> int: args = parser.parse_args() if args.skip_measured and not args.run_from_cache: print( - "error: --skip-measured resumes an AOT sweep; pass " - "--run-from-cache DIR", + "error: --skip-measured resumes an AOT sweep; pass " "--run-from-cache DIR", file=sys.stderr, ) return 2 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 b88906759043..b21a4867e952 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 @@ -1135,6 +1135,7 @@ def main() -> int: "--warmup", type=int, default=3, help="warmup iterations (default: 3)" ) add_early_stop_arg(parser) + batched_timing.add_timing_arg(parser) parser.add_argument( "--iters", type=int, default=10, help="timed iterations (default: 10)" ) @@ -1643,8 +1644,7 @@ def main() -> int: args = parser.parse_args() if args.skip_measured and not args.run_from_cache: print( - "error: --skip-measured resumes an AOT sweep; pass " - "--run-from-cache DIR", + "error: --skip-measured resumes an AOT sweep; pass " "--run-from-cache DIR", file=sys.stderr, ) return 2 @@ -4677,9 +4677,10 @@ def _run_from_cache(args, arch, target, cases, cache, directions) -> int: modules = ModuleCache(args.module_cache_bytes) # One row per case, for the comparison driver and the shard runner to # merge. Does nothing without --results. - results_out = ResultsWriter(args.results, "implicit") - # --skip-measured: cases a previous run of this sweep already recorded. - done = measured_keys(args.skip_measured, "implicit") + results_out = ResultsWriter(args.results, "implicit", args.timing) + # --skip-measured: cases a previous run of this sweep already recorded, + # timed the way this run times them. + done = measured_keys(args.skip_measured, "implicit", args.timing) n_skipped = 0 def _u8(t): @@ -5412,6 +5413,7 @@ def on_finalist(i, run): finalists=args.fast_rank_finalists, chunk=args.fast_rank_chunk, primer=args.fast_rank_primer, + timing=args.timing, on_probe=on_probe, on_finalist=on_finalist, on_drop=on_drop, @@ -5422,7 +5424,8 @@ def on_finalist(i, run): n_probed = sum(1 for r in ranked if r.probe_ms is not None) n_final = sum(1 for r in ranked if r.is_finalist) print( - f" probe: {n_probed}/{n} timed, {n_final} finalist(s) measured in full", + f" probe: {n_probed}/{n} timed, {n_final} finalist(s) measured in full " + f"({args.timing} timing)", flush=True, ) n_failed = 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py index 7b76709917cd..c9c52258c5e0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py @@ -31,6 +31,13 @@ Rows are appended and flushed per case: a shard that is killed after four hours keeps the four hours. + +Every row says how its time was measured (``timing``, one of +:data:`benchmarks.common.batched_timing.TIMINGS`; a row without the field is +from before it existed, when every sweep timed block). The protocols differ by a fixed +per-call cost, so rows of two protocols are never compared with each other: +resuming skips only cases measured the same way, and a merge takes one +protocol's rows. """ from __future__ import annotations @@ -40,6 +47,8 @@ from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Sequence +from benchmarks.common.batched_timing import BLOCK_TIMING, TIMINGS + #: The two sweeps, by the name their rows carry. ENGINES = ("direct", "implicit") @@ -76,6 +85,7 @@ def _engine_fields(engine: str) -> tuple: "winner_ms", "winner_kernel", "speedup", + "timing", ) #: Columns of the ``.csv`` beside it, in order. @@ -98,6 +108,7 @@ def _engine_fields(engine: str) -> tuple: "implicit_ms", "winner_ms", "winner_kernel", + "timing", ) @@ -148,16 +159,27 @@ def add_skip_measured_arg(parser) -> None: ) -def measured_keys(path, engine: str) -> set: - """Case keys ``engine`` already has a row for in ``path``. +def row_timing(row: dict) -> str: + """How ``row``'s time was measured; rows older than the field are block.""" + return row.get("timing") or BLOCK_TIMING + + +def measured_keys(path, engine: str, timing: str = BLOCK_TIMING) -> set: + """Case keys ``engine`` already has a row for in ``path``, timed by ``timing``. One engine's rows only: a run killed between the direct sweep and the implicit one has the case measured on one side and not the other, and - resuming has to finish the second rather than skip it. + resuming has to finish the second rather than skip it. One protocol's rows + only, for the same reason: a case measured the other way is not measured + for this run. """ if not path: return set() - return {row["key"] for row in read_rows(path) if row.get("engine") == engine} + return { + row["key"] + for row in read_rows(path) + if row.get("engine") == engine and row_timing(row) == timing + } class ResultsWriter: @@ -169,10 +191,15 @@ class ResultsWriter: no row -- its absence is the finding, and the log says why. """ - def __init__(self, path: Optional[str], engine: str) -> None: + def __init__( + self, path: Optional[str], engine: str, timing: str = BLOCK_TIMING + ) -> None: if engine not in ENGINES: raise ValueError(f"engine must be one of {ENGINES}, got {engine!r}") + if timing not in TIMINGS: + raise ValueError(f"timing must be one of {TIMINGS}, got {timing!r}") self._engine = engine + self._timing = timing self._path = Path(path) if path else None self._fh = None self._n = 0 @@ -220,6 +247,7 @@ def row( "gbps": round(float(gbps), 3), "kernel": kernel, "candidates": int(candidates), + "timing": self._timing, } if passed is not None: record["passed"] = bool(passed) @@ -378,7 +406,9 @@ def miopen_cmds_by_key( def merge( - rows: Iterable[dict], miopen_cmds: Optional[Dict[str, str]] = None + rows: Iterable[dict], + miopen_cmds: Optional[Dict[str, str]] = None, + timing: Optional[str] = None, ) -> List[dict]: """Join per-engine rows into one row per case. @@ -398,13 +428,27 @@ def merge( case can order the same near-tied kernels differently, and a mix of the two would be a ranking that never happened. + ``timing`` keeps only the rows measured that way (:func:`row_timing`). + Without it the rows must all have been measured one way, and a mix raises + ``ValueError``: the faster-row rule would otherwise pick whichever protocol + is cheaper rather than whichever kernel is faster. + Sorted by key, so two runs of the same case list produce comparable files. """ + rows = list(rows) + if timing is None: + seen = {row_timing(row) for row in rows if row.get("engine") in ENGINES} + if len(seen) > 1: + raise ValueError( + f"rows measured with different timings {sorted(seen)}; " + f"merge one protocol at a time" + ) + timing = seen.pop() if seen else BLOCK_TIMING by_key: Dict[str, dict] = {} for row in rows: key = row.get("key") engine = row.get("engine") - if not key or engine not in ENGINES: + if not key or engine not in ENGINES or row_timing(row) != timing: continue merged = by_key.setdefault( key, @@ -413,6 +457,7 @@ def merge( "shape": row.get("shape"), "dtype": row.get("dtype"), "direction": row.get("direction"), + "timing": timing, }, ) ms = row.get("ms") @@ -498,8 +543,9 @@ def summary(rows: Sequence[dict]) -> str: 1 for r in rows if r["direct_ms"] is None and r["implicit_ms"] is not None ) wins = sum(1 for r in rows if r["winner"] == "direct") + timing = f" ({rows[0].get('timing')} timing)" if rows else "" return ( - f"{len(rows)} case(s): {both} measured by both " + f"{len(rows)} case(s){timing}: {both} measured by both " f"(direct faster on {wins}), {direct_only} direct-only, " f"{implicit_only} implicit-only" ) 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 9bfe61a3daf4..c3312f25b50a 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 @@ -17,6 +17,14 @@ 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. + +What is reported and what is pruned are measured differently. The pruning +probes (the first launch, the warmup) are back-to-back block launches, cheap +and the same for every kernel. The reported time is measured with the sweep's +``--timing`` (:data:`benchmarks.common.batched_timing.TIMINGS`): by default +one stalled event pair per launch and the median, as hipDNN measures. A staged +time carries a fixed per-call cost a block probe does not, so the bound a +block probe is held to is, if anything, looser -- pruning stays safe. """ from __future__ import annotations @@ -148,7 +156,13 @@ def __init__( *, seed_ms: Optional[float] = None, record_key: Optional[str] = None, + timing: str = "block", ) -> None: + from benchmarks.common.batched_timing import TIMINGS + + if timing not in TIMINGS: + raise ValueError(f"timing must be one of {TIMINGS}, got {timing!r}") + self.timing = timing self.factor = factor self.after = after self.n_measured = 0 @@ -173,6 +187,7 @@ def for_case(cls, args, problem: object, dtype: str, direction: str) -> "EarlySt args.early_stop_after, seed_ms=seeds.get(key), record_key=key if getattr(args, "early_stop_record", None) else None, + timing=getattr(args, "timing", "block"), ) def measure( @@ -224,12 +239,12 @@ def measure( self.n_stopped += 1 return None # The probe was the warmup; only the timed loop is left. - ms = time_launches(fn, warmup=0, iters=iters, stream=stream) + ms = self._timed(fn, warmup=0, iters=iters, stream=stream) elif probing: # The probe was the first warmup launch. - ms = time_launches(fn, warmup=warmup - 1, iters=iters, stream=stream) + ms = self._timed(fn, warmup=warmup - 1, iters=iters, stream=stream) else: - ms = time_launches(fn, warmup=warmup, iters=iters, stream=stream) + ms = self._timed(fn, warmup=warmup, iters=iters, stream=stream) if passed is False: return ms self.n_measured += 1 @@ -239,6 +254,20 @@ def measure( _record(self.record_key, ms) return ms + def _timed(self, fn, *, warmup: int, iters: int, stream: int) -> float: + """The reported time of ``fn``, measured with :attr:`timing`.""" + if self.timing == "staged": + from statistics import median + + from rocke.runtime import time_launches_staged + + return median( + time_launches_staged(fn, warmup=warmup, iters=iters, stream=stream) + ) + from rocke.runtime import time_launches + + return time_launches(fn, warmup=warmup, iters=iters, stream=stream) + def note(self, ms: float, *, passed: Optional[bool] = None) -> None: """Record a time this object did not measure itself. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py index 5716e31c5e74..3ec2b0b4504b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py @@ -75,6 +75,7 @@ from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple from benchmarks.common import case_results +from benchmarks.common.batched_timing import DEFAULT_TIMING, TIMINGS _SCRIPT_DIR = Path(__file__).parent @@ -227,6 +228,24 @@ def previous_rows(log_dir: Path) -> List[dict]: return out +def forwarded_timing(forwarded: Sequence[str]) -> str: + """The finalist timing the forwarded arguments ask the sweeps for. + + The runner forwards the benchmark's flags verbatim, but it merges the rows + and carries a previous run's into ``--resume``, and rows of two timings are + not comparable (:mod:`benchmarks.common.case_results`). So it reads this + one flag back: the last ``--timing`` given, as argparse would take it, else + the sweeps' default. + """ + timing = DEFAULT_TIMING + for i, arg in enumerate(forwarded): + if arg == "--timing" and i + 1 < len(forwarded): + timing = forwarded[i + 1] + elif arg.startswith("--timing="): + timing = arg.split("=", 1)[1] + return timing + + def shard_paths(log_dir: Path, index: int, width: int) -> Dict[str, Path]: stem = f"shard-{index:0{width}d}" return { @@ -421,11 +440,17 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ) return 2 + timing = forwarded_timing(forwarded) + if timing not in TIMINGS: + print( + f"error: --timing {timing!r}: expected one of {TIMINGS}", + file=sys.stderr, + ) + return 2 + devices = resolve_devices(args.devices) shards = ( - args.shards - if args.shards is not None - else SHARDS_PER_DEVICE * len(devices) + args.shards if args.shards is not None else SHARDS_PER_DEVICE * len(devices) ) cases = read_cases(args.miopen_file) if not cases: @@ -452,6 +477,10 @@ def main(argv: Optional[Sequence[str]] = None) -> int: # does no matching of its own. carried: List[dict] = [] if args.resume: + # Every row is carried, whatever its timing, so a resume measured the + # other way loses nothing; the sweeps skip only the cases their own + # timing has (case_results.measured_keys), and the merge takes only + # those rows. carried = previous_rows(log_dir) if not carried: print(f"[info] --resume: no rows in {log_dir}; running everything") @@ -492,7 +521,7 @@ def _report(index: int, device: str, rc: int, seconds: float) -> None: # The shards' own row files, re-read on each progress tick. A shard reports # nothing until it exits, so its rows are the only live view into it. row_files = [paths["rows"] for _index, paths in plan] - baseline = {row["key"] for row in carried} + baseline = {row["key"] for row in carried if case_results.row_timing(row) == timing} def _progress(shards_done: int, busy: int) -> None: print( @@ -536,6 +565,7 @@ def _progress(shards_done: int, busy: int) -> None: # The shards' rows are keyed by geometry; the case file is what puts # each row's MIOpenDriver line back on it. case_results.miopen_cmds_by_key(miopen_file=args.miopen_file), + timing=timing, ) if rows: out = Path(args.results or (log_dir / "results.jsonl")) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py index bad7972e3530..69f86d26b290 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_batched_timing.py @@ -32,8 +32,20 @@ def __init__(self, times, *, probe_error: float = 1.0): def install(self, monkeypatch) -> None: monkeypatch.setattr(batched_timing, "time_batch", self._time_batch) + monkeypatch.setattr(batched_timing, "time_batch_staged", self._time_staged) monkeypatch.setattr(batched_timing, "prime_device", self._prime) + def _time_staged(self, runs, *, warmup, iters, stream=0): + # The finalists' round, in the default (staged) protocol. + self.batches += 1 + out = [] + for run in runs: + ms = self.times[run()] + if ms is not None: + self.measured += 1 + out.append(ms) + return out + def _prime(self, run, **_kwargs) -> None: self.primed += 1 @@ -175,9 +187,7 @@ def test_scratch_is_released_for_everything_that_cannot_rank(monkeypatch): device.install(monkeypatch) dropped = [] - ranked = _rank( - device, len(times), finalists=2, chunk=3, on_drop=dropped.append - ) + ranked = _rank(device, len(times), finalists=2, chunk=3, on_drop=dropped.append) # Everything but the two finalists: the slow ones and the one that failed. assert sorted(dropped) == [0, 2, 3, 5] @@ -238,9 +248,7 @@ def counting(runs, **kwargs): monkeypatch.setattr(batched_timing, "time_batch", counting) - batched_timing.rank( - 5, device.make_run, warmup=3, iters=10, finalists=5, chunk=2 - ) + batched_timing.rank(5, device.make_run, warmup=3, iters=10, finalists=5, chunk=2) assert batches == [2, 2, 1] @@ -262,3 +270,47 @@ def on_probe(r): assert [i for i, _, _ in seen] == [0, 1, 2, 3, 4] assert all(probe is not None for _, probe, _ in seen) assert all(ms is None for _, _, ms in seen) + + +def test_staged_timing_changes_only_how_the_finalists_are_timed(monkeypatch): + """--timing staged (the default): the probe is the same batched block round, + and only the finalists go through the stalled single-call protocol.""" + times = [4.0, 1.0, 3.0, 2.0] + device = _FakeDevice(times) + device.install(monkeypatch) + staged = [] + assert batched_timing.DEFAULT_TIMING == "staged" + + def _staged(runs, *, warmup, iters, stream=0): + staged.append((len(runs), warmup, iters)) + # A stalled sample pays the event pair the block loop amortises away. + return [times[run()] + 0.5 for run in runs] + + monkeypatch.setattr(batched_timing, "time_batch_staged", _staged) + + ranked = _rank(device, len(times), finalists=2) + + assert device.probed == len(times) + assert device.measured == 0 + assert staged == [(2, 3, 10)] + assert {r.index: r.ms for r in ranked if r.is_finalist} == {1: 1.5, 3: 2.5} + + +def test_block_timing_never_touches_the_stall_gate(monkeypatch): + device = _FakeDevice([1.0, 2.0]) + device.install(monkeypatch) + + def _staged(*_args, **_kwargs): + raise AssertionError("block timing went through the staged path") + + monkeypatch.setattr(batched_timing, "time_batch_staged", _staged) + ranked = _rank(device, 2, finalists=2, timing="block") + assert device.measured == 2 + assert all(r.is_finalist for r in ranked) + + +def test_an_unknown_timing_is_refused(monkeypatch): + device = _FakeDevice([1.0]) + device.install(monkeypatch) + with pytest.raises(ValueError, match="timing"): + _rank(device, 1, timing="events") 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 ad825509efc4..41fd49e5b33d 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 @@ -204,3 +204,41 @@ def test_failed_kernels_are_not_probed(monkeypatch): stop.measure(fast, warmup=3, iters=10) assert stop.measure(wrong, warmup=3, iters=10, passed=False) == 100.0 assert stop.best_ms == 1.0 + + +def test_staged_timing_reports_the_median_of_stalled_samples(monkeypatch): + """--timing staged: the pruning probes stay block launches; only the + reported number is measured one stalled launch at a time.""" + fast, ok = object(), object() + block = _fake_timer(monkeypatch, {fast: 1.0, ok: 3.0}) + staged = [] + + def fake_staged(fn, *, warmup, iters, stream=0): + staged.append((fn, warmup, iters)) + base = {fast: 1.5, ok: 3.5}[fn] + return [base + 1.0, base, base - 0.1] + + monkeypatch.setattr(rt_mod, "time_launches_staged", fake_staged) + stop = EarlyStop(5.0, after=1, timing="staged") + assert stop.measure(fast, warmup=3, iters=10) == 1.5 + assert block == [] and staged == [(fast, 3, 10)] + # Engaged: the probe and the rest of the warmup are block, the loop staged. + assert stop.measure(ok, warmup=3, iters=10) == 3.5 + assert block == [(ok, 0, 1), (ok, 0, 2)] + assert staged[-1] == (ok, 0, 10) + + +def test_the_sweep_flag_picks_the_timing(monkeypatch): + import argparse + + from benchmarks.common import batched_timing + from kernels.common._conv_implicit_gemm_common import ConvProblem + + _reset_handoff(monkeypatch) + parser = argparse.ArgumentParser() + batched_timing.add_timing_arg(parser) + assert parser.parse_args([]).timing == "staged" + args = parser.parse_args(["--timing", "block"]) + args.early_stop, args.early_stop_after = 5.0, 1 + problem = ConvProblem(N=1, Hi=8, Wi=8, C=8, K=8, Y=3, X=3, pH=1, pW=1) + assert EarlyStop.for_case(args, problem, "fp16", "fwd").timing == "block" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py index 5e1e53f785a4..05f97618cbab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py @@ -200,7 +200,7 @@ def test_a_jit_run_is_refused_before_any_process_starts(tmp_path, capsys): # --------------------------------------------------------------------------- merge -def _row(engine, key, ms, kernel, tflops=100.0, runners_up=None): +def _row(engine, key, ms, kernel, tflops=100.0, runners_up=None, timing=None): row = { "key": key, "engine": engine, @@ -216,6 +216,8 @@ def _row(engine, key, ms, kernel, tflops=100.0, runners_up=None): } if runners_up: row["runners_up"] = list(runners_up) + if timing is not None: + row["timing"] = timing return row @@ -246,9 +248,12 @@ def test_the_same_case_measured_twice_keeps_the_faster(): [_row("direct", "k", 0.9, "slow"), _row("direct", "k", 0.4, "fast")] ) assert rows[0]["direct_kernel"] == "fast" - assert case_results.merge( - [_row("direct", "k", 0.4, "fast"), _row("direct", "k", 0.9, "slow")] - ) == rows + assert ( + case_results.merge( + [_row("direct", "k", 0.4, "fast"), _row("direct", "k", 0.9, "slow")] + ) + == rows + ) def test_the_runner_up_kernels_get_a_column_each(): @@ -284,9 +289,7 @@ def test_the_faster_row_brings_its_whole_ranking(): def test_merge_is_ordered_by_key(): - rows = case_results.merge( - [_row("direct", k, 1.0, "x") for k in ("c", "a", "b")] - ) + rows = case_results.merge([_row("direct", k, 1.0, "x") for k in ("c", "a", "b")]) assert [r["key"] for r in rows] == ["a", "b", "c"] @@ -421,7 +424,7 @@ def test_a_writer_without_a_path_writes_nothing(tmp_path): # --------------------------------------------------------------------------- end to end -_STUB = '''\ +_STUB = """\ import argparse, json, sys ap = argparse.ArgumentParser() ap.add_argument("--miopen-file") @@ -439,10 +442,10 @@ def test_a_writer_without_a_path_writes_nothing(tmp_path): "key": f"fwd/bf16/case{n}", "engine": engine, "case": n, "shape": f"s{n}", "dtype": "bf16", "direction": "fwd", "tflops": 100.0, "ms": ms, "gbps": 1.0, - "kernel": f"{engine}-kernel-{n}", "candidates": 3, + "kernel": f"{engine}-kernel-{n}", "candidates": 3, "timing": "staged", }) + "\\n") print("stub done") -''' +""" def test_a_whole_run_splits_schedules_and_merges(tmp_path, monkeypatch, capsys): @@ -460,11 +463,17 @@ def test_a_whole_run_splits_schedules_and_merges(tmp_path, monkeypatch, capsys): rc = run_conv_shards.main( [ - "--miopen-file", str(cases), - "--devices", "2", - "--shards", "4", - "--log-dir", str(logs), - "--", "--run-from-cache", "/nonexistent", + "--miopen-file", + str(cases), + "--devices", + "2", + "--shards", + "4", + "--log-dir", + str(logs), + "--", + "--run-from-cache", + "/nonexistent", ] ) assert rc == 0, capsys.readouterr() @@ -495,8 +504,17 @@ def test_a_run_whose_shards_measured_nothing_is_a_failure(tmp_path, monkeypatch) cases = tmp_path / "cases.txt" cases.write_text("convbfp16 -n 1\n") rc = run_conv_shards.main( - ["--miopen-file", str(cases), "--devices", "1", - "--log-dir", str(tmp_path / "logs"), "--", "--run-from-cache", "/c"] + [ + "--miopen-file", + str(cases), + "--devices", + "1", + "--log-dir", + str(tmp_path / "logs"), + "--", + "--run-from-cache", + "/c", + ] ) assert rc != 0 @@ -522,8 +540,16 @@ def test_resume_after_the_separator_is_caught_here(tmp_path, capsys): f = tmp_path / "cases.txt" f.write_text("convbfp16 -n 1\n") rc = run_conv_shards.main( - ["--miopen-file", str(f), "--devices", "1", "--", - "--run-from-cache", "/c", "--resume"] + [ + "--miopen-file", + str(f), + "--devices", + "1", + "--", + "--run-from-cache", + "/c", + "--resume", + ] ) assert rc == 2 assert "goes before --" in capsys.readouterr().err @@ -543,8 +569,14 @@ def test_resume_carries_the_earlier_rows_and_skips_their_cases( cases.write_text("".join(f"convbfp16 -n {i}\n" for i in range(6))) logs = tmp_path / "logs" runner = [ - "--miopen-file", str(cases), "--devices", "1", "--shards", "2", - "--log-dir", str(logs), + "--miopen-file", + str(cases), + "--devices", + "1", + "--shards", + "2", + "--log-dir", + str(logs), ] forwarded = ["--", "--run-from-cache", "/c"] argv = runner + forwarded @@ -573,7 +605,7 @@ def test_resume_carries_the_earlier_rows_and_skips_their_cases( assert not (set(ran) & first), f"re-measured {set(ran) & first}" -_RESUME_STUB = '''\ +_RESUME_STUB = """\ import argparse, json, os, pathlib, sys ap = argparse.ArgumentParser() ap.add_argument("--miopen-file") @@ -604,12 +636,12 @@ def test_resume_carries_the_earlier_rows_and_skips_their_cases( "key": key, "engine": engine, "case": n, "shape": f"s{n}", "dtype": "bf16", "direction": "fwd", "tflops": 1.0, "ms": 1.0 if engine == "direct" else 2.0, "gbps": 1.0, - "kernel": f"{engine}-{n}", "candidates": 1, + "kernel": f"{engine}-{n}", "candidates": 1, "timing": "staged", }) + "\\n") log = pathlib.Path(args.results).parent / "ran.json" prev = json.loads(log.read_text()) if log.is_file() else [] log.write_text(json.dumps(prev + ran)) -''' +""" # --------------------------------------------------------------------------- progress @@ -824,3 +856,78 @@ def test_a_different_shape_keys_differently(): assert case_key(a, "bf16", "fwd") != case_key(b, "bf16", "fwd") assert case_key(a, "bf16", "fwd") != case_key(a, "fp16", "fwd") assert case_key(a, "bf16", "fwd") != case_key(a, "bf16", "dgrad") + + +# --------------------------------------------------------------------------- timing +def test_a_row_without_timing_is_a_block_row(): + """Rows written before the field existed were all timed back to back.""" + assert case_results.row_timing(_row("direct", "k", 1.0, "d")) == "block" + assert case_results.row_timing(_row("direct", "k", 1.0, "d", timing="staged")) == ( + "staged" + ) + + +def test_the_writer_records_how_the_row_was_timed(tmp_path): + f = tmp_path / "rows.jsonl" + with case_results.ResultsWriter(str(f), "implicit", "staged") as writer: + writer.row( + key="k", + case=1, + shape="s", + dtype="bf16", + direction="fwd", + tflops=1.0, + ms=1.0, + gbps=1.0, + kernel="x", + candidates=1, + ) + assert [r["timing"] for r in case_results.read_rows(f)] == ["staged"] + with pytest.raises(ValueError, match="timing"): + case_results.ResultsWriter(None, "direct", "events") + + +def test_resume_skips_only_cases_timed_the_same_way(tmp_path): + """A case measured back to back is not measured for a staged run.""" + f = tmp_path / "rows.jsonl" + rows = [ + _row("direct", "old", 1.0, "d"), + _row("direct", "new", 1.0, "d", timing="staged"), + ] + f.write_text("".join(json.dumps(r) + "\n" for r in rows)) + assert case_results.measured_keys(f, "direct") == {"old"} + assert case_results.measured_keys(f, "direct", "staged") == {"new"} + + +def test_a_merge_takes_one_timing(): + """The faster-row rule must not pick the cheaper protocol over the faster + kernel: a staged row of a case is never replaced by its block row.""" + rows = [ + _row("direct", "k", 1.0, "block-kernel"), + _row("direct", "k", 1.2, "staged-kernel", timing="staged"), + ] + merged = case_results.merge(rows, timing="staged") + assert merged[0]["direct_kernel"] == "staged-kernel" + assert merged[0]["timing"] == "staged" + with pytest.raises(ValueError, match="different timings"): + case_results.merge(rows) + + +def test_a_merge_of_one_timing_needs_no_argument(): + merged = case_results.merge([_row("direct", "k", 1.0, "d", timing="staged")]) + assert merged[0]["timing"] == "staged" + assert case_results.merge([_row("direct", "k", 1.0, "d")])[0]["timing"] == "block" + + +@pytest.mark.parametrize( + "forwarded, timing", + [ + ([], "staged"), + (["--fast-rank"], "staged"), + (["--timing", "block"], "block"), + (["--fast-rank", "--timing=block"], "block"), + (["--timing", "block", "--timing", "staged"], "staged"), + ], +) +def test_the_runner_reads_the_timing_it_forwards(forwarded, timing): + assert run_conv_shards.forwarded_timing(forwarded) == timing diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py index 12c5c1425566..f6c415d1584a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py @@ -69,6 +69,10 @@ create HIP events; the launchers themselves are event-free so production calls stay overhead-free. + * :func:`time_launches_staged` -- one stalled + event pair per call (:class:`StallGate`), the + way hipDNN's benchmark times a call; for + numbers compared against hipDNN or MIOpen. The previous monolithic ``launch_torch_kernel`` is now a thin back-compat shim over :class:`KernelLauncher` + :func:`time_launches`; @@ -103,6 +107,7 @@ wait_stream_and_release, ) from .packing import pack_args +from .stall_gate import StallGate, StallTimeout, time_launches_staged from .torch_interop import ( TorchLaunchSummary, empty_workspace, @@ -121,6 +126,8 @@ "LaunchSummary", "PipelineLauncher", "Runtime", + "StallGate", + "StallTimeout", "TorchLaunchSummary", "WorkspacePool", "build_hsaco_from_llvm_ir", @@ -133,5 +140,6 @@ "resolve_stream", "synchronize_and_release", "time_launches", + "time_launches_staged", "wait_stream_and_release", ] 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 aea9559d7776..ca7b159b12d1 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 @@ -171,6 +171,20 @@ def _b(name: str, *argtypes, restype=ctypes.c_int) -> _LazyFn: ctypes.c_size_t, ctypes.c_void_p, ) +_hipExtMallocWithFlags = _b( + "hipExtMallocWithFlags", + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_size_t, + ctypes.c_uint, +) +_hipStreamWaitValue32 = _b( + "hipStreamWaitValue32", + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_uint, + ctypes.c_uint32, +) _hipDeviceSynchronize = _b("hipDeviceSynchronize") _hipStreamSynchronize = _b("hipStreamSynchronize", ctypes.c_void_p) _hipEventCreate = _b("hipEventCreate", ctypes.POINTER(_HipEventHandle)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/stall_gate.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/stall_gate.py new file mode 100644 index 000000000000..194d0a0af48c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/stall_gate.py @@ -0,0 +1,286 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host-released device stall, for timing one call the way hipDNN does. + +:func:`rocke.runtime.time_launches` times N calls back to back between one +pair of events and divides. hipDNN's benchmark (``dnn-benchmark``, "staged" +timing) times every call on its own instead: it parks a wait packet on the +stream, records the start event, enqueues the work, records the stop event and +only then releases the wait from the host. The event span is gap-free device +time -- the host's submission cost is outside it -- but it does include the +fixed device cost of a start/stop event pair around a single call, which a +back-to-back loop amortises away. A sweep that is compared against hipDNN or +MIOpen numbers therefore has to measure the same way, or its small kernels look +faster than they will in that comparison. + +:class:`StallGate` is the ctypes twin of ``hipdnn_data_sdk::utilities::StallGate`` +(``hipStreamWaitValue32`` on HIP signal memory, released by a host store, with a +watchdog), and :func:`time_launches_staged` is the twin of dnn-benchmark's +``StalledRegionTimer.measure`` loop. + +The watchdog +------------ +The release is a host action that runs after the measured call returns. A call +that blocks the host on the stalled stream (a fenced launch, a synchronous +``hipMemset``) would wait forever for a stream only the host can release, so a +watchdog thread writes the release itself after ``timeout_s``. That turns the +hang into one slow, invalid sample: :meth:`StallGate.timed_out` reports it and +:func:`time_launches_staged` raises :class:`StallTimeout` rather than return a +number that includes host time. +""" + +from __future__ import annotations + +import ctypes +import threading +import time +from typing import Callable, List, Optional + +from .hip_module import ( + HipError, + Runtime, + _check, + _hipDeviceGetAttribute, + _hipExtMallocWithFlags, + _hipFree, + _hipGetDevice, + _hipStreamWaitValue32, +) + +__all__ = ["StallGate", "StallTimeout", "default_gate", "time_launches_staged"] + +# hipDeviceAttributeCanUseStreamWaitValue (AMD-specific block of the +# hipDeviceAttribute_t enum), hipMallocSignalMemory, hipStreamWaitValueGte. +_HIP_ATTR_CAN_USE_STREAM_WAIT_VALUE = 10013 +_HIP_MALLOC_SIGNAL_MEMORY = 0x2 +_HIP_STREAM_WAIT_VALUE_GTE = 0x0 + +#: Host time allowed between arm and release, as in hipDNN's StallGate. +DEFAULT_TIMEOUT_S = 2.0 + + +class StallTimeout(RuntimeError): + """The watchdog, not the host, released the gate: the sample holds host time.""" + + +class StallGate: + """Hold every later item on a stream until the host releases it. + + Usage per sample: :meth:`arm`, record the start event, enqueue the work, + record the stop event, :meth:`release`, synchronize the stop event. The + stream must have drained past the previous wait before the next + :meth:`arm` -- synchronizing the stop event is what guarantees it. + + One gate arms one stream at a time. :meth:`close` (or leaving a ``with`` + block) releases any pending wait, stops the watchdog and frees the signal. + """ + + def __init__(self, timeout_s: float = DEFAULT_TIMEOUT_S) -> None: + self._timeout_s = float(timeout_s) if timeout_s > 0 else DEFAULT_TIMEOUT_S + self._signal = 0 + self._cv = threading.Condition() + self._armed = False + self._timed_out = False + self._stop = False + self._deadline = 0.0 + self._watchdog: Optional[threading.Thread] = None + + device = ctypes.c_int(0) + _check(_hipGetDevice(ctypes.byref(device)), "hipGetDevice") + can = ctypes.c_int(0) + _check( + _hipDeviceGetAttribute( + ctypes.byref(can), _HIP_ATTR_CAN_USE_STREAM_WAIT_VALUE, device.value + ), + "hipDeviceGetAttribute(CanUseStreamWaitValue)", + ) + if not can.value: + raise HipError( + f"device {device.value} does not support hipStreamWaitValue32; " + f"staged timing needs it" + ) + # Signal memory is an 8-byte HSA signal; the 32-bit wait reads its low word. + ptr = ctypes.c_void_p(0) + _check( + _hipExtMallocWithFlags(ctypes.byref(ptr), 8, _HIP_MALLOC_SIGNAL_MEMORY), + "hipExtMallocWithFlags(hipMallocSignalMemory)", + ) + self._signal = int(ptr.value or 0) + self._write(1) + + def _write(self, value: int) -> None: + ctypes.c_uint32.from_address(self._signal).value = value + + def _watch(self) -> None: + with self._cv: + while not self._stop: + if not self._armed: + self._cv.wait() + continue + left = self._deadline - time.monotonic() + if left > 0: + self._cv.wait(left) + continue + # The blocked host is waiting on this stream; this store both + # ends the stall and unblocks it. + self._write(1) + self._armed = False + self._timed_out = True + + def arm(self, stream: int) -> None: + """Enqueue a wait on ``stream`` that holds everything queued after it.""" + if not self._signal: + raise HipError("StallGate is closed") + with self._cv: + if self._armed: + raise HipError("StallGate is already armed") + self._timed_out = False + if self._watchdog is None: + self._watchdog = threading.Thread( + target=self._watch, name="rocke-stall-gate", daemon=True + ) + self._watchdog.start() + self._write(0) + status = _hipStreamWaitValue32( + ctypes.c_void_p(int(stream)), + ctypes.c_void_p(self._signal), + 1, + _HIP_STREAM_WAIT_VALUE_GTE, + 0xFFFFFFFF, + ) + if status != 0: + # No wait was enqueued to ever raise the signal; put it back. + self._write(1) + _check(status, "hipStreamWaitValue32") + with self._cv: + self._armed = True + self._deadline = time.monotonic() + self._timeout_s + self._cv.notify_all() + + def release(self) -> None: + """Let the stalled work run. A no-op when not armed.""" + with self._cv: + if not self._armed: + return + self._write(1) + self._armed = False + self._cv.notify_all() + + def timed_out(self) -> bool: + """Whether the watchdog, not :meth:`release`, ended the last :meth:`arm`.""" + with self._cv: + return self._timed_out + + def close(self) -> None: + """Release, stop the watchdog and free the signal. Idempotent.""" + with self._cv: + if self._signal: + self._write(1) + self._armed = False + self._stop = True + self._cv.notify_all() + if self._watchdog is not None: + self._watchdog.join() + self._watchdog = None + if self._signal: + # hipFree synchronizes the device first, so the released wait + # packet has retired before its signal goes away. + _hipFree(ctypes.c_void_p(self._signal)) + self._signal = 0 + + def __enter__(self) -> "StallGate": + return self + + def __exit__(self, *_exc) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: # noqa: BLE001, S110 - HIP may be gone at teardown + pass + + +_default_gate: Optional[StallGate] = None + + +def default_gate() -> StallGate: + """The process's shared gate, created on first use and closed at exit. + + A gate costs a signal allocation, a watchdog thread and, on close, a device + drain; a sweep that times thousands of kernels one at a time shares one + rather than paying that per kernel. A gate is bound to the device current + when it is made, which is the one device a sweep process uses. + """ + global _default_gate + if _default_gate is None: + import atexit + + _default_gate = StallGate() + atexit.register(_default_gate.close) + return _default_gate + + +def time_launches_staged( + fn: Callable[[], None], + *, + warmup: int, + iters: int, + stream: int = 0, + gate: Optional[StallGate] = None, + runtime: Optional[Runtime] = None, +) -> List[float]: + """Per-call device times of ``fn`` in milliseconds, one stalled sample each. + + Every warmup and timed call is its own sample: arm the gate, record the + start event, call ``fn``, record the stop event, release, wait on the stop + event. Warmup samples are discarded. The caller summarises the returned + list (dnn-benchmark reports the median). + + ``fn`` runs under :func:`~rocke.runtime.launcher.no_fence`, so a launcher + call inside it does not synchronize with the host. ``stream`` is resolved + as :func:`~rocke.runtime.time_launches` resolves it, and ``gate`` defaults + to :func:`default_gate`. Raises :class:`StallTimeout` when ``fn`` blocked + the host on the stalled stream (the watchdog released it); the device is + drained before any exception propagates. + """ + from .launcher import no_fence + from .torch_interop import resolve_stream + + rt = runtime if runtime is not None else Runtime() + stream = resolve_stream(stream) + g = default_gate() if gate is None else gate + start = rt.event() + stop = rt.event() + out: List[float] = [] + try: + rt.sync() + with no_fence(): + for i in range(int(warmup) + int(iters)): + g.arm(stream) + try: + start.record(stream=stream) + fn() + stop.record(stream=stream) + finally: + g.release() + stop.synchronize() + if g.timed_out(): + raise StallTimeout( + "the measured call blocked the host on the stalled " + "stream; the stall watchdog released it" + ) + if i >= warmup: + out.append(start.elapsed_to(stop)) + except Exception: + try: + rt.sync() + except Exception: # noqa: BLE001, S110 - the call's error is the one to see + pass + raise + finally: + start.destroy() + stop.destroy() + rt.wait_stream(stream) + return out diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_stall_gate.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_stall_gate.py new file mode 100644 index 000000000000..d9df13ea82b7 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_stall_gate.py @@ -0,0 +1,176 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Staged timing: one stalled event pair per call, as hipDNN's benchmark times. + +The fake-gate tests need no GPU: they pin the order of one sample (arm, start +event, the call, stop event, release, wait) and what happens when the call +raises or the watchdog had to release the gate. The last tests run on a GPU +when there is one: a real wait on a real stream, released by the host and, for +a call that blocks the host, by the watchdog. +""" + +from __future__ import annotations + +import os + +import pytest +from rocke.runtime.stall_gate import StallTimeout, time_launches_staged + +_STREAM = 5 +_GPU = pytest.mark.skipif(not os.path.exists("/dev/kfd"), reason="needs an AMD GPU") + + +class _FakeEvent: + def __init__(self, log, name): + self._log = log + self._name = name + + def record(self, stream): + self._log.append(f"record:{self._name}") + + def synchronize(self): + self._log.append(f"sync:{self._name}") + + def elapsed_to(self, other): + return 2.0 + + def destroy(self): + self._log.append(f"destroy:{self._name}") + + +class _FakeRuntime: + def __init__(self): + self.log: list[str] = [] + self._events = 0 + + def sync(self): + self.log.append("device_sync") + + def wait_stream(self, stream): + self.log.append("wait_stream") + + def event(self): + self._events += 1 + return _FakeEvent(self.log, "start" if self._events == 1 else "stop") + + +class _FakeGate: + def __init__(self, log, timeout_on=None): + self._log = log + self._arms = 0 + self._timeout_on = timeout_on + + def arm(self, stream): + self._arms += 1 + self._log.append("arm") + + def release(self): + self._log.append("release") + + def timed_out(self): + return self._arms == self._timeout_on + + def close(self): + self._log.append("close") + + +def _fn(log, raise_on=None): + calls = [0] + + def fn(): + calls[0] += 1 + log.append("fn") + if calls[0] == raise_on: + raise RuntimeError("launch failed") + + return fn + + +def test_every_call_is_its_own_stalled_sample(): + rt = _FakeRuntime() + gate = _FakeGate(rt.log) + out = time_launches_staged( + _fn(rt.log), warmup=1, iters=2, stream=_STREAM, gate=gate, runtime=rt + ) + one = ["arm", "record:start", "fn", "record:stop", "release", "sync:stop"] + assert rt.log == ["device_sync"] + one * 3 + [ + "destroy:start", + "destroy:stop", + "wait_stream", + ] + # The warmup sample is discarded; a caller's gate is not closed. + assert out == [2.0, 2.0] + + +def test_a_raise_releases_and_drains_before_it_propagates(): + rt = _FakeRuntime() + gate = _FakeGate(rt.log) + with pytest.raises(RuntimeError, match="launch failed"): + time_launches_staged( + _fn(rt.log, raise_on=2), warmup=0, iters=3, gate=gate, runtime=rt + ) + tail = rt.log[rt.log.index("fn", rt.log.index("fn") + 1) :] + assert tail[:3] == ["fn", "release", "device_sync"] + assert "close" not in rt.log + + +def test_a_watchdog_release_invalidates_the_measurement(): + rt = _FakeRuntime() + gate = _FakeGate(rt.log, timeout_on=2) + with pytest.raises(StallTimeout): + time_launches_staged(_fn(rt.log), warmup=0, iters=3, gate=gate, runtime=rt) + assert rt.log.count("fn") == 2 + assert rt.log[-3:] == ["device_sync", "destroy:start", "destroy:stop"] + + +def _busy_launch(): + """A launch of a one-instruction kernel, or skip when none can be built.""" + import struct + + from rocke.runtime.hip_module import Runtime, get_device_arch + + try: + from rocke.runtime.comgr import build_hsaco_from_llvm_ir + except Exception as e: # noqa: BLE001 # pragma: no cover - environment dependent + pytest.skip(f"comgr unavailable: {e}") + arch = get_device_arch() + if not arch: + pytest.skip("no HIP device") + ir = ( + 'target triple = "amdgcn-amd-amdhsa"\n' + "define amdgpu_kernel void @nop(ptr addrspace(1) %p) {\n" + " ret void\n" + "}\n" + ) + try: + hsaco, _timings = build_hsaco_from_llvm_ir(ir, isa=f"amdgcn-amd-amdhsa--{arch}") + except Exception as e: # noqa: BLE001 # pragma: no cover - environment dependent + pytest.skip(f"cannot build a kernel here: {e}") + rt = Runtime() + fn = rt.load_module(hsaco).get_function("nop") + return rt, rt.prepare_launch(fn, (1, 1, 1), (64, 1, 1), struct.pack(" 0 for ms in out) + + +@_GPU +def test_real_watchdog_unblocks_a_call_that_waits_on_the_stalled_stream(): + from rocke.runtime.stall_gate import StallGate + + rt, run = _busy_launch() + + def blocking(): + run() + rt.stream_sync(0) + + with StallGate(timeout_s=0.2) as gate: + with pytest.raises(StallTimeout): + time_launches_staged(blocking, warmup=0, iters=1, gate=gate) + # The gate stays usable after a timeout. + assert len(time_launches_staged(run, warmup=0, iters=2, gate=gate)) == 2 From 5a1adb6f2b36a158f344b3b3c6dd2a3d34d3076e Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 28 Oct 2026 18:30:28 -0500 Subject: [PATCH 6/9] fix(rocke): address AOT sweep review findings - case_results.py: ResultsWriter no longer truncates --results in place when it names the same file as --skip-measured; measured_keys is read first and the writer appends instead of overwriting. - module_cache.py: an over-budget launcher is now released on the next miss (or at clear()) instead of leaking a HIP module per miss for the rest of the run. - benchmark_conv_compare.py: _run_script returns the child's exit status; main() now fails the shard when a sweep dies instead of always exiting 0. - CMakeLists.txt: wire the new CPU-only library tests into the conv pytest list. - batched_timing.py: time_batch_staged checks hipStreamWaitValue32 support once up front and falls back to block timing with a warning, instead of silently reporting every finalist as unranked. Co-Authored-By: Claude Sonnet 5 --- .../benchmarks/common/batched_timing.py | 22 ++++++++- .../common/benchmark_conv_compare.py | 47 ++++++++++++++----- .../common/benchmark_direct_conv.py | 16 +++++-- .../common/benchmark_implicit_gemm_conv.py | 16 +++++-- .../library/benchmarks/common/case_results.py | 27 ++++++++++- .../library/benchmarks/common/module_cache.py | 15 +++++- .../rocke/platform/CMakeLists.txt | 7 ++- 7 files changed, 125 insertions(+), 25 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py index 6450620eb39b..834fb6b96b42 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py @@ -296,11 +296,31 @@ def time_batch_staged( blocked the host on the stalled stream (the watchdog had to release it, so the sample holds host time); the second is reported, as it means the bound launch synchronizes when it should not. + + The gate itself (``hipStreamWaitValue32`` support) is checked once, before + any candidate is timed: a device that lacks it would otherwise fail every + candidate the same way, one caught exception at a time, and the case would + report no winner with nothing in the log to say why. Falls back to + :func:`time_batch` (block timing) with a warning rather than doing that. """ from statistics import median from rocke.runtime.hip_module import HipError - from rocke.runtime.stall_gate import StallTimeout, time_launches_staged + from rocke.runtime.stall_gate import ( + StallTimeout, + default_gate, + time_launches_staged, + ) + + try: + default_gate() + except HipError as e: + print( + f" [staged] {e}; this device cannot be timed with --timing staged, " + f"falling back to --timing block for this run", + flush=True, + ) + return time_batch(runs, warmup=warmup, iters=iters, stream=stream) out: List[Optional[float]] = [] for i, run in enumerate(runs): 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 551f22b2dc64..9de3ffa424bb 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 @@ -99,11 +99,14 @@ def _tee(stream, sink, lines: list) -> None: 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, output). +) -> tuple[float | None, str, str, int]: + """Run *script* with *extra_args*; return (best_tflops, kernel_name, output, + returncode). The script's stdout and stderr are echoed live, as it prints them; ``output`` - is the captured stdout followed by stderr.""" + is the captured stdout followed by stderr. ``returncode`` is the child's exit + status, negative when it was killed by a signal (including the timeout + above); the caller decides what a nonzero status means for the shard.""" import os import signal import threading @@ -153,10 +156,23 @@ def _run_script( raw_out = "".join(out_lines) output = raw_out + "".join(err_lines) + returncode = popen.returncode m = _BEST_RE.search(raw_out) if m: - return float(m.group(1)), m.group(2).strip(), output - return None, "", output + return float(m.group(1)), m.group(2).strip(), output, returncode + return None, "", output, returncode + + +def _sweep_failure_message(script: Path, returncode: int) -> str: + """One line for a sweep that did not exit 0, for the shard's stderr. + + ``returncode`` is negative when the child was killed by a signal (the + timeout path above among them); that case names the signal rather than + the raw negative status, which is a value nobody reading the log expects. + """ + if returncode < 0: + return f"[error] {script.name} killed by signal {-returncode}" + return f"[error] {script.name} exited {returncode}" def _cpg_valid_for_direct(C: int, K: int, groups: int) -> tuple[bool, str]: @@ -729,23 +745,30 @@ def main() -> int: direct_args += ["--skip-measured", args.skip_measured] implicit_args += ["--skip-measured", args.skip_measured] + overall_rc = 0 + if not args.skip_direct: - tflops, name, _ = _run_script( - _SCRIPT_DIR / "benchmark_direct_conv.py", + script = _SCRIPT_DIR / "benchmark_direct_conv.py" + tflops, name, _, rc = _run_script( + script, 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 rc != 0: + overall_rc = 1 + print(_sweep_failure_message(script, rc), file=sys.stderr) 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, - ) + script = _SCRIPT_DIR / "benchmark_implicit_gemm_conv.py" + tflops, name, output, rc = _run_script(script, implicit_args) implicit_result = (tflops, name) + if rc != 0: + overall_rc = 1 + print(_sweep_failure_message(script, rc), file=sys.stderr) m = _ALL_STOPPED_RE.search(output) if tflops is None and m: implicit_note = m.group(0) @@ -761,7 +784,7 @@ def main() -> int: ) _print_summary(direct_result, implicit_result, implicit_note) - return 0 + return overall_rc if __name__ == "__main__": 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 af878d829939..66d745257dfd 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 @@ -47,6 +47,7 @@ add_results_arg, add_skip_measured_arg, measured_keys, + same_results_file, ) from benchmarks.common.early_stop import EarlyStop, add_early_stop_arg, case_key from benchmarks.common.module_cache import add_module_cache_arg @@ -2468,12 +2469,19 @@ def _u8(t): # The binaries are shape-generic, so their modules outlive a case. modules = ModuleCache(args.module_cache_bytes) - # One row per case, for the comparison driver and the shard runner to - # merge. Does nothing without --results. - results_out = ResultsWriter(args.results, "direct", args.timing) # --skip-measured: cases a previous run of this sweep already recorded, - # timed the way this run times them. + # timed the way this run times them. Read before opening --results below: + # when the two flags name the same file, the writer must not truncate it + # out from under this read. done = measured_keys(args.skip_measured, "direct", args.timing) + # One row per case, for the comparison driver and the shard runner to + # merge. Does nothing without --results. + results_out = ResultsWriter( + args.results, + "direct", + args.timing, + append=same_results_file(args.results, args.skip_measured), + ) n_skipped = 0 from kernels.conv.direct_aot_jobs import DIRECT_DTYPES 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 b21a4867e952..f3db3d60c5d8 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 @@ -83,6 +83,7 @@ add_results_arg, add_skip_measured_arg, measured_keys, + same_results_file, ) from benchmarks.common.early_stop import ( # noqa: E402 EarlyStop, @@ -4675,12 +4676,19 @@ def _run_from_cache(args, arch, target, cases, cache, directions) -> int: # A cached binary is shape-generic: its module is loaded on the first case # that offers it and reused by every later one. modules = ModuleCache(args.module_cache_bytes) - # One row per case, for the comparison driver and the shard runner to - # merge. Does nothing without --results. - results_out = ResultsWriter(args.results, "implicit", args.timing) # --skip-measured: cases a previous run of this sweep already recorded, - # timed the way this run times them. + # timed the way this run times them. Read before opening --results below: + # when the two flags name the same file, the writer must not truncate it + # out from under this read. done = measured_keys(args.skip_measured, "implicit", args.timing) + # One row per case, for the comparison driver and the shard runner to + # merge. Does nothing without --results. + results_out = ResultsWriter( + args.results, + "implicit", + args.timing, + append=same_results_file(args.results, args.skip_measured), + ) n_skipped = 0 def _u8(t): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py index c9c52258c5e0..b012d6f21c1f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py @@ -159,6 +159,21 @@ def add_skip_measured_arg(parser) -> None: ) +def same_results_file(results_path, skip_measured_path) -> bool: + """Whether ``--results`` and ``--skip-measured`` name the same file. + + Call before constructing :class:`ResultsWriter`: a sweep resuming into its + own results file must open it for append, not truncate it out from under + the skip set it was just read for. + """ + if not results_path or not skip_measured_path: + return False + try: + return Path(results_path).resolve() == Path(skip_measured_path).resolve() + except OSError: + return Path(results_path) == Path(skip_measured_path) + + def row_timing(row: dict) -> str: """How ``row``'s time was measured; rows older than the field are block.""" return row.get("timing") or BLOCK_TIMING @@ -192,7 +207,12 @@ class ResultsWriter: """ def __init__( - self, path: Optional[str], engine: str, timing: str = BLOCK_TIMING + self, + path: Optional[str], + engine: str, + timing: str = BLOCK_TIMING, + *, + append: bool = False, ) -> None: if engine not in ENGINES: raise ValueError(f"engine must be one of {ENGINES}, got {engine!r}") @@ -205,7 +225,10 @@ def __init__( self._n = 0 if self._path is not None: self._path.parent.mkdir(parents=True, exist_ok=True) - self._fh = self._path.open("w", encoding="utf-8") + # ``append``: ``--results`` and ``--skip-measured`` name the same + # file, so the rows ``measured_keys`` already read for this run's + # skip set are the run's own history, not scratch to discard. + self._fh = self._path.open("a" if append else "w", encoding="utf-8") @property def n_rows(self) -> int: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py index 1b1d4b623abc..fc543fbf5ecd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py @@ -75,6 +75,12 @@ def __init__(self, budget_bytes: int = DEFAULT_BUDGET_BYTES) -> None: self._hits = 0 self._misses = 0 self._full = False + # The one over-budget launcher handed out most recently, kept only + # long enough to unload it on the next miss -- by then its launch has + # drained, since callers gate/time each candidate before moving on. + # Without this a long sweep past the budget leaks one HIP module per + # miss for the rest of the run. + self._unadmitted: Any = None # -- query ------------------------------------------------------------ @@ -136,6 +142,9 @@ def launcher( return held self._misses += 1 + if key is not None and self._unadmitted is not None: + self._unadmitted.unload() + self._unadmitted = None hsaco = Path(hsaco_path).read_bytes() launcher = KernelLauncher( hsaco=hsaco, kernel_name=kernel_name, signature=signature @@ -146,6 +155,7 @@ def launcher( self._bytes += len(hsaco) else: self._full = True + self._unadmitted = launcher return launcher def clear(self) -> None: @@ -155,7 +165,7 @@ def clear(self) -> None: that the launches it handed out have finished, and unloading a module under a kernel in flight faults. """ - if not self._held: + if not self._held and self._unadmitted is None: return from rocke.runtime import synchronize_and_release @@ -163,6 +173,9 @@ def clear(self) -> None: for launcher in self._held.values(): launcher.unload() self._held.clear() + if self._unadmitted is not None: + self._unadmitted.unload() + self._unadmitted = None self._bytes = 0 self._full = False diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 8b6f2a73293c..d0a0cafb8719 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -549,7 +549,12 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_direct_kernel_cache.py" "tests/library/tests/test_conv_gradient_contract.py" "tests/library/tests/test_builder_signature_contract.py" - "tests/library/tests/test_library_layering.py") + "tests/library/tests/test_library_layering.py" + "tests/library/tests/test_aot_launch_zeroing.py" + "tests/library/tests/test_batched_timing.py" + "tests/library/tests/test_capability_keys.py" + "tests/library/tests/test_module_cache.py" + "tests/library/tests/test_run_conv_shards.py") set(_ROCKE_LIBRARY_CONV_TESTS "") foreach(_rocke_conv_test ${_ROCKE_LIBRARY_CONV_TEST_RELPATHS}) string(APPEND _ROCKE_LIBRARY_CONV_TESTS From 259d0da7180c351831726e5a1070531d505507b7 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Thu, 29 Oct 2026 08:13:30 -0500 Subject: [PATCH 7/9] fix(rocke): update conv test doubles for the entry_meta/extras APIs plan_for() now calls cache.entry_meta() (a meta dict) instead of cache.get() (a (bytes, meta) tuple), and _streamk_launch() now returns an extras dict plus rt.memset_async() instead of a free() closure and rt.memset(). The test fixtures in test_grouped_conv_dispatch.py and test_streamk_sweep_launch.py still matched the old contracts, which is what failed rocke_library_conv_pytest on CI. --- .../dispatch/test_grouped_conv_dispatch.py | 4 ++-- .../tests/test_streamk_sweep_launch.py | 19 +++++++++++-------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_dispatch.py index 22e46cb867df..16eba1d76754 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_dispatch.py @@ -383,9 +383,9 @@ class _HelperCache: def __init__(self): self.looked_up: list[KernelIdentity] = [] - def get(self, identity): + def entry_meta(self, identity): self.looked_up.append(identity) - return b"", {"kernel_name": "helper"} + return {"kernel_name": "helper"} def hsaco_path(self, identity): return Path("unused.hsaco") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py index 15bc917bd8c6..b655dc41dab7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py @@ -85,6 +85,9 @@ def free(self, ptr): def memset(self, ptr, value, nbytes): self.memsets.append((ptr, value, nbytes)) + def memset_async(self, ptr, value, nbytes, stream): + self.memsets.append((ptr, value, nbytes)) + class _RecordingLauncher: def __init__(self): @@ -102,7 +105,7 @@ def _launch_host(spec, arch): stage2 = _RecordingLauncher() # Pre-seed the Stage 2 memo so the workspace reduction needs no compiler. stage2_memo = {(spec.data.dtype_d, spec.ws_replicas): stage2} - run, free, pool = _streamk_launch( + run, extras, pool = _streamk_launch( rt=rt, launcher=launcher, spec=spec, @@ -114,7 +117,7 @@ def _launch_host(spec, arch): KernelLauncher=None, LaunchConfig=LaunchConfig, ) - return rt, launcher, stage2, run, free, pool + return rt, launcher, stage2, run, extras, pool class TestStreamKSweepLaunchHost(unittest.TestCase): @@ -140,7 +143,7 @@ def test_all_dp_partition_launches(self): spec = _spec( _ALL_DP, arch=self.ARCH, reduction=reduction, dtype_d=dtype_d ) - rt, launcher, stage2, run, free, _ = _launch_host(spec, self.ARCH) + rt, launcher, stage2, run, extras, _ = _launch_host(spec, self.ARCH) # One (stub) allocation, never a zero-byte one. self.assertEqual(len(rt.allocs), 1) (ws,) = rt.allocs @@ -152,7 +155,7 @@ def test_all_dp_partition_launches(self): self.assertEqual(values["sk_flags"], ws) self.assertEqual(values["sk_partials"], ws) self.assertEqual(len(stage2.calls), reduction == "workspace") - free() + rt.free(extras["ws_ptr"]) self.assertEqual(rt.freed, [ws]) def test_split_partition_lays_out_flags_then_partials(self): @@ -167,7 +170,7 @@ def test_split_partition_lays_out_flags_then_partials(self): spec, arch=self.ARCH ) self.assertGreater(flags_bytes, 0) - rt, launcher, _, run, free, _ = _launch_host(spec, self.ARCH) + rt, launcher, _, run, extras, _ = _launch_host(spec, self.ARCH) (ws,) = rt.allocs self.assertEqual(rt.allocs[ws], flags_bytes + partials_bytes) run() @@ -176,7 +179,7 @@ def test_split_partition_lays_out_flags_then_partials(self): self.assertEqual(values["sk_partials"], ws + flags_bytes) # linear/tree reset only the flag table before each launch. self.assertEqual(rt.memsets, [(ws, 0, flags_bytes)]) - free() + rt.free(extras["ws_ptr"]) self.assertEqual(rt.freed, [ws]) @@ -240,7 +243,7 @@ def u8(t): rt.memcpy_h2d(bufs[1], u8(X), X.nbytes) # Poison dW: a reduction that leaves it alone must overwrite it. rt.memset(bufs[2], 0xFF if reduction != "atomic" else 0, dW.nbytes) - run, free, _ = _streamk_launch( + run, extras, _ = _streamk_launch( rt=rt, launcher=launcher, spec=spec, @@ -256,7 +259,7 @@ def u8(t): run() synchronize_and_release(0) finally: - free() + rt.free(extras["ws_ptr"]) rt.memcpy_d2h(u8(dW), bufs[2], dW.nbytes) finally: for b in bufs: From 45484de09731353473a0f582de5149cca08576ca Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Thu, 29 Oct 2026 08:19:52 -0500 Subject: [PATCH 8/9] style(rocke): black-format files brought in by the develop merge The merge commit only ran pre-commit on conflicted files, so these two black violations from develop slipped through. Co-Authored-By: Claude Sonnet 5 --- .../rocke/library/benchmarks/common/kernel_cache.py | 3 +-- .../rocke/library/tests/test_capability_keys.py | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) 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 c2277a1acea8..8c311a041d1f 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 @@ -456,8 +456,7 @@ def _remember(self, key: tuple, out: list) -> None: self._compatible[key] = out self._compatible_rows += len(out) while ( - self._compatible_rows > _COMPATIBLE_MEMO_ROWS - and len(self._compatible) > 1 + self._compatible_rows > _COMPATIBLE_MEMO_ROWS and len(self._compatible) > 1 ): _, dropped = self._compatible.popitem(last=False) self._compatible_rows -= len(dropped) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py index cdcce6865e41..cb05952ca8ab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_capability_keys.py @@ -171,9 +171,7 @@ def _direct_problems() -> list: return [ DirectConvProblem(N=2, H=28, W=28, groups=1, cpg=64, kpg=64), DirectConvProblem(N=1, H=56, W=56, groups=16, cpg=16, kpg=16), - DirectConvProblem( - N=4, H=14, W=14, groups=128, cpg=1, kpg=1, KH=3, KW=3, PAD=1 - ), + DirectConvProblem(N=4, H=14, W=14, groups=128, cpg=1, kpg=1, KH=3, KW=3, PAD=1), DirectConvProblem( N=2, H=32, W=32, groups=1, cpg=32, kpg=64, KH=1, KW=1, PAD=0, stride=2 ), From 2bb985a3e28740d929d03470307476ec0c78f9ef Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Thu, 29 Oct 2026 17:29:26 -0500 Subject: [PATCH 9/9] fix(rocke): address PR review findings on conv sweep benchmarks - module_cache: keep over-budget launchers alive until end_case() confirms the case's launches have drained, instead of unloading on the next miss (--fast-rank and multi-step direct plans could race an in-flight launch) - benchmark_direct_conv / case_results: exclude verification-failed rows from winner ranking and merge, so a fast-but-wrong kernel cannot become Best - batched_timing / early_stop: raise on unsupported staged-timing gate instead of silently falling back to block timing under a "staged" label - benchmark_implicit_gemm_conv: convert the remaining JIT zeroing closures to memset_async (matching the AOT path) and fix the stream-K caller that still unpacked _streamk_launch's old 3-tuple return - run_conv_shards: handle --script compare's already-merged shard rows in the final merge and in --resume's --skip-measured handling - benchmark_conv_compare: persist per-engine shard rows beside --results instead of an ephemeral tempdir, so a killed run leaves them discoverable Co-Authored-By: Claude Sonnet 5 --- .../benchmarks/common/batched_timing.py | 17 ++--- .../common/benchmark_conv_compare.py | 27 ++++++-- .../common/benchmark_direct_conv.py | 18 ++++- .../common/benchmark_implicit_gemm_conv.py | 33 +++++++--- .../library/benchmarks/common/case_results.py | 32 ++++++++- .../library/benchmarks/common/early_stop.py | 17 ++++- .../library/benchmarks/common/module_cache.py | 43 +++++++----- .../benchmarks/common/run_conv_shards.py | 51 +++++++++----- .../rocke/library/tests/test_module_cache.py | 27 ++++++++ .../library/tests/test_run_conv_shards.py | 66 ++++++++++++------- 10 files changed, 245 insertions(+), 86 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py index 834fb6b96b42..e47e38aa40bd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/batched_timing.py @@ -300,8 +300,11 @@ def time_batch_staged( The gate itself (``hipStreamWaitValue32`` support) is checked once, before any candidate is timed: a device that lacks it would otherwise fail every candidate the same way, one caught exception at a time, and the case would - report no winner with nothing in the log to say why. Falls back to - :func:`time_batch` (block timing) with a warning rather than doing that. + report no winner with nothing in the log to say why. Raises rather than + silently falling back to :func:`time_batch` (block timing): a silent + fallback would still report these rows as ``timing: staged``, and resume + or merge would then compare block averages against real staged medians -- + the mix :data:`TIMINGS`-by-row is meant to keep apart. """ from statistics import median @@ -315,12 +318,10 @@ def time_batch_staged( try: default_gate() except HipError as e: - print( - f" [staged] {e}; this device cannot be timed with --timing staged, " - f"falling back to --timing block for this run", - flush=True, - ) - return time_batch(runs, warmup=warmup, iters=iters, stream=stream) + raise HipError( + f"{e}; this device cannot be timed with --timing staged, rerun " + f"with --timing block" + ) from e out: List[Optional[float]] = [] for i, run in enumerate(runs): 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 9de3ffa424bb..cfa5c526a9b2 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 @@ -724,17 +724,32 @@ def main() -> int: with tempfile.TemporaryDirectory() as tmp: # Direct conv's best time per case, handed to the implicit-GEMM sweep - # as its starting early-stop bound. + # as its starting early-stop bound. Scratch, not an artifact anyone + # reads back: a plain tempdir is fine for it. 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] - # Each sweep writes its own per-case rows; they are joined below on - # the case key both of them record. The files are the children's, not - # the user's: --results names the merged table. - rows_direct = str(Path(tmp) / "direct_rows.jsonl") - rows_implicit = str(Path(tmp) / "implicit_rows.jsonl") + # Each sweep writes its own per-case rows, flushed per case; they are + # joined below on the case key both of them record. The files are the + # children's, not the user's: --results names the merged table. With + # --results given, they are named from it and left beside it rather + # than under the tempdir this ``with`` block deletes on exit, so a run + # killed mid-sweep -- a sharded run's SIGKILL among them -- leaves the + # completed cases' rows somewhere the caller can find and resume from, + # not inside a tempdir whose name died with the process. + if args.results: + results_path = Path(args.results) + rows_direct = str( + results_path.with_name(results_path.stem + ".direct.jsonl") + ) + rows_implicit = str( + results_path.with_name(results_path.stem + ".implicit.jsonl") + ) + else: + rows_direct = str(Path(tmp) / "direct_rows.jsonl") + rows_implicit = str(Path(tmp) / "implicit_rows.jsonl") if args.results: direct_args += ["--results", rows_direct] implicit_args += ["--results", rows_implicit] 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 a114678a28d1..fb9623b74347 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 @@ -2279,15 +2279,23 @@ def verify(plan, run) -> bool: return passed def row(ms, passed, label, idx) -> None: - """Append one measured plan's row and print its status line.""" - results.append((flop / ms / 1e9, bytes_xfer / ms / 1e6, ms, passed, label)) + """Print one measured plan's status line, and rank it unless it failed. + + A plan that failed verification is still printed (its FAIL belongs in + the log), but kept out of ``results`` -- the list the ranking and the + recorded winner are built from -- so a fast wrong answer cannot become + ``Best:``. + """ + tflops = flop / ms / 1e9 + if passed is not False: + results.append((tflops, bytes_xfer / ms / 1e6, ms, passed, label)) # One line per measured launch, as the implicit-GEMM sweep prints, so a # long run shows its progress and partial results rather than only the # ranked table at the end. verdict = "" if passed is None else (" PASS" if passed else " FAIL") print( f" [{idx:>{width}}/{n_plans}] " - f"{results[-1][0]:7.1f} TFLOPS {ms:8.3f} ms{verdict} {label}", + f"{tflops:7.1f} TFLOPS {ms:8.3f} ms{verdict} {label}", flush=True, ) @@ -2675,6 +2683,10 @@ def _u8(t): per_chunk=per_chunk, ) overall_rc = overall_rc or case_rc + # Safe only here: _measure_plans has drained the device by the + # time it returns, so no launch of this case can still be using a + # module end_case unloads. + modules.end_case() for ptr in ptrs.values(): rt.free(ptr) 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 f3db3d60c5d8..bd82c2a38d1c 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 @@ -3760,11 +3760,14 @@ def _degrees(pipeline: str, async_dma: bool) -> tuple: return 1, [] if _is_atomic_launch: + from rocke.runtime import resolve_stream + _vals_snap = dict(values) _cfg_snap = cfg + _resolved_stream = resolve_stream(stream) - def _launch_spk(_v=_vals_snap, _c=_cfg_snap): - rt.memset(dW_dev, 0, dW_t.nbytes) + def _launch_spk(_v=_vals_snap, _c=_cfg_snap, _s=_resolved_stream): + rt.memset_async(dW_dev, 0, dW_t.nbytes, _s) launcher(_v, config=_c) timed_fn = _launch_spk @@ -3982,6 +3985,9 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): s1_cfg = LaunchConfig( grid=s1_grid, block=(spec.block_size, 1, 1), stream=0 ) + from rocke.runtime import resolve_stream + + _2s_stream = resolve_stream(s1_cfg.stream) def _launch_two_stage( _s1=s1_launcher, @@ -3992,12 +3998,14 @@ def _launch_two_stage( _c2=s2_cfg, _ws=ws_dev, _ws_nb=ws_nbytes, + _s=_2s_stream, ): # Stage 1 accumulates into the scratch, so it has to start from # zero; dW is a plain store target for Stage 2 and needs no - # clearing, but zeroing it keeps a failed launch visible. - rt2.memset(_ws, 0, _ws_nb) - rt2.memset(dW_dev2, 0, dW_t.nbytes) + # clearing, but zeroing it keeps a failed launch visible. Queued + # on the launch stream, not drained for (see _streamk_launch). + rt2.memset_async(_ws, 0, _ws_nb, _s) + rt2.memset_async(dW_dev2, 0, dW_t.nbytes, _s) _s1(_v1, config=_c1) _s2(_v2, config=_c2) @@ -4121,7 +4129,7 @@ def _launch_two_stage( kernel_name=artifact.kernel_name, signature=wgrad_streamk_signature(spec), ) - run, free_ws, pool = _streamk_launch( + run, sk_extras, pool = _streamk_launch( rt=rt3, launcher=launcher, spec=spec, @@ -4168,7 +4176,7 @@ def _launch_two_stage( continue synchronize_and_release(0) finally: - free_ws() + rt3.free(sk_extras["ws_ptr"]) cur_tflops = (flop / ms) * 1e-9 cur_gbps = (bytes_xfer / ms) * 1e-6 @@ -4495,9 +4503,12 @@ def _make(*shape): return 1 if _zero_init: + from rocke.runtime import resolve_stream + + _atomic_stream = resolve_stream(stream) - def _launch_atomic(): - rt.memset(dX_dev, 0, dX_t.nbytes) + def _launch_atomic(_s=_atomic_stream): + rt.memset_async(dX_dev, 0, dX_t.nbytes, _s) launcher(values, config=cfg) timed_fn = _launch_atomic @@ -5125,6 +5136,10 @@ def _u8(t): _free_extras(rt, sk_live) synchronize_and_release() + # Safe only here: every candidate of this case has drained by now + # (the device sync above), so no launch can still be using a + # module end_case unloads. + modules.end_case() _free_extras(rt, live_extras) for buf in (A_dev, B_dev, D_dev): rt.free(buf) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py index b012d6f21c1f..6bb4b9b2a573 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/case_results.py @@ -197,6 +197,31 @@ def measured_keys(path, engine: str, timing: str = BLOCK_TIMING) -> set: } +def expand_merged_for_skip(rows: Iterable[dict]) -> List[dict]: + """Per-engine skip rows from already-merged (:data:`MERGED_FIELDS`) rows. + + A merged row carries both engines' times under ``{engine}_ms`` and no + ``engine`` field of its own, so it is not what :func:`measured_keys` (and + the sweeps' own ``--skip-measured`` reader) looks for. One output row per + engine side the merged row actually has a time for -- :func:`merge` leaves + a side ``None`` when that engine never reached the case, and that side is + not "measured" to skip. + """ + out: List[dict] = [] + for row in rows: + key = row.get("key") + if not key: + continue + for engine in ENGINES: + ms = row.get(f"{engine}_ms") + if ms is None: + continue + out.append( + {"key": key, "engine": engine, "ms": ms, "timing": row.get("timing")} + ) + return out + + class ResultsWriter: """Appends one row per case to a JSON Lines file. @@ -471,7 +496,12 @@ def merge( for row in rows: key = row.get("key") engine = row.get("engine") - if not key or engine not in ENGINES or row_timing(row) != timing: + if ( + not key + or engine not in ENGINES + or row_timing(row) != timing + or row.get("passed") is False + ): continue merged = by_key.setdefault( key, 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 c3312f25b50a..8bb28b86611f 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 @@ -176,6 +176,10 @@ def __init__( # (HOPELESS_MARGIN past the bound) rather than the full warmup. self.last_hopeless = False self.record_key = record_key + # Checked once, lazily, so a device without hipStreamWaitValue32 fails + # with one clear error instead of retrying the gate's construction (and + # printing nothing to say why) on every kernel of the case. + self._gate_checked = False @classmethod def for_case(cls, args, problem: object, dtype: str, direction: str) -> "EarlyStop": @@ -260,7 +264,18 @@ def _timed(self, fn, *, warmup: int, iters: int, stream: int) -> float: from statistics import median from rocke.runtime import time_launches_staged - + from rocke.runtime.hip_module import HipError + from rocke.runtime.stall_gate import default_gate + + if not self._gate_checked: + self._gate_checked = True + try: + default_gate() + except HipError as e: + raise HipError( + f"{e}; this device cannot be timed with --timing " + f"staged, rerun with --timing block" + ) from e return median( time_launches_staged(fn, warmup=warmup, iters=iters, stream=stream) ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py index fc543fbf5ecd..fed3d1bacc76 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/module_cache.py @@ -8,9 +8,8 @@ :class:`~rocke.runtime.launcher.KernelLauncher` -- ``read_bytes`` plus ``hipModuleLoadData`` plus ``hipModuleGetFunction`` -- inside the loop over candidates, inside the loop over cases, so a run of *C* cases over *K* -candidates loaded *C x K* modules of *K* distinct binaries. At a few hundred -microseconds each that is the single largest fixed cost per candidate in a -run that compiles nothing. +candidates loaded *C x K* modules of *K* distinct binaries. Repeating these +loads adds fixed overhead to a run that compiles nothing. :class:`ModuleCache` keys launchers by the binary and the entry point, so a blob is read and loaded once per run and every later shape reuses the loaded @@ -75,12 +74,14 @@ def __init__(self, budget_bytes: int = DEFAULT_BUDGET_BYTES) -> None: self._hits = 0 self._misses = 0 self._full = False - # The one over-budget launcher handed out most recently, kept only - # long enough to unload it on the next miss -- by then its launch has - # drained, since callers gate/time each candidate before moving on. - # Without this a long sweep past the budget leaks one HIP module per - # miss for the rest of the run. - self._unadmitted: Any = None + # Every over-budget launcher handed out since the last end_of_case, + # kept until the caller says the case's launches have drained -- + # unloading one any sooner would race a launch still in flight: + # --fast-rank binds a whole chunk before timing it, and a multi-step + # direct plan loads several steps before binding any. Without this + # list a long sweep past the budget leaks one HIP module per miss for + # the rest of the run. + self._unadmitted: list = [] # -- query ------------------------------------------------------------ @@ -142,9 +143,6 @@ def launcher( return held self._misses += 1 - if key is not None and self._unadmitted is not None: - self._unadmitted.unload() - self._unadmitted = None hsaco = Path(hsaco_path).read_bytes() launcher = KernelLauncher( hsaco=hsaco, kernel_name=kernel_name, signature=signature @@ -155,9 +153,22 @@ def launcher( self._bytes += len(hsaco) else: self._full = True - self._unadmitted = launcher + self._unadmitted.append(launcher) return launcher + def end_case(self) -> None: + """Unload this case's over-budget launchers, once its drain is done. + + The caller must not call this until every launch it handed out for + the current case has finished -- :meth:`launcher` never unloads one + itself precisely so that timing it, verifying it or ranking it against + the rest of the case's candidates cannot race an unload. A case under + budget calls this for free: there is nothing to unload. + """ + for launcher in self._unadmitted: + launcher.unload() + self._unadmitted.clear() + def clear(self) -> None: """Drain the device, then unload every module held. @@ -165,7 +176,7 @@ def clear(self) -> None: that the launches it handed out have finished, and unloading a module under a kernel in flight faults. """ - if not self._held and self._unadmitted is None: + if not self._held and not self._unadmitted: return from rocke.runtime import synchronize_and_release @@ -173,9 +184,7 @@ def clear(self) -> None: for launcher in self._held.values(): launcher.unload() self._held.clear() - if self._unadmitted is not None: - self._unadmitted.unload() - self._unadmitted = None + self.end_case() self._bytes = 0 self._full = False diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py index 3ec2b0b4504b..8c4ce2c5c3b4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/run_conv_shards.py @@ -27,9 +27,9 @@ card and a card takes the next shard off the queue when it frees up, which bounds the tail at roughly one shard rather than one quarter of the run. -The cost is that each shard re-reads the kernel cache: ~18 s for a 447k-entry -cache, so ~72 s per card over a run. Lower ``--shards`` to one per card to pay -it once, at the price of the tail. +Each shard re-reads the kernel cache, so more shards add startup overhead. +Lower ``--shards`` to one per card to pay it once, at the price of a longer +tail. Lines are dealt round-robin (line *i* to shard *i mod shards*) rather than in blocks: a case file is usually grouped by model and resolution, so contiguous @@ -486,8 +486,17 @@ def main(argv: Optional[Sequence[str]] = None) -> int: print(f"[info] --resume: no rows in {log_dir}; running everything") else: measured = log_dir / "measured.jsonl" + # --script compare's shard rows are already merged (case/winner per + # row, no "engine"), but --skip-measured is read by the per-engine + # sweeps underneath it -- expand before writing, or every row is + # silently not-measured to them and --resume re-runs everything. + skip_rows = ( + case_results.expand_merged_for_skip(carried) + if args.script == "compare" + else carried + ) measured.write_text( - "".join(json.dumps(row, sort_keys=True) + "\n" for row in carried), + "".join(json.dumps(row, sort_keys=True) + "\n" for row in skip_rows), encoding="utf-8", ) forwarded = list(forwarded) + ["--skip-measured", str(measured)] @@ -555,18 +564,28 @@ def _progress(shards_done: int, busy: int) -> None: print(f"\nshard {index} failed (rc={rc}); last lines of its log:") print(_tail(shard_paths(log_dir, index, width)["log"])) - rows = case_results.merge( - carried - + [ - row - for _index, paths in plan - for row in case_results.read_rows(paths["rows"]) - ], - # The shards' rows are keyed by geometry; the case file is what puts - # each row's MIOpenDriver line back on it. - case_results.miopen_cmds_by_key(miopen_file=args.miopen_file), - timing=timing, - ) + shard_rows = [ + row for _index, paths in plan for row in case_results.read_rows(paths["rows"]) + ] + if args.script == "compare": + # Each shard ran benchmark_conv_compare.py, whose own --results is + # already a merged row (both engines, winner, MIOpenDriver command): + # case_results.merge() expects the raw per-engine rows the "direct" + # and "implicit" scripts write, so handing it these would see no + # "engine" field on any row and silently drop them all. A fresh row + # replaces a carried one of the same key -- resume only reaches this + # case again because --skip-measured did not cover it. + by_key = {row["key"]: row for row in carried} + by_key.update({row["key"]: row for row in shard_rows}) + rows = [by_key[key] for key in sorted(by_key)] + else: + rows = case_results.merge( + carried + shard_rows, + # The shards' rows are keyed by geometry; the case file is what + # puts each row's MIOpenDriver line back on it. + case_results.miopen_cmds_by_key(miopen_file=args.miopen_file), + timing=timing, + ) if rows: out = Path(args.results or (log_dir / "results.jsonl")) csv_path = case_results.write(rows, out) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py index d4dd50059bc9..302d1cd74c94 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_module_cache.py @@ -149,6 +149,33 @@ def test_clear_unloads_every_module(fake_hip, blobs, monkeypatch): assert len(cache) == 0 and cache.nbytes == 0 +def test_over_budget_launchers_stay_alive_until_end_of_case(fake_hip, blobs): + """--fast-rank binds a whole chunk before timing it, so an over-budget + launcher must outlive every other miss of the same case; only + :meth:`ModuleCache.end_case` -- called once the case has drained -- may + unload it.""" + cache = ModuleCache(1) + first_overflow = cache.launcher( + hsaco_path=blobs[1], kernel_name="k", signature=_SIG + ) + second_overflow = cache.launcher( + hsaco_path=blobs[2], kernel_name="k", signature=_SIG + ) + assert not first_overflow.unloaded + assert not second_overflow.unloaded + cache.end_case() + assert first_overflow.unloaded + assert second_overflow.unloaded + + +def test_end_case_on_a_case_under_budget_unloads_nothing(fake_hip, blobs): + cache = ModuleCache() + held = cache.launcher(hsaco_path=blobs[0], kernel_name="k", signature=_SIG) + cache.end_case() + assert not held.unloaded + assert len(cache) == 1 + + def test_an_unhashable_signature_is_not_a_crash(fake_hip, blobs): """A signature the key cannot be built from loads per use instead of raising: the cache is an optimisation, not a dependency.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py index 05f97618cbab..1ee3262446f9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_run_conv_shards.py @@ -431,19 +431,30 @@ def test_a_writer_without_a_path_writes_nothing(tmp_path): ap.add_argument("--results") ap.add_argument("--run-from-cache") args = ap.parse_args() +# --script compare's real shards write already-merged rows (both engines, +# winner, speedup) through case_results._write_results, not one row per +# engine -- this stub mimics that so the runner's own merge of merged rows +# is what the test exercises. with open(args.results, "w") as out: for line in open(args.miopen_file): line = line.strip() if not line: continue n = int(line.split()[-1]) - for engine, ms in (("direct", 1.0 + n), ("implicit", 2.0 + n)): - out.write(json.dumps({ - "key": f"fwd/bf16/case{n}", "engine": engine, "case": n, - "shape": f"s{n}", "dtype": "bf16", "direction": "fwd", - "tflops": 100.0, "ms": ms, "gbps": 1.0, - "kernel": f"{engine}-kernel-{n}", "candidates": 3, "timing": "staged", - }) + "\\n") + direct_ms, implicit_ms = 1.0 + n, 2.0 + n + out.write(json.dumps({ + "key": f"fwd/bf16/case{n}", "miopen_cmd": None, "shape": f"s{n}", + "dtype": "bf16", "direction": "fwd", + "direct_tflops": 100.0, "direct_ms": direct_ms, + "direct_kernel": f"direct-kernel-{n}", + "direct_kernel_2": None, "direct_kernel_3": None, + "implicit_tflops": 100.0, "implicit_ms": implicit_ms, + "implicit_kernel": f"implicit-kernel-{n}", + "implicit_kernel_2": None, "implicit_kernel_3": None, + "winner": "direct", "winner_ms": direct_ms, + "winner_kernel": f"direct-kernel-{n}", + "speedup": round(implicit_ms / direct_ms, 3), "timing": "staged", + }) + "\\n") print("stub done") """ @@ -619,25 +630,30 @@ def test_resume_carries_the_earlier_rows_and_skips_their_cases( done.add(json.loads(line)["key"]) limit = int(os.environ.get("STUB_LIMIT", "999")) ran = [] +rows = [] +for n_done, line in enumerate(open(args.miopen_file)): + line = line.strip() + if not line: + continue + n = int(line.split()[-1]) + key = f"fwd/bf16/case{n}" + if key in done: + continue + if n_done >= limit: + break + ran.append(key) + rows.append({ + "key": key, "miopen_cmd": None, "shape": f"s{n}", "dtype": "bf16", + "direction": "fwd", "direct_tflops": 1.0, "direct_ms": 1.0, + "direct_kernel": f"direct-{n}", "direct_kernel_2": None, + "direct_kernel_3": None, "implicit_tflops": 1.0, "implicit_ms": 2.0, + "implicit_kernel": f"implicit-{n}", "implicit_kernel_2": None, + "implicit_kernel_3": None, "winner": "direct", "winner_ms": 1.0, + "winner_kernel": f"direct-{n}", "speedup": 2.0, "timing": "staged", + }) with open(args.results, "w") as out: - for n_done, line in enumerate(open(args.miopen_file)): - line = line.strip() - if not line: - continue - n = int(line.split()[-1]) - key = f"fwd/bf16/case{n}" - if key in done: - continue - if n_done >= limit: - break - ran.append(key) - for engine in ("direct", "implicit"): - out.write(json.dumps({ - "key": key, "engine": engine, "case": n, "shape": f"s{n}", - "dtype": "bf16", "direction": "fwd", "tflops": 1.0, - "ms": 1.0 if engine == "direct" else 2.0, "gbps": 1.0, - "kernel": f"{engine}-{n}", "candidates": 1, "timing": "staged", - }) + "\\n") + for row in rows: + out.write(json.dumps(row) + "\\n") log = pathlib.Path(args.results).parent / "ran.json" prev = json.loads(log.read_text()) if log.is_file() else [] log.write_text(json.dumps(prev + ran))