modax-solvers 0.0.3__py3-none-any.whl

modax/__init__.py

@@ -0,0 +1,5 @@
+"""Solver package initialization."""
+
+import jax
+
+jax.config.update("jax_enable_x64", True)

modax/_codegen.py

@@ -0,0 +1,37 @@
+"""Compile generated straight-line device source into a ``cuda.jit`` function.
+
+Two parts of the Rodas5P kernel are emitted as source rather than written: the
+sparse factorisation and triangular solves, with every slot a literal
+([`modax._sparse_direct`][]), and the Jacobian writer, with every colour's
+seed row a literal (``modax.rodas5P``). Both are load-bearing -- see the
+measurements in AGENTS.md -- and both need the same three things done to the
+text they produce, which is what this module does once.
+"""
+
+from __future__ import annotations
+
+import itertools
+import linecache
+
+from numba_cuda_mlir import cuda
+
+_SOURCES = itertools.count(1)
+
+
+def compile_device_source(name: str, lines: list[str], namespace: dict | None = None):
+    """Exec ``lines`` defining ``name`` and return it as a device function.
+
+    The source is registered with ``linecache`` under a filename of its own, so
+    a numba typing error inside it points at the offending line rather than at
+    nothing, and the text stays readable from a debugger as
+    ``fn._generated_source``. ``namespace`` supplies whatever the generated body
+    refers to by name.
+    """
+    source = "\n".join(lines) + "\n"
+    filename = f"<modax generated {name} {next(_SOURCES)}>"
+    linecache.cache[filename] = (len(source), None, source.splitlines(True), filename)
+    scope: dict = {}
+    exec(compile(source, filename, "exec"), dict(namespace or {}), scope)  # noqa: S102
+    generated = scope[name]
+    generated._generated_source = source
+    return cuda.jit(device=True)(generated)

modax/_jax_common.py

