kernelmeter 0.2.0__py3-none-any.whl

kernelmeter/__init__.py

@@ -0,0 +1,28 @@
+"""kernelmeter: CUDA device attributes without profiling, and kernel
+benchmarks measured against the hardware's speed of light."""
+
+from .bench import REGISTRY, BenchResult, BenchSpec, benchmark, device_peaks, run, run_registry
+from .cudadrv import CudaDriverError, CudaNotAvailableError, Driver
+from . import occupancy, roofline
+from .occupancy import Occupancy
+from .peaks import Peaks
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "BenchResult",
+    "BenchSpec",
+    "CudaDriverError",
+    "CudaNotAvailableError",
+    "Driver",
+    "Occupancy",
+    "Peaks",
+    "REGISTRY",
+    "benchmark",
+    "device_peaks",
+    "occupancy",
+    "roofline",
+    "run",
+    "run_registry",
+    "__version__",
+]

kernelmeter/attrs.py

@@ -0,0 +1,149 @@
+"""Enumerate every device attribute the driver knows about.
+
+Nsight Compute's ``device__attribute_*`` values come straight from
+``cuDeviceGetAttribute``, so we ask the driver directly. Every id from 1
+to PROBE_MAX is probed: ids the driver rejects are skipped, ids that
+succeed but aren't in the name table yet are reported as ``attribute_<id>``.
+"""
+
+from __future__ import annotations
+
+from .cudadrv import DeviceHandle, Driver
+
+PROBE_MAX = 192
+
+# CU_DEVICE_ATTRIBUTE_* names (snake_cased), keyed by enum value.
+# Deprecated/reserved slots (44, 92-94, ...) are intentionally absent.
+KNOWN_ATTRS: dict[int, str] = {
+    1: "max_threads_per_block",
+    2: "max_block_dim_x",
+    3: "max_block_dim_y",
+    4: "max_block_dim_z",
+    5: "max_grid_dim_x",
+    6: "max_grid_dim_y",
+    7: "max_grid_dim_z",
+    8: "max_shared_memory_per_block",
+    9: "total_constant_memory",
+    10: "warp_size",
+    11: "max_pitch",
+    12: "max_registers_per_block",
+    13: "clock_rate_khz",
+    14: "texture_alignment",
+    15: "gpu_overlap",
+    16: "multiprocessor_count",
+    17: "kernel_exec_timeout",
+    18: "integrated",
+    19: "can_map_host_memory",
+    20: "compute_mode",
+    21: "maximum_texture1d_width",
+    22: "maximum_texture2d_width",
+    23: "maximum_texture2d_height",
+    24: "maximum_texture3d_width",
+    25: "maximum_texture3d_height",
+    26: "maximum_texture3d_depth",
+    27: "maximum_texture2d_layered_width",
+    28: "maximum_texture2d_layered_height",
+    29: "maximum_texture2d_layered_layers",
+    30: "surface_alignment",
+    31: "concurrent_kernels",
+    32: "ecc_enabled",
+    33: "pci_bus_id",
+    34: "pci_device_id",
+    35: "tcc_driver",
+    36: "memory_clock_rate_khz",
+    37: "global_memory_bus_width_bits",
+    38: "l2_cache_size",
+    39: "max_threads_per_multiprocessor",
+    40: "async_engine_count",
+    41: "unified_addressing",
+    42: "maximum_texture1d_layered_width",
+    43: "maximum_texture1d_layered_layers",
+    45: "maximum_texture2d_gather_width",
+    46: "maximum_texture2d_gather_height",
+    47: "maximum_texture3d_width_alternate",
+    48: "maximum_texture3d_height_alternate",
+    49: "maximum_texture3d_depth_alternate",
+    50: "pci_domain_id",
+    51: "texture_pitch_alignment",
+    52: "maximum_texturecubemap_width",
+    53: "maximum_texturecubemap_layered_width",
+    54: "maximum_texturecubemap_layered_layers",
+    55: "maximum_surface1d_width",
+    56: "maximum_surface2d_width",
+    57: "maximum_surface2d_height",
+    58: "maximum_surface3d_width",
+    59: "maximum_surface3d_height",
+    60: "maximum_surface3d_depth",
+    61: "maximum_surface1d_layered_width",
+    62: "maximum_surface1d_layered_layers",
+    63: "maximum_surface2d_layered_width",
+    64: "maximum_surface2d_layered_height",
+    65: "maximum_surface2d_layered_layers",
+    66: "maximum_surfacecubemap_width",
+    67: "maximum_surfacecubemap_layered_width",
+    68: "maximum_surfacecubemap_layered_layers",
+    69: "maximum_texture1d_linear_width",
+    70: "maximum_texture2d_linear_width",
+    71: "maximum_texture2d_linear_height",
+    72: "maximum_texture2d_linear_pitch",
+    73: "maximum_texture2d_mipmapped_width",
+    74: "maximum_texture2d_mipmapped_height",
+    75: "compute_capability_major",
+    76: "compute_capability_minor",
+    77: "maximum_texture1d_mipmapped_width",
+    78: "stream_priorities_supported",
+    79: "global_l1_cache_supported",
+    80: "local_l1_cache_supported",
+    81: "max_shared_memory_per_multiprocessor",
+    82: "max_registers_per_multiprocessor",
+    83: "managed_memory",
+    84: "multi_gpu_board",
+    85: "multi_gpu_board_group_id",
+    86: "host_native_atomic_supported",
+    87: "single_to_double_precision_perf_ratio",
+    88: "pageable_memory_access",
+    89: "concurrent_managed_access",
+    90: "compute_preemption_supported",
+    91: "can_use_host_pointer_for_registered_mem",
+    95: "cooperative_launch",
+    96: "cooperative_multi_device_launch",
+    97: "max_shared_memory_per_block_optin",
+    98: "can_flush_remote_writes",
+    99: "host_register_supported",
+    100: "pageable_memory_access_uses_host_page_tables",
+    101: "direct_managed_mem_access_from_host",
+    102: "virtual_memory_management_supported",
+    103: "handle_type_posix_file_descriptor_supported",
+    104: "handle_type_win32_handle_supported",
+    105: "handle_type_win32_kmt_handle_supported",
+    106: "max_blocks_per_multiprocessor",
+    107: "generic_compression_supported",
+    108: "max_persisting_l2_cache_size",
+    109: "max_access_policy_window_size",
+    110: "gpu_direct_rdma_with_cuda_vmm_supported",
+    111: "reserved_shared_memory_per_block",
+    112: "sparse_cuda_array_supported",
+    113: "read_only_host_register_supported",
+    114: "timeline_semaphore_interop_supported",
+    115: "memory_pools_supported",
+    116: "gpu_direct_rdma_supported",
+    117: "gpu_direct_rdma_flush_writes_options",
+    118: "gpu_direct_rdma_writes_ordering",
+    119: "mempool_supported_handle_types",
+    120: "cluster_launch",
+    121: "deferred_mapping_cuda_array_supported",
+}
+
+
+def attr_name(attr_id: int) -> str:
+    return KNOWN_ATTRS.get(attr_id, f"attribute_{attr_id}")
+
+
+def query_all(driver: Driver, device: DeviceHandle, probe_max: int = PROBE_MAX) -> dict[str, int]:
+    """Return {attribute_name: value} for every attribute this driver supports."""
+    out: dict[str, int] = {}
+    for attr_id in range(1, probe_max + 1):
+        value = driver.attribute(device, attr_id)
+        if value is not None:
+            out[attr_name(attr_id)] = value
+    return out

