forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
c8ea586d0f
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>
77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
"""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",
|
|
]
|