A Primer on Just-In-Time (JIT) Compilation

With a bias towards applications to scientific computing

Dr. Kolen Cheung, Research Software Engineer

Research Software & Analytics Group, University of Exeter

August 5th, 2026

Before we start

Which language am I actually writing?

Not the one in the file extension. Who should have written it: me, or the compiler?

You already run a dozen JITs

  • Browsers. Practically every JS engine in use is a JIT.
  • Databases. PostgreSQL (jit=on, LLVM-based, since PG11); Spark whole-stage codegen.
  • CUDA itself. PTX → SASS at load time — which is why the first kernel launch is slow, and why ~/.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.
  • The Linux kernel. eBPF is JIT-compiled.
  • Your GPU driver, every time it compiles a shader.
  • Regex. PCRE2-JIT. CPython itself, since 3.13 (PEP 744, experimental, opt-in).

This is not an exotic topic you can opt out of.

Introducing Numba

A decorator, and a tower underneath it

  • Numba-jit can accelerate even a simple NumPy function that should already be quite fast — 1.87 μs → 791 ns here
  • The win is not “compiled beats interpreted”: it is operator fusion
    • memory allocation: 4 in NumPy — A @ B, α · (A @ B), β · C, and the sum
    • 2 in Numba — A @ B and the fused result
  • tower of abstractions: CPython bytecode → Numba IR → typed Numba IR → LLVM IR → assembly
  • Numba compiles from bytecode, not source — the decorator is the seam between Python-as-host-language and Python-as-mini-language

C.f. intro_numba.ipynb

def mul_numpy(C, A, B, α=True, β=False):
    return α * (A @ B) + β * C

@jit("f8[:, ::1](f8[:, ::1], f8[:, ::1], f8[:, ::1], f8, f8)", nopython=True, nogil=True)
def mul_numba(C, A, B, α=True, β=False):
    return α * (A @ B) + β * C
mul_numpy   1.87 μs
mul_numba    791 ns      2.4×

new arrays created:  numpy 4,  numba 2

Why write the loop? Fusion stops where you named a value

\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 NumPy
  • f_numba — same source, compiled — buys only 1.16×: 16M points already amortised the interpreter, and this kernel is memory-bandwidth-bound
  • Numba’s fuser works one arrayexpr node at a time, over the array notation as written — naming r2 forces it into memory: a 127 MB round trip
  • f_numba_loops is optimal: hoisted, one allocation, r2 in a register — you gave up notation and bought the top of the memory hierarchy

C.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 help

  • jit will always work, but it will not always be faster
  • When you stay inside a subset of Python (i.e. Numba’s dialect), you compile to machine code (language spec → speculation/specialisation)
  • When you don’t (fall off the guard) → interpreted: always right, but slow
  • objmode is a hole punched through the type system, drawn by hand — it lets the rest of the function keep nopython=True
  • It buys correctness at the boundary, not speed: here the brentq call was the runtime, so compiling the loop around it changes nothing

C.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 E
python:   12.63 ms
forceobj: 12.79 ms
objmode:  13.26 ms

Introducing JAX

Tracing, not typing

  • Numba type-specialises bytecode. JAX traces: it runs the function once on abstract stand-ins (Tracers) and records every array op that fires — Python only generates the recording
  • tower of abstractions: tracing → jaxpr → StableHLO → fused HLO → LLVM IR (CPU) → assembly
  • Caches on abstract shape and dtype (a ShapedArray), not values — new shape, new trace, new compile
  • JAX is not automatically faster. At 16×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-dependent

C.f. intro_jax.ipynb

@jax.jit
def mul(C, A, B, α=True, β=False):
    return α * (A @ B) + β * C
              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.

Write the maths, get the loop

  • Tracing erases the name r2 before fusion ever runs — so XLA fuses straight through it, and picks fusion boundaries by cost rather than by which Python names appeared
    • the exact thing that cost Numba a 127 MB round trip
  • Pure functional, fixed shape, static size: a tighter spec buys more aggressive optimisation
  • No low-level control (c.f. Numba’s fastest version) — you cannot hand-write the loop
  • For simple algorithms, just write the maths: four lines get what Numba needed ten hand-written ones for
  • The cost moved, it did not vanish: retracing per shape, and no data-dependent control flow or mutation inside the compiled region

C.f. jax_loop.ipynb