kernelmeter/bench.py

@@ -0,0 +1,292 @@
+"""Benchmark kernels against references and against the hardware ceiling.
+
+Workflow: decorate your implementation with @kernelmeter.benchmark, giving
+it an argument factory, a reference implementation, and (optionally) the
+bytes moved / FLOPs performed per call. Run ``kernelmeter bench yourfile.py``
+and you get latency, achieved GB/s and TFLOP/s, percent of the device's
+theoretical peak, and a numerical correctness check against the reference.
+
+torch is imported lazily so the rest of the package (``kernelmeter info``)
+works without it.
+"""
+
+from __future__ import annotations
+
+import statistics
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from . import attrs as _attrs
+from . import peaks as _peaks
+from . import roofline as _roofline
+from .cudadrv import CudaNotAvailableError, Driver
+
+
+@dataclass
+class BenchSpec:
+    fn: Callable
+    name: str
+    args: Callable[[], tuple]
+    ref: Callable | None = None
+    bytes_per_call: Callable[..., int] | int | None = None
+    flops_per_call: Callable[..., int] | int | None = None
+    peak_tflops: float | None = None  # override for tensor-core work
+    warmup: int = 10
+    iters: int = 100
+
+
+@dataclass
+class BenchResult:
+    name: str
+    ms_mean: float
+    ms_median: float
+    ms_min: float
+    gbps: float | None = None
+    tflops: float | None = None
+    intensity: float | None = None
+    bound: str | None = None
+    pct_roofline: float | None = None
+    pct_peak_bw: float | None = None
+    pct_peak_fp32: float | None = None
+    ref_ms_median: float | None = None
+    speedup_vs_ref: float | None = None
+    correct: bool | None = None
+    max_abs_err: float | None = None
+    error: str | None = None
+
+    def as_dict(self) -> dict:
+        return {k: v for k, v in self.__dict__.items()}
+
+
+REGISTRY: list[BenchSpec] = []
+
+
+def benchmark(
+    name: str | None = None,
+    *,
+    args: Callable[[], tuple],
+    ref: Callable | None = None,
+    bytes_per_call: Callable[..., int] | int | None = None,
+    flops_per_call: Callable[..., int] | int | None = None,
+    peak_tflops: float | None = None,
+    warmup: int = 10,
+    iters: int = 100,
+):
+    """Register a function for ``kernelmeter bench``. Also usable directly:
+    the decorated function is returned unchanged.
+
+    peak_tflops replaces the derived fp32 peak in the roofline when your
+    kernel runs on other units (tensor cores, fp16, fp64).
+    """
+
+    def deco(fn: Callable) -> Callable:
+        REGISTRY.append(
+            BenchSpec(
+                fn=fn,
+                name=name or fn.__name__,
+                args=args,
+                ref=ref,
+                bytes_per_call=bytes_per_call,
+                flops_per_call=flops_per_call,
+                peak_tflops=peak_tflops,
+                warmup=warmup,
+                iters=iters,
+            )
+        )
+        return fn
+
+    return deco
+
+
+# ---------------------------------------------------------------------------
+# Pure math helpers (unit-tested without a GPU)
+# ---------------------------------------------------------------------------
+
+def summarize_times(times_ms: list[float]) -> tuple[float, float, float]:
+    return (
+        statistics.fmean(times_ms),
+        statistics.median(times_ms),
+        min(times_ms),
+    )
+
+
+def achieved_gbps(bytes_per_call: int, ms: float) -> float:
+    return bytes_per_call / (ms * 1e-3) / 1e9
+
+
+def achieved_tflops(flops_per_call: int, ms: float) -> float:
+    return flops_per_call / (ms * 1e-3) / 1e12
+
+
+def pct_of_peak(achieved: float, peak: float | None) -> float | None:
+    if not peak:
+        return None
+    return 100.0 * achieved / peak
+
+
+def _resolve(metric: Callable[..., int] | int | None, args: tuple) -> int | None:
+    if metric is None:
+        return None
+    if callable(metric):
+        return metric(*args)
+    return int(metric)
+
+
+def roofline_score(
+    nbytes: int | None,
+    nflops: int | None,
+    gbps: float | None,
+    tflops: float | None,
+    peaks: _peaks.Peaks,
+    peak_tflops_override: float | None = None,
+) -> tuple[float | None, str | None, float | None]:
+    """(intensity, bound, pct of attainable) given whatever metrics exist.
+
+    With only bytes the kernel is treated as memory-bound, with only flops
+    as compute-bound. With both, the roofline decides.
+    """
+    peak_tf = peak_tflops_override or peaks.fp32_tflops
+    peak_bw = peaks.mem_bandwidth_gbs
+
+    if nbytes and nflops and peak_tf and peak_bw and tflops:
+        ai = _roofline.intensity(nflops, nbytes)
+        attainable = _roofline.attainable_tflops(ai, peak_tf, peak_bw)
+        return ai, _roofline.bound(ai, peak_tf, peak_bw), 100.0 * tflops / attainable
+    if nbytes and gbps and peak_bw:
+        return None, "mem", 100.0 * gbps / peak_bw
+    if nflops and tflops and peak_tf:
+        return None, "comp", 100.0 * tflops / peak_tf
+    return None, None, None
+
+
+def diff_results(baseline: list[dict], results: list["BenchResult"], threshold_pct: float = 5.0):
+    """Compare a run against a saved baseline. Returns (rows, regressions)
+    where rows are (name, old_ms, new_ms, delta_pct) and regressions lists
+    names that got slower by more than the threshold."""
+    old = {r["name"]: r for r in baseline}
+    rows = []
+    regressions = []
+    for r in results:
+        if r.error or r.name not in old:
+            continue
+        old_ms = old[r.name]["ms_median"]
+        delta = 100.0 * (r.ms_median - old_ms) / old_ms if old_ms else 0.0
+        rows.append((r.name, old_ms, r.ms_median, delta))
+        if delta > threshold_pct:
+            regressions.append(r.name)
+    return rows, regressions
+
+
+# ---------------------------------------------------------------------------
+# GPU execution (requires torch + an NVIDIA device)
+# ---------------------------------------------------------------------------
+
+def _time_fn(fn: Callable, args: tuple, warmup: int, iters: int, flush_l2: bool) -> list[float]:
+    import torch
+
+    for _ in range(warmup):
+        fn(*args)
+    torch.cuda.synchronize()
+
+    # Flushing L2 between iterations keeps memory-bound numbers honest:
+    # without it, small workloads get served from cache and report
+    # impossible bandwidths. Same approach as triton.testing.do_bench.
+    cache = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device="cuda") if flush_l2 else None
+
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+    for i in range(iters):
+        if cache is not None:
+            cache.zero_()
+        starts[i].record()
+        fn(*args)
+        ends[i].record()
+    torch.cuda.synchronize()
+    return [s.elapsed_time(e) for s, e in zip(starts, ends)]
+
+
+def _check_correctness(spec: BenchSpec, args: tuple) -> tuple[bool, float]:
+    import torch
+
+    got = spec.fn(*args)
+    want = spec.ref(*args)
+    max_err = (got.float() - want.float()).abs().max().item()
+    try:
+        torch.testing.assert_close(got, want, rtol=1.6e-2, atol=1e-3, check_dtype=False)
+        return True, max_err
+    except AssertionError:
+        return False, max_err
+
+
+def device_peaks() -> _peaks.Peaks:
+    drv = Driver()
+    dev = drv.device(0)
+    return _peaks.derive(_attrs.query_all(drv, dev))
+
+
+def run(spec: BenchSpec, peaks: _peaks.Peaks | None = None, flush_l2: bool = True) -> BenchResult:
+    """Execute one spec on the current CUDA device."""
+    if peaks is None:
+        try:
+            peaks = device_peaks()
+        except CudaNotAvailableError:
+            peaks = _peaks.Peaks(None, None, None)
+
+    args = spec.args()
+
+    correct = max_err = None
+    if spec.ref is not None:
+        correct, max_err = _check_correctness(spec, args)
+
+    times = _time_fn(spec.fn, args, spec.warmup, spec.iters, flush_l2)
+    ms_mean, ms_median, ms_min = summarize_times(times)
+
+    nbytes = _resolve(spec.bytes_per_call, args)
+    nflops = _resolve(spec.flops_per_call, args)
+    gbps = achieved_gbps(nbytes, ms_median) if nbytes else None
+    tflops = achieved_tflops(nflops, ms_median) if nflops else None
+    ai, kernel_bound, pct_roof = roofline_score(
+        nbytes, nflops, gbps, tflops, peaks, spec.peak_tflops
+    )
+
+    ref_ms = speedup = None
+    if spec.ref is not None:
+        ref_times = _time_fn(spec.ref, args, spec.warmup, spec.iters, flush_l2)
+        ref_ms = summarize_times(ref_times)[1]
+        speedup = ref_ms / ms_median if ms_median > 0 else None
+
+    return BenchResult(
+        name=spec.name,
+        ms_mean=ms_mean,
+        ms_median=ms_median,
+        ms_min=ms_min,
+        gbps=gbps,
+        tflops=tflops,
+        intensity=ai,
+        bound=kernel_bound,
+        pct_roofline=pct_roof,
+        pct_peak_bw=pct_of_peak(gbps, peaks.mem_bandwidth_gbs) if gbps else None,
+        pct_peak_fp32=pct_of_peak(tflops, peaks.fp32_tflops) if tflops else None,
+        ref_ms_median=ref_ms,
+        speedup_vs_ref=speedup,
+        correct=correct,
+        max_abs_err=max_err,
+    )
+
+
+def run_registry(flush_l2: bool = True) -> list[BenchResult]:
+    try:
+        peaks = device_peaks()
+    except CudaNotAvailableError:
+        peaks = _peaks.Peaks(None, None, None)
+    results = []
+    for spec in REGISTRY:
+        try:
+            results.append(run(spec, peaks=peaks, flush_l2=flush_l2))
+        except Exception as exc:  # surface per-kernel failures, keep going
+            results.append(
+                BenchResult(
+                    name=spec.name, ms_mean=0.0, ms_median=0.0, ms_min=0.0, error=str(exc)
+                )
+            )
+    return results

