With a bias towards applications to scientific computing
Research Software & Analytics Group, University of Exeter
August 5th, 2026
Not the one in the file extension. Who should have written it: me, or the compiler?
jit=on, LLVM-based, since PG11); Spark whole-stage codegen.~/.nv/ComputeCache exists.torch.compile — Dynamo rewrites bytecode → Inductor → Triton. Ubiquitous in ML; it arrives with PyTorch whether or not you went looking for a JIT.This is not an exotic topic you can opt out of.
A @ B, α · (A @ B), β · C, and the sumA @ B and the fused result\psi_\sigma(x, y) = (1 - r^2)\, e^{-r^2 / 2}, \qquad r^2 = \frac{x^2 + y^2}{\sigma^2}
f_numpy is apparently the fastest possible NumPyf_numba — same source, compiled — buys only 1.16×: 16M points already amortised the interpreter, and this kernel is memory-bandwidth-boundarrayexpr node at a time, over the array notation as written — naming r2 forces it into memory: a 127 MB round tripf_numba_loops is optimal: hoisted, one allocation, r2 in a register — you gave up notation and bought the top of the memory hierarchyC.f. numba_loop.ipynb
def f_numpy(x, y, σ):
r2 = (x.reshape(-1, 1) ** 2 + y.reshape(1, -1) ** 2) / (σ * σ)
return (1.0 - r2) * np.exp(-0.5 * r2)
@jit("f8[:, ::1](f8[::1], f8[::1], f8)", nopython=True, nogil=True)
def f_numba_loops(x, y, σ):
res = np.empty((x.shape[0], y.shape[0]))
inv = 1.0 / (σ * σ)
for i in range(x.shape[0]):
xi = x[i] * x[i] * inv # invariant in j — hoisted out of it
for j in range(y.shape[0]):
r2 = xi + y[j] * y[j] * inv # a scalar, in a register
res[i, j] = (1.0 - r2) * np.exp(-0.5 * r2)
return res ms vs np peak allocs
f_numpy broadcasting, interp. 50.26 1.00× 4.0 —
f_numba same source, compiled 43.51 1.16× 2.0 2
f_numba_inlined rewritten for fuser 46.67 1.08× 1.0 3
f_numba_loops loops, hoisted 39.67 1.27× 1.0 1
objmode: jit always compiles; it does not always helpjit will always work, but it will not always be fasterobjmode is a hole punched through the type system, drawn by hand — it lets the rest of the function keep nopython=Truebrentq call was the runtime, so compiling the loop around it changes nothingC.f. numba_objmode.ipynb
def kepler_equation(E, M, e):
return E - e * np.sin(E) - M
@jit
def kepler(M, e):
E = np.empty_like(M)
for i in range(M.shape[0]):
Mi = M[i]
E[i] = optimize.brentq(kepler_equation, Mi - 1.1, Mi + 1.1, args=(Mi, e))
return E
def kepler_objmode(M, e):
E = np.empty_like(M)
for i in range(M.shape[0]):
Mi = M[i]
with objmode(Ei="float64"):
Ei = optimize.brentq(kepler_equation, Mi - 1.1, Mi + 1.1, args=(Mi, e))
E[i] = Ei
return Epython: 12.63 ms
forceobj: 12.79 ms
objmode: 13.26 ms
Tracers) and records every array op that fires — Python only generates the recordingShapedArray), not values — new shape, new trace, new compile16×12×32 it loses to NumPy — that is per-call dispatch (pytree flatten → cache lookup → PjRt launch), not arithmetic. At 1024³ all three land within a factor of two, and which wins is run-dependentC.f. intro_jax.ipynb
16×12×32 1024³
numpy 1.86 μs 4.56 ms
numba 744 ns 2.49 ms
jax 6.17 μs ←worst 2.79 ms
XLA still finds the same fusion Numba’s arrayexpr optimiser found, from a completely different starting representation — ROOT %multiply_add_fusion. Fusion is not what is being measured here.
r2 before fusion ever runs — so XLA fuses straight through it, and picks fusion boundaries by cost rather than by which Python names appeared
C.f. jax_loop.ipynb
Two fusions for the whole kernel — one scalar, one grid:
%multiply_divide_fusion = f64[] fusion(...)
%exponential_multiply_fusion = f64[3600,4400] fusion(...)
Both squares, both broadcasts, the add, the divide, 1 - r2, -0.5 * r2, exp, the final multiply — every one of them, in a single pass. r2 never exists as an array.
w_tilde_curvature_interferometer_from in Numba vs. JAX@jit: every function is compiled to machine code, specialised on its argument types, the first time it is called with a new combinationA * B and mul! become an opaque BLAS call — the Julia compiler, including all its lowering, cannot see past itC.f. intro_julia.ipynb
impl time memory allocs
mul_naive! 633.376 ns 16.34 KiB 12
mul_fused! 277.267 ns 4.09 KiB 3
mul_native! 103.032 ns 0 bytes 0
Quantity and a Float64 (both <: Number) share identical lowered IR, and near-identical typed IR that compiles away to nothingC.f. julia_unitful.ipynb
dense: 3.614 ms
sparse: 7.946 μs 455× at N = 4000, 450 TB at production size
NTuple recursion), orC.f. julia_as_metaprogramming_lang.ipynb
implementation median time memory speed-up
runtime (Vector arg) 2204.6 μs 8000072 B 1.00x
recursive (NTuple dispatch) 414.6 μs 8000072 B 5.32x
macro (@horner) 419.4 μs 8000072 B 5.26x
The claim, in one sentence: open generic functions + multiple dispatch + parametric types + aggressive specialisation + accessible compiler infrastructure + metaprogramming makes it possible to build high-performance, composable abstraction layers over heterogeneous hardware.
@horner is a compiler pass with a name; evalpoly is that pass in the standard library.ustrip change the exponent.GPUCompiler.jl: CUDA, AMDGPU, oneAPI, Metal| Numba | JAX | Julia | |
|---|---|---|---|
| Backend | LLVM | XLA | LLVM |
| Targets | CPU; numba-cuda is a second implementation |
CPU, GPU, TPU — one source | CPU; all 4 GPU vendors, one source, via packages |
| Paradigm | C-like: loops, mutation | functional; purity enforced | a full language |
| You write | a subset of Python + NumPy | a duck-typed NumPy + SciPy | Julia |
| Closes the leak by | nothing — you hand-narrow | restricting the language | programming the compiler |
| Recompiles on | types | types and shapes | types |
| Secret weapon | smallest diff from working NumPy | jax.grad — AD as a compiler pass |
dispatch + metaprogramming |
What you can say, what is fast, and what is idiomatic are three different sets.
f_numba_loops; the dots in mul_fused!).-march=native for free on every node.NUMBA_CACHE_DIR → which storage tier?); compiling on the login node vs. the compute node.cudnn.benchmark=True, cuBLAS algorithm selection).dynamic-codesigning entitlement, historically granted only to WebKit’s JS engine.Why it worked — both missing pieces were already there
…and the tests still were not the specification
unsafe, pure functions, narrow typed interfaces) — all of which were always about making behaviour checkable without reading the code. That simply was never their main job before.Discussion
Slides, notebooks and the long-form write-up: blog.kolen.dev