Machine learning frameworks like TensorFlow, JAX, and PyTorch 2.0 all share a common architectural pattern: they trace model execution into a directed acyclic graph (DAG) of operations, then run compiler passes over that graph to optimize it before execution. This intermediate representation (IR) decouples the user-facing model API from the execution backend, allowing the same model definition to run on CPUs, GPUs, or specialized accelerators without code changes.
A Zig-based ML inference engine can implement this pattern from scratch. When a model runs, it records each operation (matrix multiply, activation function, normalization) as a node in a DAG. The edges represent data dependencies—the output tensor of one operation feeds into the input of the next. Once the full graph is traced, the engine runs a pipeline of compiler passes: constant folding (pre-computing operations on known values), operator fusion (combining adjacent operations into a single kernel), and common subexpression elimination (deduplicating redundant computations).
This approach is similar to how XLA (Accelerated Linear Algebra) compiles TensorFlow graphs. The key insight is that lazy evaluation—recording operations instead of executing them immediately—gives the compiler visibility into the entire computation, enabling optimizations that are impossible when operations execute eagerly one at a time.
Instead of executing operations immediately, tracing records them into a graph structure. When user code calls matmul(A, B), the framework doesn’t compute the result—it creates a node representing the matmul operation with edges pointing to the nodes for tensors A and B. Execution is deferred until the user requests a concrete value (e.g., calling .eval()). This lazy approach lets the compiler see the entire computation before any work happens.
The traced operations form a DAG where nodes are operations (matmul, relu, softmax) and edges are tensor data flows. The graph must be acyclic because tensor operations are pure functions—they consume input tensors and produce new output tensors without mutation. This property makes the graph safe to reorder, parallelize, and optimize without changing semantics.
A compiler pass is a function that takes a graph, analyzes or transforms it, and returns a (possibly modified) graph. Passes run in a fixed order because some passes enable others: constant folding may create new opportunities for operator fusion, which in turn may reveal common subexpressions. The pipeline is designed so each pass maintains the graph’s semantic invariants—the optimized graph must produce the same numerical results as the original.
After optimization, the graph is handed to a backend interpreter that maps each operation to a concrete implementation. A typical inference engine supports multiple backends: BLAS for CPU matrix operations, CoreML for Apple Neural Engine acceleration, and MLX for Apple GPU compute. The graph IR serves as the contract between the frontend (model definition) and the backend (hardware execution), allowing new backends to be added without changing the tracing or optimization layers.
Each node in the graph stores the operation type, input references (edges to other nodes), output tensor shape and data type, and metadata for debugging (source location, operation name). Nodes are allocated from an arena allocator for cache-friendly traversal and zero-cost deallocation.
const Op = enum {
matmul,
add,
relu,
softmax,
reshape,
transpose,
constant,
// ... more operations
};
const Node = struct {
op: Op,
inputs: [MAX_INPUTS]*Node,
num_inputs: u8,
shape: Shape,
dtype: DType,
metadata: OpMetadata,
// For constant nodes, holds the pre-computed tensor data
constant_data: ?*Tensor,
};
const Graph = struct {
nodes: ArenaAllocator(Node),
outputs: []*Node, // terminal nodes whose values are requested
pass_log: ArrayList(PassResult),
};The tracing context intercepts tensor operations and builds the graph. User code looks like normal tensor math, but each operation appends a node to the graph instead of computing a result.
fn traceForward(ctx: *TraceContext, input: *Node) *Node {
// Linear layer: y = xW + b
const weights = ctx.constant(model.linear1_weights);
const bias = ctx.constant(model.linear1_bias);
const mm = ctx.addNode(.matmul, &.{ input, weights });
const biased = ctx.addNode(.add, &.{ mm, bias });
const activated = ctx.addNode(.relu, &.{ biased });
// Second linear layer
const w2 = ctx.constant(model.linear2_weights);
const b2 = ctx.constant(model.linear2_bias);
const mm2 = ctx.addNode(.matmul, &.{ activated, w2 });
const out = ctx.addNode(.add, &.{ mm2, b2 });
return ctx.addNode(.softmax, &.{ out });
}This pass identifies operations where all inputs are constants (known at compile time) and pre-computes the result. For example, if a bias vector is added to a zero tensor, the result is just the bias vector. The pass replaces the add node with a constant node holding the pre-computed result, removing one operation from the runtime graph.
fn constantFoldPass(graph: *Graph) PassResult {
var folded: usize = 0;
for (graph.nodes.items()) |node| {
if (allInputsConstant(node)) {
const result = evaluateConstant(node);
node.op = .constant;
node.constant_data = result;
node.num_inputs = 0;
folded += 1;
}
}
return .{ .name = "constant_fold", .nodes_modified = folded };
}Operator fusion combines sequences of operations into a single fused kernel. The most common fusion pattern is matmul followed by bias add followed by activation (matmul+add+relu becomes a single fused_linear_relu operation). Fusion reduces memory bandwidth by eliminating intermediate tensor materializations—instead of writing the matmul result to memory, reading it back for the add, writing again, and reading for relu, the fused kernel does all three in a single pass through the data.
fn fusionPass(graph: *Graph) PassResult {
var fused: usize = 0;
for (graph.nodes.items()) |node| {
// Pattern: matmul -> add -> relu
if (node.op == .relu) {
if (node.inputs[0].op == .add) {
const add_node = node.inputs[0];
if (add_node.inputs[0].op == .matmul) {
node.op = .fused_linear_relu;
node.inputs = add_node.inputs[0].inputs;
// Carry bias from the add node
node.inputs[node.num_inputs] = add_node.inputs[1];
node.num_inputs += 1;
fused += 1;
}
}
}
}
return .{ .name = "op_fusion", .nodes_modified = fused };
}CSE detects when two nodes compute the same operation on the same inputs and eliminates the duplicate. This commonly occurs in attention mechanisms where the same tensor is used in multiple projections. CSE hashes each node by its operation type and input pointers, then replaces duplicates with references to the first occurrence.
Passes execute in a fixed order. Each pass returns a result indicating how many nodes it modified. The pipeline can optionally run multiple iterations until a fixed point is reached (no pass modifies any nodes), catching optimization opportunities that earlier passes create for later ones.
fn runPassPipeline(graph: *Graph) void {
const passes = [_]PassFn{
constantFoldPass,
fusionPass,
csePass,
deadCodeEliminationPass,
};
var changed = true;
var iteration: usize = 0;
while (changed and iteration < MAX_ITERATIONS) {
changed = false;
for (passes) |pass| {
const result = pass(graph);
graph.pass_log.append(result);
if (result.nodes_modified > 0) changed = true;
}
iteration += 1;
}
}Graph nodes are allocated from an arena (bump allocator) rather than individually on the heap. This gives cache-friendly sequential memory layout for graph traversal, zero-cost deallocation (free the entire arena at once when the graph is no longer needed), and no fragmentation from thousands of small allocations. Zig’s allocator interface makes this trivial to implement and swap.
Optimized graphs can be serialized to disk and reloaded on subsequent runs, skipping the tracing and optimization phases. The cache key includes the model architecture hash, input shapes, and compiler pass versions. If any of these change, the cached graph is invalidated and recompiled. This is the same strategy XLA uses with its persistent compilation cache.
In debug builds, each pass is followed by a validation step that checks the graph’s invariants: no cycles, all edges point to valid nodes, output shapes are consistent with operation semantics. This catches pass bugs early. The validation is compiled out in release builds for zero runtime cost.