zhoulin
8ace0d790c
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>
100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
"""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)
|