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>

backends/__init__.py

@@ -0,0 +1,76 @@
+"""Registry of zl_bench device backends.
+
+Aliases map common vendor spellings onto the canonical backend name. Unknown
+names fail closed rather than defaulting to CUDA: a silent CUDA fallback would
+load the cudagraph plugin on the wrong device.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Dict, Optional
+
+from .ascend import AscendBackend
+from .base import DeviceBackend, parse_device_list
+from .cuda import CudaBackend
+from .musa import MusaBackend
+from .rocm import RocmBackend
+
+ALIASES: Dict[str, str] = {
+    "nvidia": "cuda",
+    "nv": "cuda",
+    "amd": "rocm",
+    "hip": "rocm",
+    "npu": "ascend",
+    "huawei": "ascend",
+}
+
+_INSTANCES: Dict[str, DeviceBackend] = {
+    "cuda": CudaBackend(),
+    "musa": MusaBackend(),
+    "rocm": RocmBackend(),
+    "ascend": AscendBackend(),
+}
+
+
+class UnknownBackendError(ValueError):
+    pass
+
+
+def known_backends() -> tuple[str, ...]:
+    return tuple(sorted(_INSTANCES))
+
+
+def normalize_backend_name(name: str) -> str:
+    key = (name or "").strip().lower()
+    if not key:
+        return "cuda"
+    return ALIASES.get(key, key)
+
+
+def get_backend(name: Optional[str] = None) -> DeviceBackend:
+    raw = name if name is not None else os.environ.get("ZL_BACKEND", "cuda")
+    key = normalize_backend_name(raw)
+    backend = _INSTANCES.get(key)
+    if backend is None:
+        raise UnknownBackendError(
+            f"unknown backend {raw!r}; known: {', '.join(known_backends())} "
+            f"(aliases: {', '.join(sorted(ALIASES))})"
+        )
+    return backend
+
+
+def enumerate_devices(backend: Optional[DeviceBackend] = None,
+                      explicit: Optional[str] = None) -> list[str]:
+    return (backend or get_backend()).enumerate_devices(explicit)
+
+
+__all__ = [
+    "DeviceBackend",
+    "UnknownBackendError",
+    "enumerate_devices",
+    "get_backend",
+    "known_backends",
+    "normalize_backend_name",
+    "parse_device_list",
+]

backends/ascend.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class AscendBackend(DeviceBackend):
+    name = "ascend"
+    visibility_env_keys = (
+        "ASCEND_RT_VISIBLE_DEVICES",
+        "ASCEND_VISIBLE_DEVICES",
+    )
+    torch_namespace = "npu"
+    flaggems_vendor = "ascend"
+    supports_cuda_graph = False
+    supports_mm_cluster_fix = False
+    supports_triton_hooks = True
+    list_command = None

backends/base.py

@@ -0,0 +1,148 @@
+"""Device backend contract for zl_bench screening runs.
+
+This is the orchestration-layer adapter: device enumeration, process visibility
+env, and capability flags that decide which pytest plugins are safe to load.
+It does not implement kernel timing receipts or compiler-artifact identity;
+those stay with the caller (FlagGems pytest for screening, a separate harness
+for qualification).
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+# Every vendor visibility variable zl_bench knows about. apply_visibility()
+# clears the whole set before writing the active backend's keys so a parent
+# CUDA_VISIBLE_DEVICES cannot leak onto a MUSA/Ascend child.
+ALL_VISIBILITY_KEYS: Tuple[str, ...] = (
+    "CUDA_VISIBLE_DEVICES",
+    "MUSA_VISIBLE_DEVICES",
+    "HIP_VISIBLE_DEVICES",
+    "ROCR_VISIBLE_DEVICES",
+    "ASCEND_RT_VISIBLE_DEVICES",
+    "ASCEND_VISIBLE_DEVICES",
+)
+
+
+def parse_device_list(raw: Optional[str]) -> List[str]:
+    if not raw:
+        return []
+    return [tok.strip() for tok in raw.split(",") if tok.strip()]
+
+
+class DeviceBackend:
+    """One hardware runtime (CUDA, MUSA, ROCm, Ascend, ...).
+
+    Subclasses only fill class attributes and optional probe hooks. Callers
+    should not branch on ``backend.name`` for visibility or plugin policy;
+    they should read these attributes / methods.
+    """
+
+    name: str = ""
+    # First key is the canonical one written into run_identity.json.
+    visibility_env_keys: Tuple[str, ...] = ()
+    # torch.<namespace> used for optional device_count / identity probes.
+    torch_namespace: str = ""
+    # FlagGems GEMS_VENDOR value. Lives on the backend object because zl_bench
+    # still only drives FlagGems, but the mapping is owned by providers.flaggems
+    # and copied here so DeviceBackend.prepare_environment() can stay generic.
+    flaggems_vendor: str = ""
+    supports_cuda_graph: bool = False
+    supports_mm_cluster_fix: bool = False
+    supports_triton_hooks: bool = True
+    # Optional listing command; (argv, line_prefix). prefix="" means count
+    # non-empty lines. Must be a real vendor CLI; do not invent flags.
+    list_command: Optional[Tuple[Sequence[str], str]] = None
+
+    def visibility_env(self, device: str) -> Dict[str, str]:
+        return {key: device for key in self.visibility_env_keys}
+
+    def apply_visibility(self, env: Dict[str, str], device: str) -> Dict[str, str]:
+        out = dict(env)
+        for key in ALL_VISIBILITY_KEYS:
+            out.pop(key, None)
+        out.update(self.visibility_env(device))
+        return out
+
+    def enumerate_devices(self, explicit: Optional[str] = None) -> List[str]:
+        listed = parse_device_list(explicit)
+        if listed:
+            return listed
+        for key in self.visibility_env_keys:
+            listed = parse_device_list(os.environ.get(key, ""))
+            if listed:
+                return listed
+        n = self.probe_device_count()
+        if n > 0:
+            return [str(i) for i in range(n)]
+        return ["0"]
+
+    def probe_device_count(self) -> int:
+        n = self._probe_list_command()
+        if n > 0:
+            return n
+        return self._probe_torch_count()
+
+    def _probe_list_command(self) -> int:
+        if not self.list_command:
+            return 0
+        argv, prefix = self.list_command
+        try:
+            proc = subprocess.run(
+                list(argv), text=True, capture_output=True, timeout=30,
+            )
+        except Exception:
+            return 0
+        if proc.returncode != 0:
+            return 0
+        lines = [ln for ln in proc.stdout.splitlines() if ln.strip()]
+        if prefix:
+            lines = [ln for ln in lines if ln.startswith(prefix)]
+        return len(lines)
+
+    def _probe_torch_count(self) -> int:
+        if not self.torch_namespace:
+            return 0
+        try:
+            import torch
+            api = getattr(torch, self.torch_namespace, None)
+            if api is None or not hasattr(api, "device_count"):
+                return 0
+            return int(api.device_count())
+        except Exception:
+            return 0
+
+    def device_identity(self) -> Dict[str, str]:
+        """Best-effort identity for screening logs. Empty strings are allowed."""
+        info = {
+            "backend": self.name,
+            "device_name": "",
+            "device_arch": "",
+            "timer_api": f"torch.{self.torch_namespace}" if self.torch_namespace else "",
+            "flaggems_vendor": self.flaggems_vendor,
+        }
+        if not self.torch_namespace:
+            return info
+        try:
+            import torch
+            api = getattr(torch, self.torch_namespace, None)
+            if api is None:
+                return info
+            count = int(api.device_count()) if hasattr(api, "device_count") else 0
+            if count <= 0:
+                return info
+            index = int(api.current_device()) if hasattr(api, "current_device") else 0
+            if hasattr(api, "get_device_name"):
+                info["device_name"] = str(api.get_device_name(index))
+            if hasattr(api, "get_device_capability"):
+                cap = api.get_device_capability(index)
+                info["device_arch"] = f"{cap[0]}.{cap[1]}"
+            elif hasattr(api, "get_device_properties"):
+                props = api.get_device_properties(index)
+                arch = getattr(props, "name", "") or getattr(props, "arch", "")
+                info["device_arch"] = str(arch)
+        except Exception:
+            return info
+        return info

