forked from zhoulin/zl_bench
feat: make the harness backend-agnostic
Everything about which accelerator a run targets was previously spelled CUDA: device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES, the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in run_pytest.sh with no way to drop the CUDA-only members. Running the same operator corpus on another vendor's device therefore meant editing the drivers. The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with the usual aliases): - backends/ owns per-backend device enumeration, visibility env keys and capability flags, so nothing above it names a vendor; - plugin_plan.py derives the -p list from those capabilities instead of a hardcoded array, which is what lets a non-CUDA run omit the graph-capture and cluster-fix plugins rather than watch them fail; - _timing_plugin selects the timer from the same capabilities: CUDA keeps graph capture with its per-shape fallback, other backends stay on upstream do_bench and record timer=direct, because capture-then-fallback on a device without graph support only produces noise; - _device_setup_plugin replaces the NVIDIA-only vendor guard; - providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to FlagGems' own core_shapes; - _run_identity.py writes run_identity.json next to each run, recording backend, device, timer, autotune mode and the resolved test node. Two runs whose calibers differ are otherwise indistinguishable after the fact. run_batch gains --backend and --devices (--gpus kept as an alias) and carries backend/device/timer into ops_status.csv. test_backends.py covers the registry, alias resolution, visibility env and plugin planning. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
1 parent
35a928a516
commit
c8ea586d0f
24 files changed
+1448
-159
No files matched your search
@@ -214,6 +214,57 @@ def _replay_run(original):
|
||||
return runner
|
||||
|
||||
|
||||
def _write_record_meta(record_dir: str) -> None:
|
||||
"""Sidecar next to <op>.json. Replay refuses a different backend.
|
||||
|
||||
The config JSON itself stays `{kernel: {key: cfg}}` so existing REPLAY_FROM
|
||||
directories keep working.
|
||||
"""
|
||||
try:
|
||||
from backends import get_backend
|
||||
backend_name = get_backend().name
|
||||
except Exception:
|
||||
backend_name = os.environ.get("ZL_BACKEND", "")
|
||||
meta = {
|
||||
"schema_version": 1,
|
||||
"provider": "flaggems",
|
||||
"backend": backend_name,
|
||||
"op": _op_name(),
|
||||
}
|
||||
path = Path(record_dir) / "meta.json"
|
||||
try:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(meta, indent=2, ensure_ascii=False) + "\n")
|
||||
except Exception as exc:
|
||||
from _term_style import tag
|
||||
print(f"{tag('[autotune-record-plugin]')} meta.json write failed: {exc}",
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def _replay_backend_ok(replay_dir: str) -> bool:
|
||||
path = Path(replay_dir) / "meta.json"
|
||||
if not path.is_file():
|
||||
return True
|
||||
try:
|
||||
meta = json.loads(path.read_text())
|
||||
except Exception:
|
||||
return True
|
||||
if not isinstance(meta, dict):
|
||||
return True
|
||||
recorded = str(meta.get("backend") or "").strip()
|
||||
if not recorded:
|
||||
return True
|
||||
try:
|
||||
from backends import get_backend
|
||||
current = get_backend().name
|
||||
except Exception:
|
||||
current = os.environ.get("ZL_BACKEND", "").strip()
|
||||
if recorded != current:
|
||||
_emit_marker(f"backend_mismatch recorded={recorded} current={current}")
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _load_replay_dir(replay_dir: str) -> int:
|
||||
path = _record_path(replay_dir)
|
||||
if not path.is_file():
|
||||
@@ -256,29 +307,48 @@ def _dump_record(record_dir: str) -> None:
|
||||
def pytest_configure(config):
|
||||
global _wrap
|
||||
from _term_style import tag
|
||||
from _run_identity import update_identity
|
||||
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)
|
||||
update_identity(autotune_mode="uncontrolled")
|
||||
return
|
||||
if not record_dir and not replay_dir:
|
||||
update_identity(autotune_mode="uncontrolled")
|
||||
return
|
||||
|
||||
try:
|
||||
import triton.runtime.autotuner as _autotuner
|
||||
except Exception as exc:
|
||||
print(f"{tag('[autotune-record-plugin]')} Triton Autotuner unavailable "
|
||||
f"({exc}); autotune_mode=uncontrolled", file=sys.stderr, flush=True)
|
||||
update_identity(autotune_mode="uncontrolled")
|
||||
return
|
||||
|
||||
if replay_dir and not _replay_backend_ok(replay_dir):
|
||||
print(f"{tag('[autotune-record-plugin]')} refusing replay across backends",
|
||||
file=sys.stderr, flush=True)
|
||||
update_identity(autotune_mode="replay_fallback")
|
||||
return
|
||||
|
||||
import triton.runtime.autotuner as _autotuner
|
||||
_wrap = _record_run if record_dir else _replay_run
|
||||
_autotuner.Autotuner.run = _wrap(_autotuner.Autotuner.run)
|
||||
|
||||
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.
|
||||
_write_record_meta(record_dir)
|
||||
atexit.register(_dump_record, record_dir)
|
||||
update_identity(autotune_mode="record")
|
||||
print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
|
||||
f"{_record_path(record_dir)} (Autotuner)",
|
||||
file=sys.stderr, flush=True)
|
||||
else:
|
||||
loaded = _load_replay_dir(replay_dir)
|
||||
update_identity(autotune_mode="replay")
|
||||
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
|
||||
f"{_record_path(replay_dir)} (Autotuner)",
|
||||
file=sys.stderr, flush=True)
|
||||
|
||||
Reference in new issue
Block a user