@jax.jit
def f_jax(x, y, σ):
    r2 = (x.reshape(-1, 1) ** 2 + y.reshape(1, -1) ** 2) /* σ)
    return (1.0 - r2) * jnp.exp(-0.5 * r2)

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.

Numba vs. JAX on a real kernel

  • Same function, two implementations, from the Numba vs. JAX case study: w_tilde_curvature_interferometer_from in Numba vs. JAX
  • What that study actually concluded:
    • “Porting Numba to Numba” was faster in many cases — the first draft was not the opponent
    • The only fair fight is single-CPU-core
    • Which algorithm wins flips with input size, even across Numba and JAX
    • ⇒ keep all implementations → profile → pick the best, per science case per system
  • The general form, and the one thing to take from this slide: the language comparison you think you are running is usually a comparison of how hard you tried. Benchmark your own best effort, not your first draft — especially when you are about to conclude something about the language.

Introducing Julia

No decorator — the whole language works this way

  • Julia has no @jit: every function is compiled to machine code, specialised on its argument types, the first time it is called with a new combination
  • In Julia, broadcast is JAX’s fusing equivalent — lazy, better memory access pattern — but it is syntax you write, not an inference the compiler makes
  • tower of abstractions: lowered IR → typed IR → LLVM IR → assembly
    • and the fusion decision happens at lowering, before any type is known
  • See how A * B and mul! become an opaque BLAS call — the Julia compiler, including all its lowering, cannot see past it

C.f. intro_julia.ipynb

function mul_naive!(C, A, B, α = true, β = false)
    C .= α * (A * B) + β * C
    return C
end

function mul_fused!(C, A, B, α = true, β = false)
    C .= α .* (A * B) .+ β .* C
    return C
end

function mul_native!(C, A, B, α = true, β = false)
    return mul!(C, A, B, α, β)
end
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

Julia’s type system & multiple dispatch

  • Type system illustrated by Unitful:
    • the unit is encoded in the type
    • a Unitful Quantity and a Float64 (both <: Number) share identical lowered IR, and near-identical typed IR that compiles away to nothing
    • so units cost once, at compile time: 1.443 ns vs 1.443 ns per call
  • Also many array types: dense, sparse, symmetric, tridiagonal… Joy for a mathematician!
  • Julia’s type hierarchy makes unrelated libraries compose automagically — but when it stops working, it feels like magic too →

C.f. julia_unitful.ipynb

kinetic_energy(m, v) = 0.5 * m * v^2

kinetic_energy(2.0, 3.0)            # 9.0            1.443 ns
kinetic_energy(2.0u"kg", 3.0u"m/s") # 9.0 kg m^2 s^-2  1.443 ns
julia> typeof(ustrip.(u"MHz", Su))        # sparse + Unitful
SparseMatrixCSC{Float64, Int64}           # fine

julia> typeof(ustrip.(u"MHz", Symmetric(Su)))
Matrix{Float64}                           # densified!
dense:  3.614 ms
sparse: 7.946 μs      455× at N = 4000, 450 TB at production size
  • magic: methods for those types are already defined
  • broken magic → missing methods; nobody foresaw this combination
  • it changed the exponent, and nothing in the program said so
  • but the fix is always local: write the one method

Julia is its own metaprogramming language

  • In Numba and JAX, Python is the host language and the jitted subset is an embedded one. In Julia there is no such boundary: the language is its own metaprogramming language
  • All code is data, and you can manipulate it
  • Two different points in the pipeline reach the same machine code:
    • type inference resolving a chain of dispatches (NTuple recursion), or
    • a macro rewriting the source before names are even parsed into scopes
  • Both unroll the loop over coefficients away; both land ~5.3× faster than the runtime version

C.f. julia_as_metaprogramming_lang.ipynb

# baseline: the inner loop over `coeffs` survives to runtime
for c in coeffs
    acc = acc * x[i] + c
end
# 1. unrolled by type inference: Tuple{} vs Tuple are different types
horner_step(acc, xi, coeffs::Tuple{}) = acc
horner_step(acc, xi, coeffs::Tuple) =
    horner_step(acc * xi + first(coeffs), xi, Base.tail(coeffs))

out[i] = horner_step(0.0, x[i], coeffs)  # coeffs::NTuple{N,Float64}
# 2. unrolled by the macro, before lowering ever runs
macro horner(x, coeffs...)
    acc = :(0.0)
    for c in coeffs
        acc = :($acc * $x + $c)
    end
    return esc(acc)
end

@horner(x, 1.0, -3.0, 2.5)  # ⇒  ((0.0 * x + 1.0) * x + -3.0) * x + 2.5
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

Julia’s superpower — and its price

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.

  • Multiple dispatch provides composability; specialisation provides performance; metaprogramming provides syntax; compiler extensibility provides new targets.
  • Everywhere else in this talk, a leaky abstraction leaves one move: narrow by hand, at the source, once per call site. Julia lets you put the transformation where it belongs — @horner is a compiler pass with a name; evalpoly is that pass in the standard library.
  • The price: the seam between the language and the compiler is gone. You cannot tell by reading which one you are looking at — the same missing seal that let ustrip change the exponent.

Stepping back: the tower of abstractions

Numba vs. JAX vs. Julia: a high-level comparison

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.

The tower of abstractions

  • Source code → AST → IR (usually several) → assembly → machine code.
  • Note the plurality: Numba alone goes CPython bytecode → Numba IR → typed Numba IR → LLVM IR → MachineIR → x86 — six representations, none of them redundant.
  • Each layer is a self-contained abstraction. Leaky, but self-contained.
  • Each layer therefore gives you:
    • separation of concerns — the layer below doesn’t care how you got here
    • a place to verify
    • a place to optimise
    • a place to express intent — the high level carries project/author intent and correctness, in the way a mathematical proof is structured: layers built on layers
  • Hence: different languages, compilers, and libraries are tradeoffs between these abstractions, not merely differences in syntax or speed.

Lowering through the tower

  • Lowering is lossy in representation and faithful in semantics — semantic narrowing.
    • source semantics does not specify one behaviour; it specifies a set of permitted ones
    • compilation narrows that set. It does not step outside it.
  • Lossiness → room for optimisation. If lowering had to be reversible you could not constant-fold, fuse, or vectorise at all.
  • Different layers of IR → different rooms for optimisation. Some classes are only recognisable further up: in x86 you cannot see that you are in the “this is a matmul” class.
  • Each layer has a specification, and that is what makes travelling down the tower possible.

Lowering through the tower, in the case of JIT

  • With JIT you get to see the data (narrowing), so the compiler can do more aggressive partial evaluation / specialisation than an AOT compiler can.
  • All the JIT compilers we have seen have different specifications — and a narrower spec means more room to optimise (e.g. JAX).
  • Some specs are leakier than others. In Numba and Julia you often have to write the host language differently to convey the same intent in a way the compiler can act on (f_numba_loops; the dots in mul_fused!).
  • Two responses to a leak: patch it at the source, or add a rung. One is a rewrite you repeat; the other is a pass you name once — and only Julia makes the second one routine.
  • Ideally the separation of concerns between layers would be total. It isn’t, and that gap is where all the practical difficulty lives.

What “later” knows that “earlier” can’t

  • Concrete types — no boxing, no dispatch, no polymorphism to hedge against.
  • Concrete shapes — loop bounds become constants; tiling and unrolling become decidable.
  • Runtime values — things that are constant in this run but not at build time.
  • Branch frequencies — which path is actually hot.
  • The machine it is actually running on. An AOT binary shipped to a heterogeneous cluster targets the lowest common denominator. A JIT gets -march=native for free on every node.

What “later” costs

  • Warmup. You pay compile time inside the user’s wall clock, not yours.
  • Memory and shipping weight. The compiler is now a runtime dependency.
  • Budget. A JIT cannot afford the expensive passes an AOT compiler can, because the user is waiting.
  • Opacity. There is no artefact on disk to inspect, archive, or hand to someone else.
  • HPC papercuts, specifically: 512 ranks JIT-compiling the same kernel at once against a shared filesystem; where the cache lives (NUMBA_CACHE_DIR → which storage tier?); compiling on the login node vs. the compute node.

Trust

Reproducibility and provenance

  • Which binary actually ran? You cannot archive it, because it never existed as a file.
  • FMA contraction and fast-math change results; JIT decisions change which you get.
  • Autotuning is nondeterministic by construction (cudnn.benchmark=True, cuBLAS algorithm selection).
  • Different node → different ISA → different vectorisation → different reduction order → different last bits.

Signing and JIT are in irreducible tension

  • JIT’d code cannot be signed, because it does not exist until runtime.
  • iOS enforces code signing and W^X, so JIT requires a special dispensation: the dynamic-codesigning entitlement, historically granted only to WebKit’s JS engine.
  • WebKit’s “bulletproof JIT” is the mitigation: the JIT region gets a second, writable mapping at a randomised secret address, so the executable mapping is never writable. Apple Silicon adds per-thread W^X toggling.
  • Generalise: runtime code generation trades verifiability for performance.

Digression: LLM as a JIT compiler

The analogy: why an LLM is good

  • Now you can see why an LLM is good:
    • the ultimate JIT — it compiles your prompt into code;
    • the ultimate metaprogramming language — you specify intent and abstraction, not mechanism.
  • And it is genuinely the same shape: late binding, specialisation from a high-level spec, caching, and a latency/quality dial that behaves exactly like compile-time vs. runtime.

Where the analogy breaks

  • A compiler is semantics-preserving: there is a source program with a specified meaning, and correctness means the output stays inside the set of behaviours that meaning permits.
  • An LLM has no such referent. A prompt is not a specification — it does not define a set of permitted behaviours, so “preserving” it is not even a well-formed claim.
  • In the vocabulary from the last section, precisely: an LLM performs enormous narrowing — from a prompt to one program — with nothing licensing the narrowing (no spec) and nothing checking it (no guard).
  • Compare a JIT: also unsound narrowing, but bounded by a spec above and caught by a guard below.
  • An LLM is not a compiler. It is the first half of one.

The tower loses its direction

  • A compiler only goes down, because lowering destroys what it would need to climb.
  • An LLM has no preferred direction, because it is not translating — it is reconstructing from a learned prior. It moves:
    • down — prompt → code. The classic framing, and the hardest case.
    • sideways — Zig → Rust, Python → C++, COBOL → Java, JS → TS.
    • up — code → documentation, code → tests, code → a specification. And decompilation, which people are now genuinely doing.
  • It can climb precisely because it hallucinates. Going up requires inventing information that was destroyed. A decompiler cannot invent a good variable name; a model can guess one — and often guesses right, because the prior encodes what humans usually call things.
  • One mechanism, two verdicts: what makes it dangerous going down is what makes it useful going up.

Example: Bun, transpiling from Zig → Rust

  • 535,496 lines of Zig across 1,448 files → Rust in 11 days (3–14 May 2026). https://bun.com/blog/bun-in-rust
  • 64 agents in parallel (4 worktrees × 16), ~50 workflows; 6,502 commits; peak 695 commits/hour. 90%+ automated, one engineer supervising.
  • ~$165,000 of API spend, against an estimated 3 engineers × 1 year with no features shipping.
  • 128 bugs fixed, ~2–5% faster, several MB smaller binary.

Why it worked — both missing pieces were already there

  • the source program is the specification: defined semantics, so “correct” means something precise
  • the test suite is the guard — and it is written in TypeScript, so it does not depend on the implementation language: 60,624 tests, 1,386,826 assertions, zero skipped

…and the tests still were not the specification

  • 19 known semantic regressions got through all of them — each one two languages disagreeing about the spec
  • A test suite is a sample, not a specification

Why an LLM is bad

  • The input space is infinite. No input is wrong. No linting, no static analysis, no abstraction.
  • The output space is infinite. All outputs are possible — and under agentic use, so is everything that happens along the way to producing one.
    • You cannot trust the output (cf. bulletproof JIT).
    • You cannot secure the machine (cf. the lethal trifecta: private data + untrusted content + exfiltration channel).
  • A compiler’s blast radius is a process. An agent’s blast radius is your network.
    • Concrete: OpenAI ran a cybersecurity test against an unreleased model with guardrails off. Rather than solve the test, the model broke out of OpenAI’s sandbox, then found exploits to break into Hugging Face — in order to steal the answers and cheat on the test. https://simonwillison.net/2026/Jul/22/openai-cyberattack/
  • The apparent remedy is a human in the loop. But: would you put a human between the compiler and its output, and ask them to check the machine code?

Trusting and verifying

  • Once you can make a problem verifiable — at training time or at run time — an agent starts behaving like a compiler. Verifiability converts synthesis back into translation.
  • Note the inversion, and it is the whole problem: with a compiler you read the top and trust the bottom. With an agent we currently read the bottom — the largest, least reviewable representation in the tower — and trust nothing.
  • So: reduce entropy. Shrink the space of things that could come out. TDD, unit tests, property-based testing, types, contracts, enforced abstractions (Rust without 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.
  • StrongDM’s Software Factory: specs + scenarios drive agents that write code, run harnesses, and converge without human review. https://simonwillison.net/2026/Feb/7/software-factory/
    • specs and scenarios supply the missing specification; the harness supplies the missing guard — and it inherits the limit: you have validated only what the harness can see

Wrapping up

Takeaways

  • A JIT is a bet that you will know more later — types, shapes, values, profiles, and the actual machine. Sometimes the bet does not pay: JAX lost to NumPy at 16×12×32.
  • AOT vs. JIT is a dial, not a binary. The question is always when do you decide?
  • Lowering is lossy in representation and faithful in semantics. The lossiness is where all the optimisation lives. The useful question about a JIT is never “is it fast?” but “which language am I actually writing?” — and every one of them is an answer to who does the narrowing: Numba says you, JAX says the compiler (because you handed it a narrower language), Julia says you and hands you the compiler to automate yourself.
  • Every performance result here but one came from removing memory traffic, not adding FLOPs. And the language comparison you think you are running is usually a comparison of how hard you tried.
  • Runtime code generation trades verifiability for performance. Always.
  • An LLM is the extreme answer to that same question — it narrows everything, from a prompt: maximal narrowing, no specification licensing it, no guard checking it. The work is putting a layer back — and it is work we already know how to do.
  • You never delete human verification. You relocate it somewhere smaller: assembly → source → spec → a harness you can audit.

Discussion

Discussion

Slides, notebooks and the long-form write-up: blog.kolen.dev