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:
AharrypotterandClaude Opus 5 committed 2026-09-09 16:39:14 +08:00
1 parent 35a928a516
commit c8ea586d0f
24 files changed
+1448 -159

No files matched your search

+71 -1
View File
@@ -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)