kernox 1.0.1__py3-none-any.whl

kernox/__init__.py

@@ -0,0 +1,101 @@
+"""kernox: fused symbolic reductions on GPUs.
+
+Documentation: https://keckj.github.io/kernox/. Expressions are symox objects; ``compile`` produces a
+callable that validates its inputs before any native launch.
+"""
+
+from .errors import (
+    BackendUnavailableError,
+    ConstantConversionError,
+    KernoxError,
+    UnsupportedError,
+    ValidationError,
+)
+from .expr import (
+    Formula,
+    Param,
+    ReductionExpr,
+    Vi,
+    Vj,
+    abs,
+    acos,
+    asin,
+    atan,
+    atan2,
+    cat,
+    ceiling,
+    clamp,
+    constant,
+    cos,
+    cosh,
+    eq,
+    exp,
+    floor,
+    log,
+    maximum,
+    minimum,
+    ne,
+    rsqrt,
+    sign,
+    sin,
+    sinh,
+    sqrt,
+    square,
+    tan,
+    tanh,
+    where,
+)
+from .ops import (  # noqa: E402 - also attaches the PyKeOps methods to Formula
+    bspline,
+    clampint,
+    elem,
+    extract,
+    ifelse,
+    inv,
+    kron,
+    matvecmult,
+    mod,
+    normalize,
+    one_hot,
+    relu,
+    round,
+    sinc,
+    sinxdivx,
+    softdtw_sqdist,
+    grad,
+    diff,
+    grad_matrix,
+    divergence,
+    laplacian,
+    trace_operator,
+    sqdist,
+    step,
+    tensordot,
+    tensorprod,
+    vecmatmult,
+    weightedsqdist,
+    weightedsqnorm,
+    xlogx,
+)
+from .complex import ComplexFormula, angle, complex_formula, conj, exp1j, imag, imag2complex, real, real2complex
+from .linalg import LinearOperator, conjugate_gradient
+from .ranges import Ranges, ranges_from_labels
+from .runtime import RangedReduction, Reduction, compile, get_backend, register_backend
+from .tuning import export_wisdom, forget_wisdom, import_wisdom
+from .diagnostics import diagnostics  # noqa: E402 - also `python -m kernox.diagnostics`
+
+__all__ = [
+    "Vi", "Vj", "Param", "Formula", "ReductionExpr", "Reduction", "RangedReduction", "compile",
+    "Ranges", "ranges_from_labels", "LinearOperator", "conjugate_gradient",
+    "ComplexFormula", "real2complex", "imag2complex", "exp1j", "complex_formula", "conj", "angle", "real", "imag",
+    "register_backend", "get_backend", "diagnostics", "export_wisdom", "import_wisdom", "forget_wisdom", "constant", "where", "eq", "ne", "cat",
+    "exp", "log", "sqrt", "rsqrt", "sin", "cos", "tan", "asin", "acos", "atan", "atan2",
+    "sinh", "cosh", "tanh", "abs", "sign", "floor", "ceiling", "square",
+    "minimum", "maximum", "clamp",
+    "relu", "step", "sinxdivx", "sinc", "xlogx", "inv", "round", "mod", "clampint", "ifelse",
+    "normalize", "sqdist", "weightedsqnorm", "weightedsqdist", "one_hot", "elem", "extract",
+    "matvecmult", "vecmatmult", "tensorprod", "tensordot", "kron", "bspline", "softdtw_sqdist",
+    "grad", "diff", "grad_matrix", "divergence", "laplacian", "trace_operator",
+    "KernoxError", "ValidationError", "ConstantConversionError",
+    "BackendUnavailableError", "UnsupportedError",
+]

kernox/_env.py

@@ -0,0 +1,13 @@
+"""Reading of the ``KERNOX_*`` environment variables.
+
+An empty value is treated as an unset variable.
+"""
+
+from __future__ import annotations
+
+import os
+
+
+def getenv(name: str, default: str | None = None) -> str | None:
+    """Value of ``name`` if it is set and non-empty, otherwise ``default``."""
+    return os.environ.get(name) or default

kernox/_plan_key.py

