kernelport 0.1.0__py3-none-any.whl

kernelport/backends/triton/config.py

@@ -0,0 +1,41 @@
+"""Triton backend configuration."""
+
+from kernelport.backends.base import BackendConfig
+
+PORT_TEMPLATE = '''"""Triton port of <kernel>: <one line of what it computes>."""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
+    # One program per BLOCK elements (CUDA: one thread block).
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n  # CUDA: if (i < n)
+    x = tl.load(x_ptr + offs, mask=mask)
+    tl.store(y_ptr + offs, x, mask=mask)
+
+
+def run(x, y, n):
+    """Same parameter names as launch.json. Tensors are on the GPU; write outputs in place."""
+    BLOCK = 1024
+    _kernel[(triton.cdiv(n, BLOCK),)](x, y, n, BLOCK=BLOCK)
+'''
+
+# No tolerance overrides: Triton emits the same IEEE float32 add/mul/fma as CUDA,
+# and tl.exp / tl.log lower to libdevice-accurate or ex2/lg2-based sequences whose
+# error (a few ulp) is inside the float32 default (rtol 1.3e-6 ~ 11 ulp).
+# Note: tl.dot defaults to TF32 for float32 inputs, which is NOT within tolerance;
+# ports must pass input_precision="ieee" for float32 matmuls.
+CONFIG = BackendConfig(
+    name="triton",
+    extension=".py",
+    description="Triton kernels launched from Python (NVIDIA via triton / triton-windows, AMD via ROCm triton)",
+    dtypes=("float32", "float16", "bfloat16"),
+    tolerances={},
+    timing_comparable=True,
+    hardware_hint="NVIDIA GPU (CUDA) or AMD GPU (ROCm)",
+    port_template=PORT_TEMPLATE,
+    tune_space={"BLOCK": [128, 256, 512, 1024, 2048], "num_warps": [2, 4, 8]},
+)

kernelport/backends/triton/runner.py

@@ -0,0 +1,44 @@
+"""Run a Triton port.
+
+A Triton port is a Python file defining::
+
+    def run(<same parameter names as launch.json>):
+        ...launch one or more @triton.jit kernels...
+
+Tensor parameters arrive as GPU tensors (outputs pre-filled with NaN, or a
+sentinel for ints) and must be written in place. Scalars arrive as Python
+numbers. The return value is ignored.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from kernelport.backends.base import Prepared, bind_kwargs
+from kernelport.harness.inputs import CaseData
+from kernelport.harness.reference import load_module
+from kernelport.harness.spec import LaunchSpec
+
+
+def load_port(path: Path):
+    """Import the port and return its run() function."""
+    fn = getattr(load_module(path, "kp_triton"), "run", None)
+    if fn is None:
+        raise AttributeError(f"{Path(path).name} does not define run(...)")
+    return fn
+
+
+def prepare(path: Path, spec: LaunchSpec, data: CaseData) -> Prepared:
+    """Load the port and bind it to one case's tensors on the GPU."""
+    fn = load_port(path)
+    dev = {k: v.to("cuda") for k, v in data.tensors.items()}
+    kwargs = bind_kwargs(fn, spec, dev, data.scalars)
+    outputs = [p.name for p in spec.params if p.is_output]
+    return Prepared(
+        launch=lambda: fn(**kwargs),
+        sync=torch.cuda.synchronize,
+        fetch=lambda: {n: dev[n].cpu() for n in outputs},
+        timer="cuda_event",
+    )

kernelport/backends/webgpu/__init__.py

@@ -0,0 +1 @@
+"""WebGPU backend: WGSL compute shaders run through wgpu-native."""

kernelport/backends/webgpu/capability.py

@@ -0,0 +1,34 @@
+"""Capability check for the webgpu backend (wgpu-native via the `wgpu` package)."""
+
+from __future__ import annotations
+
+import importlib.util
+
+
+def is_available() -> tuple[bool, str]:
+    """(available, reason)."""
+    if importlib.util.find_spec("wgpu") is None:
+        return False, "wgpu not installed (pip install kernelport[webgpu])"
+    if importlib.util.find_spec("torch") is None:
+        return False, "torch not installed (needed to build inputs and run references)"
+    try:
+        from kernelport.backends.webgpu import runner
+
+        name = runner.adapter_name()
+    except Exception as exc:
+        return False, f"no usable WebGPU adapter: {type(exc).__name__}: {exc}"
+    return True, f"wgpu-native on {name}"
+
+
+def supported_dtypes() -> tuple[str, ...]:
+    """float16 only when the adapter offers the standard shader-f16 feature."""
+    from kernelport.backends.webgpu import runner
+
+    return ("float32", "float16") if runner.has_f16() else ("float32",)
+
+
+def hardware_name() -> str:
+    """Adapter name and native API, e.g. 'NVIDIA GeForce RTX 5060 Laptop GPU (WebGPU/Vulkan)'."""
+    from kernelport.backends.webgpu import runner
+
+    return runner.adapter_name()

