[Feature] Extract reusable event-based benchmarking helper

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.
 

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."""
 

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]

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