diff --git a/docs/testing_benchmarking_guide.md b/docs/testing_benchmarking_guide.md index 3ceb8d7eb..cb971b714 100644 --- a/docs/testing_benchmarking_guide.md +++ b/docs/testing_benchmarking_guide.md @@ -194,7 +194,27 @@ def insert_point(ctx): ## 4. Performance Measurement -### 4.1 `tests/test_common.py` +### 4.1 `flydsl.profiling.do_bench` + +Use the package-level event timer when a benchmark or tuning workflow only +needs GPU latency: + +```python +from flydsl.profiling import do_bench + +latency_ms = do_bench(lambda: launch_kernel(...), warmup=5, rep=25) +``` + +`warmup` and `rep` are iteration counts. The default result is the +upper-middle latency in milliseconds; pass a non-empty `quantiles` sequence to +return selected sorted samples. The helper uses PyTorch's CUDA-compatible event +interface, so the same call works with PyTorch CUDA and HIP builds. The measured +callable must enqueue its work on PyTorch's current stream. Use a non-negative +`warmup`, a positive `rep`, and quantiles in the inclusive range `[0, 1]`. +Pass `setup` to restore or reset inputs before every iteration without including +that work in the timed event interval. + +### 4.2 `tests/test_common.py` Core performance testing utilities (adapted from AIter). @@ -224,7 +244,7 @@ verify_output(c_out, c_ref, atol=1e-2, rtol=1e-2, msg='') ``` High-level validation wrapper around `checkAllclose`. -### 4.2 `tests/kernels/benchmark_common.py` +### 4.3 `tests/kernels/benchmark_common.py` Shared benchmark harness for performance comparison. diff --git a/python/flydsl/autotune.py b/python/flydsl/autotune.py index b243d528b..5783953ae 100644 --- a/python/flydsl/autotune.py +++ b/python/flydsl/autotune.py @@ -9,10 +9,7 @@ from pathlib import Path from typing import Callable, Dict, List -try: - import torch -except ImportError: - torch = None +from .profiling import do_bench as do_bench def _env_fingerprint() -> tuple: @@ -129,26 +126,6 @@ def from_dict(cls, d): ) -def do_bench(fn, warmup=5, rep=25, quantiles=None): - """Benchmark a GPU kernel using CUDA/HIP events. Returns median ms.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - times = [] - for _ in range(rep): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - torch.cuda.synchronize() - times.append(start.elapsed_time(end)) - times.sort() - if quantiles: - return [times[min(int(q * len(times)), len(times) - 1)] for q in quantiles] - return times[len(times) // 2] - - class Autotuner: """Wraps a @jit function, benchmarks configs, caches best.""" diff --git a/python/flydsl/profiling.py b/python/flydsl/profiling.py new file mode 100644 index 000000000..0bf70d395 --- /dev/null +++ b/python/flydsl/profiling.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""GPU event timing shared across FlyDSL compile backends.""" + +from typing import Callable, List, Optional, Sequence, Union + +__all__ = ["do_bench"] + + +def _get_torch_cuda(): + """Return PyTorch's CUDA-compatible namespace lazily. + + PyTorch exposes the same ``torch.cuda`` event API on CUDA and HIP builds, + so callers do not need to branch on the FlyDSL compile backend. + """ + try: + import torch + except ImportError as exc: + raise RuntimeError("GPU profiling requires PyTorch with CUDA or HIP support") from exc + if not torch.cuda.is_available(): + raise RuntimeError("GPU profiling requires an available CUDA or HIP device") + return torch.cuda + + +def do_bench( + fn: Callable[[], object], + warmup: int = 5, + rep: int = 25, + quantiles: Optional[Sequence[float]] = None, + setup: Optional[Callable[[], object]] = None, +) -> Union[float, List[float]]: + """Benchmark a GPU callable with CUDA/HIP events. + + ``warmup`` and ``rep`` are iteration counts. Timings are returned in + milliseconds. By default the upper-middle sample is returned; when a + non-empty ``quantiles`` sequence is provided, the corresponding sorted + samples are returned. ``warmup`` must be non-negative, ``rep`` must be + positive, and quantiles must be in the inclusive range ``[0, 1]``. + + ``fn`` must enqueue the measured work on PyTorch's current CUDA/HIP stream. + Work submitted only to another stream is outside the event interval unless + ``fn`` synchronizes that stream itself. + + When provided, ``setup`` runs before every warmup and measured iteration. + For measured iterations it runs before the start event, so restore/reset + work is not included in the reported kernel latency. + """ + if warmup < 0: + raise ValueError("warmup must be non-negative") + if rep <= 0: + raise ValueError("rep must be positive") + if quantiles is not None and any(not 0.0 <= q <= 1.0 for q in quantiles): + raise ValueError("quantiles must be between 0 and 1") + + device = _get_torch_cuda() + + for _ in range(warmup): + if setup is not None: + setup() + fn() + device.synchronize() + + times = [] + for _ in range(rep): + if setup is not None: + setup() + start = device.Event(enable_timing=True) + end = device.Event(enable_timing=True) + start.record() + fn() + end.record() + device.synchronize() + times.append(start.elapsed_time(end)) + + times.sort() + if quantiles: + return [times[min(int(q * len(times)), len(times) - 1)] for q in quantiles] + return times[len(times) // 2] diff --git a/tests/unit/test_profiling.py b/tests/unit/test_profiling.py new file mode 100644 index 000000000..423a589d9 --- /dev/null +++ b/tests/unit/test_profiling.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""GPU-free tests for the shared profiling helpers.""" + +import sys +from types import SimpleNamespace + +import pytest + +import flydsl.profiling as profiling +from flydsl.autotune import do_bench as autotune_do_bench + +pytestmark = pytest.mark.l0_backend_agnostic + + +class FakeEvent: + def __init__(self, device): + self.device = device + self.timestamp = None + + def record(self): + self.timestamp = self.device.clock + + def elapsed_time(self, end): + return end.timestamp - self.timestamp + + +class FakeDeviceInterface: + def __init__(self): + self.clock = 0.0 + self.event_count = 0 + self.synchronize_count = 0 + + def Event(self, *, enable_timing): + assert enable_timing + self.event_count += 1 + return FakeEvent(self) + + def synchronize(self): + self.synchronize_count += 1 + + +def _install_fake(monkeypatch): + device = FakeDeviceInterface() + monkeypatch.setattr(profiling, "_get_torch_cuda", lambda: device) + return device + + +def test_do_bench_warmup_repetitions_and_median(monkeypatch): + device = _install_fake(monkeypatch) + durations = iter([99.0, 3.0, 1.0, 5.0, 2.0, 4.0]) + call_count = 0 + + def fn(): + nonlocal call_count + call_count += 1 + device.clock += next(durations) + + assert profiling.do_bench(fn, warmup=1, rep=5) == 3.0 + assert call_count == 6 + assert device.event_count == 10 + assert device.synchronize_count == 6 + + +def test_do_bench_preserves_upper_middle_for_even_repetitions(monkeypatch): + device = _install_fake(monkeypatch) + durations = iter([4.0, 1.0, 3.0, 2.0]) + + def fn(): + device.clock += next(durations) + + assert profiling.do_bench(fn, warmup=0, rep=4) == 3.0 + + +def test_do_bench_quantiles(monkeypatch): + device = _install_fake(monkeypatch) + durations = iter([3.0, 1.0, 5.0, 2.0, 4.0]) + + def fn(): + device.clock += next(durations) + + assert profiling.do_bench(fn, warmup=0, rep=5, quantiles=[0.0, 0.5, 0.9, 1.0]) == [1.0, 3.0, 5.0, 5.0] + + +def test_do_bench_runs_setup_before_each_iteration_and_outside_timing(monkeypatch): + device = _install_fake(monkeypatch) + durations = iter([99.0, 1.0, 2.0]) + order = [] + + def setup(): + order.append("setup") + device.clock += 100.0 + + def fn(): + order.append("kernel") + device.clock += next(durations) + + assert profiling.do_bench(fn, warmup=1, rep=2, setup=setup) == 2.0 + assert order == ["setup", "kernel", "setup", "kernel", "setup", "kernel"] + + +def test_do_bench_reports_missing_pytorch(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + with pytest.raises(RuntimeError, match="requires PyTorch"): + profiling._get_torch_cuda() + + +def test_do_bench_reports_unavailable_device(monkeypatch): + cuda = SimpleNamespace(is_available=lambda: False) + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(cuda=cuda)) + with pytest.raises(RuntimeError, match="available CUDA or HIP device"): + profiling._get_torch_cuda() + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"warmup": -1, "rep": 1}, "warmup must be non-negative"), + ({"warmup": 0, "rep": 0}, "rep must be positive"), + ({"warmup": 0, "rep": 1, "quantiles": [-0.1]}, "quantiles must be between 0 and 1"), + ({"warmup": 0, "rep": 1, "quantiles": [1.1]}, "quantiles must be between 0 and 1"), + ], +) +def test_do_bench_rejects_invalid_arguments(kwargs, message): + with pytest.raises(ValueError, match=message): + profiling.do_bench(lambda: None, **kwargs) + + +def test_autotune_keeps_compatibility_alias(): + assert autotune_do_bench is profiling.do_bench