backends/cuda.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class CudaBackend(DeviceBackend):
+    name = "cuda"
+    visibility_env_keys = ("CUDA_VISIBLE_DEVICES",)
+    torch_namespace = "cuda"
+    flaggems_vendor = "nvidia"
+    supports_cuda_graph = True
+    supports_mm_cluster_fix = True
+    supports_triton_hooks = True
+    list_command = (("nvidia-smi", "--list-gpus"), "GPU ")

backends/musa.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class MusaBackend(DeviceBackend):
+    name = "musa"
+    visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
+    torch_namespace = "musa"
+    flaggems_vendor = "musa"
+    supports_cuda_graph = False
+    supports_mm_cluster_fix = False
+    supports_triton_hooks = True
+    # Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
+    # a vendor CLI here; confirm the real listing tool before adding it.
+    list_command = None

backends/rocm.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class RocmBackend(DeviceBackend):
+    name = "rocm"
+    # HIP runtimes honor both; write both so a parent CUDA/HIP mix cannot
+    # leave the child looking at the wrong cards.
+    visibility_env_keys = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES")
+    torch_namespace = "cuda"  # ROCm PyTorch still exposes torch.cuda
+    flaggems_vendor = "amd"
+    supports_cuda_graph = False
+    supports_mm_cluster_fix = False
+    supports_triton_hooks = True
+    list_command = None