kernelport/backends/webgpu/config.py

@@ -0,0 +1,38 @@
+"""webgpu (WGSL) backend configuration."""
+
+from kernelport.backends.base import BackendConfig
+
+PORT_TEMPLATE = '''"""WebGPU (WGSL) port of <kernel>: <one line of what it computes>."""
+
+# The runner appends: alias T, one storage buffer per tensor param (named like the
+# launch.json param), and `params` (a uniform struct holding every scalar param).
+WGSL = """
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;                      // CUDA: blockIdx.x * blockDim.x + threadIdx.x
+    if (i < u32(params.n)) {            // CUDA: if (i < n)
+        y[i] = x[i];
+    }
+}
+"""
+
+
+def grid(n):
+    """Workgroups per dimension (CUDA: the grid). Each must be <= 65535."""
+    return ((n + 255) // 256,)
+'''
+
+# No tolerance overrides: WGSL f32 add/mul are IEEE single precision. exp/log/sqrt
+# precision is implementation-defined in WGSL (exp: 3 + 2*|x| ulp), so ports that
+# need tighter results compute those in a numerically safe form; nothing is loosened.
+CONFIG = BackendConfig(
+    name="webgpu",
+    extension=".py",
+    description="WGSL compute shaders, run natively through wgpu-native (core WebGPU features only)",
+    dtypes=("float32", "float16"),
+    tolerances={},
+    timing_comparable=False,
+    timing_note="wall-clock time including WebGPU submission; not comparable to native kernel timings",
+    hardware_hint="any GPU with Vulkan, Metal or D3D12",
+    port_template=PORT_TEMPLATE,
+)

kernelport/backends/webgpu/runner.py

