kernelport 0.1.0__py3-none-any.whl

kernelport/harness/device_source.py

@@ -0,0 +1,273 @@
+"""Device-only view of a real-world CUDA file, for NVRTC.
+
+Real .cu files mix kernels with host code: CPU reference functions, main(),
+cudaMalloc, <<<...>>> launches, cuBLAS handles, <stdio.h>. NVRTC compiles device
+code only and has no host headers, so compiling such a file as-is fails.
+
+``device_source`` walks the file's top level and keeps what device code can need:
+
+- ``__global__`` / ``__device__`` functions (also via macros such as
+  ``#define DEVICE __device__ __forceinline__``)
+- types (struct/class/union/enum), typedefs, ``using``, templates of those
+- ``__device__`` / ``__constant__`` / const / constexpr variables
+- preprocessor directives, except host-only ``#include <...>`` headers
+- local ``#include "x.h"`` files, inlined (once each) with ``#line`` markers
+
+and blanks out host functions, host prototypes and host globals. Blanked text is
+replaced by spaces with every newline kept, so NVRTC's line numbers still point at
+the user's file. The original file is what gets hashed and frozen; this view is
+only what NVRTC compiles.
+
+This is lexical, not a C++ front end: kernels generated by macros, or a top-level
+macro call with no trailing ``;`` followed by a kernel, can confuse it.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from kernelport.analyzer.cuda_parser import strip_comments_and_strings
+
+# Headers NVRTC (or hipRTC) can provide; every other <...> include is host-only.
+DEVICE_HEADER_PREFIXES = ("cuda_fp16", "cuda_bf16", "cuda_fp8", "cuda_fp4", "cuda_fp6", "cooperative_groups",
+                          "mma.h", "cuda/", "cub/", "curand_kernel", "curand_philox", "cuda_pipeline",
+                          "cuda_awbarrier", "sm_", "crt/", "hip/", "vector_types.h", "vector_functions.h")
+
+_CUDA_HOST_HEADERS = ("cuda_runtime", "cuda.h", "cublas", "cudnn", "cusparse", "curand.h", "cufft", "nccl",
+                      "cusolver", "nvToolsExt", "nvtx")
+
+# Constants host headers would have supplied (math.h, float.h, limits.h).
+SHIM = """#ifndef M_PI
+#define M_PI 3.14159265358979323846
+#endif
+#ifndef M_SQRT2
+#define M_SQRT2 1.41421356237309504880
+#endif
+#ifndef M_SQRT1_2
+#define M_SQRT1_2 0.70710678118654752440
+#endif
+#ifndef M_E
+#define M_E 2.7182818284590452354
+#endif
+#ifndef M_LN2
+#define M_LN2 0.69314718055994530942
+#endif
+#ifndef FLT_MAX
+#define FLT_MAX 3.402823466e+38F
+#endif
+#ifndef FLT_MIN
+#define FLT_MIN 1.175494351e-38F
+#endif
+#ifndef FLT_EPSILON
+#define FLT_EPSILON 1.192092896e-07F
+#endif
+#ifndef DBL_MAX
+#define DBL_MAX 1.7976931348623158e+308
+#endif
+#ifndef INT_MAX
+#define INT_MAX 2147483647
+#endif
+#ifndef INT_MIN
+#define INT_MIN (-2147483647 - 1)
+#endif
+#ifndef UINT_MAX
+#define UINT_MAX 0xffffffffU
+#endif
+"""
+
+_DEVICE_WORDS = {"__global__", "__device__"}
+_KEEP_DECL_WORDS = {"__device__", "__constant__", "__shared__", "__managed__", "constexpr", "static_assert",
+                    "typedef", "using", "const", "struct", "class", "union", "enum", "template"}
+_TYPE_WORDS = {"struct", "class", "union", "enum"}
+_INCLUDE = re.compile(r'#\s*include\s*([<"])([^>"]+)[>"]')
+_OBJECT_MACRO = re.compile(r"^[ \t]*#[ \t]*define[ \t]+(\w+)[ \t]+(.*)$", re.M)
+
+
+def _strip_templates(s: str) -> str:
+    """Remove <...> template argument lists (innermost first)."""
+    prev = None
+    while prev != s:
+        prev, s = s, re.sub(r"<[^<>;{}()]*>", "", s)
+    return s
+
+
+def _flatten(header: str) -> str:
+    """Header with operator names, comparisons, template args and (...)/[...] groups removed."""
+    s = re.sub(r"\boperator\s*[^\w\s(]+", "operator", header)
+    s = re.sub(r"[<>!=]=", "", s)
+    s = _strip_templates(s)
+    prev = None
+    while prev != s:
+        prev, s = s, re.sub(r"\([^()]*\)|\[[^\[\]]*\]", "", s)
+    return s
+
+
+class SliceError(ValueError):
+    """The file could not be sliced (e.g. an include escapes the repo)."""
+
+
+def _blank(chars: list[str], start: int, end: int) -> None:
+    for k in range(start, end):
+        if chars[k] != "\n":
+            chars[k] = " "
+
+
+def _device_markers(clean: str) -> set[str]:
+    """Words that mark device code: __global__/__device__ and object-like macros expanding to them."""
+    marks = set(_DEVICE_WORDS)
+    for name, body in _OBJECT_MACRO.findall(clean):
+        if set(re.findall(r"\w+", body)) & marks:
+            marks.add(name)
+    return marks
+
+
+def _match(clean: str, i: int, open_ch: str, close_ch: str) -> int:
+    depth = 0
+    for k in range(i, len(clean)):
+        if clean[k] == open_ch:
+            depth += 1
+        elif clean[k] == close_ch:
+            depth -= 1
+            if depth == 0:
+                return k + 1
+    return len(clean)
+
+
+def _directive_end(clean: str, i: int) -> int:
+    """End (exclusive, before the newline) of the preprocessor line starting at i, with continuations."""
+    while True:
+        j = clean.find("\n", i)
+        if j < 0:
+            return len(clean)
+        if clean[:j].rstrip(" \t\r").endswith("\\"):
+            i = j + 1
+            continue
+        return j
+
+
+def _at_line_start(clean: str, i: int) -> bool:
+    k = clean.rfind("\n", 0, i) + 1
+    return clean[k:i].strip() == ""
+
+
+class _Slicer:
+    def __init__(self, root: Path | None):
+        self.root = root.resolve() if root else None
+        self.seen: set[Path] = set()
+        self.dropped_cuda_host = False
+
+    def file(self, path: Path) -> str:
+        path = path.resolve()
+        self.seen.add(path)
+        text = path.read_text(encoding="utf-8", errors="replace")
+        clean = strip_comments_and_strings(text)
+        chars = list(text)
+        self.marks = getattr(self, "marks", set()) | _device_markers(clean)  # includes add theirs
+        inserts: dict[int, str] = {}
+        self._region(text, clean, chars, 0, len(text), path, inserts)
+        out, last = [], 0
+        for pos in sorted(inserts):
+            out.append("".join(chars[last:pos]))
+            out.append(inserts[pos])
+            last = pos
+        out.append("".join(chars[last:]))
+        return "".join(out)
+
+    def _region(self, text, clean, chars, i, end, path, inserts) -> None:
+        while i < end:
+            while i < end and clean[i] in " \t\r\n":
+                i += 1
+            if i >= end:
+                return
+            if clean[i] == "#" and _at_line_start(clean, i):
+                j = _directive_end(clean, i)
+                self._directive(text, clean, chars, i, j, path, inserts)
+                i = j
+                continue
+            if clean[i] == ";":  # stray semicolon
+                i += 1
+                continue
+            j, body = self._item_end(clean, i, end)
+            header = clean[i:body[0]] if body else clean[i:j]
+            words = set(re.findall(r"\w+", header))
+            if body and (header.split()[:1] == ["namespace"] or re.fullmatch(r'\s*extern\s*"\s*"\s*', header)):
+                self._region(text, clean, chars, body[0] + 1, body[1] - 1, path, inserts)
+            elif body and not self._is_type_or_init(header):
+                if not words & self.marks:  # host function definition
+                    _blank(chars, i, j)
+            elif not body and not self._keep_declaration(header, words):
+                _blank(chars, i, j)
+            i = j
+
+    def _item_end(self, clean, i, end) -> tuple[int, tuple[int, int] | None]:
+        """(end of the item, (body_start, body_end) if it has a top-level { } body)."""
+        k = i
+        while k < end:
+            c = clean[k]
+            if c in "([":
+                k = _match(clean, k, c, ")" if c == "(" else "]")
+                continue
+            if c == ";":
+                return k + 1, None
+            if c == "{":
+                b_end = _match(clean, k, "{", "}")
+                header = clean[i:k]
+                if self._is_type_or_init(header):  # struct S {...}; / int a[] = {...};
+                    semi = clean.find(";", b_end, end)
+                    return (end if semi < 0 else semi + 1), (k, b_end)
+                return b_end, (k, b_end)
+            k += 1
+        return end, None
+
+    @staticmethod
+    def _is_type_or_init(header: str) -> bool:
+        """True for `struct S {` / `enum class E : int {` / `T a[] = {` headers; False for functions."""
+        h = _flatten(header)
+        if "=" in h:
+            return True
+        return bool(set(re.findall(r"\w+", h)) & _TYPE_WORDS) and not header.rstrip().endswith(")")
+
+    def _keep_declaration(self, header: str, words: set[str]) -> bool:
+        if words & self.marks or words & {"__constant__", "__shared__", "__managed__"}:
+            return True
+        if words & {"typedef", "using", "static_assert"}:
+            return True
+        if "(" in re.sub(r"=.*", "", _strip_templates(header), flags=re.S):
+            return False  # host function prototype, or a host object built by a constructor
+        return bool(words & _KEEP_DECL_WORDS)
+
+    def _directive(self, text, clean, chars, i, j, path, inserts) -> None:
+        m = _INCLUDE.match(text[i:j].strip())
+        if not m:
+            return
+        kind, name = m.groups()
+        if kind == "<":
+            if not name.startswith(DEVICE_HEADER_PREFIXES):
+                _blank(chars, i, j)
+                if name.startswith(_CUDA_HOST_HEADERS):
+                    self.dropped_cuda_host = True
+            return
+        inc = (path.parent / name).resolve()
+        if not inc.is_file():
+            return  # leave it to the compiler's -I search (and its error message)
+        if self.root and self.root not in inc.parents:
+            raise SliceError(f"{path.name} includes {name}, which is outside the repo")
+        _blank(chars, i, j)
+        if inc in self.seen:
+            return
+        line = text.count("\n", 0, j) + 2
+        body = self.file(inc)
+        at = j + 1 if j < len(text) else j
+        inserts[at] = (("" if at > j else "\n") +
+                       f'#line 1 "{inc.name}"\n{body}\n#line {line} "{path.name}"\n')
+
+
+def device_source(path: Path, root: Path | None = None) -> str:
+    """The device-only view of ``path`` (see module docstring), with a small constants shim."""
+    slicer = _Slicer(root)
+    sliced = slicer.file(Path(path))
+    # Under nvcc, cuda_runtime.h / cublas headers pull in the half and bfloat16 types; code
+    # that relied on that still needs them.
+    fp = "#include <cuda_fp16.h>\n#include <cuda_bf16.h>\n" if slicer.dropped_cuda_host else ""
+    return f'{SHIM}{fp}#line 1 "{Path(path).name}"\n{sliced}'

kernelport/harness/expr.py

@@ -0,0 +1,80 @@
+"""Tiny, safe integer expression evaluator for launch specs.
+
+Shapes, grid sizes and scalar values in launch.json are written as strings
+like ``"cdiv(n, 256)"`` or ``"rows * cols"``. They are evaluated against the
+case's dims. Only arithmetic on names and numbers plus a few helper functions
+is allowed; anything else raises ValueError.
+"""
+
+from __future__ import annotations
+
+import ast
+import math
+import operator
+from typing import Any
+
+_BINOPS = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.Div: operator.truediv,
+    ast.FloorDiv: operator.floordiv,
+    ast.Mod: operator.mod,
+    ast.Pow: operator.pow,
+    ast.LShift: operator.lshift,
+    ast.RShift: operator.rshift,
+}
+
+
+def _next_pow2(x: int) -> int:
+    return 1 if x <= 1 else 1 << (int(x) - 1).bit_length()
+
+
+_FUNCS = {
+    "cdiv": lambda a, b: -(-a // b),
+    "min": min,
+    "max": max,
+    "next_pow2": _next_pow2,
+    "ceil": math.ceil,
+    "floor": math.floor,
+    "int": int,
+}
+
+
+def evaluate(expr: Any, dims: dict[str, Any]) -> Any:
+    """Evaluate ``expr`` (a number or expression string) with ``dims`` as variables."""
+    if isinstance(expr, (int, float)):
+        return expr
+    if not isinstance(expr, str):
+        raise ValueError(f"expression must be a number or string, got {expr!r}")
+    tree = ast.parse(expr, mode="eval")
+    return _eval(tree.body, dims)
+
+
+def _eval(node: ast.AST, dims: dict[str, Any]) -> Any:
+    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
+        return node.value
+    if isinstance(node, ast.Name):
+        if node.id not in dims:
+            raise ValueError(f"unknown dimension {node.id!r}")
+        return dims[node.id]
+    if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS:
+        return _BINOPS[type(node.op)](_eval(node.left, dims), _eval(node.right, dims))
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
+        v = _eval(node.operand, dims)
+        return -v if isinstance(node.op, ast.USub) else v
+    if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id in _FUNCS:
+        if node.keywords:
+            raise ValueError("keyword arguments are not allowed")
+        return _FUNCS[node.func.id](*[_eval(a, dims) for a in node.args])
+    raise ValueError(f"unsupported expression element: {ast.dump(node)}")
+
+
+def evaluate_int(expr: Any, dims: dict[str, Any]) -> int:
+    """Evaluate and require an integer result."""
+    v = evaluate(expr, dims)
+    if isinstance(v, float):
+        if not v.is_integer():
+            raise ValueError(f"{expr!r} evaluated to non-integer {v}")
+        v = int(v)
+    return int(v)

kernelport/harness/inputs.py

@@ -0,0 +1,194 @@
+"""Deterministic input generation from a launch spec.
+
+Values come from NumPy's PCG64 ``random_raw`` stream, whose bit output is
+guaranteed stable across NumPy versions and platforms. Uniform and normal
+values are derived from those raw bits here (not via NumPy's distribution
+methods, which are allowed to change), so any machine regenerates
+bit-identical inputs from the same spec.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import zlib
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any
+
+import numpy as np
+import torch
+
+from kernelport.harness.expr import evaluate
+from kernelport.harness.spec import FLOAT_DTYPES, CaseSpec, LaunchSpec, ParamSpec
+
+TORCH_DTYPES: dict[str, torch.dtype] = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "float64": torch.float64,
+    "int8": torch.int8,
+    "uint8": torch.uint8,
+    "int16": torch.int16,
+    "int32": torch.int32,
+    "int64": torch.int64,
+}
+
+INT_SENTINEL = -0x5A5A5A5  # fills int outputs so unwritten elements are visible
+DEFAULT_EDGE_BIG = 1000.0
+NUMERIC_INIT_KEYS = ("low", "high", "mean", "std", "edge_big", "step")
+
+
+@dataclass
+class CaseData:
+    """Concrete tensors and scalars for one (case, dtype)."""
+
+    case: str
+    dtype: str
+    dims: dict[str, Any]
+    tensors: dict[str, torch.Tensor]  # CPU tensors, inputs and pre-filled outputs
+    scalars: dict[str, Any]
+    input_hash: str
+
+    @property
+    def label(self) -> str:
+        """Short name such as ``odd_1000/float16``."""
+        return f"{self.case}/{self.dtype}"
+
+
+def _uniform01(gen: np.random.Generator, n: int) -> np.ndarray:
+    """n doubles in [0, 1): 52 raw bits placed in the mantissa of [1, 2), minus 1. Exact."""
+    raw = gen.bit_generator.random_raw(n)
+    raw >>= np.uint64(12)
+    raw |= np.uint64(0x3FF0000000000000)
+    out = raw.view(np.float64)
+    out -= 1.0
+    return out
+
+
+def _normal(gen: np.random.Generator, n: int) -> np.ndarray:
+    """n values with mean 0 and std 1, bell-shaped (Irwin-Hall: sum of 4 uniforms).
+
+    Only uses IEEE add/multiply, which are exactly rounded on every platform.
+    Box-Muller would need log/sin/cos, whose last bit differs between libms.
+    Values are bounded to about (-3.46, 3.46), which is fine for test inputs.
+    """
+    s = _uniform01(gen, n)
+    for _ in range(3):  # elementwise adds in a fixed order
+        s += _uniform01(gen, n)
+    s -= 2.0
+    s *= 1.7320508075688772  # sqrt(3): var of the sum of 4 uniforms is 1/3
+    return s
+
+
+def _seed_for(spec: LaunchSpec, case: CaseSpec, param: str) -> int:
+    # Same values for every dtype of a case, so dtypes are comparable.
+    return (spec.seed * 1_000_003 + zlib.crc32(f"{case.name}|{param}".encode())) & 0xFFFFFFFF
+
+
+def _edge_values(big: float) -> list[float]:
+    return [0.0, -0.0, 1.0, -1.0, big, -big, 1.0 / big, -1.0 / big]
+
+
+def _generate(init: dict, shape: tuple[int, ...], dtype: str, seed: int) -> torch.Tensor:
+    is_float = dtype in FLOAT_DTYPES
+    vals = _base_values(json.dumps(init, sort_keys=True), shape, is_float, seed)
+    return torch.from_numpy(vals.copy()).to(TORCH_DTYPES[dtype])
+
+
+@lru_cache(maxsize=4)  # dtypes of one case share values; keep the last few
+def _base_values(init_json: str, shape: tuple[int, ...], is_float: bool, seed: int) -> np.ndarray:
+    """float64 values before the dtype cast."""
+    init = json.loads(init_json)
+    n = int(np.prod(shape)) if shape else 1
+    gen = np.random.Generator(np.random.PCG64(seed))
+    dist = init.get("dist", "normal" if is_float else "randint")
+    if dist == "normal":
+        vals = init.get("mean", 0.0) + init.get("std", 1.0) * _normal(gen, n)
+    elif dist == "uniform":
+        lo, hi = init.get("low", -1.0), init.get("high", 1.0)
+        vals = lo + (hi - lo) * _uniform01(gen, n)
+    elif dist == "randint":
+        lo, hi = int(init.get("low", 0)), int(init.get("high", 100))
+        vals = np.floor(lo + (hi - lo) * _uniform01(gen, n))
+    elif dist == "zeros":
+        vals = np.zeros(n)
+    elif dist == "ones":
+        vals = np.ones(n)
+    elif dist == "arange":
+        vals = np.arange(n, dtype=np.float64)
+    else:
+        raise ValueError(f"unknown dist {dist!r}")
+    if init.get("edges") and is_float:
+        edges = _edge_values(float(init.get("edge_big", DEFAULT_EDGE_BIG)))[:n]
+        vals[: len(edges)] = edges
+    step = init.get("step")
+    if step and is_float:
+        # Snap to multiples of a power-of-two step. Sums of such values are exact
+        # in float32 (while partial sums stay below 2**24 * step), so reductions
+        # give the same answer in any order and strict tolerances still apply.
+        if not math.log2(step).is_integer():
+            raise ValueError(f"step must be a power of two, got {step}")
+        vals = np.round(vals / step) * step
+    vals = np.ascontiguousarray(vals.reshape(shape))
+    vals.flags.writeable = False
+    return vals
+
+
+def _sentinel(shape: tuple[int, ...], dtype: str) -> torch.Tensor:
+    if dtype in FLOAT_DTYPES:
+        return torch.full(shape, float("nan"), dtype=TORCH_DTYPES[dtype])
+    return torch.full(shape, INT_SENTINEL, dtype=torch.int64).to(TORCH_DTYPES[dtype])
+
+
+def tensor_bytes(t: torch.Tensor) -> bytes:
+    """Raw little-endian bytes of a CPU tensor, for hashing."""
+    t = t.detach().cpu().contiguous()
+    if t.dim() == 0:
+        t = t.reshape(1)
+    return t.view(torch.uint8).numpy().tobytes()
+
+
+def _as_scalar_dtype(value, dtype: str):
+    """The value the kernel actually receives: 0.999 as a float32 param is 0.99900001...
+
+    References, originals and ports all get this rounded value, so a reference never
+    disagrees with the kernel just because it used the unrounded Python float.
+    """
+    if dtype in ("float32", "float16"):
+        return float(np.dtype(dtype).type(value))
+    if dtype in ("int32", "int64", "uint32", "int8", "uint8", "int16"):
+        return int(value)
+    return value
+
+
+def make_case(spec: LaunchSpec, case: CaseSpec, dtype: str) -> CaseData:
+    """Build CPU tensors for one case/dtype. Outputs are pre-filled with a sentinel."""
+    tensors: dict[str, torch.Tensor] = {}
+    scalars: dict[str, Any] = {}
+    h = hashlib.sha256()
+    for p in spec.params:
+        if not p.is_tensor:
+            scalars[p.name] = _as_scalar_dtype(p.resolved_value(case.dims), p.dtype)
+            continue
+        shape = p.resolved_shape(case.dims)
+        pdtype = p.resolved_dtype(dtype)
+        if p.kind == "output":
+            t = _sentinel(shape, pdtype)
+        else:
+            init = {**spec.init.get(p.name, {}), **case.init.get(p.name, {})}
+            # Numeric fields may be expressions over dims, e.g. "high": "max_pos".
+            init = {k: (evaluate(v, case.dims) if k in NUMERIC_INIT_KEYS and isinstance(v, str) else v)
+                    for k, v in init.items()}
+            t = _generate(init, shape, pdtype, _seed_for(spec, case, p.name))
+        tensors[p.name] = t
+        h.update(f"{p.name}:{pdtype}:{shape}".encode())
+        h.update(tensor_bytes(t))
+    h.update(repr(sorted(scalars.items())).encode())
+    return CaseData(case.name, dtype, dict(case.dims), tensors, scalars, h.hexdigest())
+
+
+def output_params(spec: LaunchSpec) -> list[ParamSpec]:
+    """Params the kernel writes."""
+    return [p for p in spec.params if p.is_output]

