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,275 @@
|
||||
"""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.
|
||||
print(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", 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.
|
||||
print(f"[autotune-record-plugin] dump failed: {exc}", file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
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("[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"[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"[autotune-record-plugin] replaying {loaded} recorded entries from "
|
||||
f"{_record_path(replay_dir)} ({'+'.join(patched)})",
|
||||
file=sys.stderr, flush=True)
|
||||
Reference in New Issue
Block a user