# SPDX-FileCopyrightText: 2026 John C. Travers and contributors
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""The protected operations of the ``mlx`` backend, and the probes that establish them.
SEP 33 section 6.3's arithmetic policy. Each protected operation has a mechanism and its
evidence:
* **Matrix products** are float32 only with ``MLX_ENABLE_TF32=0``, which the package
sets when it is imported and MLX has not been initialised; :func:`gemm_probe` measures
every dispatch family a run uses against a discriminating bound (``2**-17``) and
refuses a device that computes them in TF32.
* **The stage propagator and every combination of stage terms** are formed in
double-float arithmetic from binary64 terms held as pairs of binary32 numbers
(:class:`DoubleFloatTerms`), then rounded once: on the Metal device by one custom
kernel compiled in ``"safe"`` math mode, whose two-sum and two-product cannot be
reassociated; on MLX's CPU device, available for local tests, in binary64 on
the host with NumPy (the memory is shared). The phase is range-reduced in double-float
before its binary32 cosine and sine, and the exponential of the real part is
``exp(hi) (1 + lo)``.
Nothing outside :mod:`solprop.backends.mlx_gpu` imports MLX.
"""
import math
from dataclasses import dataclass
from functools import cache
import mlx.core as mx # pyright: ignore[reportMissingImports] - macOS arm64 only
import numpy as np
from solprop.backends.exponential import _SERIES_TERMS, PHI_JOIN
from solprop.exceptions import BackendError
__all__ = [
"GEMM_PROBE_BOUND",
"DoubleFloatTerms",
"fft_bound_32",
"fft_probe",
"gemm_probe",
"on_gpu",
"population",
"population_parameters",
"prefix_scan",
"rate_kernel",
]
GEMM_PROBE_BOUND = 2.0**-17
"""The normwise relative error a float32 product must stay below (SEP 33 7.1).
Float32 products of inner dimension up to ``10**3`` are expected near
``sqrt(n) u <= 2e-6`` and TF32 near its own unit roundoff ``2**-11``: the bound
separates the two by more than a factor of four on each side.
"""
# 2 pi as a double-float pair, and its reciprocal: the reduction of a phase of up to
# 1e7 rad keeps the residual to about 2**-48 of the phase.
_TWO_PI_HI = float(np.float32(2 * math.pi))
_TWO_PI_LO = float(np.float32(2 * math.pi - _TWO_PI_HI))
_HEADER = """
inline float2 two_sum(float a, float b) {
float s = a + b;
float bb = s - a;
float e = (a - (s - bb)) + (b - bb);
return float2(s, e);
}
inline float2 quick_two_sum(float a, float b) {
float s = a + b;
float e = b - (s - a);
return float2(s, e);
}
inline float2 two_prod(float a, float b) {
float p = a * b;
float e = metal::fma(a, b, -p);
return float2(p, e);
}
inline float2 df_add(float2 x, float2 y) {
float2 s = two_sum(x.x, y.x);
float e = s.y + x.y + y.y;
return quick_two_sum(s.x, e);
}
inline float2 df_mul(float2 x, float2 y) {
float2 p = two_prod(x.x, y.x);
p.y = metal::fma(x.x, y.y, metal::fma(x.y, y.x, p.y));
return quick_two_sum(p.x, p.y);
}
"""
_COMBINE = """
uint i = thread_position_in_grid.x;
if (i >= N) { return; }
float2 re = float2(0.0f, 0.0f);
float2 im = float2(0.0f, 0.0f);
for (int t = 0; t < TERMS; ++t) {
float2 w = float2(w_hi[t], w_lo[t]);
uint k = t * N + i;
re = df_add(re, df_mul(w, float2(re_hi[k], re_lo[k])));
im = df_add(im, df_mul(w, float2(im_hi[k], im_lo[k])));
}
if (EXPONENTIAL == 0) {
out_re[i] = re.x + re.y;
out_im[i] = im.x + im.y;
return;
}
float sign = SIGN > 0 ? 1.0f : -1.0f;
float turns = metal::rint(im.x * INV_TWO_PI);
float2 period = float2(TWO_PI_HI, TWO_PI_LO);
float2 phase = df_add(im, df_mul(float2(-turns, 0.0f), period));
float angle = sign * (phase.x + phase.y);
// exp(hi + lo) = exp(hi) (1 + lo) to binary64-level accuracy, lo being below
// hi's unit in the last place.
float magnitude = metal::precise::exp(sign * re.x) * (1.0f + sign * re.y);
float c = metal::precise::cos(angle);
float s = metal::precise::sin(angle);
// Binary32 cosine and sine leave c^2 + s^2 off one by a few u with a bias that
// a long march accumulates in the energy; one Newton step on the double-float
// squared modulus returns the phasor to unit modulus. The step is applied as
// c + c delta by a fused multiply-add: a factor 1 + delta rounded to binary32
// would itself be biased, its spacing being u below one and 2u above.
float2 norm = df_add(two_prod(c, c), two_prod(s, s));
float delta = -0.5f * ((norm.x - 1.0f) + norm.y);
out_re[i] = magnitude * metal::fma(c, delta, c);
out_im[i] = magnitude * metal::fma(s, delta, s);
"""
_PAIR = """
uint i = thread_position_in_grid.x;
if (i >= N) { return; }
float2 re = float2(0.0f, 0.0f);
float2 im = float2(0.0f, 0.0f);
for (int t = 0; t < TERMS; ++t) {
float2 w = float2(w_hi[t], w_lo[t]);
uint k = t * N + i;
re = df_add(re, df_mul(w, float2(re_hi[k], re_lo[k])));
im = df_add(im, df_mul(w, float2(im_hi[k], im_lo[k])));
}
// The combination kernel's exponential at both signs from one combination: the
// inverse's magnitude is the same expression at sign -1, and its phasor is the
// forward's conjugate, cosine and sine being even and odd exactly.
float turns = metal::rint(im.x * INV_TWO_PI);
float2 period = float2(TWO_PI_HI, TWO_PI_LO);
float2 phase = df_add(im, df_mul(float2(-turns, 0.0f), period));
float angle = phase.x + phase.y;
float c = metal::precise::cos(angle);
float s = metal::precise::sin(angle);
float2 norm = df_add(two_prod(c, c), two_prod(s, s));
float delta = -0.5f * ((norm.x - 1.0f) + norm.y);
float cc = metal::fma(c, delta, c);
float ss = metal::fma(s, delta, s);
float forward = metal::precise::exp(re.x) * (1.0f + re.y);
float back = metal::precise::exp(-re.x) * (1.0f + -re.y);
forward_re[i] = forward * cc;
forward_im[i] = forward * ss;
back_re[i] = back * cc;
back_im[i] = -(back * ss);
"""
_MOMENTS_HEADER = """
struct dfc { float2 re; float2 im; };
inline float2 df_neg(float2 a) { return float2(-a.x, -a.y); }
inline float2 df_sub(float2 a, float2 b) { return df_add(a, df_neg(b)); }
inline float2 df_div_f(float2 a, float b) {
float q1 = a.x / b;
float2 p = two_prod(q1, b);
float r = ((a.x - p.x) - p.y) + a.y;
return quick_two_sum(q1, r / b);
}
inline float2 df_div(float2 a, float2 b) {
float q1 = a.x / b.x;
float2 r = df_sub(a, df_mul(float2(q1, 0.0f), b));
float q2 = r.x / b.x;
r = df_sub(r, df_mul(float2(q2, 0.0f), b));
float q3 = r.x / b.x;
return df_add(quick_two_sum(q1, q2), float2(q3, 0.0f));
}
inline dfc c_add(dfc a, dfc b) { return {df_add(a.re, b.re), df_add(a.im, b.im)}; }
inline dfc c_sub(dfc a, dfc b) { return {df_sub(a.re, b.re), df_sub(a.im, b.im)}; }
inline dfc c_mul(dfc a, dfc b) {
return {df_sub(df_mul(a.re, b.re), df_mul(a.im, b.im)),
df_add(df_mul(a.re, b.im), df_mul(a.im, b.re))};
}
inline dfc c_real(dfc a, float2 s) { return {df_mul(a.re, s), df_mul(a.im, s)}; }
inline dfc c_div_f(dfc a, float b) { return {df_div_f(a.re, b), df_div_f(a.im, b)}; }
inline dfc c_ldexp(dfc a, int e) {
return {float2(metal::ldexp(a.re.x, e), metal::ldexp(a.re.y, e)),
float2(metal::ldexp(a.im.x, e), metal::ldexp(a.im.y, e))};
}
inline dfc c_div(dfc a, dfc b) {
float2 d = df_add(df_mul(b.re, b.re), df_mul(b.im, b.im));
dfc n = {df_add(df_mul(a.re, b.re), df_mul(a.im, b.im)),
df_sub(df_mul(a.im, b.re), df_mul(a.re, b.im))};
return {df_div(n.re, d), df_div(n.im, d)};
}
// exp(y) for |y| <= 1/2 by its Taylor series, Horner's form, 18 terms (the tail is
// below 2^-60 relative).
inline dfc c_exp_small(dfc y) {
dfc e = {float2(1.0f, 0.0f), float2(0.0f, 0.0f)};
for (int m = 18; m >= 1; --m) {
e = c_div_f(c_mul(e, y), float(m));
e.re = df_add(e.re, float2(1.0f, 0.0f));
}
return e;
}
// exp(x): the phase reduced by whole turns of a double-float 2 pi, then scaled to at
// most 1/2, summed and squared back.
inline dfc c_exp(dfc x) {
float turns = metal::rint(x.im.x * INV_TWO_PI);
dfc r = {x.re, df_add(x.im, df_mul(float2(-turns, 0.0f),
float2(TWO_PI_HI, TWO_PI_LO)))};
float size = metal::sqrt(r.re.x * r.re.x + r.im.x * r.im.x);
int s = size > 0.5f ? int(metal::ceil(metal::log2(size / 0.5f))) : 0;
dfc e = c_exp_small(c_ldexp(r, -s));
for (int level = 0; level < s; ++level) { e = c_mul(e, e); }
return e;
}
// One combination sum_t w_ct T_t of the terms, in double-float.
#define COMBINE(c, out) \
out = {float2(0.0f, 0.0f), float2(0.0f, 0.0f)}; \
for (int t = 0; t < TERMS; ++t) { \
float2 w = float2(w_hi[(c) * TERMS + t], w_lo[(c) * TERMS + t]); \
uint k = t * N + i; \
out.re = df_add(out.re, df_mul(w, float2(re_hi[k], re_lo[k]))); \
out.im = df_add(out.im, df_mul(w, float2(im_hi[k], im_lo[k]))); \
}
"""
_MOMENTS = """
uint i = thread_position_in_grid.x;
if (i >= N) { return; }
// x = sum_t w_t T_t in double-float, as the propagator kernels combine it.
dfc x;
COMBINE(0, x)
dfc phi[P + 1];
float magnitude = metal::sqrt(x.re.x * x.re.x + x.im.x * x.im.x);
if (magnitude > JOIN) {
// exp(x) with the phase reduced by whole turns of a double-float 2 pi, then
// scaled to at most 1/2, summed and squared back; recurred up beyond the join,
// where the recurrence is stable (solprop.backends.exponential._phi_far).
phi[0] = c_exp(x);
for (int k = 1; k <= P; ++k) {
dfc shifted = phi[k - 1];
shifted.re = df_sub(shifted.re, float2(inv_fact[2 * (k - 1)],
inv_fact[2 * (k - 1) + 1]));
phi[k] = c_div(shifted, x);
}
} else {
// Near zero: the reference's scaling, highest series, downward recurrence and
// doubling (solprop.backends.exponential._phi_near), in double-float.
int s = magnitude > 0.5f ? int(metal::ceil(metal::log2(magnitude / 0.5f))) : 0;
dfc y = c_ldexp(x, -s);
dfc term = {float2(inv_fact[2 * P], inv_fact[2 * P + 1]), float2(0.0f, 0.0f)};
dfc total = term;
for (int m = 1; m < SERIES; ++m) {
term = c_div_f(c_mul(term, y), float(m + P));
total = c_add(total, term);
}
phi[P] = total;
for (int k = P - 1; k >= 1; --k) {
phi[k] = c_mul(y, phi[k + 1]);
phi[k].re = df_add(phi[k].re, float2(inv_fact[2 * k], inv_fact[2 * k + 1]));
}
phi[0] = c_exp_small(y);
for (int level = 0; level < s; ++level) {
dfc old[P + 1];
for (int k = 0; k <= P; ++k) { old[k] = phi[k]; }
phi[0] = c_mul(old[0], old[0]);
for (int k = 1; k <= P; ++k) {
dfc sum = c_mul(old[0], old[k]);
for (int j = 1; j <= k; ++j) {
sum = c_add(sum, c_real(old[j], float2(inv_fact[2 * (k - j)],
inv_fact[2 * (k - j) + 1])));
}
phi[k] = c_ldexp(sum, -k);
}
}
}
flow_re[i] = phi[0].re.x + phi[0].re.y;
flow_im[i] = phi[0].im.x + phi[0].im.y;
for (int k = 0; k < P; ++k) {
// k! phi_{k+1}, the factorial exact in binary32 (at most 7!).
float f = 1.0f;
for (int j = 2; j <= k; ++j) { f *= float(j); }
dfc m = c_real(phi[k + 1], float2(f, 0.0f));
moments_re[k * N + i] = m.re.x + m.re.y;
moments_im[k * N + i] = m.im.x + m.im.y;
}
// A varying operator's compensation factors (SEP 16 section 4.2): per history
// sample exp(s_j x) (D(z_j) - x / h), the rates' difference being combination
// 1 + j, each product rounded once; then the endpoint's D(z + h) - x / h.
for (int j = 0; j < SLOTS; ++j) {
dfc d;
COMBINE(1 + j, d)
dfc f = c_mul(c_exp(c_real(x, float2(spans[2 * j], spans[2 * j + 1]))), d);
factors_re[j * N + i] = f.re.x + f.re.y;
factors_im[j * N + i] = f.im.x + f.im.y;
}
if (ENDPOINT) {
dfc d;
COMBINE(1 + SLOTS, d)
factors_re[SLOTS * N + i] = d.re.x + d.re.y;
factors_im[SLOTS * N + i] = d.im.x + d.im.y;
}
"""
@cache
def _moments_kernel() -> object:
"""Return the compiled kernel of an exponential Adams trial's flow and moments.
Returns
-------
object
The ``mlx.core.fast.metal_kernel`` callable.
"""
return mx.fast.metal_kernel(
name="solprop_double_float_moments",
input_names=[
"re_hi",
"re_lo",
"im_hi",
"im_lo",
"w_hi",
"w_lo",
"inv_fact",
"spans",
],
output_names=[
"flow_re",
"flow_im",
"moments_re",
"moments_im",
"factors_re",
"factors_im",
],
source=_source(_MOMENTS)
.replace("JOIN", f"{PHI_JOIN!r}f")
.replace("SERIES", str(_SERIES_TERMS)),
header=_HEADER + _source(_MOMENTS_HEADER),
compile_options={"math_mode": "safe"},
)
def _source(body: str) -> str:
"""Return a double-float kernel's source with its constants in place.
Parameters
----------
body : str
The kernel's body.
Returns
-------
str
The source.
"""
return (
body.replace("INV_TWO_PI", f"{1 / (2 * math.pi)!r}f")
.replace("TWO_PI_HI", f"{_TWO_PI_HI!r}f")
.replace("TWO_PI_LO", f"{_TWO_PI_LO!r}f")
)
@cache
def _pair_kernel() -> object:
"""Return the compiled kernel of a propagator and its inverse.
Returns
-------
object
The ``mlx.core.fast.metal_kernel`` callable.
"""
return mx.fast.metal_kernel(
name="solprop_double_float_pair",
input_names=["re_hi", "re_lo", "im_hi", "im_lo", "w_hi", "w_lo"],
output_names=["forward_re", "forward_im", "back_re", "back_im"],
source=_source(_PAIR),
header=_HEADER,
compile_options={"math_mode": "safe"},
)
@cache
def _kernel() -> object:
"""Return the compiled double-float combination kernel.
Returns
-------
object
The ``mlx.core.fast.metal_kernel`` callable.
"""
return mx.fast.metal_kernel(
name="solprop_double_float_combine",
input_names=["re_hi", "re_lo", "im_hi", "im_lo", "w_hi", "w_lo"],
output_names=["out_re", "out_im"],
source=_source(_COMBINE),
header=_HEADER,
compile_options={"math_mode": "safe"},
)
[docs]
def on_gpu() -> bool:
"""Return whether MLX's default device is the Metal GPU.
Returns
-------
bool
True on the GPU, False on MLX's CPU device.
"""
return mx.default_device() == mx.gpu
def _split(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Split binary64 values into a binary32 head and tail.
Parameters
----------
values : numpy.ndarray
Real binary64 values.
Returns
-------
tuple of numpy.ndarray
``hi = float32(v)`` and ``lo = float32(v - hi)``; ``hi + lo`` holds ``v`` to
about 48 bits.
"""
hi = values.astype(np.float32)
lo = (values - hi.astype(np.float64)).astype(np.float32)
return hi, lo
[docs]
@dataclass(frozen=True)
class DoubleFloatTerms:
"""Complex binary64 terms held on the device for double-float combination.
On the Metal device each term is four float32 arrays (the real and imaginary heads
and tails); on MLX's CPU device it is one complex128-equivalent pair of float64
arrays. :meth:`combine` returns ``sum_t w_t T_t`` rounded once to complex64, and
:meth:`exponential` ``exp(sign sum_t w_t T_t)`` with the phase range-reduced in
double-float before the binary32 exponential, cosine and sine (SEP 33 4.2).
Attributes
----------
shape : tuple of int
The shape of one term.
n_terms : int
The number of terms.
parts : tuple of mlx.core.array
The device arrays, each of shape ``(n_terms * prod(shape),)``; empty on MLX's
CPU device.
host : numpy.ndarray or None
The binary64 terms, kept on MLX's CPU device only.
real_bound : tuple of float
Each term's largest ``|Re T_t|``, so that ``sum_t |w_t| real_bound_t`` bounds
a combination's real parts without forming it.
"""
shape: tuple[int, ...]
n_terms: int
parts: tuple[mx.array, ...]
host: np.ndarray | None = None
real_bound: tuple[float, ...] = ()
[docs]
@classmethod
def upload(cls, terms: np.ndarray) -> "DoubleFloatTerms": # noqa: UP037 - the runtime type hook resolves this while the class body runs
"""Upload binary64 terms of shape ``(n_terms, *shape)``.
Parameters
----------
terms : numpy.ndarray
Complex binary64 terms, finite.
Returns
-------
DoubleFloatTerms
The terms on the device.
Raises
------
BackendError
If a term is not finite.
"""
values = np.ascontiguousarray(terms, dtype=np.complex128)
if not np.all(np.isfinite(values)):
raise BackendError("a stage term must be finite to reach the device")
shape = values.shape[1:]
flat = values.reshape(-1)
bound = tuple(float(np.max(np.abs(term.real), initial=0.0)) for term in values)
if on_gpu():
re_hi, re_lo = _split(flat.real)
im_hi, im_lo = _split(flat.imag)
parts = tuple(mx.array(part) for part in (re_hi, re_lo, im_hi, im_lo))
return cls(shape, int(values.shape[0]), parts, real_bound=bound)
return cls(shape, int(values.shape[0]), (), values, real_bound=bound)
[docs]
def real_reach(self, weights: np.ndarray) -> float:
"""Return a bound on ``max |Re sum_t w_t T_t|`` from the terms' own bounds.
Parameters
----------
weights : numpy.ndarray
One binary64 weight per term.
Returns
-------
float
``sum_t |w_t| max |Re T_t|``.
"""
w = np.abs(np.asarray(weights, dtype=np.float64).reshape(-1))
return float(np.dot(w, np.asarray(self.real_bound)))
[docs]
@property
def nbytes(self) -> int:
"""Return the device bytes the terms hold.
Returns
-------
int
Their size.
"""
host = 0 if self.host is None else int(self.host.nbytes)
return int(sum(part.nbytes for part in self.parts)) + host
[docs]
def device_weights(self, weights: np.ndarray) -> tuple[mx.array, mx.array]:
"""Split binary64 weights into the kernel's float32 heads and tails.
Parameters
----------
weights : numpy.ndarray
One binary64 weight per term.
Returns
-------
tuple of mlx.core.array
The heads and the tails, ``(n_terms,)`` each.
Raises
------
BackendError
If the weights do not match the terms.
"""
w = np.asarray(weights, dtype=np.float64).reshape(-1)
if w.size != self.n_terms:
msg = f"{w.size} weights for {self.n_terms} stage terms"
raise BackendError(msg)
w_hi, w_lo = _split(w)
return mx.array(w_hi), mx.array(w_lo)
[docs]
def pair(self, w_hi: mx.array, w_lo: mx.array) -> tuple[mx.array, mx.array]:
"""Return ``exp(sum_t w_t T_t)`` and its inverse from one combination.
On the Metal device only; each equals :meth:`exponential` at its sign. Its
inputs are device arrays, so a compiled graph may call it.
Parameters
----------
w_hi, w_lo : mlx.core.array
The weights' heads and tails (:meth:`device_weights`).
Returns
-------
tuple of mlx.core.array
The propagator and its inverse, complex64.
"""
n = math.prod(self.shape)
forward_re, forward_im, back_re, back_im = _pair_kernel()( # pyright: ignore[reportCallIssue]
inputs=[*self.parts, w_hi, w_lo],
template=[("N", n), ("TERMS", self.n_terms)],
grid=(n, 1, 1),
threadgroup=(min(256, n), 1, 1),
output_shapes=[(n,)] * 4,
output_dtypes=[mx.float32] * 4,
)
forward = (forward_re + 1j * forward_im).astype(mx.complex64)
back = (back_re + 1j * back_im).astype(mx.complex64)
return forward.reshape(self.shape), back.reshape(self.shape)
[docs]
def combination(self, w_hi: mx.array, w_lo: mx.array) -> mx.array:
"""Return ``sum_t w_t T_t`` rounded once to complex64, from device weights.
:meth:`combine` with weights already on the device (:meth:`device_weights`).
Parameters
----------
w_hi, w_lo : mlx.core.array
The weights' heads and tails.
Returns
-------
mlx.core.array
The combination.
"""
n = math.prod(self.shape)
out_re, out_im = _kernel()( # pyright: ignore[reportCallIssue]
inputs=[*self.parts, w_hi, w_lo],
template=[
("N", n),
("TERMS", self.n_terms),
("EXPONENTIAL", 0),
("SIGN", 1),
],
grid=(n, 1, 1),
threadgroup=(min(256, n), 1, 1),
output_shapes=[(n,), (n,)],
output_dtypes=[mx.float32, mx.float32],
)
return (out_re + 1j * out_im).astype(mx.complex64).reshape(self.shape)
[docs]
def moments(
self, w_hi: mx.array, w_lo: mx.array, count: int
) -> tuple[mx.array, mx.array]:
"""Return ``exp(x)`` and ``k! phi_{k+1}(x)`` for ``k < count``, ``x = sum w T``.
An exponential Adams trial's flow and moments (SEP 16) in double-float, each
rounded once to complex64 (:meth:`trial` without a compensation). On the
Metal device only.
Parameters
----------
w_hi, w_lo : mlx.core.array
The weights' heads and tails (:meth:`device_weights`).
count : int
The moments, at most :data:`~solprop.backends.exponential.MAX_PHI_ORDER`.
Returns
-------
tuple of mlx.core.array
The flow, of :attr:`shape`, and the moments, ``(count, *shape)``.
"""
flow, moments, _ = self._trial(w_hi, w_lo, count, np.zeros(0), endpoint=False)
return flow, moments
[docs]
def trial(
self,
weights: np.ndarray,
count: int,
spans: np.ndarray,
*,
endpoint: bool,
) -> tuple[mx.array, mx.array, list[mx.array]]:
"""Return an exponential Adams trial's flow, moments and compensation factors.
One pass over the terms, in double-float, every output rounded once:
``exp(x)`` and ``k! phi_{k+1}(x)`` -- the reference's far recurrence beyond
its join and its scaled series and doubling within it
(:func:`solprop.backends.exponential.phi_functions`), ``exp`` reduced by whole
turns of a double-float ``2 pi`` -- with ``x`` the first row's combination;
then for each span ``s_j`` the factor ``exp(s_j x) R_j``, ``R_j`` the
combination of row ``1 + j``; then, with `endpoint`, the last row's
combination alone. On the Metal device only.
Parameters
----------
weights : numpy.ndarray
Binary64 weights over the terms, ``(n_rows, n_terms)``: ``x``'s, each
history sample's difference of rates, and the endpoint's.
count : int
The moments, at most :data:`~solprop.backends.exponential.MAX_PHI_ORDER`.
spans : numpy.ndarray
``s_j``, binary64, one per history sample.
endpoint : bool
Whether the last row is the endpoint's.
Returns
-------
tuple
The flow, of :attr:`shape`; the moments, ``(count, *shape)``; and the
samples' factors then the endpoint's, each of :attr:`shape`.
Raises
------
BackendError
If the weights do not match the terms or the rows.
"""
rows = np.asarray(weights, dtype=np.float64)
if rows.shape != (1 + len(spans) + int(endpoint), self.n_terms):
msg = f"trial weights of shape {rows.shape} for {self.n_terms} terms"
raise BackendError(msg)
hi, lo = _split(rows.reshape(-1))
return self._trial(
mx.array(hi), mx.array(lo), count, np.asarray(spans), endpoint=endpoint
)
def _trial(
self,
w_hi: mx.array,
w_lo: mx.array,
count: int,
spans: np.ndarray,
*,
endpoint: bool,
) -> tuple[mx.array, mx.array, list[mx.array]]:
"""Launch the trial kernel (:meth:`trial`) on weights already on the device.
Parameters
----------
w_hi, w_lo : mlx.core.array
The rows' heads and tails, flattened.
count : int
The moments.
spans : numpy.ndarray
``s_j``, binary64.
endpoint : bool
Whether the last row is the endpoint's.
Returns
-------
tuple
As :meth:`trial`.
"""
n = math.prod(self.shape)
factorials = np.array(
[1.0 / math.factorial(k) for k in range(count + 1)], dtype=np.float64
)
hi, lo = _split(factorials)
span_hi, span_lo = _split(np.asarray(spans, dtype=np.float64))
n_factors = len(spans) + int(endpoint)
span_pairs = (
np.stack([span_hi, span_lo], axis=1)
if len(spans)
else np.zeros((1, 2), np.float32)
)
outputs = _moments_kernel()( # pyright: ignore[reportCallIssue]
inputs=[
*self.parts,
w_hi,
w_lo,
mx.array(np.stack([hi, lo], axis=1)),
mx.array(span_pairs),
],
template=[
("N", n),
("TERMS", self.n_terms),
("P", count),
("SLOTS", len(spans)),
("ENDPOINT", int(endpoint)),
],
grid=(n, 1, 1),
threadgroup=(min(256, n), 1, 1),
output_shapes=[(n,), (n,), (count * n,), (count * n,)]
+ [(max(n_factors, 1) * n,)] * 2,
output_dtypes=[mx.float32] * 6,
)
flow_re, flow_im, moments_re, moments_im, factors_re, factors_im = outputs
flow = (flow_re + 1j * flow_im).astype(mx.complex64).reshape(self.shape)
moments = (moments_re + 1j * moments_im).astype(mx.complex64)
factors = (factors_re + 1j * factors_im).astype(mx.complex64)
return (
flow,
moments.reshape(count, *self.shape),
[
factors[j * n : (j + 1) * n].reshape(self.shape)
for j in range(n_factors)
],
)
def _evaluate(
self, weights: np.ndarray, *, exponential: int, sign: int
) -> mx.array:
"""Combine the terms over binary64 weights on the device.
Parameters
----------
weights : numpy.ndarray
One binary64 weight per term.
exponential : int
1 to return the exponential, 0 the combination.
sign : int
The exponential's sign, +1 or -1.
Returns
-------
mlx.core.array
Complex64 of shape :attr:`shape`.
Raises
------
BackendError
If the weights do not match the terms.
"""
w = np.asarray(weights, dtype=np.float64).reshape(-1)
if w.size != self.n_terms:
msg = f"{w.size} weights for {self.n_terms} stage terms"
raise BackendError(msg)
if self.host is not None:
total = np.tensordot(w, self.host, axes=1)
if exponential:
total = np.exp(sign * total)
return mx.array(total.astype(np.complex64))
n = math.prod(self.shape)
w_hi, w_lo = _split(w)
out_re, out_im = _kernel()( # pyright: ignore[reportCallIssue]
inputs=[*self.parts, mx.array(w_hi), mx.array(w_lo)],
template=[
("N", n),
("TERMS", self.n_terms),
("EXPONENTIAL", exponential),
("SIGN", sign),
],
grid=(n, 1, 1),
threadgroup=(min(256, n), 1, 1),
output_shapes=[(n,), (n,)],
output_dtypes=[mx.float32, mx.float32],
)
return (out_re + 1j * out_im).astype(mx.complex64).reshape(self.shape)
[docs]
def combine(self, weights: np.ndarray) -> mx.array:
"""Return ``sum_t w_t T_t`` rounded once to complex64.
Parameters
----------
weights : numpy.ndarray
One binary64 weight per term.
Returns
-------
mlx.core.array
The combination.
"""
return self._evaluate(weights, exponential=0, sign=1)
[docs]
def exponential(self, weights: np.ndarray, sign: int = 1) -> mx.array:
"""Return ``exp(sign sum_t w_t T_t)`` to binary64 accuracy, rounded once.
Parameters
----------
weights : numpy.ndarray
One binary64 weight per term.
sign : int, optional
+1 for the propagator, -1 for its inverse.
Returns
-------
mlx.core.array
The complex64 propagator.
"""
return self._evaluate(weights, exponential=1, sign=sign)
_POPULATION = """
uint i = thread_position_in_grid.x;
if (i >= N) { return; }
float x = exposure[i];
float keep = constants[0];
float initial = constants[1];
// Kahan's expm1(-x) from the precise exponential and logarithm: a few ulp at every
// argument, where MLX's own GPU expm1 loses up to 1e-4 relative at small ones.
float u = metal::precise::exp(-x);
float m;
if (u == 1.0f) {
m = -x;
} else if (u - 1.0f == -1.0f) {
m = -1.0f;
} else {
m = (u - 1.0f) * (-x) / metal::precise::log(u);
}
fraction[i] = initial + keep * (-m);
survival[i] = keep * u;
"""
@cache
def _population_kernel() -> object:
"""Return the compiled population kernel.
Returns
-------
object
The ``mlx.core.fast.metal_kernel`` callable.
"""
return mx.fast.metal_kernel(
name="solprop_stable_population",
input_names=["exposure", "constants"],
output_names=["fraction", "survival"],
source=_POPULATION,
compile_options={"math_mode": "safe"},
)
[docs]
def population_parameters(preionisation: float) -> np.ndarray:
"""Return the population kernel's parameters, ``(1 - p, p)`` rounded to binary32.
Parameters
----------
preionisation : float
``p`` in ``[0, 1]``.
Returns
-------
numpy.ndarray
Float32, ``(2,)``.
"""
return np.array([1.0 - preionisation, preionisation], dtype=np.float32)
[docs]
def population(
exposure: mx.array, preionisation: float | mx.array
) -> tuple[mx.array, mx.array]:
"""Return the ionised fraction and the neutral survival in their stable forms.
The oracle's ``stable_population`` (SEP 33 section 4.4): ``f = p + (1 - p)
(-expm1(-X))`` and ``sigma = (1 - p) exp(-X)``. On the Metal device it is a custom
kernel in ``"safe"`` math mode with Kahan's ``expm1``; on MLX's CPU device, whose
``exp`` errs by up to 3e-6 relative at large arguments, NumPy's binary64 forms
rounded once.
Parameters
----------
exposure : mlx.core.array
``X``, float32, non-negative.
preionisation : float or mlx.core.array
``p`` in ``[0, 1]``, or on the Metal device the kernel's parameters
(:func:`population_parameters`) already on it, as a compiled stage passes them.
Returns
-------
tuple of mlx.core.array
The fraction and the survival, float32.
"""
if isinstance(preionisation, mx.array):
parameters = preionisation
else:
parameters = mx.array(population_parameters(preionisation))
if not on_gpu():
keep, initial = (float(value) for value in np.array(parameters))
# MLX's CPU exp errs by up to 3e-6 relative near exp(-80); on the CPU device
# the host's binary64 forms, rounded once, are the oracle's own.
wide = np.array(exposure, dtype=np.float64)
fraction = initial + keep * -np.expm1(-wide)
survival = keep * np.exp(-wide)
return mx.array(fraction.astype(np.float32)), mx.array(
survival.astype(np.float32)
)
flat = exposure.reshape(-1)
n = flat.size
fraction, survival = _population_kernel()( # pyright: ignore[reportCallIssue]
inputs=[flat, parameters],
template=[("N", n)],
grid=(n, 1, 1),
threadgroup=(min(256, n), 1, 1),
output_shapes=[(n,), (n,)],
output_dtypes=[mx.float32, mx.float32],
)
return fraction.reshape(exposure.shape), survival.reshape(exposure.shape)
_SCAN = """
// One threadgroup per row: the doubling scan of the single-precision oracle
// (solprop.translate.precision.prefix_sum), level by level, in place, every level
// the same additions in the same order, so the result is the oracle's exactly.
uint row = threadgroup_position_in_grid.x;
uint tid = thread_position_in_threadgroup.x;
uint width = threads_per_threadgroup.x;
device float* a = out + row * N;
device float* b = scratch + row * N;
for (uint i = tid; i < N; i += width) { a[i] = values[row * N + i]; }
threadgroup_barrier(metal::mem_flags::mem_device);
for (uint shift = 1; shift < N; shift *= 2) {
for (uint i = tid; i < N; i += width) {
b[i] = i >= shift ? a[i] + a[i - shift] : a[i];
}
threadgroup_barrier(metal::mem_flags::mem_device);
device float* swap = a; a = b; b = swap;
}
if (a != out + row * N) {
for (uint i = tid; i < N; i += width) { out[row * N + i] = a[i]; }
}
"""
_SCAN_THREADS = (512, 1024)
"""A row's threads in the scan: 512 when there are many rows, 1024 for a few.
On an Apple M5 Max (MLX 0.32) against the 256 before, chained scans: one row of 16383
took 62 us against 105, sixteen of 8191 43 against 65, and 1024 rows of 2047 114
against 116. Holding the row in threadgroup memory was slower at every shape.
"""
@cache
def _scan_kernel() -> object:
"""Return the compiled doubling-scan kernel.
Returns
-------
object
The ``mlx.core.fast.metal_kernel`` callable.
"""
return mx.fast.metal_kernel(
name="solprop_doubling_scan",
input_names=["values"],
output_names=["out", "scratch"],
source=_SCAN,
compile_options={"math_mode": "safe"},
)
[docs]
def prefix_scan(values: mx.array) -> mx.array:
"""Return the oracle's doubling scan along the last axis, in one launch.
Parameters
----------
values : mlx.core.array
Float32 increments.
Returns
-------
mlx.core.array
The running sums, the oracle's to the bit.
"""
n = values.shape[-1]
rows = values.size // n
# Enough threads for the row in whole SIMD groups, up to the shape's width.
width = min(_SCAN_THREADS[rows < 64], 32 * math.ceil(n / 32))
out, _ = _scan_kernel()( # pyright: ignore[reportCallIssue]
inputs=[values.reshape(rows, n)],
template=[("N", n)],
grid=(rows * width, 1, 1),
threadgroup=(width, 1, 1),
output_shapes=[(rows, n), (rows, n)],
output_dtypes=[mx.float32, mx.float32],
)
return out.reshape(values.shape)
_RATE = """
uint i = thread_position_in_grid.x;
if (i >= N) { return; }
float e = magnitude[i];
float log_min = constants[0];
float inverse_width = constants[1];
float field_min = constants[2];
float field_max = constants[3];
if (e < field_min) { rate[i] = 0.0f; return; }
if (e > field_max) { rate[i] = constants[4]; return; }
int k = int((metal::precise::log(e) - log_min) * inverse_width);
k = metal::clamp(k, 0, INTERVALS - 1);
// E / E_k as an exact two-product, so the local coordinate carries binary32's
// rounding relative to itself.
float p = e * inverse_edges[k];
float error = metal::fma(e, inverse_edges[k], -p) + e * inverse_tail[k];
float x = (p - 1.0f) + error;
float u = 1.0f + x;
float local = u == 1.0f ? x : metal::precise::log(u) * x / (u - 1.0f);
float c0 = coefficients[k];
float c1 = coefficients[INTERVALS + k];
float c2 = coefficients[2 * INTERVALS + k];
float c3 = coefficients[3 * INTERVALS + k];
float remainder = ((c0 * local + c1) * local + c2) * local;
float head = c3 + remainder;
float shifted = head - c3;
float tail = (c3 - (head - shifted)) + (remainder - shifted) + constant_tail[k];
rate[i] = metal::precise::exp(head) * (1.0f + tail);
"""
@cache
def _rate_kernel() -> object:
"""Return the compiled rate-table kernel.
Returns
-------
object
The ``mlx.core.fast.metal_kernel`` callable.
"""
return mx.fast.metal_kernel(
name="solprop_rate_table",
input_names=[
"magnitude",
"coefficients",
"constant_tail",
"inverse_edges",
"inverse_tail",
"constants",
],
output_names=["rate"],
source=_RATE,
compile_options={"math_mode": "safe"},
)
[docs]
def rate_kernel(
magnitude: mx.array,
coefficients: mx.array,
constant_tail: mx.array,
inverse_edges: mx.array,
inverse_tail: mx.array,
constants: mx.array,
) -> mx.array:
"""Evaluate a rate table's spline at field magnitudes in one launch.
Parameters
----------
magnitude : mlx.core.array
``|E|``, float32.
coefficients : mlx.core.array
The spline's ``(4, n_intervals)`` coefficients flattened row by row, the
constant term's head last.
constant_tail : mlx.core.array
The constant term's tail.
inverse_edges : mlx.core.array
``1 / E_k``'s head.
inverse_tail : mlx.core.array
``1 / E_k``'s tail.
constants : mlx.core.array
``ln E_0``, ``1 / d ln E``, ``E_min``, ``E_max``, ``W_max``, float32.
Returns
-------
mlx.core.array
The rate, float32, of `magnitude`'s shape.
"""
flat = magnitude.reshape(-1)
n = flat.size
(rate,) = _rate_kernel()( # pyright: ignore[reportCallIssue]
inputs=[
flat,
coefficients,
constant_tail,
inverse_edges,
inverse_tail,
constants,
],
template=[("N", n), ("INTERVALS", constant_tail.size)],
grid=(n, 1, 1),
threadgroup=(min(256, n), 1, 1),
output_shapes=[(n,)],
output_dtypes=[mx.float32],
)
return rate.reshape(magnitude.shape)
[docs]
def fft_bound_32(samples: int) -> float:
"""Return the established transform bound at binary32, ``512 eps32 ceil(log2 n)``.
The ladder's per-transform envelope (``tests/bounds.py``, ``fft_bound``) evaluated
at binary32's machine epsilon (SEP 33 7.1).
Parameters
----------
samples : int
The transform length.
Returns
-------
float
The bound, relative and normwise.
"""
return 512 * 2.0**-23 * math.ceil(math.log2(max(samples, 2)))
[docs]
def fft_probe(lengths: list[tuple[str, int]]) -> float:
"""Measure the device's transforms at a run's own kinds and lengths (SEP 33 6.3).
A Gaussian random signal of each length is transformed forward and back on the
device and compared with binary64: the forward transform normwise, and the round
trip, whose scale bias MLX's GPU transform carries, against the identity.
Parameters
----------
lengths : list of tuple
``(kind, n)``, ``kind`` ``"field"`` (real transforms) or ``"envelope"``.
Returns
-------
float
The largest relative error found.
Raises
------
BackendError
If a transform exceeds its bound at binary32.
"""
rng = np.random.default_rng(33)
worst = 0.0
for kind, n in lengths:
if kind == "field":
x = rng.standard_normal((2, n)).astype(np.float32)
forward = np.array(mx.fft.rfft(mx.array(x), axis=-1), dtype=np.complex128)
exact = np.fft.rfft(x.astype(np.float64), axis=-1)
back = np.array(
mx.fft.irfft(mx.fft.rfft(mx.array(x), axis=-1), n=n, axis=-1),
dtype=np.float64,
)
else:
x = (rng.standard_normal((2, n)) + 1j * rng.standard_normal((2, n))).astype(
np.complex64
)
forward = np.array(mx.fft.fft(mx.array(x), axis=-1), dtype=np.complex128)
exact = np.fft.fft(x.astype(np.complex128), axis=-1)
back = np.array(
mx.fft.ifft(mx.fft.fft(mx.array(x), axis=-1), axis=-1),
dtype=np.complex128,
)
errors = (
float(np.linalg.norm(forward - exact) / np.linalg.norm(exact)),
float(np.linalg.norm(back - x) / np.linalg.norm(x)),
)
worst = max(worst, *errors)
if max(errors) > fft_bound_32(n):
msg = (
f"MLX's {kind} transform of {n} samples errs by {max(errors):.2e}, "
f"beyond its binary32 bound {fft_bound_32(n):.1e}"
)
raise BackendError(msg)
return worst
[docs]
def gemm_probe(shapes: list[tuple[int, int, int]]) -> float:
"""Measure float32 matrix products of each dispatch family a run uses.
Gaussian random factors of each ``(m, k, n)`` are multiplied on the device and
compared with binary64 (SEP 33 6.3, 7.1).
Parameters
----------
shapes : list of tuple of int
The products' shapes.
Returns
-------
float
The largest normwise relative error.
Raises
------
BackendError
If any product exceeds :data:`GEMM_PROBE_BOUND`: the device computes in TF32
or lower, which SEP 25 forbids.
"""
rng = np.random.default_rng(33)
worst = 0.0
for m, k, n in shapes:
a = rng.standard_normal((m, k)).astype(np.float32)
b = rng.standard_normal((k, n)).astype(np.float32)
product = np.asarray(mx.array(a) @ mx.array(b), dtype=np.float64)
exact = a.astype(np.float64) @ b.astype(np.float64)
error = float(np.linalg.norm(product - exact) / np.linalg.norm(exact))
worst = max(worst, error)
if error > GEMM_PROBE_BOUND:
msg = (
f"MLX computed a {m}x{k} by {k}x{n} float32 product to {error:.2e}, "
f"beyond {GEMM_PROBE_BOUND:.1e}: its matrix products run in TF32. Set "
"MLX_ENABLE_TF32=0 before MLX is first used in the process (in the "
"shell, or before any import of mlx)"
)
raise BackendError(msg)
return worst