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.
281 lines
10 KiB
Python
281 lines
10 KiB
Python
"""pytest plugin: record or replay autotune configs across benchmark runs.
|
|
|
|
When comparing two runs (e.g. two compiler builds), replay the first run's
|
|
recorded configs in the second: the autotuner's choice is pinned, so the
|
|
latency delta reflects the change under test rather than the autotuner picking
|
|
a different config under different IR.
|
|
|
|
Modes (mutually exclusive, selected by env var; op identity from
|
|
FLAGGEMS_PERF_CURRENT_OP):
|
|
- FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR: run full autotune, dump the chosen Config
|
|
per (kernel, key) to <dir>/<op>.json at session end.
|
|
- FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR: pre-populate Autotuner.cache from
|
|
<dir>/<op>.json, skipping the sweep. Missing key / non-compiling config /
|
|
missing json fall back to full autotune and emit an AUTOTUNE_REPLAY_FALLBACK
|
|
marker into the log.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
|
|
_RECORD_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR"
|
|
_REPLAY_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR"
|
|
_OP_ENV = "FLAGGEMS_PERF_CURRENT_OP"
|
|
|
|
# Module-global aggregates, accumulated as run() is called and dumped at exit.
|
|
# Lock guards merges in case a kernel uses threads internally.
|
|
_record_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
|
_record_lock = threading.Lock()
|
|
_replay_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
|
_dump_done = False
|
|
|
|
|
|
def _op_name() -> str:
|
|
return os.environ.get(_OP_ENV, "").strip() or "default"
|
|
|
|
|
|
def _record_path(record_dir: str) -> Path:
|
|
return Path(record_dir) / f"{_op_name()}.json"
|
|
|
|
|
|
def _kernel_id(tuner: Any) -> str:
|
|
# Stable across runs: source module + qualified name. Shapes/keys split
|
|
# by the inner key dict.
|
|
name = (
|
|
getattr(tuner, "__name__", None)
|
|
or getattr(getattr(tuner, "base_fn", None), "__name__", None)
|
|
or getattr(getattr(tuner, "fn", None), "__name__", None)
|
|
or "unknown_kernel"
|
|
)
|
|
module = getattr(getattr(tuner, "fn", None), "__module__", "") or ""
|
|
return f"{module}::{name}"
|
|
|
|
|
|
def _serialize_key(key: Tuple[Any, ...]) -> str:
|
|
# Per-element repr, JSON-encoded into a flat string usable as a JSON key;
|
|
# stable across runs given seeded RNG + yaml-fixed shapes.
|
|
return json.dumps([repr(x) for x in key])
|
|
|
|
|
|
def _serialize_config(cfg: Any) -> Optional[Dict[str, Any]]:
|
|
# Drop configs with a pre_hook (un-serializable callable) rather than replay
|
|
# without it and compute incorrectly.
|
|
if getattr(cfg, "pre_hook", None) is not None:
|
|
return None
|
|
try:
|
|
return {
|
|
"kwargs": dict(cfg.kwargs),
|
|
"num_warps": cfg.num_warps,
|
|
"num_stages": cfg.num_stages,
|
|
"num_ctas": cfg.num_ctas,
|
|
"maxnreg": cfg.maxnreg,
|
|
}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _deserialize_config(entry: Dict[str, Any]) -> Any:
|
|
# Late import (triton loads during collection, after plugin registration).
|
|
from triton import Config
|
|
return Config(
|
|
kwargs=dict(entry["kwargs"]),
|
|
num_warps=entry.get("num_warps", 4),
|
|
num_stages=entry.get("num_stages", 3),
|
|
num_ctas=entry.get("num_ctas", 1),
|
|
maxnreg=entry.get("maxnreg"),
|
|
)
|
|
|
|
|
|
def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Optional[Tuple[Any, ...]]:
|
|
# Reproduce Autotuner/LibTuner key computation; return None on any internals
|
|
# shift (caller then just loses replay for that call).
|
|
try:
|
|
arg_names = tuner.arg_names
|
|
nargs = dict(zip(arg_names, args))
|
|
all_args = {**nargs, **kwargs}
|
|
_args = {k: v for k, v in all_args.items() if k in arg_names}
|
|
if hasattr(tuner, "get_key"):
|
|
return tuner.get_key(_args)
|
|
key = [_args[name] for name in tuner.keys if name in _args]
|
|
for _, arg in _args.items():
|
|
if hasattr(arg, "dtype"):
|
|
key.append(str(arg.dtype))
|
|
return tuple(key)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _emit_marker(reason: str) -> None:
|
|
# Marker lands in run.log (grep it to audit replay coverage); flush so it
|
|
# appears before the kernel runs. Yellow on the live terminal only —
|
|
# run_pytest.sh strips ANSI from run.log.
|
|
from _term_style import YELLOW, paint
|
|
print(paint(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", YELLOW), flush=True)
|
|
|
|
|
|
def _record_run(original):
|
|
# Snapshot self.cache before/after run() to capture the chosen config (works
|
|
# for both Autotuner and LibTuner). Single-config kernels skip the cache
|
|
# write, so record configs[0] explicitly for uniform replay.
|
|
def runner(self, *args, **kwargs):
|
|
try:
|
|
keys_before = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
|
|
except Exception:
|
|
keys_before = set()
|
|
result = original(self, *args, **kwargs)
|
|
try:
|
|
keys_after = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
|
|
except Exception:
|
|
keys_after = set()
|
|
new_keys = keys_after - keys_before
|
|
if not new_keys:
|
|
# No new entry: record configs[0] for single-config kernels (cache
|
|
# write bypassed); otherwise nothing to record (disk-cache hit).
|
|
if len(getattr(self, "configs", []) or []) == 1:
|
|
key = _compute_key(self, args, kwargs)
|
|
if key is not None and key not in self.cache:
|
|
cfg = self.configs[0]
|
|
entry = _serialize_config(cfg)
|
|
if entry is not None:
|
|
kid = _kernel_id(self)
|
|
with _record_lock:
|
|
bucket = _record_map.setdefault(kid, {})
|
|
bucket[_serialize_key(key)] = entry
|
|
return result
|
|
kid = _kernel_id(self)
|
|
with _record_lock:
|
|
bucket = _record_map.setdefault(kid, {})
|
|
for k in new_keys:
|
|
try:
|
|
cfg = self.cache[k]
|
|
except Exception:
|
|
continue
|
|
entry = _serialize_config(cfg)
|
|
if entry is None:
|
|
continue
|
|
bucket[_serialize_key(k)] = entry
|
|
return result
|
|
return runner
|
|
|
|
|
|
def _replay_run(original):
|
|
# Inject recorded configs before benchmarking; on key-missing or a recorded
|
|
# config failing to compile/launch, fall back to full autotune (+ marker).
|
|
def runner(self, *args, **kwargs):
|
|
injected_key = None
|
|
if len(getattr(self, "configs", []) or []) > 1:
|
|
kid = _kernel_id(self)
|
|
bucket = _replay_map.get(kid)
|
|
if bucket:
|
|
key = _compute_key(self, args, kwargs)
|
|
if key is not None and key not in self.cache:
|
|
rec = bucket.get(_serialize_key(key))
|
|
if rec is None:
|
|
_emit_marker("key_missing")
|
|
else:
|
|
try:
|
|
cfg = _deserialize_config(rec)
|
|
self.cache[key] = cfg
|
|
injected_key = key
|
|
except Exception as exc:
|
|
_emit_marker(f"deserialize_{type(exc).__name__}")
|
|
try:
|
|
return original(self, *args, **kwargs)
|
|
except Exception as exc:
|
|
if injected_key is not None:
|
|
# Recorded config failed under the current build: drop it and
|
|
# let run() autotune.
|
|
try:
|
|
del self.cache[injected_key]
|
|
except Exception:
|
|
pass
|
|
_emit_marker(f"compile_{type(exc).__name__}")
|
|
return original(self, *args, **kwargs)
|
|
raise
|
|
return runner
|
|
|
|
|
|
def _load_replay_dir(replay_dir: str) -> int:
|
|
path = _record_path(replay_dir)
|
|
if not path.is_file():
|
|
_emit_marker("record_missing")
|
|
return 0
|
|
try:
|
|
data = json.loads(path.read_text())
|
|
except Exception as exc:
|
|
_emit_marker(f"load_{type(exc).__name__}")
|
|
return 0
|
|
if not isinstance(data, dict):
|
|
return 0
|
|
count = 0
|
|
for kid, bucket in data.items():
|
|
if not isinstance(bucket, dict):
|
|
continue
|
|
_replay_map[kid] = dict(bucket)
|
|
count += len(bucket)
|
|
return count
|
|
|
|
|
|
def _dump_record(record_dir: str) -> None:
|
|
global _dump_done
|
|
if _dump_done:
|
|
return
|
|
_dump_done = True
|
|
path = _record_path(record_dir)
|
|
try:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with _record_lock:
|
|
payload = {kid: dict(bucket) for kid, bucket in _record_map.items()}
|
|
path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n")
|
|
except Exception as exc:
|
|
# Don't break the run if dump fails; the perf row is still produced.
|
|
from _term_style import tag
|
|
print(f"{tag('[autotune-record-plugin]')} dump failed: {exc}",
|
|
file=sys.stderr, flush=True)
|
|
|
|
|
|
def pytest_configure(config):
|
|
from _term_style import tag
|
|
record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
|
|
replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
|
|
if record_dir and replay_dir:
|
|
# Mutually exclusive: recording while replaying is pointless. Surface loudly.
|
|
print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
|
|
"REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
|
|
return
|
|
if not record_dir and not replay_dir:
|
|
return
|
|
|
|
import triton.runtime.autotuner as _autotuner
|
|
wrap = _record_run if record_dir else _replay_run
|
|
_autotuner.Autotuner.run = wrap(_autotuner.Autotuner.run)
|
|
patched = ["Autotuner"]
|
|
try:
|
|
from flag_gems.utils.libentry import LibTuner
|
|
if "run" in LibTuner.__dict__:
|
|
LibTuner.run = wrap(LibTuner.__dict__["run"])
|
|
patched.append("LibTuner")
|
|
except Exception:
|
|
pass
|
|
|
|
if record_dir:
|
|
# atexit (not sessionfinish): persist whatever was recorded even if an op
|
|
# crash kills the session; a later replay run falls back for missing keys.
|
|
atexit.register(_dump_record, record_dir)
|
|
print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
|
|
f"{_record_path(record_dir)} ({'+'.join(patched)})",
|
|
file=sys.stderr, flush=True)
|
|
else:
|
|
loaded = _load_replay_dir(replay_dir)
|
|
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
|
|
f"{_record_path(replay_dir)} ({'+'.join(patched)})",
|
|
file=sys.stderr, flush=True)
|