Skip to main content

The Graph IR

The graph IR is poot's central abstraction. It is a flat, typed, topologically ordered list of primitive tensor equations in SSA (static single assignment) form. Every other subsystem (the tracer, the fusion pass, the backends, the executor) is either a producer or a consumer of this one representation.

Why a fine-grained primitive IR​

ML compilers from XLA 1 to TVM 2 to JAX 3 share a common design insight: general kernel fusion requires the IR to be finer-grained than the kernels it produces. poot's IR follows the same principle. Attention, RMSNorm, RoPE, and SwiGLU are not opcodes. They are compositions of the actual atomic operations (matmul, reduce, elementwise, broadcast). The fusion pass can then synthesize kernels for supported regions of those primitives.

A practical consequence: flash attention 4 is not a fixed primitive in the IR. compile recognizes supported attention subgraphs and rewrites them to a single fused dispatch. A model author composes attention with ops::attention_masked(...); the rewrite depends on the shape, target and fusion policy.

Graph structure​

A simplified view of the main tables (all private; a Graph is read through accessors and changed only through a checked edit):

struct Graph {
    values:  ValueTable<ValueMeta>,        // shape + dtype + storage class per id
    inputs:  Vec<ValueId>,                 // all input binders: consts, slots, state
    consts:  Vec<ValueId>,                 // the subset of inputs that are weights/tables
    slots:   Vec<(ValueId, Slot)>,         // per-step inputs and their kind
    eqns:    Vec<Eqn>,                     // topologically ordered SSA equations
    outputs: Vec<GraphOutput>,             // one list: the result, witnesses, exports, state next values
}

enum OutputKind { Result, Witness { .. }, Export(BoundaryId), StateNext { input: ValueId } }

A Graph is always well formed: Builder::finish and GraphEdit::commit check the whole graph. A pass opens graph.edit(), changes the draft and commits it; a rewrite that loses a witness, an export or a carried-state pair, or changes the result's type, is refused with a typed error.

struct Eqn {
    op:     OpKind,                        // the operation, with its parameters inline
    inputs: Vec<Operand>,                  // Value(ValueId) or Lit(Scalar)
    out:    ValueId,                       // single output
}

struct ValueMeta {
    aval:     TensorType,                  // shape + dtype - no data
    storage:  Storage,                     // Device | Const | Slot(Slot) | State | Computed
    identity: Identity,                    // Anonymous | Named(String) | Weight(WeightRef) | Slot(SlotKey)
}

ValueId is a u32 newtype that indexes a ValueTable; a bare usize is rejected by the compiler. Each computed value is defined by one Eqn; graph inputs are binders. There are no separate def-use chains; a use is a ValueId into values, and an equation result is defined by the unique Eqn whose out is that id. This makes all transforms (fusion, substitution, dead-code elimination) straightforward.

Storage classes​

Device. An ordinary intermediate tensor. The planner derives its storage requirements from ValueMeta.aval; the executor allocates or aliases buffers and retains them for cached execution. Activation tensors such as attention scores, MLP outputs and the residual stream use this class.

Const. A weight or constant tensor (weight matrices, embedding tables, RoPE tables). Bound from the input map and uploaded for execution. A mapped weight is identified by its WeightRef, not by a name; weight caches use payload identity and storage representation, so different bytes never reuse a cached weight. Graph::consts() lists the const binders.

Slot(Slot). A per-step input routed through a stable device buffer that the host rewrites in place. The Slot kind says what it carries: Token, Pos, SlotMap and WriteMap (paged KV reads and writes; both declare the range of the words they hold, which the bind checks), and others. A role only a pass or a runtime feature declares (the adapter index, the hosted embedding row, the expert pool and its map, the sampler inputs) is defined by that owner, not by the IR: the IR knows its label and the one type rule it declares inputs by. The attention mask is not a slot: it is computed in-graph from Pos.

State. A persistent buffer carried across decode steps: the KV cache. Pre-allocated to cap capacity per layer. Graph::state() pairs each state_in binder with a state_out value of identical type. The executor feeds the buffer in and overwrites it with state_out (the JAX donate_argnums contract), which also orders the KV read before the KV write.

Computed. A compiler-defined constant, such as a folded Iota (an F32 range, or an exact I32 range). Its value comes from the graph rather than a checkpoint weight binder.

Primitive operations​

