Add FlagGems single-op perf benchmark harness
pytest plugin set + driver script for compiler A/B perf comparison: - reproducible runs: fixed seed, yaml-driven shapes, autotune record/replay - per-shape ttgir dump of actually-used variants with readable naming - cudagraph-based timing with documented fallback semantics Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,99 @@
|
||||
"""pytest plugin: make our shape yaml win over subclasses that override
|
||||
`set_shapes()` with hardcoded shapes (~9% of upstream Benchmark subclasses).
|
||||
|
||||
Wraps `Benchmark.init_user_config` (runs right after `set_shapes()`): if the
|
||||
op_name has a yaml entry, overwrite `self.shapes` with the yaml shapes. Ops with
|
||||
no yaml entry keep their own shapes. An arity guard skips ops whose subclass
|
||||
normalizes shape arity after reading yaml (e.g. BLAS (B,M,N,K)↔(M,N,K)).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import yaml as _yaml
|
||||
|
||||
|
||||
_YAML_CACHE: dict | None = None # filled lazily from Config.shape_file
|
||||
_PATCHED = False
|
||||
|
||||
|
||||
def _load_yaml(path: str) -> dict:
|
||||
if not path or not os.path.isfile(path):
|
||||
return {}
|
||||
try:
|
||||
with open(path, "r") as f:
|
||||
return _yaml.safe_load(f) or {}
|
||||
except Exception as exc:
|
||||
print(f"[shape-inject-plugin] failed to load {path}: {exc}", file=sys.stderr)
|
||||
return {}
|
||||
|
||||
|
||||
def _normalize_shapes(raw):
|
||||
"""Mirror base.py's `[tuple(shape) for shape in self.shapes]`, defensively
|
||||
handling scalar entries (`5` → `(5,)`)."""
|
||||
out = []
|
||||
for shape in raw:
|
||||
if isinstance(shape, (list, tuple)):
|
||||
out.append(tuple(shape))
|
||||
else:
|
||||
out.append((shape,))
|
||||
return out
|
||||
|
||||
|
||||
def _patch_init_user_config():
|
||||
global _PATCHED
|
||||
if _PATCHED:
|
||||
return
|
||||
from benchmark import base as fg_base
|
||||
from benchmark.conftest import Config as fg_Config
|
||||
|
||||
original = fg_base.Benchmark.init_user_config
|
||||
|
||||
def _shape_arity(shape):
|
||||
if isinstance(shape, (list, tuple)):
|
||||
return len(shape)
|
||||
return 1
|
||||
|
||||
def patched(self):
|
||||
original(self)
|
||||
global _YAML_CACHE
|
||||
if _YAML_CACHE is None:
|
||||
_YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
|
||||
entry = _YAML_CACHE.get(getattr(self, "op_name", None))
|
||||
if not isinstance(entry, dict):
|
||||
return
|
||||
shapes = entry.get("shapes")
|
||||
if not shapes:
|
||||
return
|
||||
# Skip ops whose subclass normalizes shape arity after reading yaml (e.g.
|
||||
# BLAS (B,M,N,K)↔(M,N,K)): detected as both sides being homogeneous in
|
||||
# arity but differing. Heterogeneous subclass shapes are safe to override.
|
||||
if getattr(self, "shapes", None) and len(self.shapes) > 0:
|
||||
cur_arities = {_shape_arity(s) for s in self.shapes}
|
||||
yaml_arities = {_shape_arity(s) for s in shapes}
|
||||
cur_homogeneous = len(cur_arities) == 1
|
||||
yaml_homogeneous = len(yaml_arities) == 1
|
||||
if cur_homogeneous and yaml_homogeneous and cur_arities != yaml_arities:
|
||||
print(
|
||||
f"[shape-inject-plugin] skip override for op_name={self.op_name!r}: "
|
||||
f"subclass produced homogeneous arity {next(iter(cur_arities))}, "
|
||||
f"yaml has homogeneous arity {next(iter(yaml_arities))} "
|
||||
f"— subclass normalization preserved",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return
|
||||
self.shapes = _normalize_shapes(shapes)
|
||||
if "shape_desc" in entry:
|
||||
self.shape_desc = entry["shape_desc"]
|
||||
|
||||
fg_base.Benchmark.init_user_config = patched
|
||||
_PATCHED = True
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
_patch_init_user_config()
|
||||
print("[shape-inject-plugin] Benchmark.init_user_config patched: "
|
||||
"yaml shapes now win over hardcoded subclass set_shapes()",
|
||||
file=sys.stderr, flush=True)
|
||||
Reference in New Issue
Block a user