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>
70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
"""Set FlagGems vendor + record backend identity before flag_gems is imported.
|
|
|
|
Replaces the NVIDIA-only `_device_guard_plugin`. Setting GEMS_VENDOR from the
|
|
zl_bench backend name skips FlagGems' hang-prone nvidia-smi subprocess probe
|
|
on CUDA, and maps musa/rocm/ascend onto the matching FlagGems vendor.
|
|
|
|
Import-time side effect is intentional: FlagGems reads GEMS_VENDOR while the
|
|
benchmark module is being collected. Device name / arch are filled later in
|
|
pytest_configure so import does not initialize a device context.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
from backends import get_backend
|
|
from providers.flaggems import apply_vendor_env
|
|
from _run_identity import update_identity
|
|
|
|
|
|
def _setup_at_import() -> None:
|
|
try:
|
|
backend = get_backend()
|
|
except Exception as exc:
|
|
print(f"[device-setup-plugin] cannot resolve backend ({exc}); "
|
|
"leaving GEMS_VENDOR to flag_gems",
|
|
file=sys.stderr, flush=True)
|
|
return
|
|
vendor = apply_vendor_env(backend)
|
|
os.environ["ZL_BACKEND"] = backend.name
|
|
device = os.environ.get("ZL_DEVICE", "").strip()
|
|
if not device:
|
|
for key in backend.visibility_env_keys:
|
|
raw = os.environ.get(key, "").strip()
|
|
if raw:
|
|
device = raw.split(",")[0].strip()
|
|
break
|
|
update_identity(
|
|
backend=backend.name,
|
|
device=device,
|
|
flaggems_vendor=vendor or "",
|
|
visibility_env=backend.visibility_env_keys[0] if backend.visibility_env_keys else "",
|
|
)
|
|
if vendor:
|
|
from _term_style import tag
|
|
print(f"{tag('[device-setup-plugin]')} backend={backend.name} "
|
|
f"GEMS_VENDOR={vendor}",
|
|
file=sys.stderr, flush=True)
|
|
|
|
|
|
def pytest_configure(config):
|
|
try:
|
|
backend = get_backend()
|
|
ident = backend.device_identity()
|
|
update_identity(
|
|
backend=backend.name,
|
|
device_name=ident.get("device_name", ""),
|
|
device_arch=ident.get("device_arch", ""),
|
|
timer_api=ident.get("timer_api", ""),
|
|
flaggems_vendor=ident.get("flaggems_vendor", ""),
|
|
)
|
|
except Exception as exc:
|
|
from _term_style import tag
|
|
print(f"{tag('[device-setup-plugin]')} identity probe failed ({exc})",
|
|
file=sys.stderr, flush=True)
|
|
|
|
|
|
_setup_at_import()
|