Suggestion: auto CPU/GPU dispatch heuristic + GPU-side extended-precision ("decimal") arithmetic

[bahaehmimdi]

Two suggestions from a benchmarking pass I did on mlx.core CPU/GPU
dispatch, both aimed at making it easier to get good performance (and
more numeric precision) out of MLX without the caller having to hand-manage
device placement.

Setup: Apple Silicon Mac, 34GB unified memory, MLX 0.31.2.

1. Built-in adaptive CPU/GPU dispatch based on residency + fusion

Right now getting a real GPU speedup requires the caller to already know
two non-obvious things: keep data resident on the GPU stream across
calls (don't convert per-call), and fuse multiple ops into one
@mx.compile kernel. Get either wrong and GPU can actually lose to
plain NumPy.

What I measured (payload: sqrt(a*b + a/b) - (a-b)**2/(a+b+1), 5
ops/element, timed with mx.eval() forced before stopping the clock):

• Single GPU op, converting NumPy→mx.array every call: ~1.2x over
  NumPy at 16000x16000, loses at smaller sizes.
• Same op fused into one compiled kernel, still converting per call:
  ~4.3-4.5x at 4000x8000.
• Same fused kernel, data already resident on GPU: 16.7-23.7x at
  the same sizes — this was the single biggest lever, bigger than
  fusion or size scaling combined. Kept growing with size, 12.9x →
  46.2x from 4000x4000 up to 28000x28000 (784M elements).
• The same resident+fused pattern holds for mx.where-based
  conditional logic too (22-29.5x), not just arithmetic.
• Below roughly 5×10^5 elements, CPU (numba.njit, which itself beats
  plain NumPy by 1.5-4.6x at small scale) wins outright — GPU
  dispatch/conversion overhead isn't worth paying.

Suggestion: could mlx.core expose (or default to, where safe) a
dispatch heuristic that keeps ops on whichever stream currently holds
their operands, and warns or auto-recommends mx.compile when it
detects an unfused chain of ops repeatedly hitting the same tensors?
Even a documented rule of thumb in the perf guide (data residency +
fusion matter far more than raw op count) would help — right now this
is easy to get backwards and quietly eat all the GPU gain.

2. GPU-side extended-precision ("decimal") arithmetic

Apple GPUs don't do hardware double precision, so mlx.core has no
float64 on the GPU stream today (only CPU, via stream=mx.cpu). For
work that needs more than float32 precision but doesn't want to fall
back to CPU entirely, I prototyped a digit-decomposition approach
that stays on GPU:

• Represent a number as an array of digits and reconstruct it via a
  single fused kernel: value = Σ digit_i × 10^-i (parallel across all
  digits in one @mx.compile call).
• Digit-by-digit addition with explicit carry/overflow detection via
  mx.where(sum > 9, ...), so multi-digit addition can run as one
  vectorized GPU pass instead of a serial float64 add.
• A fractional-part filter (x - floor(x), masked by
  mx.where(mask, ...)) for isolating and processing the sub-integer
  part separately from the integer part, which is where most precision
  loss in float32 actually happens.

This is essentially software-emulated extended precision (comparable in
spirit to "double-single"/compensated-arithmetic techniques used
elsewhere in GPGPU code), built entirely from ops MLX already has —
mx.where, mx.floor, mx.compile. It got real precision back for
addition/reconstruction workloads without ever leaving the GPU stream.

Suggestion: would there be interest in mlx.core shipping something
like this as a first-class extended-precision dtype/helper (e.g. a
mx.decimal array type backed by a digit array under the hood), rather
than every user who needs float32-beyond precision on GPU reinventing
digit-decomposition themselves? Happy to share the prototype code if
useful as a starting point.

---

Benchmark details (rerun just now, same machine, MLX 0.31.2)

A. Single op (a * b), converting NumPy→mx.array every call:

Size	CPU	GPU	Winner
4000x4000	7.94 ms	9.18 ms	CPU x1.16
8000x8000	35.66 ms	30.11 ms	GPU x1.18
12000x12000	80.85 ms	63.63 ms	GPU x1.27
16000x16000	124.26 ms	141.93 ms	CPU x1.14

B. 5-op fused formula (sqrt(a*b + a/b) - (a-b)**2/(a+b+1)), still converting every call:

Size	CPU	GPU	Winner
4000x4000	49.49 ms	6.99 ms	GPU x7.08
8000x8000	221.43 ms	22.47 ms	GPU x9.86

C. Same fused formula, data already resident on GPU (no per-call conversion):

Size	CPU	GPU (resident)	Winner
4000x4000	50.12 ms	3.79 ms	GPU x13.2
8000x8000	207.26 ms	9.14 ms	GPU x22.7
12000x12000	577.9 ms	19.8 ms	GPU x29.2
20000x20000	1459.9 ms	54.1 ms	GPU x27.0

D. Same resident+fused pattern applied to mx.where-based conditional logic (diff = a>b ? a-b : (b-a)*2, then clamp diff>50 -> 50):

Size	CPU	GPU (resident)	Winner
4000x4000	63.34 ms	2.21 ms	GPU x28.6
12000x12000	605.94 ms	19.39 ms	GPU x31.3
20000x20000	1648.98 ms	55.65 ms	GPU x29.6

E. Task-level hybrid (8 independent 3000x3000 formula-evaluations, half dispatched to CPU + half to GPU concurrently via ThreadPoolExecutor, vs. running the whole batch on one backend):

Strategy	Time
CPU only	264.47 ms
GPU only	50.88 ms
Hybrid (parallel CPU+GPU)	131.84 ms

GPU-only wins outright — splitting work across both backends is slower than just using the faster one, because coordination overhead isn't amortized once GPU alone is already far ahead.

F. Small-data CPU comparison (same formula, numba.njit(cache=True) vs plain NumPy):

Size	NumPy	Numba
scalar	0.296 µs	0.168 µs
50 elements	2.915 µs	0.654 µs
5,000 elements	15.219 µs	9.625 µs

Code

Suggestion 1 — the adaptive dispatcher (routes on element count; small → Numba, large → GPU-resident fused kernel):

import numpy as np
import numba
import mlx.core as mx

@numba.njit(cache=True)
def _numba_scalar(a, b):
    return np.sqrt(a * b + a / b) - (a - b) ** 2 / (a + b + 1)

@numba.njit(cache=True)
def _numba_array(a, b):
    out = np.empty_like(a)
    for i in range(a.shape[0]):
        out[i] = np.sqrt(a[i] * b[i] + a[i] / b[i]) - (a[i] - b[i]) ** 2 / (a[i] + b[i] + 1)
    return out

@mx.compile
def _mlx_kernel(a, b):
    return mx.sqrt(a * b + a / b) - (a - b) ** 2 / (a + b + 1)

SEUIL_GPU = 500_000  # element-count threshold

def calcul_adaptatif(a, b):
    if np.isscalar(a) and np.isscalar(b):
        return "numba-scalar", _numba_scalar(float(a), float(b))

    n = np.asarray(a).size if not isinstance(a, mx.array) else a.size
    if n < SEUIL_GPU:
        a_np = np.asarray(a, dtype=np.float64)
        b_np = np.asarray(b, dtype=np.float64)
        return "numba-array", _numba_array(a_np.ravel(), b_np.ravel()).reshape(a_np.shape)

    ax = a if isinstance(a, mx.array) else mx.array(a, dtype=mx.float32)
    bx = b if isinstance(b, mx.array) else mx.array(b, dtype=mx.float32)
    r = _mlx_kernel(ax, bx)
    mx.eval(r)
    return "mlx-gpu", r

Suggestion 2 — GPU-side digit-decomposition ("decimal") arithmetic prototype:

import mlx.core as mx

# reconstruct a number from its digits: value = sum(digit_i * 10^-i),
# all products computed in parallel in one fused kernel
@mx.compile
def _kernel_decimal(chiffres, facteurs):
    return mx.sum(chiffres * facteurs)

def mlx_decimal_from_digits(chiffres):
    x = mx.array(chiffres, dtype=mx.float32)
    n = x.shape[0]
    facteurs = 0.1 ** mx.arange(n, dtype=mx.float32)
    mx.eval(facteurs)
    total = _kernel_decimal(x, facteurs)
    mx.eval(total)
    return total  # e.g. [3,1,4] -> 3.14

# digit-by-digit addition with carry/overflow detection, vectorized
@mx.compile
def _kernel_retenue(a, b):
    somme = a + b
    retenue = mx.where(somme > 9, mx.array(1, dtype=a.dtype), mx.array(0, dtype=a.dtype))
    exces = mx.where(somme > 9, somme - 10, somme)
    return somme, retenue, exces

def mlx_addition_avec_retenue(a_vals, b_vals):
    a = mx.array(a_vals, dtype=mx.int32)
    b = mx.array(b_vals, dtype=mx.int32)
    somme, retenue, exces = _kernel_retenue(a, b)
    mx.eval(somme, retenue, exces)
    return somme, retenue, exces  # carry propagation across digits stays CPU-side (sequential dependency)

[bahaehmimdi]

Follow-up: I built a small from-scratch PyTorch-API-compatible layer on top of mlx.core (torch-mlx — import torch backed by MLX instead of ATen), and used it as a testbed to actually measure the residency/fusion findings above against real PyTorch, not just synthetic benchmarks.

Speed comparison, with torch-mlx's own MLX-backed timings shown explicitly (not just ratios) alongside real PyTorch CPU and MPS/Apple GPU, median of 8 iterations, full sync forced before timing stops on every leg:

Operation	torch-mlx (MLX)	real PyTorch CPU	real PyTorch MPS	MLX vs CPU	MLX vs MPS
matmul (2048x2048 @ 2048x2048)	4.35 ms	10.86 ms	4.27 ms	2.5x faster	tied (0.98x)
nn.Linear(1024,1024), batch 256, fwd+bwd	1.49 ms	1.60 ms	1.16 ms	tied (1.08x)	0.78x (MPS faster)
nn.Conv2d(32,64,k=3), (16,32,64,64) input, fwd+bwd	4.55 ms	54.94 ms	11.14 ms	12.1x faster	2.45x faster
nn.LSTM(256,512,layers=2), (32,50,256) input, fwd+bwd	65.33 ms	127.52 ms	22.17 ms	1.95x faster	0.34x (MPS ~3x faster)
nn.MultiheadAttention(256,heads=8), (16,128,256) input, fwd+bwd	3.75 ms	10.85 ms	3.78 ms	2.9x faster	tied (1.01x)
3-layer MLP train step (+Adam)	1.40 ms	1.45 ms	2.07 ms	tied (1.03x)	1.47x faster
nn.BatchNorm2d(128), (32,128,32,32) input, fwd+bwd	5.27 ms	5.09 ms	1.57 ms	tied (0.97x)	0.30x (MPS ~3.3x faster)

All computation in the "torch-mlx (MLX)" column runs on mlx.core — that's the actual MLX number being compared, not derived. Reported honestly including the losses: LSTM and BatchNorm2d lose clearly to MPS, because real PyTorch has hand-fused MPS kernels for both that this project's ops-composed-from-MLX-primitives implementations don't match yet. Conv2d and the full training step win clearly, even against Apple's own official PyTorch GPU backend.

The wins are consistent with the residency+fusion finding above: this project's core never calls mx.eval() per-op anywhere in its forward/backward/optimizer path, so MLX's lazy graph accumulates an entire training step naturally before materializing — which turns out to matter a lot (a separate, unrelated project, LucasSte/MlxTorch, a real fork of PyTorch's C++ ATen source with MLX as a backend, had to add an experimental deferred-execution branch just to match MPS on a training loop, for exactly this reason).

Benchmark code, methodology, and full results: https://github.com/bahaehmimdi/torch-mlx/tree/speed-benchmarks (see benchmarks/PERFORMANCE.md)