Numba occupies a unique spot in the Python performance world: it is a compiler that ships as a pip package. You put a decorator on a plain Python function, call the function once, and by the time that call returns, your Python bytecode has been type-inferred, translated into LLVM IR, optimized, and emitted as machine code for your exact CPU. Every later call skips Python entirely and jumps straight into the native function pointer. There is no separate build step, no C toolchain to drive, and no Ahead-of-Time ceremony — the compiler is embedded in the library and runs lazily, at import time of the call, not of the module.
In this source tour we read the numba/numba repository the way its authors layered it. The story runs through five territories: the decorator entry points (numba/core/decorators.py, numba/core/dispatcher.py), the interpreter that lifts bytecode into Numba’s own IR (numba/core/interpreter.py), the type inference engine (numba/core/typing/, numba/core/typeinfer.py), the lowering and code generation layers (numba/core/lowering.py, numba/core/codegen.py), and the accelerators — numba/parfors/ for automatic CPU parallelism and numba/cuda/ for NVIDIA GPUs. Along the way we will point at exact classes and line numbers so you can follow the same path in the code.
Why care about a numeric JIT on an AI blog? Because the techniques here are the same ones underneath every serious ML runtime: PyTorch’s compilation stack, tinygrad’s kernel generation, and the C metaprogramming in llm.c all solve the problem “turn high-level array math into tight machine code” in related ways. Numba is the cleanest open-source classroom for that subject because each stage is a separate, readable Python module.
The overview: a call enters through the decorators and the Dispatcher, becomes Numba IR, gets typed by the constraint engine, is rewritten for auto-parallelism, then lowered to LLVM and machine code — with the CUDA target reusing the same skeleton.
Reading the overview from left to right:
@jit, L27 — the single public entry point; it accepts a function or a signature and composes target, locals, and options into a dispatcher factory.@njit, L250 — a thin alias that pinsnopython=True, the mode you should always use; the old object-mode fallback is gone from modern versions.Dispatcher, L757 — the callable object your decorated name actually refers to; it keeps one compiled specialization per argument-type signature.Interpreter, L1423 — walks CPython bytecode and rebuilds it as Numba IR, a typed-node SSA-friendly representation.TypeInferer, L952 — the constraint solver that decides every variable’s Numba type before any code is generated.Lower, L346 — walks the typed IR one instruction at a time and emits equivalent LLVM IR through llvmlite.JITCPUCodegen, L1364 — hands the LLVM modules to the optimization and back-end layers and produces executable machine code.ParforPass, L2975 — the auto-parallelizer that recognizes loop and NumPy patterns and rewrites them intoParfornodes.CUDACompiler, L119 — proof that the whole skeleton is target-agnostic: the GPU path subclasses the same base and swaps in an NVVM back end.
From Install to Your First Model
Installation is deliberately boring: pip install numba pulls in llvmlite (the LLVM binding layer) plus a small set of dependencies — no compilers to install, no CUDA toolkit required unless you want the GPU path. The classic first model is a numerical kernel; here is a logistic-regression-style update, which is exactly the shape of code Numba loves:
import numpy as np
from numba import njit
@njit(cache=True, fastmath=True)
def train_step(w, X, y, lr):
n = X.shape[0]
for i in range(n):
pred = 1.0 / (1.0 + np.exp(-(X[i] @ w)))
err = y[i] - pred
w += lr * err * X[i]
return w
The first call is where everything interesting happens. train_step(X, y, ...) is not a Python function anymore — it is a CPUDispatcher, L44. Its call path enters _compile_for_args, L342, which maps each Python argument to a Numba type using machinery like resolve_value_type, L441. A float64 NumPy array becomes array(float64, 1d, C); that tuple of types is the cache key. If a compiled specialization for this key already exists — on disk too, thanks to cache=True — the call is a pointer jump. If not, the dispatcher blocks and runs the whole compile stack described below, then stores the result. This is why the first call can take hundreds of milliseconds while every subsequent call runs 30–100x faster than CPython. In a serving loop you warm the kernel once and pay native prices forever after; the same lazy-compile discipline shows up in other runtimes we have covered, from PyTorch’s compilation features to tinygrad’s JIT.
The Stage List: How a Compile Actually Proceeds
Every compilation is a sequence of small passes threaded through one mutable state object. compile_extra, L713 is the target-facing entry: it builds a CompilerBase, L386 subclass — Compiler for the CPU — whose state carries the function, the flags, the typing context, and later the IR and the LLVM module. DefaultPassBuilder, L528 assembles the nopython compilation stages by concatenating three lists: the untyped passes, the typed passes, and the lowering passes.
The executor is deliberately simple. PassManager, L169 holds an ordered list of stage descriptors, and its run, L342 method walks that list, fetching each registered stage from the PassRegistry, L408 and executing it. Stages subclass CompilerPass, L31; most are FunctionPass, L116 instances that mutate the IR in place. That is the entire control flow — there is no scheduler, no planner, just an ordered list you can print with NUMBA_DEBUG_ARRAY_OPT or the -O debug flags.
The untyped half of the sequence lives in numba/core/untyped_passes.py. TranslateByteCode, L73 invokes the interpreter to build func_ir; IRProcessing, L112 performs post-processing on the freshly built IR; and ReconstructSSA, L1745 rewrites the IR into static single assignment form, which is what makes the later rewrites and the type solver well-behaved. In between sit dozens of small rewrites — closure inlining, literal unrolling, dead branch pruning — each one an independent, testable class.
The Detail View: Inside Each Stage
The detail diagram spreads the same story across all six territories, including the stage machinery and the two accelerators, with the exact files each participant lives in.
The detail view: the Dispatcher-to-compile_extra handshake on the left, the untyped and typed stage lists in the middle, the LLVM code generation layers on the right, and the parfors plus CUDA accelerators along the bottom.
Type Inference: The Constraint Engine
Typing is where Numba earns its speed. The frontend builds a TypeInferer over the SSA IR; build_constraint, L1033 inspects each IR instruction and registers constraints — assignments propagate types between variables, calls register a CallConstraint, and so on. Then propagate, L1064 iterates the constraint network until it reaches a fixpoint: every variable must end up with exactly one Numba type, and inference failures raise the famous long TypingError diagnostics rather than falling back to slow code.
Resolving what a call means is delegated to the typing context. BaseContext, L210 is the registry of everything the compiler knows how to type, and resolve_function_type, L268 is the hot entry every CallConstraint eventually hits. It consults typing templates: AbstractTemplate, L349 subclasses implement a generic() method that, given argument types, returns a candidate Signature, L27 — return type plus argument types. NumPy functions, math functions, array indexing, @overload implementations: all of them register templates into this context. When inference succeeds, the IR is annotated and ready for lowering; when it fails, the error message lists every candidate signature that was tried, which is the single most useful debugging surface in the whole project.
Lowering to LLVM and Machine Code
The typed stage list ends with the lowering passes. NativeLowering, L517 constructs the lowering driver — BaseLower, L25 sets up the LLVM function declaration, the code library, and the variable-to-LLVM-value map — and then Lower, L346 runs the main lower, L188 loop over the IR. Each instruction is handed to lower_inst, L457, which dispatches on the node kind: an assignment lowers through the operation implementations registered in numba/core/imputils.py, a control-flow node lowers to LLVM branches, a Parfor node gets its own dedicated lowering path in numba/parfors/parfor_lowering.py.
The final stage is NoPythonBackend, L569: it asks the target context to optimize the finished LLVM module and produce the callable. On the CPU path that work happens in numba/core/codegen.py, where the Codegen, L1130 base defines the contract and JITCodeLibrary, L984 accumulates modules until JITCPUCodegen, L1364 runs the LLVM optimization passes and the native back end, then hands back a raw function pointer wrapped into the dispatcher. Environment arrays, exception state, and the boxing/unboxing glue between Python objects and native values are all woven in at this layer. It is code generation in the same family as tinygrad’s kernel emitter or llm.c’s hand-rolled GGML loops — just with LLVM doing the instruction selection.
Auto-Parallelism and the GPU Path
Pass parallel=True to @njit and the typed stage list grows a pair of extra stages. PreParforPass, L1514 rewrites convenient syntax — prange loops, in-place operators, reductions — into canonical forms, then ParforPass, L2975 pattern-matches loops and NumPy expressions and replaces them with a Parfor, L633 IR node: an explicit, analyzable parallel loop carrying its own loop nests, array analysis, and scheduling information. Array analysis proves bounds and independence so fusion and privatization are legal. At lowering time the choice in DefaultPassBuilder between NativeParforLowering, L528 and the scalar path decides whether those nodes become OpenMP-parallel loops. The same subsystem powers @vectorize, L64 and guvectorize, L139, which broadcast a scalar kernel across arrays exactly like NumPy ufuncs — a natural companion to the NumPy compatibility story.
The GPU path reuses the whole skeleton with different concrete classes. @cuda.jit, L14 decorates a device kernel; the resulting CUDADispatcher, L580 compiles the kernel through CUDACompiler, L119, whose typed back end is CUDABackend, L68, and emits PTX via JITCUDACodegen, L351 driving NVVM instead of the CPU back end. Thread indexing (threadIdx, shared memory, barriers) is expressed inside the kernel body, so the frontend and typing layers barely change — a nice demonstration that the decorator-to-LLVM design generalizes.
Honest Limits
The first-call compile latency is real and unavoidable: warm every specialization before serving traffic, or the cache file pays for it. The nopython subset excludes large parts of dynamic Python — arbitrary class hierarchies, mixed-type dicts, most reflection — and the failure mode is a Enjoyed this post? Never miss out on future posts by following us TypingError whose wall of text you must learn to read from the bottom up. parallel=True needs a working threading layer (OpenMP or TBB) and its speedups depend on array analysis succeeding; when the analyzer cannot prove independence it silently stays serial. The CUDA target requires an NVIDIA GPU, the CUDA toolkit, and driver versions in a supported window, and it compiles for your specific compute capability. Caching interacts poorly with closures and with functions defined in notebooks, where the source-to-cache-key mapping changes between sessions. And while numba composes beautifully with NumPy, it is not a general ML graph compiler: no autodiff, no dynamic shapes, no kernel-fusion planners at the scale of the deep-learning runtimes. Its sweet spot is exactly one thing — making a numerical Python function run at native speed — and within that scope it remains one of the best-engineered compilers you can read.