@@ -0,0 +1,118 @@
+"""Stable key of a :class:`kernox.plan.Plan` (contract v0).
+
+The key is the hexadecimal sha256 of a canonical JSON serialization of:
+
+- the key format version (``PLAN_KEY_FORMAT``);
+- for each argument, in ``plan.args`` order: name, role, dimension and the
+  symox ``srepr`` of each component symbol (name + assumptions);
+- the symox ``srepr`` of each output, in order; if the unfolded tree of an
+  output exceeds ``DAG_THRESHOLD`` nodes (deep shared DAG: ``srepr`` would be
+  exponential), the output is serialized as a graph: ``["dag", nodes]``, the
+  list of distinct nodes in postfix order, each one being ``srepr`` for an atom
+  or ``[class, head, child indices]`` otherwise. Keys of existing plans
+  (smaller trees) are unchanged;
+- reduction (kind, K), ``dtype``, ``acc_dtype``, ``index_dtype``, ``fast_math``, ``fp_contract``;
+- ``sum_scheme`` (amendment v0.7) only if it differs from ``"direct"``: the
+  key of an existing plan is unchanged.
+- ``batch_rank`` (amendment v0.6) only if nonzero: the key of a plan
+  without batch is unchanged;
+- ``ranges`` (amendment v0.6 §2) only if true.
+
+``plan.meta`` is not part of the key (diagnostic information only).
+The key depends neither on ``PYTHONHASHSEED`` nor on object creation order
+(tested in ``tests/test_plan_key.py``). Backends add the generator version,
+compiler, options and architecture to it (see the contract).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+
+PLAN_KEY_FORMAT = "kernox-plan-key-v1"
+
+
+#: unfolded tree size beyond which an output is serialized as a graph
+DAG_THRESHOLD = 200_000
+
+
+def _srepr(expr) -> str:
+    from symox import srepr
+
+    return srepr(expr)
+
+
+def _tree_size(expr, limit):
+    """Node count of the unfolded tree (memoized per subexpression), capped at ``limit``."""
+    memo = {}
+
+    def size(e):
+        hit = memo.get(e)
+        if hit is None:
+            hit = min(limit, 1 + sum(size(a) for a in e.args))
+            memo[e] = hit
+        return hit
+
+    return size(expr)
+
+
+def _dag(expr) -> list:
+    """Canonical graph serialization: distinct nodes in postfix order."""
+    from symox import srepr
+
+    index, nodes = {}, []
+
+    def visit(e):
+        hit = index.get(e)
+        if hit is not None:
+            return hit
+        if not e.args:
+            node = srepr(e)
+        else:
+            children = [visit(a) for a in e.args]
+            node = [type(e).__name__, str(e.func), children]
+        index[e] = len(nodes)
+        nodes.append(node)
+        return index[e]
+
+    visit(expr)
+    return ["dag", nodes]
+
+
+def _output(expr):
+    if _tree_size(expr, DAG_THRESHOLD + 1) > DAG_THRESHOLD:
+        return _dag(expr)
+    return _srepr(expr)
+
+
+def plan_payload(plan) -> dict:
+    """Canonical JSON structure hashed by :func:`plan_key` (useful for debugging)."""
+    payload = {
+        "format": PLAN_KEY_FORMAT,
+        "args": [
+            [a.name, a.role, int(a.dim), [_srepr(s) for s in a.symbols]]
+            for a in plan.args
+        ],
+        "outputs": [_output(e) for e in plan.outputs],
+        "reduction": [plan.reduction.kind.value, plan.reduction.k],
+        "dtype": plan.dtype,
+        "acc_dtype": plan.acc_dtype,
+        "index_dtype": plan.index_dtype,
+        "fast_math": bool(plan.fast_math),
+        "fp_contract": bool(plan.fp_contract),
+    }
+    scheme = getattr(plan, "sum_scheme", "direct")
+    if scheme != "direct":
+        payload["sum_scheme"] = scheme
+    rank = int(getattr(plan, "batch_rank", 0))
+    if rank:
+        payload["batch_rank"] = rank
+    if getattr(plan, "ranges", False):
+        payload["ranges"] = True
+    return payload
+
+
+def plan_key(plan) -> str:
+    """Hex sha256 of :func:`plan_payload`, canonically serialized."""
+    text = json.dumps(plan_payload(plan), sort_keys=True, separators=(",", ":"), ensure_ascii=True)
+    return hashlib.sha256(text.encode("ascii")).hexdigest()

kernox/_symox_compat.py

@@ -0,0 +1,128 @@
+"""Detection of the symox capabilities used by Kernox.
+
+Kernox targets symox 0.1.1 from PyPI (built from ``53d9576e`` or later),
+which contains the codegen extensions below, and remains compatible with the
+older development revision ``2e623e66``, which lacks them:
+
+- ``cudacode``: logical operators ``And/Or/Xor/Not``;
+- ``cudacode``: ``reciprocal_sqrt`` option (``1/sqrt(x)``, opt-in);
+- ``symox.extra.cudacode_function``: generic CUDA/HIP device function
+  (pointer arguments, named output struct, no math macro);
+- native NumPy ``lambdify`` with NumPy broadcasting.
+
+Detection tests the **behavior** (not a version number: both revisions
+carry version 0.1.1). Each probe is cached; the environment variable
+``KERNOX_SYMOX_COMPAT=pinned`` forces the compatibility code path, without
+the extensions (regression tests).
+"""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+
+from ._env import getenv
+
+PINNED_SYMOX = "53d9576e9f39db16e9ca8b3483d754385ebaf986"  # https://github.com/keckj/symox (main, printer SYCL sycl::half)
+# Oldest revision still supported (compatibility path, without the
+# codegen extensions); KERNOX_SYMOX_COMPAT=pinned forces this path.
+MIN_SUPPORTED_SYMOX = "2e623e66bb31882e31f6dbfff6f77ae4bd945408"
+
+
+def _forced_pinned() -> bool:
+    return getenv("KERNOX_SYMOX_COMPAT", "") == "pinned"
+
+
+@lru_cache(maxsize=None)
+def _cuda_logic() -> bool:
+    import symox
+    from symox.extra import cudacode
+
+    x, y = symox.symbols("x y", real=True)
+    try:
+        return "&&" in cudacode(symox.And(x < 1, y < 1))
+    except NotImplementedError:
+        return False
+
+
+@lru_cache(maxsize=None)
+def _cuda_reciprocal_sqrt() -> bool:
+    import symox
+    from symox.extra import cudacode
+
+    x = symox.Symbol("x", real=True)
+    try:
+        return cudacode(1 / symox.sqrt(x) + 1, use_rsqrt=False, reciprocal_sqrt=True) == "1 + 1.0/sqrt(x)"
+    except TypeError:
+        return False
+
+
+@lru_cache(maxsize=None)
+def _cudacode_function() -> bool:
+    import symox.extra
+
+    return "cudacode_function" in getattr(symox.extra, "__all__", ())
+
+
+@lru_cache(maxsize=None)
+def _lambdify_broadcasts() -> bool:
+    import numpy as np
+    import symox
+
+    x, y = symox.symbols("x y", real=True)
+    f = symox.lambdify((x, y), symox.Piecewise((x, x < y), (y, True)), modules="numpy")
+    try:
+        return np.shape(f(np.zeros((2, 1)), np.ones((1, 3)))) == (2, 3)
+    except ValueError:
+        return False
+
+
+@lru_cache(maxsize=None)
+def _sycl_half() -> bool:
+    import symox
+    from symox import _core
+
+    try:
+        return "sycl::half" in _core._syclcode(symox.Symbol("x") + 1, dtype="float16")
+    except (ValueError, TypeError):
+        return False
+
+
+def sycl_half() -> bool:
+    """``_syclcode(..., dtype="float16")`` prints ``sycl::half`` (sycl backend)."""
+    return not _forced_pinned() and _sycl_half()
+
+
+def cuda_logic() -> bool:
+    """``cudacode`` prints ``And/Or/Xor/Not``."""
+    return not _forced_pinned() and _cuda_logic()
+
+
+def cuda_reciprocal_sqrt() -> bool:
+    """``cudacode(..., reciprocal_sqrt=True)`` is available."""
+    return not _forced_pinned() and _cuda_reciprocal_sqrt()
+
+
+def cudacode_function() -> bool:
+    """``symox.extra.cudacode_function`` is available (implies ``cuda_logic``)."""
+    return not _forced_pinned() and _cudacode_function() and _cuda_logic()
+
+
+def lambdify_broadcasts() -> bool:
+    """The native NumPy ``lambdify`` broadcasts shapes like NumPy."""
+    return not _forced_pinned() and _lambdify_broadcasts()
+
+
+def capabilities() -> dict:
+    """Summary (diagnostics, benchmark reports)."""
+    import symox
+
+    return {
+        "symox_version": getattr(symox, "__version__", "?"),
+        "forced_pinned": _forced_pinned(),
+        "cuda_logic": cuda_logic(),
+        "cuda_reciprocal_sqrt": cuda_reciprocal_sqrt(),
+        "cudacode_function": cudacode_function(),
+        "lambdify_broadcasts": lambdify_broadcasts(),
+        "sycl_half": sycl_half(),
+    }