The OpKind vocabulary. Tracers emit the primitives (elementwise, reduce, movement, matmul, gather/scatter, quantization). The optimizer introduces the composite ops (Fused, FusedRow, ContractionEpilogue, Attention, Rope, BlockTopK, SelectedAttention and the packed/dense contraction ops), each with the same semantics as the primitive chain it replaces; a tracer cannot emit them. The one exception is Scan, which states a recurrence rather than a schedule: tracers emit it as a source composite with a primitive definition the CPU oracle evaluates. A sparse-attention indexer writes its block selection and the attention over the selected blocks as primitive chains, and the selection pass of compile forms BlockTopK and SelectedAttention from them; a device that cannot plan them refuses the composite, because the chains have no device lowering. No op records a kernel choice: which kernel runs an eqn (a tiled GEMM, the attention kernel, an imported kernel) is a property of its plan.

OpCategoryInputsOutput shapeNotes
Unary(op)Elementwisex: [...][...] (same)Neg, Exp, Log, Sqrt, Recip, Tanh, Erf; Not, Clz on I32
Binary(op)Elementwisea: [...], b: [...]broadcast([a], [b])Add, Sub, Mul, Div, Max, Ge (indicator); GeU, And, Or, Xor, Shl, Shr on I32. Numpy broadcast rules
SelectElementwisecond, if_true, if_false[...] (same)cond is I32, nonzero picks if_true; branches share one dtype. A native hardware select that moves the chosen element bit for bit, so a NaN or Inf in the unselected branch never reaches the output. I32, F32, F16 and BF16 branches; ROCm and PTX lower all four, while SPIR-V (wgpu, raw Vulkan) refuses F16 and BF16 with a typed DtypeLowering
Cast(to)Elementwisex: [...][...] (same)Storage dtype change; f32 -> bf16 rounds to bf16 precision
BroadcastShapex: [...], shapeshapeExpand dims
ReshapeShapex: [...], shapeshapeSame element count; lowers to a buffer alias (no dispatch)
TransposeShapex: [...], permpermuted [...]Reorder axes; may lower to a strided view
SliceShapex: [...], axis, start, endsliced [...]Contiguous subrange on one axis
ConcatShapexs: [[...]], axisconcatenated [...]Fusion boundary
GatherDatasrc: [...], idx, axisindexed [...]Embedding lookup, RoPE table row, paged KV read
ScatterUpdateDatabase: [P, ...], src: [N, ...], slots: [N]same as basePaged KV write: row t of src goes to pool slot slots[t] (-1 writes nothing), in place on the donated pool; cost follows N, not P
DynamicUpdateSliceDataoperand, update, indexsame as operandDense KV-cache slot write at a scalar index
ArgTopK(k)Datarank: [..., E][..., k]Expert ids from a stable rank tensor (MoE routing)
Reduce(op, axis, keepdim)Reducex: [...][...] (axis reduced)Sum or Max over one axis; keepdim preserves rank
MatMulComputea: [..., M, K], b: [..., K, N][..., M, N]Batched; B=1 (GEMV) and B>1 (GEMM) are one op; specialization at plan time
IndexedMatMulComputex: [M, K], W: [E, K, N], idx: [M][M, N]Each row contracts against its own expert; no weight copy
PackedDequantComputepacked sources by descriptor role[out, K]Packed-weight decode primitive; format and shape from poot-quant
PackedContraction / PackedRowGatherComputeactivation or row ids, packed sourcescontraction or selected rowsCompiler-only; decode packed values while consuming them
DenseContraction / DenseRowGatherComputebf16, f32 or f16 weight, or bf16 tablef32Compiler-only; contract a weight in checkpoint [N, K] order (bf16, f32, or f16 stored as two-byte words packed in u32) and gather rows of a bf16 table, with no transpose or widened copy
GroupedDenseContractionComputeactivation [G, M, K] f32, bf16 weight [G*N, K][G, M, N] f32Compiler-only; contract a checkpoint [G*N, K] weight read as G blocks of N rows against a [G, M, K] activation, so a grouped projection keeps its bf16 weight resident with no transpose or widened copy
Rope(rot)Computex, cos, sinsame as xIntroduced by rope_fusion from the rotate-half chain
AttentionComputeq, k, v, masksame as qIntroduced by the attention rewrite in compile, for decode and prefill; online softmax, no score matrix. The planner generates its kernel
Fused / FusedRowComputeregion inputsregion rootPointwise and row-wise (reduction-rooted) fused regions
ContractionEpilogueComputehead operands (a, b), then the epilogue's extra inputs (bias, residual, ...)as the head contractionIntroduced by the epilogue fusion from a dense contraction and the single-use pointwise chain that consumes it. The head's result keeps its dtype (a BF16 head rounds to BF16 before the first step reads it) and every step rounds through it
AllReduce / AllGatherComputex: [...]same shapeTensor-parallel collectives. AllGather takes the number of ranks and multiplies the gather axis by it. A collective is a cut between stages: compiling one inside a stage graph is refused, and the reference evaluator refuses an AllGather over more than one rank

