forked from zhoulin/zl_bench
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:
1 parent
35a928a516
commit
c8ea586d0f
24 files changed
+1448
-159
No files matched your search
+107
@@ -0,0 +1,107 @@
|
||||
"""Decide which pytest plugins a zl_bench run should load.
|
||||
|
||||
CUDA-specific plugins (cudagraph patch, Hopper mm cluster workaround) are not
|
||||
placed on the default list; `_timing_plugin` enables graph capture only when
|
||||
the backend declares `supports_cuda_graph`. IR / autotune plugins stay on for
|
||||
Triton-compatible compilers and can be turned off with env flags.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from typing import List, Optional, Sequence
|
||||
|
||||
from backends import get_backend
|
||||
|
||||
CORE_PLUGINS: tuple[str, ...] = (
|
||||
"_device_setup_plugin",
|
||||
"_parallel_warmup_plugin",
|
||||
"_seed_plugin",
|
||||
"_shape_inject_plugin",
|
||||
"_shape_iter_inject_plugin",
|
||||
"_bespoke_shape_plugin",
|
||||
)
|
||||
|
||||
ALWAYS_AFTER: tuple[str, ...] = (
|
||||
"_pretty_report_plugin",
|
||||
)
|
||||
|
||||
|
||||
def _env_flag(name: str, default: str = "1") -> bool:
|
||||
raw = os.environ.get(name, default).strip().lower()
|
||||
return raw not in ("0", "false", "no", "off", "")
|
||||
|
||||
|
||||
def plan_plugins(
|
||||
backend_name: Optional[str] = None,
|
||||
extra: Optional[Sequence[str]] = None,
|
||||
enable_autotune: Optional[bool] = None,
|
||||
enable_ir: Optional[bool] = None,
|
||||
enable_mm_cluster: Optional[bool] = None,
|
||||
) -> List[str]:
|
||||
backend = get_backend(backend_name)
|
||||
plugins: List[str] = list(CORE_PLUGINS)
|
||||
|
||||
autotune = (
|
||||
backend.supports_triton_hooks
|
||||
if enable_autotune is None
|
||||
else enable_autotune
|
||||
)
|
||||
if enable_autotune is None:
|
||||
autotune = autotune and _env_flag("ZL_ENABLE_AUTOTUNE", "1")
|
||||
if autotune:
|
||||
plugins.append("_autotune_record_plugin")
|
||||
|
||||
plugins.append("_timing_plugin")
|
||||
plugins.extend(ALWAYS_AFTER)
|
||||
|
||||
ir = backend.supports_triton_hooks if enable_ir is None else enable_ir
|
||||
if enable_ir is None:
|
||||
ir = ir and _env_flag("ZL_ENABLE_IR", "1")
|
||||
if ir:
|
||||
plugins.append("_ir_meta_plugin")
|
||||
|
||||
mm = enable_mm_cluster
|
||||
if mm is None:
|
||||
mm = backend.supports_mm_cluster_fix and _env_flag("ZL_ENABLE_MM_CLUSTER", "0")
|
||||
elif mm:
|
||||
mm = backend.supports_mm_cluster_fix
|
||||
if mm:
|
||||
plugins.append("_mm_cluster_fix_plugin")
|
||||
|
||||
if extra is None:
|
||||
extras = [p for p in os.environ.get("ZL_EXTRA_PLUGINS", "").split() if p]
|
||||
else:
|
||||
extras = [p for p in extra if p]
|
||||
for name in extras:
|
||||
if name not in plugins:
|
||||
plugins.append(name)
|
||||
return plugins
|
||||
|
||||
|
||||
def pytest_args(plugins: Optional[Sequence[str]] = None, **kwargs) -> List[str]:
|
||||
names = list(plugins) if plugins is not None else plan_plugins(**kwargs)
|
||||
args: List[str] = []
|
||||
for name in names:
|
||||
args.extend(("-p", name))
|
||||
return args
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
p = argparse.ArgumentParser(description="Print zl_bench pytest -p args")
|
||||
p.add_argument("--backend", default=None)
|
||||
p.add_argument("--pytest-args", action="store_true",
|
||||
help="print `-p name` tokens on one line")
|
||||
p.add_argument("--one-per-line", action="store_true")
|
||||
args = p.parse_args(argv)
|
||||
names = plan_plugins(backend_name=args.backend)
|
||||
if args.one_per_line:
|
||||
print("\n".join(names))
|
||||
else:
|
||||
print(" ".join(pytest_args(names)))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in new issue
Block a user