forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
92dae82c7f
run_pytest.sh hardcoded --level core, so the only shapes ever reachable were each operator's core_shapes entry. FlagGems merges a subclass's set_more_shapes() on top of that at --level comprehensive, and for BLAS that addition is where the large GEMMs live: mm goes from 5 shapes to 29, taking in four large-K tuples plus models_shapes/*.yaml, with N/K up to 152064. The comprehensive path also yields each shape a second time (transposed, for BLAS), so the case count doubles. None of that was reachable without editing the driver. BENCH_LEVEL defaults to core, so existing runs are unchanged, and rejects anything other than core/comprehensive rather than passing it through to pytest. The chosen level is echoed in the run header, written into the archived shapes.yaml, and recorded in run_identity.json, because an A/B whose two sides used different levels is not comparable and is otherwise invisible afterwards. Sweep cost scales with the shape count, so comprehensive on a BLAS operator wants PARALLEL_WARMUP_GPUS. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
"""Merge-update `$OUT_DIR/run_identity.json` so plugins can record backend facts.
|
|
|
|
Shell writes a stub before pytest; plugins add device_name / timer / vendor
|
|
without clobbering earlier keys. Missing OUT_DIR is a silent no-op so unit
|
|
tests and ad-hoc pytest invocations do not fail.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
|
|
IDENTITY_NAME = "run_identity.json"
|
|
|
|
|
|
def identity_path(out_dir: Optional[str] = None) -> Optional[Path]:
|
|
root = out_dir or os.environ.get("OUT_DIR", "").strip()
|
|
if not root:
|
|
return None
|
|
return Path(root) / IDENTITY_NAME
|
|
|
|
|
|
def load_identity(out_dir: Optional[str] = None) -> Dict[str, Any]:
|
|
path = identity_path(out_dir)
|
|
if path is None or not path.is_file():
|
|
return {}
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def update_identity(out_dir: Optional[str] = None, **fields: Any) -> Dict[str, Any]:
|
|
path = identity_path(out_dir)
|
|
if path is None:
|
|
return dict(fields)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
current = load_identity(out_dir)
|
|
for key, value in fields.items():
|
|
if value is None:
|
|
continue
|
|
current[key] = value
|
|
tmp = path.with_name(path.name + ".tmp")
|
|
tmp.write_text(json.dumps(current, indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8")
|
|
os.replace(tmp, path)
|
|
return current
|
|
|
|
|
|
def init_identity(out_dir: str, **fields: Any) -> Dict[str, Any]:
|
|
payload = {"schema_version": 1, "provider": "flaggems"}
|
|
payload.update(fields)
|
|
os.environ["OUT_DIR"] = str(out_dir)
|
|
return update_identity(out_dir, **payload)
|
|
|
|
|
|
def main(argv=None) -> int:
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("out_dir")
|
|
p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
|
|
p.add_argument("--op", default=os.environ.get("OP", ""))
|
|
p.add_argument("--op-file", default=os.environ.get("OP_FILE", ""))
|
|
p.add_argument("--flaggems-dir", default=os.environ.get("FLAGGEMS_DIR", ""))
|
|
p.add_argument("--autotune-mode", default="")
|
|
p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
|
|
p.add_argument("--test-file", default=os.environ.get("TEST_FILE", ""))
|
|
p.add_argument("--shape-mode", default=os.environ.get("SHAPE_MODE", ""))
|
|
p.add_argument("--bench-level", default=os.environ.get("BENCH_LEVEL", "core"))
|
|
args = p.parse_args(argv)
|
|
init_identity(
|
|
args.out_dir,
|
|
backend=args.backend,
|
|
op=args.op,
|
|
op_file=args.op_file,
|
|
flaggems_dir=args.flaggems_dir,
|
|
autotune_mode=args.autotune_mode,
|
|
device=args.device,
|
|
test_file=args.test_file,
|
|
shape_mode=args.shape_mode,
|
|
bench_level=args.bench_level,
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|