What composite ops decompose to​

Higher-level operations are not opcodes; they are sequences of primitives recorded by the tracer. Here are three examples.

RMSNorm 5, seven eqns, fused into one reduction kernel:

sq  = Mul(x, x)                            # [L, D]
ss  = Reduce(Sum, sq, axis=-1, keepdim)    # [L, 1]
ms  = Mul(ss, 1/D)                         # [L, 1]
mse = Add(ms, eps)                         # [L, 1]
den = Unary(Sqrt, mse)                     # [L, 1]
xn  = Div(x, den)                          # [L, D], den broadcasts
out = Mul(xn, weight)                      # [L, D]

Numerically stable softmax. Five eqns (binary ops broadcast implicitly, numpy style), fused into one reduction kernel:

m   = Reduce(Max, scores, axis=-1, keepdim=True)
xs  = Binary(Sub, scores, m)
ex  = Unary(Exp, xs)
s   = Reduce(Sum, ex, axis=-1, keepdim=True)
out = Binary(Div, ex, s)

The Max and Sum reductions share the same traversal. This is the two-pass safe pattern (subtract max, then normalize). The streaming "online" algorithm used inside flash attention is different, see Attention. After flash attention rewrite, this entire pattern (plus the surrounding matmuls) disappears into one fused dispatch.

RoPE 6. The tracer emits the rotate-half chain (ops::rope_partial), with cos/sin gathered from constant tables by position:

x1, x2 = Slice(x, 0..D/2), Slice(x, D/2..D)
out = Concat([Sub(Mul(x1,cos), Mul(x2,sin)),
              Add(Mul(x1,sin), Mul(x2,cos))], axis=-1)

compile recognizes this chain and replaces it with one Rope { rot } eqn with inputs [x, cos, sin], lowered to one synthesized kernel instead of ~6-8 dispatches.

What an attention layer looks like as a DAG​

The diagram below shows the primitive eqns for one decode-step attention layer. Hover a node to highlight its edges.

Primitive eqns for one attention layer (decode step): hover a node for details
q[H,D]k_cache[H,cap,D]v_cache[H,cap,D]causal_mask[cap]MatMul[H,cap]Mul / sqrt(D)[H,cap]Add mask[H,cap]Reduce Max[H,1]Sub[H,cap]Exp[H,cap]Reduce Sum[H,1]Div[H,cap]MatMul[H,D]
MatmulReduceElementwiseInput / const

Host-side shape inference​

Every OpKind implements infer(input_types) -> output_type. Shape inference runs entirely on the host at trace time. No shape computation ever happens on the device.

Two consequences. First, the executor knows every buffer size before the first dispatch and can allocate everything upfront at stable addresses, a requirement for recorded replay (CUDA graphs, HSA packets, Vulkan command buffers). Second, because all shapes are static, the planner can select kernel specializations (GEMV vs GEMM, the attention schedule) at compile time rather than conditionally at runtime.

The exception: per-step inputs (token id, position) are Storage::Slot inputs (Slot::Token, Slot::Pos), small device buffers at stable addresses written by the host each token. The causal/windowed mask is built in-graph from Slot::Pos against an Iota; no slot carries it. Output shapes do not depend on slot values (the KV capacity is pre-fixed), so the captured graph remains valid across all positions.

Quantization in the IR​

PackedDequant is the single packed-weight primitive. Its poot-quant::PackedWeight descriptor names the WeightFormat, shape and source roles. The compiler recognizes dense projections as PackedContraction and embedding lookups as PackedRowGather; generated kernels decode only the weight values they consume.

GPTQ, AWQ, supported GGUF quants and the registered FP8/E2M1 layouts use the descriptor-driven lowering. Support still depends on the consuming graph: a packed MoE expert stack is rewritten when its weights are bound into one indexed packed contraction that decodes only the selected experts' blocks, and the packed loader keeps the weights packed.

References​

Footnotes​

  1. OpenXLA / Google, "XLA: Compiling Machine Learning for Peak Performance" (engineering docs / talks). https://openxla.org/xla ↩

  2. Chen, T. et al., "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning," OSDI, 2018. ↩

  3. Bradbury, J. et al., "JAX: Composable Transformations of Python+NumPy Programs," 2018. ↩

  4. Dao, T. et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS, 2022. https://arxiv.org/abs/2205.14135 ↩

  5. Zhang, B. & Sennrich, R., "Root Mean Square Layer Normalization," NeurIPS, 2019. ↩

  6. Su, J. et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding," arXiv:2104.09864, 2021. https://arxiv.org/abs/2104.09864 ↩