kernelmeter/ceiling.py

@@ -0,0 +1,111 @@
+"""Measure what the card actually delivers, not what the spec sheet says.
+
+Theoretical peaks are computed from max boost clocks and are never fully
+reachable. This module runs the four STREAM kernels (copy, scale, add,
+triad) through torch to find the real bandwidth ceiling, and a large
+TF32-disabled matmul to find the real FP32 ceiling. Judge your kernels
+against these numbers and 100% actually means 100%.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from . import bench as _bench
+from . import peaks as _peaks
+
+
+@dataclass
+class CeilingResult:
+    name: str
+    ms: float
+    gbps: float | None = None
+    tflops: float | None = None
+    pct_theoretical: float | None = None
+
+    def as_dict(self) -> dict:
+        return dict(self.__dict__)
+
+
+def _stream_specs(n: int):
+    """The four STREAM kernels and the bytes each one moves."""
+    import torch
+
+    a = torch.randn(n, device="cuda")
+    b = torch.empty_like(a)
+    c = torch.empty_like(a)
+    q = 2.5
+    elt = a.element_size()
+    return [
+        ("copy", lambda: b.copy_(a), 2 * n * elt),
+        ("scale", lambda: torch.mul(a, q, out=b), 2 * n * elt),
+        ("add", lambda: torch.add(a, b, out=c), 3 * n * elt),
+        ("triad", lambda: torch.add(b, c, alpha=q, out=a), 3 * n * elt),
+    ]
+
+
+def measure(
+    mb: int = 256, matmul_n: int = 4096, warmup: int = 10, iters: int = 50
+) -> list[CeilingResult]:
+    import torch
+
+    try:
+        device_peaks = _bench.device_peaks()
+    except Exception:
+        device_peaks = _peaks.Peaks(None, None, None)
+
+    n = mb * 1024 * 1024 // 4  # fp32 elements
+    results = []
+    for name, fn, nbytes in _stream_specs(n):
+        times = _bench._time_fn(lambda *_: fn(), (), warmup, iters, flush_l2=True)
+        ms = _bench.summarize_times(times)[1]
+        gbps = _bench.achieved_gbps(nbytes, ms)
+        results.append(
+            CeilingResult(
+                name=name,
+                ms=ms,
+                gbps=gbps,
+                pct_theoretical=_bench.pct_of_peak(gbps, device_peaks.mem_bandwidth_gbs),
+            )
+        )
+
+    allow_tf32 = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False  # keep it on the fp32 units
+    try:
+        m = torch.randn(matmul_n, matmul_n, device="cuda")
+        out = torch.empty_like(m)
+        times = _bench._time_fn(
+            lambda *_: torch.mm(m, m, out=out), (), warmup, iters, flush_l2=True
+        )
+        ms = _bench.summarize_times(times)[1]
+        tflops = _bench.achieved_tflops(2 * matmul_n**3, ms)
+        results.append(
+            CeilingResult(
+                name="fp32 matmul",
+                ms=ms,
+                tflops=tflops,
+                pct_theoretical=_bench.pct_of_peak(tflops, device_peaks.fp32_tflops),
+            )
+        )
+    finally:
+        torch.backends.cuda.matmul.allow_tf32 = allow_tf32
+
+    return results
+
+
+def format_table(results: list[CeilingResult]) -> list[str]:
+    header = f"{'test':<14} {'median ms':>10} {'GB/s':>9} {'TFLOP/s':>9} {'% of theoretical':>17}"
+    lines = [header, "-" * len(header)]
+    for r in results:
+        gbps = f"{r.gbps:.1f}" if r.gbps is not None else "-"
+        tflops = f"{r.tflops:.2f}" if r.tflops is not None else "-"
+        pct = f"{r.pct_theoretical:.1f}%" if r.pct_theoretical is not None else "-"
+        lines.append(f"{r.name:<14} {r.ms:>10.4f} {gbps:>9} {tflops:>9} {pct:>17}")
+    best_bw = max((r.gbps for r in results if r.gbps), default=None)
+    if best_bw:
+        lines.append("")
+        lines.append(
+            f"measured bandwidth ceiling: {best_bw:.1f} GB/s "
+            "(use this as the honest 100% for memory-bound kernels)"
+        )
+    return lines