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.5 KiB
Python
104 lines
3.5 KiB
Python
"""pytest plugin: make our yaml shapes win over subclasses that hardcode their
|
|
shape list inside an overridden `get_input_iter` (instead of `self.shapes`),
|
|
which `_shape_inject_plugin` can't reach (conv padding / pool / nll /
|
|
scaled_softmax families, ~17 ops).
|
|
|
|
At `pytest_collection_finish` (all Benchmark subclasses defined), redirect
|
|
`get_input_iter` ONLY on classes whose own implementation does not reference
|
|
`self.shapes` (the source-level guard that protects well-behaved BLAS/generic
|
|
classes). The redirect yields one input per yaml shape via the op's standard
|
|
`input_fn(shape, dtype, device)`; ops with no yaml entry or a non-standard
|
|
input_fn fall through unchanged. This also bounds compile cost vs the large
|
|
hardcoded lists.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import os
|
|
import sys
|
|
|
|
import yaml as _yaml
|
|
|
|
|
|
_YAML_CACHE: dict | None = None
|
|
|
|
|
|
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-iter-inject-plugin]')} failed to load {path}: {exc}",
|
|
file=sys.stderr)
|
|
return {}
|
|
|
|
|
|
def _reads_self_shapes(func) -> bool:
|
|
"""True if the method references self.shapes (already covered by
|
|
_shape_inject_plugin, so don't redirect). Defaults True on introspection
|
|
failure (safe: leave the method alone)."""
|
|
try:
|
|
return "self.shapes" in inspect.getsource(func)
|
|
except (OSError, TypeError):
|
|
return True
|
|
|
|
|
|
def _accepts_shape_dtype_device(fn) -> bool:
|
|
"""True if `fn` can be called as `fn(shape, dtype, device)` (the standard
|
|
input_fn signature). Guards against blas-style `fn(b,m,n,k,dtype,device,...)`.
|
|
Defaults False on introspection failure (safe: fall back to original)."""
|
|
if not callable(fn):
|
|
return False
|
|
try:
|
|
inspect.signature(fn).bind(None, None, None)
|
|
except (TypeError, ValueError):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _make_patched(original):
|
|
def patched(self, dtype):
|
|
entry = (_YAML_CACHE or {}).get(getattr(self, "op_name", None))
|
|
shapes = entry.get("shapes") if isinstance(entry, dict) else None
|
|
input_fn = getattr(self, "input_fn", None)
|
|
if shapes and _accepts_shape_dtype_device(input_fn):
|
|
for shape in shapes:
|
|
yield from input_fn(tuple(shape), dtype, self.device)
|
|
else:
|
|
yield from original(self, dtype)
|
|
return patched
|
|
|
|
|
|
def pytest_collection_finish(session):
|
|
global _YAML_CACHE
|
|
from benchmark import base as fg_base
|
|
from benchmark.conftest import Config as fg_Config
|
|
|
|
_YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
|
|
|
|
def _subclasses(cls):
|
|
for sub in cls.__subclasses__():
|
|
yield sub
|
|
yield from _subclasses(sub)
|
|
|
|
patched = 0
|
|
seen = set()
|
|
for cls in _subclasses(fg_base.Benchmark):
|
|
if cls in seen:
|
|
continue
|
|
seen.add(cls)
|
|
own = cls.__dict__.get("get_input_iter")
|
|
# Redirect only classes with their own get_input_iter that ignore self.shapes.
|
|
if own is not None and not _reads_self_shapes(own):
|
|
cls.get_input_iter = _make_patched(own)
|
|
patched += 1
|
|
|
|
from _term_style import tag
|
|
print(f"{tag('[shape-iter-inject-plugin]')} get_input_iter redirected on {patched} "
|
|
f"hardcoded-shape Benchmark classes; yaml shapes now win",
|
|
file=sys.stderr, flush=True)
|