forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
29af59be9b
Three pieces from the fork's feature/musa-benchmark-support branch that never reached main, rebased onto the backend abstraction rather than copied verbatim. _musa_benchmark_compat_plugin drops FlagGems benchmark skip guards written as `torch.cuda.is_available()` where the code under test dispatches through flag_gems.device. On MThreads those turn a working operator into "not measured": test_fft.py skipped entirely, and unskipped on an S5000 it produces 25 rows. The plugin pins each override to one file and one exact skip reason and only acts when flag_gems has dispatched to the mthreads vendor; two guards were probed on device and deliberately left in place, which is why they are named in its docstring rather than silently absent. It is wired through a new DeviceBackend.compat_plugins rather than a name test in the planner, so the policy stays where every other backend difference lives. Compat plugins are planned first because they rewrite collection-time markers and nothing may read the item list before them; a test pins that position and another pins that they do not leak to other backends. batch_summary gains --ab: the speedup column compares gems against torch, not one build against another, so an A/B needs the two run.logs paired by operator, dtype and size detail. Only rows that succeeded on both sides with a positive latency count, so an incomplete pairing cannot look like a complete range, and the range is annotated when either side carries a BENCHMARK_DIRECT_NO_CUDAGRAPH marker, since a fallback row was timed under a different regime. ab_rlc_test.sh is the same caliber as ab_fold_test.sh with the comparison axis moved to FLAGTREE_<VENDOR>_RLC_ENHANCE. Both are caliber examples, not regression scripts: if the switch under test no longer exists the script still runs every round and prints a full table showing no difference, so confirm the axis is live before reusing either. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
"""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)
|
|
# Compat first: these override collection-time markers, so they have to be
|
|
# in place before anything else inspects the collected items.
|
|
plugins: List[str] = list(backend.compat_plugins) + 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())
|