zhoulin
26b071c6e1
- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture (internal 5-iter warmup is too short), tag fallback markers with the failing phase; first-run latency no longer jitters run-to-run. - _device_guard_plugin (new): set GEMS_VENDOR via torch probe before importing flag_gems, avoiding its timeout-less nvidia-smi subprocess probe that can hang import in fork-broken environments. - _pretty_report_plugin (new) + _term_style (new): fold inputs identical across all result rows into a legend line, color status/plugin tags/markers on the live terminal; run.log is ANSI-stripped and keeps upstream SUCCESS/column wording for grep compatibility. - run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a config section with Chinese comments, add start/end banners. - README: document the warmup semantics, new plugins/env vars, and the A/B rule that both sides must use the same USE_FLAGTUNE.
104 lines
3.6 KiB
Python
104 lines
3.6 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:
|
|
from _term_style import tag
|
|
print(f"{tag('[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:
|
|
from _term_style import tag
|
|
print(
|
|
f"{tag('[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):
|
|
from _term_style import tag
|
|
_patch_init_user_config()
|
|
print(f"{tag('[shape-inject-plugin]')} Benchmark.init_user_config patched: "
|
|
"yaml shapes now win over hardcoded subclass set_shapes()",
|
|
file=sys.stderr, flush=True)
|