@@ -0,0 +1,111 @@
+"""Shared scaffolding for exposing the numba-cuda ensemble solvers to JAX."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+from jax.custom_batching import custom_vmap
+
+
+def normalize_y0_params(y0, params):
+    """Broadcast ``y0`` / ``params`` to a consistent ``(N, …)`` ensemble layout.
+
+    Accepts either 1-D (``(n_vars,)`` / ``(n_params,)``) or 2-D
+    (``(N, n_vars)`` / ``(N, n_params)``) inputs and returns 2-D arrays with a
+    common leading axis, so every numba-cuda solver shares one calling
+    convention.
+    """
+    y0_arr = jnp.asarray(y0, dtype=jnp.float64)
+    params_arr = jnp.asarray(params, dtype=jnp.float64)
+
+    if y0_arr.ndim not in (1, 2) or params_arr.ndim not in (1, 2):
+        raise ValueError(
+            "y0 must have shape (n_vars,) or (N, n_vars) and params shape "
+            f"(n_params,) or (N, n_params); got y0.shape={y0_arr.shape} and "
+            f"params.shape={params_arr.shape}"
+        )
+    if y0_arr.ndim == 2:
+        n = y0_arr.shape[0]
+        if params_arr.ndim == 2 and params_arr.shape[0] != n:
+            raise ValueError(
+                "params must have shape (n_params,) or (N, n_params) when y0 has "
+                f"shape (N, n_vars); got y0.shape={y0_arr.shape} and "
+                f"params.shape={params_arr.shape}"
+            )
+    elif params_arr.ndim == 2:
+        n = params_arr.shape[0]
+    else:
+        n = 1
+
+    if y0_arr.ndim == 1:
+        y0_arr = jnp.broadcast_to(y0_arr, (n, y0_arr.shape[0]))
+    if params_arr.ndim == 1:
+        params_arr = jnp.broadcast_to(params_arr, (n, params_arr.shape[0]))
+    return y0_arr, params_arr, n, y0_arr.shape[1]
+
+
+def _broadcast_for_vmap(arg, is_batched: bool, axis_size: int, name: str):
+    arr = jnp.asarray(arg)
+    if is_batched:
+        if arr.ndim != 2:
+            raise NotImplementedError(
+                f"vmap over an already-ensembled {name} is not supported; "
+                "call the solver with batched y0/params directly instead."
+            )
+        if arr.shape[0] != axis_size:
+            raise ValueError(
+                f"batched {name} has leading axis {arr.shape[0]}, expected {axis_size}"
+            )
+        return arr
+    if arr.ndim != 1:
+        raise NotImplementedError(
+            f"vmap with unbatched ensemble-shaped {name} is not supported; "
+            "call the solver with batched y0/params directly instead."
+        )
+    return jnp.broadcast_to(arr, (axis_size,) + arr.shape)
+
+
+def make_custom_vmap_solver(solve_impl: Callable, *, return_stats: bool):
+    """Wrap a solver implementation so outer ``jax.vmap`` becomes one ensemble call.
+
+    ``solve_impl`` must accept ``(y0, t_span, params)`` and return the normal
+    public solver result for those arrays.  The custom batching rule supports
+    vmapping scalar solves over ``y0`` and/or ``params`` and lowers that vmap to
+    a single native ensemble solve with a leading trajectory axis.  Every stats
+    field the kernels emit is a per-trajectory counter, so the stats pytree
+    only needs a trailing solve axis added.
+    """
+
+    @custom_vmap
+    def _solve(y0, t_span, params):
+        return solve_impl(y0, t_span, params)
+
+    @_solve.def_vmap
+    def _solve_vmap(axis_size, in_batched, y0, t_span, params):
+        y0_batched, t_span_batched, params_batched = in_batched
+        if t_span_batched:
+            t_span_arr = jnp.asarray(t_span)
+            if t_span_arr.ndim != 2:
+                raise NotImplementedError(
+                    "vmap over nested t_span values is not supported; use a shared "
+                    "t_span and vmap over y0 and/or params, or call the solver directly."
+                )
+            # JAX can mark closed-over constant save times as batched inside
+            # a larger vmapped function.  Treat that as a shared time grid.
+            t_span = t_span_arr[0]
+
+        y0_arr = _broadcast_for_vmap(y0, y0_batched, axis_size, "y0")
+        params_arr = _broadcast_for_vmap(params, params_batched, axis_size, "params")
+        result = solve_impl(y0_arr, t_span, params_arr)
+
+        if not return_stats:
+            return result[:, None, :, :], True
+
+        sol, stats = result
+        stats_out = jax.tree_util.tree_map(lambda x: x[:, None], stats)
+        stats_batched = jax.tree_util.tree_map(lambda _: True, stats_out)
+        return (sol[:, None, :, :], stats_out), (True, stats_batched)
+
+    return _solve

modax/_jax_numba_custom_call.py

@@ -0,0 +1,371 @@
+"""JAX custom-call bridge for launching numba-cuda kernels.
+
+The public surface in this module is intentionally small: compile a CUDA
+kernel, register a typed XLA FFI launcher for it, and call it from JAX.  The C++ FFI shim is built lazily into ``/tmp`` so the project can
+keep using plain ``uv run python`` without a package build step.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import hashlib
+import subprocess
+import sysconfig
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import jax
+import numpy as np
+
+_CAPSULE_NAME = b"xla._CUSTOM_CALL_TARGET"
+_TARGET_NAME = "modax_numba_cuda_abi_launch"
+_CUSTOM_CALL_API_VERSION = 4
+_REGISTERED = False
+_LOADED_LIB: ctypes.CDLL | None = None
+
+ABI_ARRAY = 0
+ABI_SCALAR_F64 = 1
+ABI_SCALAR_I32 = 2
+
+
+@dataclass(frozen=True)
+class CudaLaunch:
+    """Compiled CUDA kernel launch metadata for an XLA FFI call."""
+
+    function: int
+    grid: tuple[int, int, int]
+    block: tuple[int, int, int]
+    shared_mem: int = 0
+
+
+def _as_3d(value: int | Sequence[int]) -> tuple[int, int, int]:
+    if isinstance(value, int):
+        return (int(value), 1, 1)
+    parts = tuple(int(v) for v in value)
+    if len(parts) == 1:
+        return (parts[0], 1, 1)
+    if len(parts) == 2:
+        return (parts[0], parts[1], 1)
+    if len(parts) == 3:
+        return parts
+    raise ValueError(f"launch dimensions must have rank 1, 2, or 3; got {value!r}")
+
+
+def _pycapsule_new(ptr: int, name: bytes = _CAPSULE_NAME) -> object:
+    ctypes.pythonapi.PyCapsule_New.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_char_p,
+        ctypes.c_void_p,
+    ]
+    ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object
+    return ctypes.pythonapi.PyCapsule_New(ctypes.c_void_p(ptr), name, None)
+
+
+def _source() -> str:
+    return r"""
+#include <cstdint>
+#include <dlfcn.h>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include "xla/ffi/api/ffi.h"
+
+namespace ffi = xla::ffi;
+
+using CuLaunchKernel = int (*)(void*, unsigned int, unsigned int, unsigned int,
+                               unsigned int, unsigned int, unsigned int,
+                               unsigned int, void*, void**, void**);
+
+static CuLaunchKernel LoadCuLaunchKernel() {
+  static std::once_flag once;
+  static CuLaunchKernel fn = nullptr;
+  std::call_once(once, []() {
+    void* lib = dlopen("libcuda.so.1", RTLD_NOW | RTLD_LOCAL);
+    if (lib == nullptr) return;
+    fn = reinterpret_cast<CuLaunchKernel>(dlsym(lib, "cuLaunchKernel"));
+  });
+  return fn;
+}
+
+// numba-cuda-mlir lowers an array parameter to an MLIR MemRef descriptor,
+// passed flattened as {allocated, aligned, offset, sizes..., strides...} --
+// 3 + 2*rank kernel parameters.  The offset and strides count ELEMENTS, unlike
+// the byte strides of the numba-cuda array ABI this replaced (which also led
+// with a meminfo/parent pair, for 5 + 2*rank parameters).
+struct ArrayArg {
+  void* allocated = nullptr;
+  void* aligned = nullptr;
+  int64_t offset = 0;
+  std::vector<int64_t> sizes;
+  std::vector<int64_t> strides;
+};
+
+struct KernelArgStorage {
+  ArrayArg array;
+  double f64 = 0.0;
+  int32_t i32 = 0;
+};
+
+static void AddArrayParams(ffi::AnyBuffer buf, KernelArgStorage& storage,
+                           std::vector<void*>& params) {
+  storage.array.allocated = buf.untyped_data();
+  storage.array.aligned = buf.untyped_data();
+  storage.array.offset = 0;
+  auto dims = buf.dimensions();
+  storage.array.sizes.assign(dims.begin(), dims.end());
+  storage.array.strides.resize(storage.array.sizes.size());
+  // Row-major element counts, so the innermost stride is 1.
+  int64_t stride = 1;
+  for (int64_t i = static_cast<int64_t>(storage.array.sizes.size()) - 1; i >= 0; --i) {
+    storage.array.strides[static_cast<size_t>(i)] = stride;
+    stride *= storage.array.sizes[static_cast<size_t>(i)];
+  }
+
+  params.push_back(&storage.array.allocated);
+  params.push_back(&storage.array.aligned);
+  params.push_back(&storage.array.offset);
+  for (int64_t& size : storage.array.sizes) params.push_back(&size);
+  for (int64_t& stride_value : storage.array.strides) params.push_back(&stride_value);
+}
+
+static ffi::Error AddBufferParam(ffi::AnyBuffer buf, int64_t kind,
+                                 KernelArgStorage& storage,
+                                 std::vector<void*>& params) {
+  if (kind != 0) {
+    return ffi::Error(ffi::ErrorCode::kInvalidArgument,
+                      "unknown Numba CUDA ABI argument kind");
+  }
+  AddArrayParams(buf, storage, params);
+  return ffi::Error::Success();
+}
+
+static ffi::Error LaunchNumbaCudaAbi(
+    void* stream, int64_t function, int64_t grid_x, int64_t grid_y,
+    int64_t grid_z, int64_t block_x, int64_t block_y, int64_t block_z,
+    int64_t shared_mem, ffi::Span<const int64_t> arg_kinds,
+    ffi::Span<const double> scalar_f64_values,
+    ffi::Span<const int32_t> scalar_i32_values,
+    ffi::RemainingArgs args, ffi::RemainingRets rets) {
+  CuLaunchKernel cuLaunchKernel = LoadCuLaunchKernel();
+  if (cuLaunchKernel == nullptr) {
+    return ffi::Error(ffi::ErrorCode::kInternal,
+                      "could not load cuLaunchKernel from libcuda.so.1");
+  }
+  std::vector<KernelArgStorage> storage(arg_kinds.size());
+  std::vector<void*> params;
+  params.reserve(arg_kinds.size() * 12);
+  size_t arg_idx = 0;
+  size_t ret_idx = 0;
+  size_t f64_idx = 0;
+  size_t i32_idx = 0;
+
+  for (size_t i = 0; i < arg_kinds.size(); ++i) {
+    const int64_t kind = arg_kinds[i];
+    if (kind == 1) {
+      if (f64_idx >= scalar_f64_values.size()) {
+        return ffi::Error(ffi::ErrorCode::kInvalidArgument,
+                          "not enough f64 scalar values");
+      }
+      storage[i].f64 = scalar_f64_values[f64_idx++];
+      params.push_back(&storage[i].f64);
+    } else if (kind == 2) {
+      if (i32_idx >= scalar_i32_values.size()) {
+        return ffi::Error(ffi::ErrorCode::kInvalidArgument,
+                          "not enough i32 scalar values");
+      }
+      storage[i].i32 = scalar_i32_values[i32_idx++];
+      params.push_back(&storage[i].i32);
+    } else if (arg_idx < args.size()) {
+      auto arg = args.get<ffi::AnyBuffer>(arg_idx++);
+      if (!arg.has_value()) return arg.error();
+      ffi::Error err = AddBufferParam(arg.value(), kind, storage[i], params);
+      if (!err.success()) return err;
+    } else {
+      auto ret = rets.get<ffi::AnyBuffer>(ret_idx++);
+      if (!ret.has_value()) return ret.error();
+      ffi::Error err = AddBufferParam(*ret.value(), kind, storage[i], params);
+      if (!err.success()) return err;
+    }
+  }
+  if (arg_idx != args.size() || ret_idx != rets.size()) {
+    return ffi::Error(ffi::ErrorCode::kInvalidArgument,
+                      "kernel ABI kinds did not consume all buffers");
+  }
+
+  int err = cuLaunchKernel(reinterpret_cast<void*>(function),
+                           static_cast<unsigned int>(grid_x),
+                           static_cast<unsigned int>(grid_y),
+                           static_cast<unsigned int>(grid_z),
+                           static_cast<unsigned int>(block_x),
+                           static_cast<unsigned int>(block_y),
+                           static_cast<unsigned int>(block_z),
+                           static_cast<unsigned int>(shared_mem),
+                           stream, params.data(), nullptr);
+  if (err != 0) {
+    return ffi::Error(ffi::ErrorCode::kInternal,
+                      "cuLaunchKernel failed with CUDA driver error " +
+                          std::to_string(err));
+  }
+  return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    modax_numba_cuda_abi_launch, LaunchNumbaCudaAbi,
+    ffi::Ffi::Bind()
+        .Ctx<ffi::PlatformStream<void*>>()
+        .Attr<int64_t>("function")
+        .Attr<int64_t>("grid_x")
+        .Attr<int64_t>("grid_y")
+        .Attr<int64_t>("grid_z")
+        .Attr<int64_t>("block_x")
+        .Attr<int64_t>("block_y")
+        .Attr<int64_t>("block_z")
+        .Attr<int64_t>("shared_mem")
+        .Attr<ffi::Span<const int64_t>>("arg_kinds")
+        .Attr<ffi::Span<const double>>("scalar_f64_values")
+        .Attr<ffi::Span<const int32_t>>("scalar_i32_values")
+        .RemainingArgs()
+        .RemainingRets());
+"""
+
+
+def _build_bridge() -> Path:
+    include_dir = Path(jax.ffi.include_dir())
+    build_dir = Path(tempfile.gettempdir()) / "modax_jax_numba_cuda_bridge"
+    build_dir.mkdir(parents=True, exist_ok=True)
+    source = _source()
+    digest = hashlib.sha256(source.encode()).hexdigest()[:16]
+    src_path = build_dir / f"bridge-{digest}.cc"
+    so_path = build_dir / f"bridge-{digest}{sysconfig.get_config_var('EXT_SUFFIX')}"
+    if so_path.exists():
+        return so_path
+    src_path.write_text(source)
+    cmd = [
+        "g++",
+        "-std=c++17",
+        "-shared",
+        "-fPIC",
+        "-O2",
+        f"-I{include_dir}",
+        str(src_path),
+        "-ldl",
+        "-o",
+        str(so_path),
+    ]
+    subprocess.run(cmd, check=True, capture_output=True, text=True)
+    return so_path
+
+
+def register_target() -> None:
+    """Register the generic CUDA launcher with JAX once per process."""
+
+    global _LOADED_LIB, _REGISTERED
+    if _REGISTERED:
+        return
+    so_path = _build_bridge()
+    _LOADED_LIB = ctypes.CDLL(str(so_path))
+    symbol = getattr(_LOADED_LIB, _TARGET_NAME)
+    address = ctypes.cast(symbol, ctypes.c_void_p).value
+    # A symbol resolved out of a loaded library always has one; the
+    # `None` is what `c_void_p` carries for a null pointer.
+    assert address is not None
+    capsule = _pycapsule_new(address)
+    jax.ffi.register_ffi_target(_TARGET_NAME, capsule, platform="CUDA", api_version=1)
+    _REGISTERED = True
+
+
+# Loaded modules are held for the process lifetime: the ``CUfunction`` handed to
+# the FFI target stays valid only while its module is loaded, and these kernels
+# live as long as the JAX primitives that launch them.
+_LOADED_MODULES: list[Any] = []
+
+
+def compile_kernel(kernel: Any, argtypes: Sequence[Any]) -> int:
+    """Compile a ``cuda.jit`` kernel and return its ``CUfunction`` pointer.
+
+    numba-cuda-mlir exposes no ``get_cufunc()`` -- its ``MLIRLibrary`` only
+    offers the textual IR.  The linked cubin and the mangled entry name are on
+    the compile result's metadata instead, so load the module through the driver
+    and look the function up by name.
+    """
+
+    cres = kernel.compile(tuple(argtypes)).cres
+    cubin = cres.metadata["cubin"]
+    func_name = cres.metadata["func_name"]
+
+    libcuda = ctypes.CDLL("libcuda.so.1")
+    module = ctypes.c_void_p()
+    err = libcuda.cuModuleLoadData(ctypes.byref(module), ctypes.c_char_p(cubin))
+    if err != 0:
+        raise RuntimeError(f"cuModuleLoadData failed with CUDA driver error {err}")
+    _LOADED_MODULES.append(module)
+
+    function = ctypes.c_void_p()
+    err = libcuda.cuModuleGetFunction(
+        ctypes.byref(function), module, func_name.encode()
+    )
+    if err != 0:
+        raise RuntimeError(f"cuModuleGetFunction failed with CUDA driver error {err}")
+    if function.value is None:
+        raise RuntimeError("cuModuleGetFunction returned a null function pointer")
+    return int(function.value)
+
+
+def make_launch(
+    kernel: Any,
+    argtypes: Sequence[Any],
+    *,
+    grid: int | Sequence[int],
+    block: int | Sequence[int],
+    shared_mem: int = 0,
+) -> CudaLaunch:
+    return CudaLaunch(
+        function=compile_kernel(kernel, argtypes),
+        grid=_as_3d(grid),
+        block=_as_3d(block),
+        shared_mem=int(shared_mem),
+    )
+
+
+def ffi_abi_call(
+    launch: CudaLaunch,
+    inputs: Sequence[Any],
+    output_specs: Sequence[jax.ShapeDtypeStruct],
+    *,
+    input_kinds: Sequence[int],
+    scalar_f64_values: Sequence[float] = (),
+    scalar_i32_values: Sequence[int] = (),
+) -> tuple[Any, ...]:
+    """Launch a Numba CUDA kernel using Numba's normal array/scalar ABI.
+
+    Outputs are always arrays, so only the input kinds need spelling out.
+    """
+
+    register_target()
+    attrs = {
+        "function": np.int64(launch.function),
+        "grid_x": np.int64(launch.grid[0]),
+        "grid_y": np.int64(launch.grid[1]),
+        "grid_z": np.int64(launch.grid[2]),
+        "block_x": np.int64(launch.block[0]),
+        "block_y": np.int64(launch.block[1]),
+        "block_z": np.int64(launch.block[2]),
+        "shared_mem": np.int64(launch.shared_mem),
+        "arg_kinds": np.asarray(
+            tuple(input_kinds) + (ABI_ARRAY,) * len(output_specs), dtype=np.int64
+        ),
+        "scalar_f64_values": np.asarray(tuple(scalar_f64_values), dtype=np.float64),
+        "scalar_i32_values": np.asarray(tuple(scalar_i32_values), dtype=np.int32),
+    }
+    result = jax.ffi.ffi_call(
+        _TARGET_NAME,
+        tuple(output_specs),
+        has_side_effect=False,
+        custom_call_api_version=_CUSTOM_CALL_API_VERSION,
+    )(*inputs, **attrs)
+    if not isinstance(result, tuple):
+        return (result,)
+    return result