|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Prepare native Criterion executables, enumerate cases, or run one exact case. |
| 3 | +
|
| 4 | +Full characterization is opt-in: one case per job retains Criterion's configured |
| 5 | +sample count while excluding compilation from the 30/45 minute measurement budget. |
| 6 | +""" |
| 7 | +import argparse |
| 8 | +import json |
| 9 | +import os |
| 10 | +from pathlib import Path |
| 11 | +import re |
| 12 | +import shutil |
| 13 | +import subprocess |
| 14 | + |
| 15 | +BENCHES = ("matmul_bench", "attention_bench", "ane_bench") |
| 16 | + |
| 17 | + |
| 18 | +def case_names(binary): |
| 19 | + result = subprocess.run([str(binary.resolve()), "--list"], check=True, |
| 20 | + capture_output=True, text=True) |
| 21 | + names = [line.removesuffix(": benchmark") for line in result.stdout.splitlines() |
| 22 | + if line.endswith(": benchmark")] |
| 23 | + if not names or len(names) != len(set(names)) or len(names) > 256: |
| 24 | + raise ValueError(f"Invalid case inventory for {binary}: {len(names)} (limit 256)") |
| 25 | + return names |
| 26 | + |
| 27 | + |
| 28 | +def prepare(features): |
| 29 | + benches = BENCHES if features else BENCHES[:2] |
| 30 | + command = ["cargo", "bench", "--locked", "-p", "ruvllm", "--no-run", |
| 31 | + "--message-format=json"] |
| 32 | + for bench in benches: |
| 33 | + command.extend(["--bench", bench]) |
| 34 | + if features: |
| 35 | + command.extend(["--features", features]) |
| 36 | + result = subprocess.run(command, check=True, stdout=subprocess.PIPE, text=True) |
| 37 | + directory = Path("benchmark-binaries") |
| 38 | + directory.mkdir(exist_ok=True) |
| 39 | + executables = {} |
| 40 | + for line in result.stdout.splitlines(): |
| 41 | + message = json.loads(line) |
| 42 | + if message.get("reason") == "compiler-artifact" and message.get("executable"): |
| 43 | + name = message["target"]["name"] |
| 44 | + if name in benches: |
| 45 | + executables[name] = message["executable"] |
| 46 | + if set(executables) != set(benches): |
| 47 | + raise ValueError(f"Missing benchmark executables: {executables}") |
| 48 | + with open(os.environ["GITHUB_OUTPUT"], "a", encoding="utf-8") as output: |
| 49 | + for bench in benches: |
| 50 | + binary = directory / bench |
| 51 | + shutil.copy2(executables[bench], binary) |
| 52 | + names = case_names(binary) |
| 53 | + (directory / f"{bench}.cases.json").write_text(json.dumps(names, indent=2)) |
| 54 | + matrix = {"include": [{"case": name} for name in names]} |
| 55 | + output.write(f"{bench}={json.dumps(matrix, separators=(',', ':'))}\n") |
| 56 | + print(f"{bench}: {len(names)} independently budgeted cases") |
| 57 | + |
| 58 | + |
| 59 | +def main(): |
| 60 | + parser = argparse.ArgumentParser(description=__doc__) |
| 61 | + parser.add_argument("mode", choices=("prepare", "run")) |
| 62 | + parser.add_argument("--features", default="") |
| 63 | + parser.add_argument("--bench", choices=BENCHES) |
| 64 | + args = parser.parse_args() |
| 65 | + if args.mode == "prepare": |
| 66 | + prepare(args.features) |
| 67 | + else: |
| 68 | + if not args.bench: |
| 69 | + parser.error("run requires --bench") |
| 70 | + binary = Path("benchmark-binaries") / args.bench |
| 71 | + name = os.environ["BENCHMARK_CASE"] |
| 72 | + inventory = json.loads((binary.parent / f"{args.bench}.cases.json").read_text()) |
| 73 | + if name not in inventory: |
| 74 | + raise ValueError("Requested case is absent from the built inventory") |
| 75 | + binary.chmod(binary.stat().st_mode | 0o100) |
| 76 | + subprocess.run([str(binary.resolve()), "--bench", f"^{re.escape(name)}$", |
| 77 | + "--output-format", "bencher"], check=True) |
| 78 | + |
| 79 | + |
| 80 | +if __name__ == "__main__": |
| 81 | + main() |
0 commit comments