Turning Python into a compilable language
Machine learning libraries like JAX and Triton popularized a pattern where a Python function decorated with jit is not executed by the Python interpreter in the usual way. Instead, the function body acts as a domain-specific language that a library-specific compiler processes. Python serves as a meta-language for describing computations, and the decorated function's code is transformed into an internal representation, compiled, and executed natively. This article explores several implementation strategies behind such decorators using a simplified educational example.
The educational jit decorator follows a three-step pipeline:
- Translate the Python function into an expression IR (
Expr). - Convert the
Exprto LLVM IR. - JIT-execute the LLVM IR using
llvmlite.
The Expr IR is deliberately minimal — it only supports functions that return a single arithmetic expression. This simplification keeps the demonstration focused while still showing a complete compilation path from Python source to native execution. The full Expr structures and the code generator that lowers them to LLVM IR are available in the referenced implementation; the generator walks the Expr tree and emits equivalent LLVM instructions.
This design mirrors how real libraries work without overcomplicating the example. The point is not to evaluate the expression directly but to show that arbitrary compilation complexity can be hidden behind a simple decorator interface. With this foundation, we can examine how different JIT strategies convert a Python function into an Expr.
AST-based JIT: leveraging Python's parser
Python's built-in introspection makes AST-based translation straightforward. The astjit decorator is a standard decorator that uses functools.wraps to preserve function metadata:
def astjit(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
if kwargs:
raise ASTJITError("Keyword arguments are not supported")
source = inspect.getsource(func)
tree = ast.parse(source)
emitter = _ExprCodeEmitter()
emitter.visit(tree)
return llvm_jit_evaluate(emitter.return_expr, *args)
return wrapper
When the wrapped function is called, the wrapper fetches the original function's AST and feeds it to an _ExprCodeEmitter visitor. This visitor walks the AST nodes and builds the Expr tree representing the function's return value:
class _ExprCodeEmitter(ast.NodeVisitor):
def __init__(self):
self.args = []
self.return_expr = None
self.op_map = {
ast.Add: Op.ADD,
ast.Sub: Op.SUB,
ast.Mult: Op.MUL,
ast.Div: Op.DIV,
}
def visit_FunctionDef(self, node):
self.args = [arg.arg for arg in node.args.args]
if len(node.body) != 1 or not isinstance(node.body[0], ast.Return):
raise ASTJITError("Function must consist of a single return statement")
self.visit(node.body[0])
def visit_Return(self, node):
self.return_expr = self.visit(node.value)
def visit_Name(self, node):
try:
idx = self.args.index(node.id)
except ValueError:
raise ASTJITError(f"Unknown variable {node.id}")
return VarExpr(node.id, idx)
def visit_Constant(self, node):
return ConstantExpr(node.value)
def visit_BinOp(self, node):
left = self.visit(node.left)
right = self.visit(node.right)
try:
op = self.op_map[type(node.op)]
return BinOpExpr(left, right, op)
except KeyError:
raise ASTJITError(f"Unsupported operator {node.op}")
After the AST traversal, the emitter's return_expr field contains the full expression, which is then handed to llvm_jit_evaluate for native execution. The key insight is that the decorator intercepts normal Python execution: instead of compiling to bytecode and running in the VM, the function body is translated to LLVM IR and JIT-compiled.
Triton uses essentially this approach. A function decorated with @triton.jit has its body parsed to a Python AST, converted through several internal IRs to LLVM IR, and finally lowered to PTX by the NVPTX backend for GPU execution. Triton also exposes intrinsics — special calls from the triton.language package that the compiler handles directly. The supported Python subset is limited, but sufficient for writing performant GPU kernels alongside regular "host" Python code.
Bytecode-based JIT: standing on Python's compiler
Python is a large language with rich semantics. A JIT that needs to support a substantial portion of those semantics might prefer to start from Python's own bytecode rather than the AST. The bytecodejit decorator takes this path:
def bytecodejit(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
if kwargs:
raise BytecodeJITError("Keyword arguments are not supported")
expr = _emit_exprcode(func)
return llvm_jit_evaluate(expr, *args)
return wrapper
def _emit_exprcode(func):
bc = func.__code__
stack = []
for inst in dis.get_instructions(func):
match inst.opname:
case "LOAD_FAST":
idx = inst.arg
stack.append(VarExpr(bc.co_varnames[idx], idx))
case "LOAD_CONST":
stack.append(ConstantExpr(inst.argval))
case "BINARY_OP":
right = stack.pop()
left = stack.pop()
match inst.argrepr:
case "+":
stack.append(BinOpExpr(left, right, Op.ADD))
case "-":
stack.append(BinOpExpr(left, right, Op.SUB))
case "*":
stack.append(BinOpExpr(left, right, Op.MUL))
case "/":
stack.append(BinOpExpr(left, right, Op.DIV))
case _:
raise BytecodeJITError(f"Unsupported operator {inst.argval}")
case "RETURN_VALUE":
if len(stack) != 1:
raise BytecodeJITError("Invalid stack state")
return stack.pop()
case "RESUME" | "CACHE":
# Skip nops
pass
case _:
raise BytecodeJITError(f"Unsupported opcode {inst.opname}")
Because the Python VM is stack-based, the translator emulates a stack to convert the function's bytecode into the Expr IR — similar to how a reverse-Polish-notation evaluator works. The resulting Expr is then lowered to LLVM IR and JIT-executed just as before. Swapping astjit for bytecodejit requires no other changes.
Numba is the prominent real-world example here. Its numba.njit decorator also relies on Python's bytecode as a starting point. Numba compiles the bytecode into its own IR and then to LLVM using llvmlite. Starting from bytecode saves Numba from reimplementing Python's frontend, but it comes at a cost: by the time code is in bytecode, semantic information from higher-level constructs is already lost, forcing Numba to do extra work to recover control flow information via a specialized interpreter.
The bottom line
Whether you start from the AST or the bytecode, the underlying pattern is the same: a decorator converts Python source into a library-specific IR, which is then compiled and executed natively. AST-based strategies keep more source-level information but must handle Python's full grammar; bytecode-based strategies reuse Python's compiler but have to reconstruct structure from a stack-machine encoding. Both approaches enable Python to act as a frontend for highly optimized, domain-specific execution engines.
Tracing: A Different Angle on JIT
Where the AST and bytecode approaches both rely on introspection to lower the function’s source into an intermediate representation, tracing takes a fundamentally different route. Instead of analyzing code, it executes the wrapped function with specially-boxed arguments that capture the flow of operations via overloaded operators. The result is an IR trace of the computation, which can then be lowered and JIT-compiled.
The implementation for the smile demo is remarkably concise:
def tracejit(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
if kwargs:
raise TraceJITError("Keyword arguments are not supported")
argspec = inspect.getfullargspec(func)
argboxes = []
for i, arg in enumerate(args):
if i >= len(argspec.args):
raise TraceJITError("Too many arguments")
argboxes.append(_Box(VarExpr(argspec.args[i], i)))
out_box = func(*argboxes)
return llvm_jit_evaluate(out_box.expr, *args)
return wrapper
Each runtime argument is wrapped in a _Box containing a VarExpr, and the box’s overloaded operators build up an Expr tree as the traced function executes:
@dataclass
class _Box:
expr: Expr
_Box.__add__ = _Box.__radd__ = _register_binary_op(Op.ADD)
_Box.__sub__ = _register_binary_op(Op.SUB)
_Box.__rsub__ = _register_binary_op(Op.SUB, reverse=True)
_Box.__mul__ = _Box.__rmul__ = _register_binary_op(Op.MUL)
_Box.__truediv__ = _register_binary_op(Op.DIV)
_Box.__rtruediv__ = _register_binary_op(Op.DIV, reverse=True)
The core mechanism lives in _register_binary_op:
def _register_binary_op(opcode, reverse=False):
"""Registers a binary opcode for Boxes.
If reverse is True, the operation is registered as arg2 <op> arg1,
instead of arg1 <op> arg2.
"""
def _op(arg1, arg2):
if reverse:
arg1, arg2 = arg2, arg1
box1 = arg1 if isinstance(arg1, _Box) else _Box(ConstantExpr(arg1))
box2 = arg2 if isinstance(arg2, _Box) else _Box(ConstantExpr(arg2))
return _Box(BinOpExpr(box1.expr, box2.expr, opcode))
return _op
To see this in action, consider a simple decorated addition function:
@tracejit
def add(a, b):
return a + b
print(add(1, 2))
After decoration, add points to the wrapper defined inside tracejit. When you call add(1, 2), the wrapper goes through these steps:
- It creates a fresh
_Boxholding aVarExprfor each of the function’s parameters (aandb). - The wrapped function is invoked with those boxes as arguments.
- Inside the wrapped function, the expression
a + btriggers the box’s overloaded__add__, producing a newBinOpExprwith the twoVarExprs as children. - The wrapper unboxes the resulting
Exprand passes it tollvm_jit_evaluate, which emits LLVM IR and JIT-compiles it with the original runtime values (1,2) to produce the final result.
This design introduces two distinct execution phases:
- Tracing step: the wrapped function runs normally under the Python interpreter, but the boxed arguments force it to build an
ExprIR rather than compute anything concrete. - Execution step: the captured IR is lowered to LLVM IR and JIT-executed with the actual argument values.
Because tracing never inspects the function’s source, it transparently supports much richer control flow. For instance, code with intermediate local variables “just works”, since the tracer follows the flow of values and is oblivious to how they’re bound:
@tracejit
def use_locals(a, b, c):
x = a + 2
y = b - a
z = c * x
return y / x - z
print(use_locals(2, 8, 11))
Data-independent loops also work without special handling. Here, the resulting Expr becomes a long chain of BinExpr additions that accumulate the loop-variable values with b * c:
@tracejit
def use_loop(a, b, c):
result = 0
for i in range(1, 11):
result += i
return result + b * c
print(use_loop(10, 2, 3))
That example reveals tracing’s main limitation: the control flow cannot depend on runtime argument values. The tracer has no notion of concrete values, so it would have no way to know how many loop iterations to trace through—unless you’re willing to re-trace on every invocation.
Tracing is especially well-suited to automatic differentiation (AD); the author’s radgrad project offers a deeper look.
Case Study: JAX
The JAX ML framework uses a tracing strategy that closely mirrors the one described here. JAX wraps Numpy with its own traced variant (its boxes are called “tracers”), letting users write familiar Numpy code that can be JIT-compiled and executed on accelerators via XLA. Its tracer produces an IR called jaxpr, which is then lowered into XLA operations.
JAX also inherits the data-dependent-control-flow limitation. Code like the following fails because count is a runtime value:
import jax
@jax.jit
def sum_datadep(a, b, count):
total = a
for i in range(count):
total += b
return total
print(sum_datadep(10, 3, 3))
The error is explicit:
This concrete value was not available in Python because it depends on the value of the argument count.
JAX provides jax.lax intrinsics as a workaround. The equivalent fori_loop version compiles fine, because JAX can trace through it and lower it to XLA’s While operations:
import jax
from jax import lax
@jax.jit
def sum_datadep_fori(a, b, count):
def body(i, total):
return total + b
return lax.fori_loop(0, count, body, a)
The payoff is that tracing handles arbitrarily complex Python—closures, metaprogramming, and other value-flow patterns—without requiring source analysis. The full sample code is available on GitHub, and the autodidax doc provides a thorough deep-dive into JAX’s internals.
