feat(rocke): improve conv sweep benchmarks

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,21 @@
 ``--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.
+
+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)
@@ -65,6 +80,8 @@
 import tempfile
 from pathlib import Path
 
+from benchmarks.common import batched_timing, case_results
+
 _SCRIPT_DIR = Path(__file__).parent
 
 # Regex patterns that match the "Best:" summary printed by both benchmarks.
@@ -82,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
@@ -136,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]:
@@ -214,6 +247,24 @@ 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
+    # 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 +293,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
 
 
@@ -272,6 +328,40 @@ 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,
+    *,
+    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``).
+
+    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.
+    ``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)
+        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],
@@ -355,6 +445,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",
@@ -396,6 +487,88 @@ 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",
+    )
+
+    case_results.add_results_arg(parser, merged=True)
+    case_results.add_skip_measured_arg(parser)
+
+    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 +674,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)
@@ -541,35 +724,82 @@ 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, 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]
+        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]
+
+        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)
 
+        if args.results:
+            _write_results(
+                args.results,
+                rows_direct,
+                rows_implicit,
+                miopen_file=args.miopen_file,
+                miopen_cmd=args.miopen_cmd,
+                timing=args.timing,
+            )
+
     _print_summary(direct_result, implicit_result, implicit_note)
-    return 0
+    return overall_rc
 
 
 if __name__ == "__main__":