zl_bench/_autotune_record_plugin.py

"""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):
    # Read back this call's own key after run(), rather than diffing cache keys.
    # LibTuner.cache is backed by a persistent sqlite DB (flag_gems libcache), so
    # on a warm DB the key is already present and a before/after diff comes up
    # empty -- which silently recorded nothing for every libtuner kernel.
    # Single-config kernels bypass the cache write, so fall back to configs[0].
    def runner(self, *args, **kwargs):
        result = original(self, *args, **kwargs)
        key = _compute_key(self, args, kwargs)
        if key is None:
            return result
        cfg = None
        try:
            cfg = self.cache[key]
        except Exception:
            cfg = None
        if cfg is None and len(getattr(self, "configs", []) or []) == 1:
            cfg = self.configs[0]
        if cfg is None:
            return result
        entry = _serialize_config(cfg)
        if entry is not None:
            kid = _kernel_id(self)
            with _record_lock:
                _record_map.setdefault(kid, {})[_serialize_key(key)] = 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)
                # Overwrite unconditionally: LibTuner.cache is backed by a
                # persistent sqlite DB, so gating on `key not in self.cache`
                # would skip injection whenever that DB is warm -- leaving the
                # run silently self-tuned instead of replaying.
                if key is not None:
                    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)