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>
15 lines
392 B
Python
15 lines
392 B
Python
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 ")
|