kernelport/harness/reference.py

@@ -0,0 +1,76 @@
+"""Load and run a plain-PyTorch reference.
+
+A reference file defines::
+
+    def reference(a, b, n):          # inputs/inouts and scalars, by name
+        return a + b                 # one output, or a dict {name: tensor}
+
+Floating-point inputs are passed as float64 CPU tensors so the reference is
+as close to the exact answer as possible. Integer inputs keep their dtype.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import inspect
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import torch
+
+from kernelport.harness.inputs import CaseData
+from kernelport.harness.spec import LaunchSpec
+from kernelport.hashing import sha256_file
+
+
+def file_hash(path: Path) -> str:
+    """sha256 of a file's contents (line-ending insensitive)."""
+    return sha256_file(path)
+
+
+def load_module(path: Path, prefix: str = "kp_mod") -> ModuleType:
+    """Import a Python file under a unique module name."""
+    path = Path(path).resolve()
+    name = f"{prefix}_{hashlib.sha1(str(path).encode()).hexdigest()[:12]}_{file_hash(path)[:8]}"
+    spec = importlib.util.spec_from_file_location(name, path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"cannot import {path}")
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[name] = mod  # some libraries (triton) need to find the module again
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def run_reference(path: Path, spec: LaunchSpec, data: CaseData) -> dict[str, torch.Tensor]:
+    """Run the reference for one case and return {output_name: float64/int CPU tensor}."""
+    fn = getattr(load_module(path, "kp_ref"), "reference", None)
+    if fn is None:
+        raise AttributeError(f"{path} does not define reference(...)")
+    kwargs = {}
+    for p in spec.params:
+        if p.kind in ("input", "inout"):
+            t = data.tensors[p.name].clone()
+            kwargs[p.name] = t.double() if t.is_floating_point() else t
+        elif p.kind == "scalar":
+            kwargs[p.name] = data.scalars[p.name]
+    accepted = inspect.signature(fn).parameters
+    if not any(v.kind == v.VAR_KEYWORD for v in accepted.values()):
+        missing = [k for k in accepted if k not in kwargs]
+        if missing:
+            raise TypeError(f"reference() asks for unknown params {missing}; available: {list(kwargs)}")
+        kwargs = {k: v for k, v in kwargs.items() if k in accepted}
+    with torch.no_grad():
+        out = fn(**kwargs)
+    outs = [p for p in spec.params if p.is_output]
+    if isinstance(out, torch.Tensor):
+        if len(outs) != 1:
+            raise TypeError(f"reference returned one tensor but kernel has outputs {[p.name for p in outs]}")
+        out = {outs[0].name: out}
+    if not isinstance(out, dict):
+        raise TypeError("reference must return a tensor or a dict of tensors")
+    missing = [p.name for p in outs if p.name not in out]
+    if missing:
+        raise TypeError(f"reference did not return outputs {missing}")
+    return {k: v.detach().cpu() for k, v in out.items()}