From 9b73823ac56013e92477af9a1cbb930f06ffb9ae Mon Sep 17 00:00:00 2001 From: NCIC-AlphaSparse Date: Sun, 27 Sep 2026 21:21:36 +0800 Subject: [PATCH] dcu spmv csr accu and bandwidth --- src/flagsparse/sparse_operations/spmv_csr.py | 22 +- tests/ci/test_baseline_bound.py | 355 ++++++++++ tools/baseline_bound.py | 643 +++++++++++++++++++ 3 files changed, 1001 insertions(+), 19 deletions(-) create mode 100644 tests/ci/test_baseline_bound.py create mode 100644 tools/baseline_bound.py diff --git a/src/flagsparse/sparse_operations/spmv_csr.py b/src/flagsparse/sparse_operations/spmv_csr.py index adb8abb..309051f 100644 --- a/src/flagsparse/sparse_operations/spmv_csr.py +++ b/src/flagsparse/sparse_operations/spmv_csr.py @@ -951,25 +951,9 @@ def _prepare_spmv_csr_matrix( raise IndexError("indices out of range for n_cols") kernel_indices = indices kernel_indptr = indptr - # gfx936 has substantially cheaper address arithmetic for 32-bit index - # tensors. Delivery inputs commonly arrive as int64 even though their - # complete index domain is safely representable in int32; keep the public - # input dtype unchanged but use compact kernel metadata when provably safe. - if _is_rocm_runtime(): - indices_fit_i32 = ( - indices.dtype == torch.int64 - and n_cols <= _INDEX_LIMIT_INT32 - and (nnz == 0 or int(indices.max().item()) <= _INDEX_LIMIT_INT32) - ) - indptr_fit_i32 = ( - indptr.dtype == torch.int64 - and n_rows <= _INDEX_LIMIT_INT32 - and nnz <= _INDEX_LIMIT_INT32 - ) - if indices_fit_i32: - kernel_indices = indices.to(torch.int32).contiguous() - if indptr_fit_i32: - kernel_indptr = indptr.to(torch.int32).contiguous() + # No int64 -> int32 compaction here. A ROCm-only one (37e16a9) failed the DCU + # accuracy tests and was removed on 2026-09-27; int64 input therefore keeps + # int64 kernel indices, which also leaves supports_opt False for it. row_lengths = kernel_indptr[1:] - kernel_indptr[:-1] max_row_nnz = int(row_lengths.max().item()) if n_rows > 0 else 0 return ( diff --git a/tests/ci/test_baseline_bound.py b/tests/ci/test_baseline_bound.py new file mode 100644 index 0000000..4137c12 --- /dev/null +++ b/tests/ci/test_baseline_bound.py @@ -0,0 +1,355 @@ +# Copyright 2026 FlagOS Contributors +# SPDX-License-Identifier: Apache-2.0 +"""tools/baseline_bound.py: T' <= T_h800 * (P_h800 / P_vendor) / ratio.""" + +import csv +import json +import re +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +TOOL = ROOT / "tools" / "baseline_bound.py" +sys.path.insert(0, str(ROOT)) +from tools.baseline_bound import REFERENCE_PEAKS # noqa: E402 + +H800_BW = REFERENCE_PEAKS["h800-sxm"]["mem_bw_gbs"] +HALF_BW = str(H800_BW / 2) + +PEAKS = { + "reference": { + "mem_bw_gbs": 1000.0, + "cuda_tflops": {"fp32": 40.0, "fp64": 1.0}, + "tensor_tflops": {"fp16": 200.0}, + }, + "vendor": { + "mem_bw_gbs": 500.0, + "cuda_tflops": {"fp32": 10.0, "fp64": 0.5}, + "tensor_tflops": {"fp16": 100.0}, + }, +} + + +def _write_csv(path, rows, fields=None): + path.parent.mkdir(parents=True, exist_ok=True) + if fields is None: + fields = [] + for row in rows: + fields += [k for k in row if k not in fields] + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + + +def _ref(matrix="a.mtx", dtype="float32", t=1.0, status="PASS", **extra): + # An H800 row: FlagSparse time, the cuSPARSE baseline it was accepted against. + return { + "matrix": matrix, + "value_dtype": dtype, + "triton_ms": t, + "cusparse_ms": t, + "status": status, + **extra, + } + + +def _ven(matrix="a.mtx", dtype="float32", t=1.0, status="PASS", **extra): + # A vendor row with no library baseline: cusparse_ms is present but empty. + return { + "matrix": matrix, + "value_dtype": dtype, + "triton_ms": t, + "cusparse_ms": "", + "status": status, + **extra, + } + + +def _call(tmp_path, *args): + cmd = [sys.executable, str(TOOL), *map(str, args), "--csv", tmp_path / "out.csv"] + proc = subprocess.run(cmd, capture_output=True, text=True, cwd=ROOT) + out = tmp_path / "out.csv" + rows = list(csv.DictReader(out.open(encoding="utf-8"))) if out.exists() else [] + return proc, rows + + +def _run(tmp_path, ref_rows, vendor_rows=None, *extra): + _write_csv(tmp_path / "ref.csv", ref_rows) + args = [tmp_path / "ref.csv", *extra] + if not {"--peaks", "--vendor-bw-gbs", "--vendor-card"} & set(extra): + (tmp_path / "peaks.json").write_text(json.dumps(PEAKS), encoding="utf-8") + args += ["--peaks", tmp_path / "peaks.json"] + if vendor_rows is not None: + _write_csv(tmp_path / "vendor.csv", vendor_rows) + args += ["--vendor", tmp_path / "vendor.csv"] + return _call(tmp_path, *args) + + +def test_h800_sxm_peaks_match_the_datasheet(): + h800 = REFERENCE_PEAKS["h800-sxm"] + assert h800["mem_bw_gbs"] == 3050.0 # measured, not the 3350 datasheet + assert h800["cuda_tflops"]["fp32"] == 67.0 and h800["cuda_tflops"]["fp64"] == 1.0 + pcie = REFERENCE_PEAKS["h800-pcie"] + assert pcie["mem_bw_gbs"] == 2000.0 and pcie["cuda_tflops"]["fp32"] == 51.0 + + +def test_template_carries_h800_and_leaves_the_vendor_blank(): + proc = subprocess.run( + [sys.executable, str(TOOL), "--print-template"], + capture_output=True, + text=True, + cwd=ROOT, + ) + assert proc.returncode == 0, proc.stderr + tpl = json.loads(proc.stdout) + assert tpl["reference"]["mem_bw_gbs"] == H800_BW + assert tpl["vendor"]["mem_bw_gbs"] is None + assert all(v is None for v in tpl["vendor"]["cuda_tflops"].values()) + + +def test_half_the_h800_bandwidth_allows_twice_the_time_over_0_8(tmp_path): + # 10 ms on H800; vendor has half the bandwidth -> 20 ms expected, 25 ms allowed. + ref = [_ref("a.mtx", t=10.0), _ref("b.mtx", t=10.0), _ref("c.mtx", t=10.0)] + vendor = [_ven("a.mtx", t=24.0), _ven("b.mtx", t=25.0), _ven("c.mtx", t=26.0)] + proc, rows = _run(tmp_path, ref, vendor, "--vendor-bw-gbs", HALF_BW) + assert proc.returncode == 1 + by = {r["key"].split("|")[0]: r for r in rows} + assert {k: float(r["bound_ms"]) for k, r in by.items()} == { + "a.mtx": 25.0, + "b.mtx": 25.0, + "c.mtx": 25.0, + } + # exactly at the bound counts + assert {k: r["verdict"] for k, r in by.items()} == { + "a.mtx": "PASS", + "b.mtx": "PASS", + "c.mtx": "FAIL", + } + assert by["a.mtx"]["resource"] == "mem" and by["a.mtx"]["source"] == "assumed" + assert "H800 SXM" in proc.stderr and "assumed" in proc.stderr + + +def test_bound_only_mode_needs_no_vendor_run(tmp_path): + proc, rows = _run(tmp_path, [_ref(t=10.0)], None, "--vendor-bw-gbs", HALF_BW) + assert proc.returncode == 0, proc.stderr + assert rows[0]["verdict"] == "BOUND" and float(rows[0]["bound_ms"]) == 25.0 + + +def test_vendor_card_uses_its_measured_bandwidth(tmp_path): + # MetaX C550 1440 GB/s against H800 3050: 1.0 * 3050 / 1440 / 0.8 + proc, rows = _run(tmp_path, [_ref(t=1.0)], None, "--vendor-card", "maca-c550") + assert proc.returncode == 0, proc.stderr + assert abs(float(rows[0]["bound_ms"]) - 3050 / 1440 / 0.8) < 1e-3 + # an explicit bandwidth wins over the card table + proc, rows = _run( + tmp_path, + [_ref(t=1.0)], + None, + "--vendor-card", + "maca-c550", + "--vendor-bw-gbs", + "3050", + ) + assert float(rows[0]["bound_ms"]) == 1.25 + + +def test_pcie_reference_changes_the_ratio(tmp_path): + proc, rows = _run( + tmp_path, + [_ref(t=1.0)], + None, + "--reference", + "h800-pcie", + "--vendor-bw-gbs", + 1000, + ) + assert proc.returncode == 0, proc.stderr + assert float(rows[0]["bound_ms"]) == 2.5 # 1.0 * 2000 / 1000 / 0.8 + + +def test_a_row_with_its_own_vendor_baseline_is_not_judged(tmp_path): + vendor = [_ven("a.mtx", t=100.0, cusparse_ms=90.0)] + proc, rows = _run(tmp_path, [_ref("a.mtx", t=1.0)], vendor) + assert proc.returncode == 0, proc.stderr + assert rows[0]["verdict"] == "HAS_BASELINE" and rows[0]["bound_ms"] == "" + proc, rows = _run(tmp_path, [_ref("a.mtx", t=1.0)], vendor, "--all-rows") + assert proc.returncode == 1 and rows[0]["verdict"] == "FAIL" + + +def test_the_reference_row_must_itself_have_passed(tmp_path): + proc, rows = _run(tmp_path, [_ref(status="FAIL")], [_ven()]) + assert proc.returncode == 1 + assert rows[0]["verdict"] == "N/A" and "not PASS" in rows[0]["note"] + + +def test_a_failed_vendor_row_fails_even_when_fast(tmp_path): + proc, rows = _run(tmp_path, [_ref(t=1.0)], [_ven(t=0.1, status="FAIL")]) + assert proc.returncode == 1 and rows[0]["verdict"] == "FAIL" + + +def test_a_vendor_row_the_reference_lacks_is_no_reference(tmp_path): + ref = [_ref("a.mtx"), _ref("wide-sweep-only.mtx")] + vendor = [_ven("a.mtx", t=1.0), _ven("b.mtx", t=1.0)] + proc, rows = _run(tmp_path, ref, vendor) + assert proc.returncode == 1 + # Driven by the vendor run: the reference's extra row is not reported. + assert {r["key"].split("|")[0]: r["verdict"] for r in rows} == { + "a.mtx": "PASS", + "b.mtx": "NO_REFERENCE", + } + + +def test_time_columns_are_picked_per_side_from_the_runner_schemas(tmp_path): + # CUDA spmv_csr writes triton_ms/cusparse_ms, ROCm writes ms/vendor_ms. + ref = [_ref(t=1.0, pytorch_ms=0.1)] + vendor = [ + { + "matrix": "a.mtx", + "dtype": "float32", + "ms": 2.4, + "vendor_ms": "", + "status": "PASS", + } + ] + proc, rows = _run(tmp_path, ref, vendor) + assert proc.returncode == 0, proc.stderr + assert float(rows[0]["t_ref_ms"]) == 1.0 and float(rows[0]["t_vendor_ms"]) == 2.4 + # spmv_coo: the delivered kernel is opt_ms, not base_ms. + ref = [ + { + "matrix": "a.mtx", + "value_dtype": "float32", + "base_ms": 9.0, + "opt_ms": 1.0, + "cusparse_ms": 1.0, + "status": "PASS", + } + ] + proc, rows = _run(tmp_path, ref, None) + assert float(rows[0]["t_ref_ms"]) == 1.0 + + +def test_spsv_columns_are_recognised(tmp_path): + ref = [ + { + "matrix": "a.mtx", + "value_dtype": "float32", + "opA": "NON", + "FlagSparse_ms": 1.0, + "CuPy/cuSPARSE_ms": 1.0, + "status": "PASS", + } + ] + vendor = [ + { + "matrix": "a.mtx", + "value_dtype": "float32", + "opA": "NON", + "FlagSparse_ms": 2.0, + "hipSPARSE_ms": "", + "status": "PASS", + } + ] + proc, rows = _run(tmp_path, ref, vendor) + assert proc.returncode == 0, proc.stderr + assert rows[0]["verdict"] == "PASS" and rows[0]["key"] == "a.mtx|float32|non" + + +def test_bottleneck_follows_the_largest_nsight_utilisation(tmp_path): + rows_in = [ + _ref("mem.mtx", u_cuda=30, u_tensor=5, u_mem=70), + _ref("cuda.mtx", u_cuda=80, u_tensor=1, u_mem=20), + _ref("named.mtx", bottleneck="mem", u_cuda=99), + ] + proc, rows = _run(tmp_path, rows_in) + assert proc.returncode == 0, proc.stderr + by = {r["key"].split("|")[0]: r for r in rows} + assert (by["mem.mtx"]["resource"], by["mem.mtx"]["source"]) == ("mem", "profile") + assert float(by["mem.mtx"]["bound_ms"]) == 2.5 + # fp32 vector peaks 40 vs 10: 1.0 * 4 / 0.8 = 5.0 + assert (by["cuda.mtx"]["resource"], by["cuda.mtx"]["source"]) == ("cuda", "profile") + assert float(by["cuda.mtx"]["bound_ms"]) == 5.0 + assert (by["named.mtx"]["resource"], by["named.mtx"]["source"]) == ("mem", "column") + + +def test_complex_dtype_uses_the_real_unit_of_the_same_precision(tmp_path): + # complex128 -> fp64 vector peaks 1.0 vs 0.5: 1.0 * 2 / 0.8 = 2.5 + proc, rows = _run(tmp_path, [_ref(dtype="complex128")], None, "--resource", "cuda") + assert proc.returncode == 0, proc.stderr + assert float(rows[0]["bound_ms"]) == 2.5 + + +def test_ratio_is_configurable(tmp_path): + proc, rows = _run(tmp_path, [_ref()], None, "--ratio", "0.5") + assert proc.returncode == 0, proc.stderr + assert float(rows[0]["bound_ms"]) == 4.0 + + +def test_missing_peak_is_reported_not_guessed(tmp_path): + proc, rows = _run(tmp_path, [_ref()], None, "--resource", "tensor") + assert proc.returncode == 1 + assert rows[0]["verdict"] == "N/A" and "tensor_tflops[fp32]" in rows[0]["note"] + + +def test_no_vendor_peaks_is_an_error(tmp_path): + _write_csv(tmp_path / "ref.csv", [_ref()]) + proc = subprocess.run( + [sys.executable, str(TOOL), str(tmp_path / "ref.csv")], + capture_output=True, + text=True, + cwd=ROOT, + ) + assert proc.returncode != 0 and "--vendor-bw-gbs" in proc.stderr + + +def test_results_directories_pair_up_by_operator(tmp_path): + ref_dir, ven_dir = tmp_path / "h800", tmp_path / "vendor" + _write_csv(ref_dir / "gather" / "performance.csv", [_ref(t=10.0)]) + _write_csv(ref_dir / "spmm_csr" / "performance.csv", [_ref(t=1.0)]) + _write_csv(ref_dir / "spmv_csr" / "performance.csv", [_ref(t=1.0)]) + _write_csv(ven_dir / "gather" / "performance.csv", [_ven(t=20.0)]) + _write_csv(ven_dir / "spmv_csr" / "performance.csv", [], fields=["matrix", "ms"]) + _write_csv(ven_dir / "sddmm_csr" / "performance.csv", [_ven()]) + proc, rows = _call( + tmp_path, ref_dir, "--vendor", ven_dir, "--vendor-bw-gbs", HALF_BW + ) + assert proc.returncode == 1 + assert {r["op"]: r["verdict"] for r in rows} == { + "gather": "PASS", + "sddmm_csr": "NO_REFERENCE", + "spmv_csr": "EMPTY", # header-only CSV must not vanish from the table + } + + +def test_duplicate_keys_are_an_error_until_disambiguated(tmp_path): + rows_in = [_ref(alg="1"), _ref(alg="2", t=2.0)] + proc, rows = _run(tmp_path, rows_in) + assert proc.returncode == 1 + assert rows[0]["verdict"] == "ERROR" and "--extra-key" in rows[0]["note"] + proc, rows = _run(tmp_path, rows_in, None, "--extra-key", "alg") + assert proc.returncode == 0, proc.stderr + assert sorted(float(r["bound_ms"]) for r in rows) == [2.5, 5.0] + + +def test_a_case_axis_both_sides_carry_joins_the_key(tmp_path): + ref = [_ref(alg="1", t=1.0), _ref(alg="2", t=2.0)] + vendor = [_ven(alg="1", t=2.5), _ven(alg="2", t=2.5)] + proc, rows = _run(tmp_path, ref, vendor) + assert proc.returncode == 0, proc.stderr + # alg 1: 1.0 * 2 / 0.8 = 2.5 (at the bound); alg 2: bound 5.0 + assert sorted((r["key"], r["bound_ms"], r["verdict"]) for r in rows) == [ + ("a.mtx|float32|1", "2.5", "PASS"), + ("a.mtx|float32|2", "5.0", "PASS"), + ] + + +def test_markdown_table_is_well_formed(tmp_path): + proc, _ = _run(tmp_path, [_ref()], None, "--markdown") + lines = proc.stdout.strip().splitlines() + assert lines[0].startswith("| op | case |") and set(lines[1]) <= {"|", "-"} + # "|" inside a cell (case labels are joined with it) must be escaped, or the row + # grows extra columns. + assert len({len(re.split(r"(? expected 20 ms, bound +20 / 0.8 = 25 ms. + +Inputs + * the reference run: a results directory (``//performance.csv``) or + one ``performance.csv``. Only rows whose ``status`` is PASS are used -- the + reference must itself be a verified result; + * the reference peaks: built in for ``--reference h800-sxm`` (default) and + ``h800-pcie``, or the ``reference`` side of a ``--peaks`` JSON; + * the vendor peaks: ``--vendor-card`` (built-in measured bandwidth for C550, + S5000, BW1000, BI-V150, Ascend 910B/910C), ``--vendor-bw-gbs`` (either is enough for + the default memory-bound unit), or the ``vendor`` side of ``--peaks``; + * optionally the vendor run (``--vendor``, same layout as the reference). Each + vendor row with no usable vendor-library baseline is judged PASS/FAIL; a row + that has one is HAS_BASELINE and left to its measured speedup (``--all-rows`` + judges it anyway); a vendor row the reference run lacks is NO_REFERENCE. + +The FlagSparse time column is picked per CSV from the runner's +``PERFORMANCE_SPEEDUP_SCHEMAS``: the first vendor-library schema whose FlagSparse +column the CSV carries (``triton_ms``, ``ms``, ``opt_ms``, ...). The two sides may +differ -- CUDA spmv_csr writes ``triton_ms``, ROCm writes ``ms``. + +The bottleneck unit of a row is, in order: its ``bottleneck`` column +(``mem``/``cuda``/``tensor``), the largest of its ``u_cuda``/``u_tensor``/``u_mem`` +columns (Nsight Compute utilisation, %), else ``--default-resource`` (``mem``). +The last case is an assumption -- sparse kernels are usually memory-bound, but +nothing measured it -- and is labelled ``assumed`` in the output. + +Exit status is 0 when every row has a bound (and, with ``--vendor``, passes or +has its own baseline), 1 otherwise. + +Usage: + python3 tools/baseline_bound.py h800_results/ --vendor-bw-gbs 1600 + python3 tools/baseline_bound.py h800_results/ --vendor-card dcu-bw1000 \ + --vendor vendor_results/ [--markdown] [--csv out.csv] + python3 tools/baseline_bound.py --print-template > peaks.json +""" + +from __future__ import annotations + +import argparse +import ast +import csv +import json +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +RUNNER = ROOT / "run_flagsparse_pytest.py" +RESOURCES = ("mem", "cuda", "tensor") + +# NVIDIA H800 peaks. Memory bandwidth on h800-sxm is MEASURED (3050 GB/s; the +# datasheet says 3350) so that it divides the vendors' measured figures below +# like for like -- a datasheet over a measurement would loosen every bound by +# ~10%. h800-pcie has no measurement and keeps its datasheet 2000. The compute +# figures are datasheet (dense; the datasheet's Tensor figures are quoted +# with 2:4 sparsity and are halved here). H800 is H100 with FP64 cut to ~1 +# TFLOPS -- a compute-bound fp64 row therefore gets a very tight bound, since +# almost any vendor card out-runs it. The datasheet gives no non-Tensor FP16/BF16 +# figure; Hopper runs them at 2x FP32 (H100 whitepaper: 133.8 vs 66.9 TFLOPS), +# so those two entries are derived, not quoted. +REFERENCE_PEAKS = { + "h800-sxm": { + "name": "NVIDIA H800 SXM5 80GB HBM3", + "mem_bw_gbs": 3050.0, + "cuda_tflops": {"fp64": 1.0, "fp32": 67.0, "fp16": 134.0, "bf16": 134.0}, + "tensor_tflops": {"fp64": 1.0, "tf32": 494.5, "fp16": 989.5, "bf16": 989.5}, + }, + "h800-pcie": { + "name": "NVIDIA H800 PCIe 80GB HBM2e", + "mem_bw_gbs": 2000.0, + "cuda_tflops": {"fp64": 0.8, "fp32": 51.0, "fp16": 102.0, "bf16": 102.0}, + "tensor_tflops": {"fp64": 0.8, "tf32": 378.0, "fp16": 756.5, "bf16": 756.5}, + }, +} + +# Measured memory bandwidth per vendor card, GB/s (2026-09-27, same method as the +# H800 3050). --vendor-bw-gbs overrides. +VENDOR_BANDWIDTH_GBS = { + "maca-c550": 1440.0, # MetaX C550 + "musa-s5000": 1370.0, # Moore Threads S5000 + "dcu-bw1000": 1530.0, # Hygon BW1000 + "iluvatar-biv150": 1150.0, # Iluvatar CoreX BI-V150 + # Ascend delivery tests run on a 910B. No measurement exists for it, so this is + # the PUBLISHED 1.6 TB/s -- the one entry that is not measured; a published peak + # usually exceeds a measured one, so the bound is somewhat tighter than the rest. + "ascend-910b": 1600.0, + # Measured on an Atlas 800T A3 (910C, two dies per package). Not the tested chip. + "ascend-910c": 3070.0, +} + +# A column holding a vendor sparse library's time: cusparse_ms, vendor_ms, and the +# SpSV/SpSM spellings cuSPARSE_ms, hipSPARSE_ms, CuPy/cuSPARSE_ms. PyTorch is not +# one -- on Ascend it is the same library FlagSparse falls back to. +_BASELINE_EXACT = {"vendor_ms"} +_BASELINE_SUFFIX = "sparse_ms" +# SpSV/SpSM name their FlagSparse column outside the runner's schema table. +_EXTRA_TIME_COLUMNS = ("FlagSparse_ms",) + +# Complex values are computed on the real unit of the matching precision. +_UNIT_DTYPE = { + "float16": "fp16", + "half": "fp16", + "bfloat16": "bf16", + "float32": "fp32", + "float": "fp32", + "float64": "fp64", + "double": "fp64", + "complex64": "fp32", + "complex128": "fp64", + "tf32": "tf32", +} +_MATRIX_KEYS = ("matrix", "path", "name", "case_id", "case") +_DTYPE_KEYS = ( + "dtype", + "value_dtype", + "value_dtype_req", + "value_dtype_compute", + "data_dtype", +) +_OP_KEYS = ("op", "opa", "transpose") +# Case axes joined into the key when BOTH sides carry the column; one side alone +# (ROCm spmv_csr has `alg`, CUDA does not) would make every key miss. +_OPTIONAL_KEYS = ("alg", "layout", "mode", "dense_cols", "k") +_PASS = {"PASS", "PASSED", "OK", "SUCCESS"} +_TEMPLATE_SIDE = { + "name": None, + "mem_bw_gbs": None, + "cuda_tflops": {"fp16": None, "fp32": None, "fp64": None}, + "tensor_tflops": {"fp16": None, "bf16": None, "tf32": None}, +} + + +def _template() -> dict: + return { + "_note": ( + "Fill in measured peaks. Units only need to agree between the reference " + "and vendor sides of the same field; the ratio is what is used. Delete " + "dtype entries you do not need. cuda_tflops on the vendor side is the " + "vector-unit peak. Omit 'reference' to use the built-in --reference." + ), + "reference": json.loads(json.dumps(REFERENCE_PEAKS["h800-sxm"])), + "vendor": json.loads(json.dumps(_TEMPLATE_SIDE)), + } + + +def _schemas() -> list[tuple[str, str | None, str | None]]: + """The runner's (speedup, baseline, FlagSparse) column table, read statically. + + Importing the runner costs seconds and pulls in torch; the literal is enough. + """ + + tree = ast.parse(RUNNER.read_text(encoding="utf-8")) + for node in tree.body: + if isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + if any( + getattr(t, "id", "") == "PERFORMANCE_SPEEDUP_SCHEMAS" for t in targets + ): + return list(ast.literal_eval(node.value)) + raise SystemExit(f"PERFORMANCE_SPEEDUP_SCHEMAS not found in {RUNNER}") + + +def _is_baseline_column(name: str | None) -> bool: + lowered = str(name or "").lower() + if lowered == "flagsparse_ms": # ends in "sparse_ms" but is FlagSparse itself + return False + return lowered in _BASELINE_EXACT or lowered.endswith(_BASELINE_SUFFIX) + + +def _time_column(fields: list[str]) -> str | None: + for _, base, latency in _schemas(): + if latency and latency in fields and _is_baseline_column(base): + return latency + for name in _EXTRA_TIME_COLUMNS: + if name in fields: + return name + return None + + +def _read_csv(path: Path) -> tuple[list[str], list[dict[str, str]]]: + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle) + return list(reader.fieldnames or []), list(reader) + + +def _collect(target: str) -> dict[str, Path]: + """Map operator -> performance.csv for a results directory or a single CSV.""" + + path = Path(target) + if path.is_file(): + return {path.parent.name or path.stem: path} + if path.is_dir(): + found = {p.parent.name: p for p in sorted(path.glob("*/performance.csv"))} + if (path / "performance.csv").is_file(): + found.setdefault(path.name, path / "performance.csv") + if found: + return found + raise SystemExit(f"no */performance.csv under {path}") + raise SystemExit(f"no such file or directory: {path}") + + +def _num(value: object) -> float | None: + try: + out = float(str(value).strip()) + except (TypeError, ValueError): + return None + return out if out == out and out not in (float("inf"), float("-inf")) else None + + +def _lower_fields(row: dict[str, str]) -> dict[str, str]: + return {str(k).lower(): v for k, v in row.items() if k is not None} + + +def _matrix_key(row: dict[str, str]) -> str: + for key in _MATRIX_KEYS: + value = str(row.get(key) or "").strip() + if value: + return Path(value).name.lower() + return "" + + +def _dtype(row: dict[str, str]) -> str: + for key in _DTYPE_KEYS: + value = str(row.get(key) or "").strip() + if value: + return value.replace("torch.", "").lower() + return "unknown" + + +def _op(row: dict[str, str]) -> str: + low = _lower_fields(row) + for key in _OP_KEYS: + value = str(low.get(key) or "").strip().lower() + if value: + return value + return "" + + +def _row_key(row: dict[str, str], extra: list[str]) -> tuple[str, ...]: + parts = [ + _matrix_key(row), + _dtype(row), + str(row.get("index_dtype") or "").strip().lower(), + _op(row), + ] + parts.extend(str(row.get(col) or "").strip().lower() for col in extra) + return tuple(parts) + + +def _label(key: tuple[str, ...]) -> str: + return "|".join(part for part in key if part) + + +def _index_rows( + rows: list[dict[str, str]], extra: list[str], source: str +) -> dict[tuple[str, ...], dict[str, str]]: + indexed: dict[tuple[str, ...], dict[str, str]] = {} + for row in rows: + key = _row_key(row, extra) + if key in indexed: + raise ValueError( + f"{source}: two rows share the key {_label(key)!r}; add a column that " + "tells them apart with --extra-key (e.g. --extra-key alg)" + ) + indexed[key] = row + return indexed + + +def _passed(row: dict[str, str]) -> bool: + status = str(row.get("status") or row.get("matrix_status") or "").strip().upper() + return status in _PASS + + +def _baseline_ms(row: dict[str, str]) -> float | None: + for key, value in row.items(): + if _is_baseline_column(key): + number = _num(value) + if number is not None and number > 0: + return number + return None + + +def _load_peaks(args) -> dict: + peaks: dict = {} + if args.peaks: + try: + peaks = json.loads(Path(args.peaks).read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise SystemExit(f"cannot read peaks file {args.peaks}: {exc}") + reference = peaks.get("reference") or peaks.get("nvidia") + if not isinstance(reference, dict): + reference = json.loads(json.dumps(REFERENCE_PEAKS[args.reference])) + vendor = peaks.get("vendor") + vendor = dict(vendor) if isinstance(vendor, dict) else {} + if args.vendor_card: + vendor["name"] = args.vendor_card + vendor["mem_bw_gbs"] = VENDOR_BANDWIDTH_GBS[args.vendor_card] + if args.vendor_bw_gbs is not None: + vendor["mem_bw_gbs"] = args.vendor_bw_gbs + if _num(vendor.get("mem_bw_gbs")) is None and not any( + isinstance(vendor.get(f), dict) for f in ("cuda_tflops", "tensor_tflops") + ): + raise SystemExit("no vendor peaks: pass --vendor-bw-gbs or a --peaks JSON") + return {"reference": reference, "vendor": vendor} + + +def _peak(side: dict, resource: str, dtype: str) -> float: + unit = _UNIT_DTYPE.get(dtype, dtype) + if resource == "mem": + value, what = side.get("mem_bw_gbs"), "mem_bw_gbs" + else: + field = "cuda_tflops" if resource == "cuda" else "tensor_tflops" + table = side.get(field) + value = table.get(unit) if isinstance(table, dict) else None + what = f"{field}[{unit}]" + number = _num(value) + if number is None or number <= 0: + raise KeyError(what) + return number + + +def _resource_for(row: dict[str, str], forced: str, default: str) -> tuple[str, str]: + if forced != "auto": + return forced, "forced" + named = str(row.get("bottleneck") or "").strip().lower() + if named in RESOURCES: + return named, "column" + util = {r: _num(row.get(f"u_{r}")) for r in RESOURCES} + util = {r: u for r, u in util.items() if u is not None} + if util: + return max(util, key=util.get), "profile" + return default, "assumed" + + +def _fmt(value: float | None) -> str: + return "-" if value is None else f"{value:.4g}" + + +def _result(op: str, key: str, **fields) -> dict: + out = { + "op": op, + "key": key, + "dtype": "", + "resource": "", + "source": "", + "t_ref_ms": None, + "bound_ms": None, + "t_vendor_ms": None, + "margin": None, + "verdict": "", + "note": "", + } + out.update(fields) + return out + + +def _pick(fields: list[str], wanted: str | None, side: str, path: Path) -> str: + if wanted: + if wanted not in fields: + raise ValueError(f"{side} {path}: column {wanted!r} is not in the header") + return wanted + column = _time_column(fields) + if column is None: + times = [f for f in fields if f.lower().endswith("ms")] + raise ValueError( + f"{side} {path}: no FlagSparse time column found; pass --ref-column / " + f"--vendor-column. Time-like columns here: {', '.join(times) or 'none'}" + ) + return column + + +def _bound(out: dict, row: dict[str, str], ref_col: str, args, peaks) -> bool: + """Fill ``out`` with the reference time and its bound; False when there is none.""" + + t_ref = _num(row.get(ref_col)) + out["t_ref_ms"] = t_ref + if not _passed(row): + out["verdict"] = "N/A" + out["note"] = f"reference row is not PASS ({row.get('status') or 'no status'})" + return False + if t_ref is None or t_ref <= 0: + out["verdict"], out["note"] = "N/A", f"no usable reference {ref_col}" + return False + try: + p_ref = _peak(peaks["reference"], out["resource"], out["dtype"]) + p_v = _peak(peaks["vendor"], out["resource"], out["dtype"]) + except KeyError as exc: + out["verdict"], out["note"] = "N/A", f"peak missing: {exc.args[0]}" + return False + out["bound_ms"] = t_ref * (p_ref / p_v) / args.ratio + return True + + +def _evaluate_op(op: str, ref_path: Path, vendor_path: Path | None, args, peaks): + ref_fields, ref_rows = _read_csv(ref_path) + ref_col = _pick(ref_fields, args.ref_column, "reference", ref_path) + vendor_fields, vendor_rows = _read_csv(vendor_path) if vendor_path else ([], []) + # A header-only CSV (a benchmark that recorded nothing) would otherwise add no + # rows and vanish from the table. + if not ref_rows or (vendor_path and not vendor_rows): + empty = ref_path if not ref_rows else vendor_path + return [_result(op, "*", verdict="EMPTY", note=f"{empty} has no rows")] + if vendor_path is None: + results = [] + for key, row in _index_rows(ref_rows, args.extra_key, str(ref_path)).items(): + resource, source = _resource_for(row, args.resource, args.default_resource) + out = _result( + op, _label(key), dtype=_dtype(row), resource=resource, source=source + ) + if _bound(out, row, ref_col, args, peaks): + out["verdict"] = "BOUND" + results.append(out) + return results + + vendor_col = _pick(vendor_fields, args.vendor_column, "vendor", vendor_path) + extra = list(args.extra_key) + [ + c + for c in _OPTIONAL_KEYS + if c in ref_fields and c in vendor_fields and c not in args.extra_key + ] + ref_index = _index_rows(ref_rows, extra, str(ref_path)) + vendor_index = _index_rows(vendor_rows, extra, str(vendor_path)) + + # Driven by the vendor run: its rows are what is being accepted, and a delivery + # run sweeps fewer axes than a full reference run (int32/non only), so the + # reference rows it skipped are not failures. Completeness of the vendor run + # itself is tools/delivery_table.py's job. + results = [] + for key, vrow in vendor_index.items(): + out = _result(op, _label(key), dtype=_dtype(vrow)) + results.append(out) + if not args.all_rows and _baseline_ms(vrow) is not None: + out["verdict"] = "HAS_BASELINE" + out["note"] = "vendor row has its own library baseline" + continue + row = ref_index.get(key) + if row is None: + out["verdict"], out["note"] = "NO_REFERENCE", "no reference row" + continue + out["resource"], out["source"] = _resource_for( + row, args.resource, args.default_resource + ) + if not _bound(out, row, ref_col, args, peaks): + continue + t_v = _num(vrow.get(vendor_col)) + if not _passed(vrow): + out["verdict"] = "FAIL" + out["note"] = ( + f"vendor row is not PASS ({vrow.get('status') or 'no status'})" + ) + continue + if t_v is None or t_v <= 0: + out["verdict"], out["note"] = "FAIL", f"no usable vendor {vendor_col}" + continue + out["t_vendor_ms"] = t_v + out["margin"] = out["bound_ms"] / t_v + out["verdict"] = "PASS" if t_v <= out["bound_ms"] else "FAIL" + return results + + +def _evaluate(args, peaks: dict) -> list[dict]: + ref_ops = _collect(args.reference_run) + vendor_ops = _collect(args.vendor) if args.vendor else None + wanted = {o.strip() for o in (args.ops or "").split(",") if o.strip()} + single = len(ref_ops) == 1 and vendor_ops is not None and len(vendor_ops) == 1 + driver = vendor_ops if vendor_ops is not None else ref_ops + results = [] + for op in driver: + if wanted and op not in wanted: + continue + # Two bare CSVs pair up whatever their parent directories are called. + ref_path = next(iter(ref_ops.values())) if single else ref_ops.get(op) + if ref_path is None: + results.append( + _result( + op, "*", verdict="NO_REFERENCE", note="reference run has no such op" + ) + ) + continue + vendor_path = vendor_ops[op] if vendor_ops is not None else None + try: + results.extend(_evaluate_op(op, ref_path, vendor_path, args, peaks)) + except ValueError as exc: + results.append(_result(op, "*", verdict="ERROR", note=str(exc))) + return results + + +_HEADERS = ( + ("op", "op"), + ("key", "case"), + ("dtype", "dtype"), + ("resource", "unit"), + ("source", "unit from"), + ("t_ref_ms", "T_ref ms"), + ("bound_ms", "T' max ms"), + ("t_vendor_ms", "T' ms"), + ("margin", "margin"), + ("verdict", "verdict"), + ("note", "note"), +) + + +def _cell(row: dict, name: str) -> str: + value = row[name] + return _fmt(value) if isinstance(value, float) or value is None else str(value) + + +def _render(results: list[dict], markdown: bool, with_vendor: bool) -> str: + columns = [ + (n, h) for n, h in _HEADERS if with_vendor or n not in ("t_vendor_ms", "margin") + ] + table = [[h for _, h in columns]] + table += [[_cell(r, n) for n, _ in columns] for r in results] + if markdown: + # Case labels are joined with "|" (and gather case ids contain it), which + # would otherwise add columns to the row. + table = [[cell.replace("|", "\\|") for cell in row] for row in table] + lines = ["| " + " | ".join(table[0]) + " |"] + lines.append("|" + "|".join("---" for _ in columns) + "|") + lines += ["| " + " | ".join(r) + " |" for r in table[1:]] + return "\n".join(lines) + widths = [max(len(r[i]) for r in table) for i in range(len(columns))] + return "\n".join( + " ".join(cell.ljust(widths[i]) for i, cell in enumerate(r)).rstrip() + for r in table + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description=__doc__.split("\n\n")[0], + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "reference_run", nargs="?", help="H800 results dir or performance.csv" + ) + parser.add_argument( + "--reference", + choices=sorted(REFERENCE_PEAKS), + default="h800-sxm", + help="built-in reference card peaks (default h800-sxm)", + ) + parser.add_argument( + "--vendor-bw-gbs", type=float, help="vendor peak memory bandwidth, GB/s" + ) + parser.add_argument( + "--vendor-card", + choices=sorted(VENDOR_BANDWIDTH_GBS), + help="use this card's built-in measured bandwidth", + ) + parser.add_argument("--peaks", help="peaks JSON (see --print-template)") + parser.add_argument("--print-template", action="store_true") + parser.add_argument("--vendor", help="vendor results dir or performance.csv") + parser.add_argument( + "--all-rows", + action="store_true", + help="also judge vendor rows that have their own library baseline", + ) + parser.add_argument("--ops", help="comma-separated operators to evaluate") + parser.add_argument("--ref-column", help="reference time column (default: auto)") + parser.add_argument("--vendor-column", help="vendor time column (default: auto)") + parser.add_argument("--ratio", type=float, default=0.8) + parser.add_argument( + "--resource", + choices=("auto", *RESOURCES), + default="auto", + help="force the bottleneck unit for every row (default: per row, see above)", + ) + parser.add_argument( + "--default-resource", + choices=RESOURCES, + default="mem", + help="unit for rows nothing identifies a bottleneck for (default mem)", + ) + parser.add_argument( + "--extra-key", + action="append", + default=[], + metavar="COL", + help="extra column that makes rows unique (repeatable)", + ) + parser.add_argument("--markdown", action="store_true") + parser.add_argument("--csv", help="also write the result table here") + args = parser.parse_args(argv) + + if args.print_template: + print(json.dumps(_template(), indent=2, ensure_ascii=False)) + return 0 + if not args.reference_run: + parser.error("the reference run is required (or use --print-template)") + if not 0 < args.ratio <= 1: + parser.error("--ratio must be in (0, 1]") + + peaks = _load_peaks(args) + results = _evaluate(args, peaks) + print(_render(results, args.markdown, bool(args.vendor))) + + counts: dict[str, int] = {} + for r in results: + counts[r["verdict"]] = counts.get(r["verdict"], 0) + 1 + summary = ", ".join(f"{k}={v}" for k, v in sorted(counts.items())) or "no rows" + ref_name = peaks["reference"].get("name") or args.reference + print( + f"\n{len(results)} rows: {summary}; reference={ref_name}; ratio={args.ratio}", + file=sys.stderr, + ) + assumed = sum(1 for r in results if r["source"] == "assumed" and r["bound_ms"]) + if assumed: + print( + f"note: {assumed} row(s) use the assumed '{args.default_resource}' " + "bottleneck; add u_cuda/u_tensor/u_mem or a bottleneck column to measure it", + file=sys.stderr, + ) + + if args.csv: + with Path(args.csv).open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=[n for n, _ in _HEADERS]) + writer.writeheader() + writer.writerows(results) + bad = [r for r in results if r["verdict"] not in ("BOUND", "PASS", "HAS_BASELINE")] + return 1 if bad else 0 + + +if __name__ == "__main__": + raise SystemExit(main())