@@ -0,0 +1,285 @@
+"""Run a webgpu (WGSL) port through wgpu-native.
+
+A webgpu port is a Python file defining::
+
+    WGSL = '''
+    @compute @workgroup_size(256)
+    fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+        let i = gid.x;
+        if (i < u32(params.n)) { y[i] = max(x[i], T(0.0)); }
+    }
+    '''
+
+    def grid(n):                      # launch.json names (scalars, or arrays for .shape)
+        return (cdiv(n, 256),)         # workgroups per dimension, each <= 65535
+
+The runner appends the binding declarations, so the port never writes them:
+
+- ``alias T = f32;`` (or ``f16`` for float16 cases; ``enable f16;`` is added)
+- one storage buffer per tensor param, in launch.json order, named like the
+  param: ``var<storage, read>`` for inputs, ``read_write`` for outputs/inouts.
+  Element type: ``T`` for dtype "T", else f32 / f16 / i32 / u32; int64 arrives
+  as ``vec2<u32>`` (low word in ``.x``), since core WGSL has no 64-bit integers.
+- ``var<uniform> params: Params`` holding every scalar param (i32 / u32 / f32).
+
+A port may override a buffer's declared type, e.g. for atomics:
+``BINDINGS = {"out": "array<atomic<u32>>"}``.
+
+Only core WebGPU features are enabled (plus shader-f16 when the adapter has it),
+so a port that passes here uses WGSL a browser also accepts.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from kernelport.backends.base import Prepared, bind_kwargs
+from kernelport.harness.inputs import CaseData
+from kernelport.harness.reference import file_hash, load_module
+from kernelport.harness.results import CaseNotRunnable
+from kernelport.harness.spec import LaunchSpec
+
+MAX_WORKGROUPS = 65535
+PARAVIRTUAL_MAX_BOUND = 500 * 2**20
+_ELEM = {"float32": ("f32", np.float32), "float16": ("f16", np.float16), "int32": ("i32", np.int32),
+         "uint32": ("u32", np.uint32), "int64": ("vec2<u32>", np.int64)}
+_SCALAR = {"int32": ("i32", np.int32), "int64": ("i32", np.int32), "uint32": ("u32", np.uint32),
+           "float32": ("f32", np.float32), "float64": ("f32", np.float32)}
+
+_device = None
+_pipelines: dict[str, tuple] = {}
+_loaded: dict[tuple[str, str], object] = {}
+
+
+def device():
+    """The shared wgpu device (high-performance adapter, core features + shader-f16)."""
+    global _device
+    if _device is None:
+        import wgpu
+
+        adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance")
+        if adapter is None:
+            raise RuntimeError("no WebGPU adapter found")
+        lim = adapter.limits
+        cap = 2**31 - 4  # what desktop browsers allow at most
+        limits = {"max-storage-buffer-binding-size": min(lim["max-storage-buffer-binding-size"], cap),
+                  "max-buffer-size": min(lim["max-buffer-size"], cap)}
+        feats = ["shader-f16"] if "shader-f16" in adapter.features else []
+        _device = adapter.request_device_sync(required_features=feats, required_limits=limits)
+        _device._kp_adapter_info = adapter.info
+    return _device
+
+
+_fence = None
+
+
+def sync() -> None:
+    """Wait for all submitted work. Reading back a tiny buffer waits for the queue.
+
+    (queue.on_submitted_work_done_sync is broken in wgpu 0.32: wrong C callback signature.)
+    """
+    global _fence
+    import wgpu
+
+    dev = device()
+    if _fence is None:
+        _fence = dev.create_buffer(size=4, usage=wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST)
+    dev.queue.read_buffer(_fence)
+
+
+def has_f16() -> bool:
+    return "shader-f16" in device().features
+
+
+def cdiv(a: int, b: int) -> int:
+    return -(-a // b)
+
+
+def load_port(path: Path):
+    """Import the port once per content hash; returns the module."""
+    path = Path(path).resolve()
+    key = (str(path), file_hash(path))
+    if key not in _loaded:
+        mod = load_module(path, "kp_wgpu")
+        if not isinstance(getattr(mod, "WGSL", None), str):
+            raise AttributeError(f"{path.name} must define WGSL = '''...''' (the compute shader)")
+        if not callable(getattr(mod, "grid", None)):
+            raise AttributeError(f"{path.name} must define grid(...) returning workgroup counts")
+        _loaded[key] = mod
+    return _loaded[key]
+
+
+def _elem_type(p, case_dtype: str) -> tuple[str, type]:
+    dt = p.resolved_dtype(case_dtype)
+    if dt not in _ELEM:
+        raise TypeError(f"{p.name}: dtype {dt} is not supported by the webgpu backend")
+    if dt == "float16" and not has_f16():
+        raise TypeError("this WebGPU adapter has no shader-f16 support")
+    wgsl, np_t = _ELEM[dt]
+    return ("T" if p.dtype == "T" else wgsl), np_t
+
+
+def bindings_source(spec: LaunchSpec, case_dtype: str, overrides: dict[str, str]) -> str:
+    """WGSL declarations the runner appends to the port's shader."""
+    t = "f16" if case_dtype == "float16" else "f32"
+    lines = [f"alias T = {t};"]
+    b = 0
+    for p in spec.params:
+        if not p.is_tensor:
+            continue
+        elem, _ = _elem_type(p, case_dtype)
+        access = "read" if p.kind == "input" else "read_write"
+        ty = overrides.get(p.name, f"array<{elem}>")
+        lines.append(f"@group(0) @binding({b}) var<storage, {access}> {p.name}: {ty};")
+        b += 1
+    scalars = [p for p in spec.params if p.kind == "scalar"]
+    if scalars:
+        fields = ", ".join(f"{p.name}: {_SCALAR[p.dtype][0]}" for p in scalars)
+        lines.append(f"struct Params {{ {fields} }};")
+        lines.append(f"@group(0) @binding({b}) var<uniform> params: Params;")
+    return "\n".join(lines) + "\n"
+
+
+def _pipeline(code: str, kinds: list[str]):
+    """Compile (cached) shader + explicit bind group layout."""
+    import wgpu
+
+    key = code + "|" + ",".join(kinds)
+    if key not in _pipelines:
+        dev = device()
+        types = {"input": wgpu.BufferBindingType.read_only_storage,
+                 "rw": wgpu.BufferBindingType.storage,
+                 "uniform": wgpu.BufferBindingType.uniform}
+        bgl = dev.create_bind_group_layout(entries=[
+            {"binding": i, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": types[k]}}
+            for i, k in enumerate(kinds)])
+        layout = dev.create_pipeline_layout(bind_group_layouts=[bgl])
+        module = dev.create_shader_module(code=code)
+        pipe = dev.create_compute_pipeline(layout=layout, compute={"module": module, "entry_point": "main"})
+        _pipelines[key] = (pipe, bgl)
+    return _pipelines[key]
+
+
+def _padded(raw: bytes, multiple: int = 4) -> bytes:
+    extra = (-len(raw)) % multiple
+    return raw + b"\0" * extra if extra or not raw else raw
+
+
+def binding_limits() -> tuple[int, int | None]:
+    """(largest single storage binding, largest total bound per dispatch or None) in bytes."""
+    dev = device()
+    per = min(dev.limits["max-storage-buffer-binding-size"], dev.limits["max-buffer-size"])
+    total = None
+    if "Paravirtual" in str(dev._kp_adapter_info.get("device", "")):
+        # Apple's paravirtual GPU (macOS VMs such as GitHub's macos-14 runners) reports 3.5 GiB
+        # per binding, but a dispatch binding 628 MiB in total silently does not run (outputs
+        # untouched), while 500 MiB was measured to work. Stay within what was measured.
+        per = total = min(per, PARAVIRTUAL_MAX_BOUND)
+    return per, total
+
+
+def _check_binding_sizes(spec: LaunchSpec, data: CaseData, per: int, total: int | None) -> None:
+    """Raise CaseNotRunnable if the case's tensors exceed what this adapter can bind."""
+    sizes = {}
+    for p in spec.params:
+        if p.is_tensor:
+            t = data.tensors[p.name]
+            itemsize = np.dtype(_elem_type(p, data.dtype)[1]).itemsize if t.dtype != torch.bfloat16 else 2
+            sizes[p.name] = t.numel() * itemsize
+    mib = 2**20
+    for name, n in sizes.items():
+        if n > per:
+            raise CaseNotRunnable(f"{name} needs a {n / mib:.0f} MiB storage binding; "
+                                  f"this adapter allows {per / mib:.0f} MiB")
+    if total is not None and sum(sizes.values()) > total:
+        raise CaseNotRunnable(f"the case binds {sum(sizes.values()) / mib:.0f} MiB in total; "
+                              f"this adapter runs at most {total / mib:.0f} MiB per dispatch")
+
+
+def prepare(path: Path, spec: LaunchSpec, data: CaseData) -> Prepared:
+    """Compile the port for this case's dtype and upload the case's buffers."""
+    import wgpu
+
+    mod = load_port(path)
+    dev = device()
+    enable = "enable f16; " if data.dtype == "float16" or any(
+        p.resolved_dtype(data.dtype) == "float16" for p in spec.params if p.is_tensor) else ""
+    overrides = dict(getattr(mod, "BINDINGS", {}) or {})
+    unknown = set(overrides) - {p.name for p in spec.params if p.is_tensor}
+    if unknown:
+        raise KeyError(f"BINDINGS names unknown tensor params {sorted(unknown)}")
+    code = enable + mod.WGSL + "\n" + bindings_source(spec, data.dtype, overrides)
+
+    _check_binding_sizes(spec, data, *binding_limits())
+
+    kinds, buffers, outputs = [], [], {}
+    arrays = {}
+    for p in spec.params:
+        if not p.is_tensor:
+            continue
+        t = data.tensors[p.name]
+        if t.dtype == torch.bfloat16:
+            raise TypeError(f"{p.name}: bfloat16 is not a WGSL type; the webgpu backend skips bfloat16 cases")
+        _, np_t = _elem_type(p, data.dtype)
+        arr = t.contiguous().numpy().astype(np_t, copy=False)
+        arrays[p.name] = arr
+        usage = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST
+        buf = dev.create_buffer_with_data(data=_padded(arr.tobytes()), usage=usage)
+        buffers.append(buf)
+        kinds.append("input" if p.kind == "input" else "rw")
+        if p.is_output:
+            outputs[p.name] = (buf, arr)
+    scalars = [p for p in spec.params if p.kind == "scalar"]
+    if scalars:
+        vals = []
+        for p in scalars:
+            v = data.scalars[p.name]
+            if p.dtype in ("int64", "int32") and not (-2**31 <= int(v) < 2**31):
+                raise OverflowError(f"scalar {p.name}={v} does not fit in i32")
+            vals.append(np.array(v, dtype=_SCALAR[p.dtype][1]).tobytes())
+        ubuf = dev.create_buffer_with_data(data=_padded(b"".join(vals), 16), usage=wgpu.BufferUsage.UNIFORM)
+        buffers.append(ubuf)
+        kinds.append("uniform")
+
+    pipe, bgl = _pipeline(code, kinds)
+    group = dev.create_bind_group(layout=bgl, entries=[
+        {"binding": i, "resource": {"buffer": b, "offset": 0, "size": b.size}} for i, b in enumerate(buffers)])
+
+    kwargs = bind_kwargs(mod.grid, spec, arrays, data.scalars)
+    g = mod.grid(**kwargs)
+    g = (g,) if isinstance(g, int) else tuple(int(x) for x in g)
+    if not 1 <= len(g) <= 3 or any(x < 0 for x in g):
+        raise ValueError(f"grid() must return 1 to 3 non-negative ints, got {g}")
+    if any(x > MAX_WORKGROUPS for x in g):
+        raise ValueError(f"grid() returned {g}: WebGPU allows at most {MAX_WORKGROUPS} workgroups per "
+                         "dimension; use a 2-D grid or loop inside the shader")
+    g = g + (1,) * (3 - len(g))
+
+    def launch() -> None:
+        enc = dev.create_command_encoder()
+        cp = enc.begin_compute_pass()
+        cp.set_pipeline(pipe)
+        cp.set_bind_group(0, group)
+        cp.dispatch_workgroups(*g)
+        cp.end()
+        dev.queue.submit([enc.finish()])
+
+    def fetch() -> dict:
+        out = {}
+        for name, (buf, arr) in outputs.items():
+            raw = dev.queue.read_buffer(buf)
+            got = np.frombuffer(raw, dtype=arr.dtype, count=arr.size).reshape(arr.shape).copy()
+            out[name] = torch.from_numpy(got)
+        return out
+
+    return Prepared(launch=launch, sync=sync, fetch=fetch, timer="wall_clock")
+
+
+def adapter_name() -> str:
+    info = device()._kp_adapter_info
+    name = info.get("device") or "WebGPU adapter"
+    return f"{name} (WebGPU/{info.get('backend_type', '?')})"
+

kernelport/bench/timer.py

@@ -0,0 +1,76 @@
+"""Device-synchronised kernel timing.
+
+Warm up, then time each launch individually with device events and report
+the median. On NVIDIA and ROCm we use torch.cuda.Event (ROCm maps it to HIP
+events). CPU-side timers are only used for backends with no device events,
+and those numbers are labelled as not comparable.
+"""
+
+from __future__ import annotations
+
+import statistics
+import time
+from dataclasses import asdict, dataclass
+from typing import Callable
+
+DEFAULT_WARMUP = 10
+DEFAULT_ITERS = 50
+GPU_HOLD_CYCLES = 50_000_000  # ~20-40 ms of GPU time: covers queueing 50 launches from Python
+
+
+@dataclass
+class Timing:
+    """Timing summary in milliseconds."""
+
+    median_ms: float
+    min_ms: float
+    max_ms: float
+    iters: int
+    method: str  # "cuda_event" or "wall_clock"
+
+    def to_dict(self) -> dict:
+        """JSON-serialisable dict."""
+        return asdict(self)
+
+
+def time_cuda(launch: Callable[[], None], warmup: int = DEFAULT_WARMUP,
+              iters: int = DEFAULT_ITERS) -> Timing:
+    """Median time of ``launch()`` measured with CUDA/HIP events."""
+    import torch
+
+    for _ in range(warmup):
+        launch()
+    torch.cuda.synchronize()
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+    # Keep the GPU busy while the CPU queues every timed launch. Otherwise a short kernel
+    # waits for the next Python launch (tens of microseconds for Triton) and the events
+    # measure CPU launch overhead instead of the kernel.
+    sleep = getattr(torch.cuda, "_sleep", None)
+    if sleep is not None:
+        try:
+            sleep(GPU_HOLD_CYCLES)
+        except Exception:  # private API; ROCm or future builds may lack it. Timing still works.
+            pass
+    for s, e in zip(starts, ends):
+        s.record()
+        launch()
+        e.record()
+    torch.cuda.synchronize()
+    times = [s.elapsed_time(e) for s, e in zip(starts, ends)]
+    return Timing(statistics.median(times), min(times), max(times), iters, "cuda_event")
+
+
+def time_wall(launch: Callable[[], None], sync: Callable[[], None] = lambda: None,
+              warmup: int = DEFAULT_WARMUP, iters: int = DEFAULT_ITERS) -> Timing:
+    """Median wall-clock time of ``launch(); sync()``."""
+    for _ in range(warmup):
+        launch()
+    sync()
+    times = []
+    for _ in range(iters):
+        t0 = time.perf_counter()
+        launch()
+        sync()
+        times.append((time.perf_counter() - t0) * 1e3)
+    return Timing(statistics.median(times), min(times), max(times), iters, "wall_clock")