kernelport 0.1.0__py3-none-any.whl

kernelport/doctor.py

@@ -0,0 +1,285 @@
+"""Environment report: OS, Python, torch, GPU, installed extras, and a smoke
+kernel for every backend that is available on this machine.
+
+Every check is isolated: a missing package or a failing backend is reported
+with a reason, never raised.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import platform
+import sys
+import traceback
+from dataclasses import asdict, dataclass, field
+from typing import Callable
+
+# Package that marks each optional extra as installed.
+EXTRA_MARKERS: dict[str, list[str]] = {
+    "ref": ["torch"],
+    "cuda": ["torch", "cupy"],
+    "triton": ["torch", "triton"],
+    "metal": ["mlx"],
+    "webgpu": ["wgpu"],
+    "cpu": ["numba"],
+    "agent": ["anthropic"],
+}
+
+
+@dataclass
+class BackendCheck:
+    """Result of checking one backend on this machine."""
+
+    name: str
+    available: bool
+    reason: str
+    smoke: str = "skipped"  # "pass", "fail", or "skipped"
+    detail: str = ""
+
+
+@dataclass
+class DoctorReport:
+    """Full environment report."""
+
+    os: str
+    python: str
+    torch: str
+    cuda_available: bool
+    gpu_name: str
+    compute_capability: str
+    extras: dict[str, bool] = field(default_factory=dict)
+    backends: list[BackendCheck] = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        """Return a JSON-serialisable dict."""
+        return asdict(self)
+
+    @property
+    def ok(self) -> bool:
+        """True when no available backend failed its smoke test."""
+        return all(b.smoke != "fail" for b in self.backends)
+
+
+def _has(module: str) -> bool:
+    return importlib.util.find_spec(module) is not None
+
+
+def _torch_info() -> tuple[str, bool, str, str]:
+    if not _has("torch"):
+        return "not installed", False, "", ""
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            cc = torch.cuda.get_device_capability(0)
+            return torch.__version__, True, torch.cuda.get_device_name(0), f"sm_{cc[0]}{cc[1]}"
+        return torch.__version__, False, "", ""
+    except Exception as exc:  # pragma: no cover - defensive
+        return f"error: {exc}", False, "", ""
+
+
+# ---------------------------------------------------------------- smoke tests
+# Each returns None on success or raises. Inputs are tiny and fixed.
+
+def _smoke_cuda() -> None:
+    import cupy as cp
+    import numpy as np
+
+    src = r"""extern "C" __global__ void add(const float* a, const float* b, float* c, int n) {
+        int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i < n) c[i] = a[i] + b[i];
+    }"""
+    kern = cp.RawKernel(src, "add")
+    n = 1000
+    a = cp.arange(n, dtype=cp.float32)
+    b = cp.full(n, 2.0, dtype=cp.float32)
+    c = cp.empty_like(a)
+    kern(((n + 255) // 256,), (256,), (a, b, c, np.int32(n)))
+    if not bool(cp.allclose(c, a + 2.0)):
+        raise AssertionError("CuPy RawKernel produced wrong result")
+
+
+def _smoke_triton() -> None:
+    import torch
+    import triton
+    import triton.language as tl
+
+    @triton.jit
+    def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
+        offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+        mask = offs < n
+        tl.store(out_ptr + offs, tl.load(x_ptr + offs, mask=mask) + tl.load(y_ptr + offs, mask=mask), mask=mask)
+
+    n = 1000
+    x = torch.arange(n, device="cuda", dtype=torch.float32)
+    y = torch.full((n,), 2.0, device="cuda")
+    out = torch.empty_like(x)
+    add_kernel[(triton.cdiv(n, 256),)](x, y, out, n, BLOCK=256)
+    if not torch.equal(out, x + 2.0):
+        raise AssertionError("Triton kernel produced wrong result")
+
+
+def _smoke_cpu() -> None:
+    import numpy as np
+    from numba import njit, prange
+
+    @njit(parallel=True, cache=False)
+    def add(a, b, c):
+        for i in prange(a.shape[0]):
+            c[i] = a[i] + b[i]
+
+    a = np.arange(1000, dtype=np.float32)
+    c = np.empty_like(a)
+    add(a, np.full_like(a, 2.0), c)
+    if not np.array_equal(c, a + 2.0):
+        raise AssertionError("Numba kernel produced wrong result")
+
+
+def _smoke_metal() -> None:
+    import mlx.core as mx
+
+    kern = mx.fast.metal_kernel(
+        name="add",
+        input_names=["a"],
+        output_names=["out"],
+        source="uint i = thread_position_in_grid.x; out[i] = a[i] + 2.0f;",
+    )
+    a = mx.arange(1000, dtype=mx.float32)
+    (out,) = kern(
+        inputs=[a], grid=(1000, 1, 1), threadgroup=(256, 1, 1),
+        output_shapes=[a.shape], output_dtypes=[a.dtype],
+    )
+    if not bool(mx.array_equal(out, a + 2.0)):
+        raise AssertionError("MLX metal kernel produced wrong result")
+
+
+def _smoke_webgpu() -> None:
+    import numpy as np
+    import wgpu
+
+    from kernelport.backends.webgpu.runner import device
+
+    dev = device()
+    a = np.arange(1000, dtype=np.float32)
+    buf = dev.create_buffer_with_data(data=a.tobytes(), usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC)
+    module = dev.create_shader_module(code="""
+@group(0) @binding(0) var<storage, read_write> a: array<f32>;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    if (gid.x < 1000u) { a[gid.x] = a[gid.x] + 2.0; }
+}""")
+    pipe = dev.create_compute_pipeline(layout="auto", compute={"module": module, "entry_point": "main"})
+    group = dev.create_bind_group(layout=pipe.get_bind_group_layout(0),
+                                  entries=[{"binding": 0, "resource": {"buffer": buf}}])
+    enc = dev.create_command_encoder()
+    cp = enc.begin_compute_pass()
+    cp.set_pipeline(pipe)
+    cp.set_bind_group(0, group)
+    cp.dispatch_workgroups(4)
+    cp.end()
+    dev.queue.submit([enc.finish()])
+    out = np.frombuffer(dev.queue.read_buffer(buf), dtype=np.float32)
+    if not np.array_equal(out, a + 2.0):
+        raise AssertionError("WGSL kernel produced wrong result")
+
+
+# --------------------------------------------------------------- capabilities
+
+def _cuda_gpu() -> tuple[bool, str]:
+    if not _has("torch"):
+        return False, "torch not installed (pip install kernelport[cuda])"
+    import torch
+
+    if not torch.cuda.is_available():
+        return False, "torch.cuda.is_available() is False (CPU-only torch or no NVIDIA GPU)"
+    if getattr(torch.version, "hip", None):
+        return False, "ROCm GPU, not NVIDIA CUDA"
+    return True, "ok"
+
+
+def _check_cuda() -> tuple[bool, str]:
+    ok, why = _cuda_gpu()
+    if not ok:
+        return ok, why
+    if not _has("cupy"):
+        return False, "cupy not installed (pip install kernelport[cuda])"
+    return True, "CuPy + NVRTC on NVIDIA GPU"
+
+
+def _registry_check(name: str) -> Callable[[], tuple[bool, str]]:
+    def check() -> tuple[bool, str]:
+        from kernelport.backends.base import get_backend
+
+        return get_backend(name).is_available()
+    return check
+
+
+# The original CUDA runner is not a port backend; every port backend comes from the registry.
+SMOKES: dict[str, Callable[[], None]] = {"triton": _smoke_triton, "cpu": _smoke_cpu, "metal": _smoke_metal,
+                                         "webgpu": _smoke_webgpu}
+BACKENDS: list[tuple[str, Callable[[], tuple[bool, str]], Callable[[], None] | None]] = [
+    ("cuda (original)", _check_cuda, _smoke_cuda),
+    *[(n, _registry_check(n), SMOKES.get(n)) for n in ("triton", "cpu", "metal", "webgpu", "hip")],
+]
+
+
+def _run_backend(name: str, check: Callable, smoke: Callable | None) -> BackendCheck:
+    try:
+        ok, reason = check()
+    except Exception as exc:
+        return BackendCheck(name, False, f"capability check crashed: {exc}")
+    result = BackendCheck(name, ok, reason)
+    if ok and smoke is not None:
+        try:
+            smoke()
+            result.smoke = "pass"
+        except Exception as exc:
+            result.smoke = "fail"
+            result.detail = f"{type(exc).__name__}: {exc}\n" + traceback.format_exc(limit=3)
+    return result
+
+
+def run_doctor() -> DoctorReport:
+    """Collect environment info and smoke-test every available backend."""
+    torch_ver, cuda_ok, gpu, cc = _torch_info()
+    report = DoctorReport(
+        os=f"{platform.system()} {platform.release()} ({platform.machine()})",
+        python=sys.version.split()[0],
+        torch=torch_ver,
+        cuda_available=cuda_ok,
+        gpu_name=gpu,
+        compute_capability=cc,
+        extras={k: all(_has(m) for m in mods) for k, mods in EXTRA_MARKERS.items()},
+    )
+    report.backends = [_run_backend(*b) for b in BACKENDS]
+    return report
+
+
+def print_report(report: DoctorReport) -> None:
+    """Pretty-print a doctor report with rich."""
+    from rich.console import Console
+    from rich.markup import escape
+    from rich.table import Table
+
+    con = Console()
+    con.print("[bold]kernelport doctor[/bold]")
+    con.print(f"OS:       {report.os}")
+    con.print(f"Python:   {report.python}")
+    con.print(f"torch:    {report.torch}")
+    con.print(f"CUDA:     {report.cuda_available}")
+    if report.gpu_name:
+        con.print(f"GPU:      {report.gpu_name} ({report.compute_capability})")
+    con.print("Extras:   " + ", ".join(f"{k}={'yes' if v else 'no'}" for k, v in report.extras.items()))
+
+    table = Table(title="Backends")
+    for col in ("backend", "available", "smoke", "reason"):
+        table.add_column(col)
+    colors = {"pass": "green", "fail": "red", "skipped": "dim"}
+    for b in report.backends:
+        table.add_row(b.name, "yes" if b.available else "no",
+                      f"[{colors[b.smoke]}]{b.smoke}[/]", escape(b.reason))
+    con.print(table)
+    for b in report.backends:
+        if b.smoke == "fail":
+            con.print(f"[red]{b.name} smoke failed:[/red]\n{b.detail}")
+    con.print("[green]doctor: OK[/green]" if report.ok else "[red]doctor: FAILED[/red]")

kernelport/hardware.py

@@ -0,0 +1,60 @@
+"""Human-readable hardware names, recorded with every result."""
+
+from __future__ import annotations
+
+import os
+import platform
+import subprocess
+import sys
+from functools import lru_cache
+
+
+def _cpu_model() -> str:
+    if sys.platform == "win32":
+        try:
+            import winreg
+
+            with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
+                                r"HARDWARE\DESCRIPTION\System\CentralProcessor\0") as key:
+                return str(winreg.QueryValueEx(key, "ProcessorNameString")[0]).strip()
+        except OSError:
+            pass
+    elif sys.platform == "darwin":
+        try:
+            return subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True,
+                                  text=True, timeout=5).stdout.strip()
+        except (OSError, subprocess.SubprocessError):
+            pass
+    else:
+        try:
+            with open("/proc/cpuinfo", encoding="utf-8") as f:
+                for line in f:
+                    if line.startswith("model name"):
+                        return line.split(":", 1)[1].strip()
+        except OSError:
+            pass
+    return platform.processor() or platform.machine()
+
+
+@lru_cache(maxsize=1)
+def cpu_name() -> str:
+    """CPU model with logical thread count, e.g. 'AMD Ryzen 7 7840HS (16 threads)'."""
+    model = " ".join(_cpu_model().split()) or "unknown CPU"
+    threads = os.cpu_count()
+    return f"{model} ({threads} threads)" if threads else model
+
+
+def gpu_name() -> str:
+    """Name of GPU 0 as seen by torch, with arch, or '' if none."""
+    try:
+        import torch
+    except ImportError:
+        return ""
+    if not torch.cuda.is_available():
+        return ""
+    name = torch.cuda.get_device_name(0)
+    if getattr(torch.version, "hip", None):
+        arch = getattr(torch.cuda.get_device_properties(0), "gcnArchName", "")
+        return f"{name} ({arch.split(':')[0]})" if arch else name
+    major, minor = torch.cuda.get_device_capability(0)
+    return f"{name} (sm_{major}{minor})"

kernelport/harness/compare.py

@@ -0,0 +1,152 @@
+"""Compare a backend's output against the frozen reference.
+
+Checks, in order: shape, dtype, NaN placement, Inf placement (with sign),
+then element-wise tolerance. The failure report is compact so it fits in an
+agent's context: one line of numbers plus small expected/actual slices.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import torch
+
+from kernelport.harness.tolerances import Tolerance
+
+SLICE_RADIUS = 4  # elements shown either side of the worst index
+
+
+@dataclass
+class TensorCompare:
+    """Result of comparing one output tensor."""
+
+    name: str
+    passed: bool
+    shape: list[int]
+    dtype: str
+    reason: str = ""  # empty on pass
+    max_abs_err: float = 0.0
+    max_rel_err: float = 0.0
+    worst_index: list[int] | None = None
+    num_mismatched: int = 0
+    num_elements: int = 0
+    expected_slice: list[Any] = field(default_factory=list)
+    actual_slice: list[Any] = field(default_factory=list)
+    slice_start: int = 0  # flat index of the first slice element
+
+    def to_dict(self) -> dict:
+        """JSON-serialisable dict."""
+        return asdict(self)
+
+    def summary(self) -> str:
+        """One-line human summary."""
+        if self.passed:
+            return f"{self.name}: ok (max abs {self.max_abs_err:.3g}, max rel {self.max_rel_err:.3g})"
+        return f"{self.name}: {self.reason}"
+
+
+def _dtype_name(t: torch.dtype) -> str:
+    return str(t).replace("torch.", "")
+
+
+def _fmt(v: float) -> Any:
+    if v != v or v in (float("inf"), float("-inf")):
+        return str(v)
+    return float(f"{v:.6g}")
+
+
+def _slices(expected: torch.Tensor, actual: torch.Tensor, flat_idx: int) -> tuple[list, list, int]:
+    e, a = expected.reshape(-1), actual.reshape(-1)
+    lo = max(0, flat_idx - SLICE_RADIUS)
+    hi = min(e.numel(), flat_idx + SLICE_RADIUS + 1)
+    return [_fmt(x) for x in e[lo:hi].tolist()], [_fmt(x) for x in a[lo:hi].tolist()], lo
+
+
+def _unravel(flat: int, shape: tuple[int, ...]) -> list[int]:
+    idx = []
+    for s in reversed(shape):
+        idx.append(flat % s if s else 0)
+        flat //= s if s else 1
+    return list(reversed(idx))
+
+
+def compare_tensor(name: str, actual: torch.Tensor, expected: torch.Tensor,
+                   out_dtype: torch.dtype, tol: Tolerance) -> TensorCompare:
+    """Compare ``actual`` (backend output) to ``expected`` (reference, any precision)."""
+    res = TensorCompare(name, False, list(expected.shape), _dtype_name(out_dtype),
+                        num_elements=expected.numel())
+    actual = actual.detach().cpu()
+    if tuple(actual.shape) != tuple(expected.shape):
+        res.reason = f"shape mismatch: expected {tuple(expected.shape)}, got {tuple(actual.shape)}"
+        return res
+    if actual.dtype != out_dtype:
+        res.reason = f"dtype mismatch: expected {_dtype_name(out_dtype)}, got {_dtype_name(actual.dtype)}"
+        return res
+    if expected.numel() == 0:
+        res.passed = True
+        return res
+
+    if not out_dtype.is_floating_point:
+        exp_i = expected.to(out_dtype).to(torch.int64)
+        act_i = actual.to(torch.int64)
+        bad = exp_i != act_i
+        diff = (exp_i - act_i).abs()
+        res.max_abs_err = float(diff.max())
+        res.num_mismatched = int(bad.sum())
+        if res.num_mismatched:
+            flat = int(torch.argmax(bad.reshape(-1).to(torch.int8)))
+            res.worst_index = _unravel(flat, tuple(expected.shape))
+            res.expected_slice, res.actual_slice, res.slice_start = _slices(exp_i, act_i, flat)
+            res.reason = (f"{res.num_mismatched}/{res.num_elements} integer values differ; "
+                          f"first at {res.worst_index}")
+            return res
+        res.passed = True
+        return res
+
+    # Round the exact reference to the output dtype, then compare in float64.
+    exp = expected.to(out_dtype).double()
+    act = actual.double()
+
+    exp_nan, act_nan = torch.isnan(exp), torch.isnan(act)
+    if not torch.equal(exp_nan, act_nan):
+        bad = exp_nan != act_nan
+        flat = int(torch.argmax(bad.reshape(-1).to(torch.int8)))
+        res.num_mismatched = int(bad.sum())
+        res.worst_index = _unravel(flat, tuple(exp.shape))
+        res.expected_slice, res.actual_slice, res.slice_start = _slices(exp, act, flat)
+        kind = "unexpected NaN" if bool(act_nan.reshape(-1)[flat]) else "missing NaN"
+        res.reason = (f"{kind}: {res.num_mismatched}/{res.num_elements} NaN positions differ; "
+                      f"first at {res.worst_index} (NaN may mean the element was never written)")
+        return res
+
+    exp_inf, act_inf = torch.isinf(exp), torch.isinf(act)
+    inf_bad = (exp_inf | act_inf) & ~((exp_inf & act_inf) & (torch.sign(exp) == torch.sign(act)))
+    if bool(inf_bad.any()):
+        flat = int(torch.argmax(inf_bad.reshape(-1).to(torch.int8)))
+        res.num_mismatched = int(inf_bad.sum())
+        res.worst_index = _unravel(flat, tuple(exp.shape))
+        res.expected_slice, res.actual_slice, res.slice_start = _slices(exp, act, flat)
+        res.reason = f"Inf mismatch: {res.num_mismatched}/{res.num_elements} positions; first at {res.worst_index}"
+        return res
+
+    finite = torch.isfinite(exp)
+    e, a = torch.where(finite, exp, 0.0), torch.where(finite, act, 0.0)
+    abs_err = (a - e).abs()
+    rel_err = abs_err / e.abs().clamp_min(torch.finfo(torch.float64).tiny)
+    allowed = tol.atol + tol.rtol * e.abs()
+    excess = abs_err - allowed
+    res.max_abs_err = float(abs_err.max())
+    res.max_rel_err = float(torch.where(abs_err > 0, rel_err, 0.0).max())
+    bad = excess > 0
+    res.num_mismatched = int(bad.sum())
+    if res.num_mismatched:
+        flat = int(torch.argmax(excess.reshape(-1)))
+        res.worst_index = _unravel(flat, tuple(exp.shape))
+        res.expected_slice, res.actual_slice, res.slice_start = _slices(exp, act, flat)
+        res.reason = (f"{res.num_mismatched}/{res.num_elements} values outside tolerance "
+                      f"(rtol={tol.rtol:g}, atol={tol.atol:g}); max abs err {res.max_abs_err:.4g}, "
+                      f"max rel err {res.max_rel_err:.4g}, worst at {res.worst_index}")
+        return res
+    res.passed = True
+    return res