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.
@@ -214,6 +231,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 +277,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 +312,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 +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",
@@ -396,6 +471,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 +658,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)
@@ -547,6 +714,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",
@@ -568,6 +750,16 @@ 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,
+                miopen_file=args.miopen_file,
+                miopen_cmd=args.miopen_cmd,
+                timing=args.timing,
+            )
+
     _print_summary(direct_result, implicit_result, implicit_note)
     return 0