Lowering and fast CPU

In short: the first benchmark found rml slow wherever GEMM dominated, single-threaded outside GEMM, and slowed by JAX and TensorFlow graphs that spell out layer norm, softmax and GELU as chains of small ops. Four pieces answer it. rml/graph/fusion turns a run of elementwise nodes into one program that streams 512-element tiles through L1 and writes only the final value. rml/graph/kernels gives elementwise math, softmax, layer norm and row reductions loops compiled per instruction set with branch-free polynomials, split over threads only in chunks large enough to pay for waking them. The GEMM gained AVX-512 and AVX2 micro-kernels, a matrix-vector path and a few-row path. And the exporters fuse the spelled-out patterns back into single ops before rml ever sees them.

The levels above are on The inference stack; the executor that calls these pieces is on Executing a graph. Why fusing helps at all is the memory-traffic argument on GPU execution, which applies to CPU caches the same way.

The gaps

The first framework benchmark, on rml's 4-vCPU Linux container, set the agenda (progress.md L128–133): rml beat every framework only on the batch-1 CNN, and was 2–20× slower wherever GEMM dominated. Only GEMM was threaded, and the JAX and TensorFlow graphs arrived with layer norm, softmax and GELU decomposed. The plan that followed (L157–173) has one item per gap:

Fused programs

fusion.h calls this the kernel-level layer below rml-graph. A run of consecutive elementwise nodes with one output shape (a GELU written out as tanh, pow, mul and add; a layer norm's subtract, scale and shift; a bias add followed by an activation) becomes one program. The program streams the output in tiles small enough to stay in L1, keeps every intermediate in a tile buffer and writes only the values read outside it, which removes the intermediate tensors and all but one pass over memory, whichever framework the graph came from.

Members need not feed each other (fusion.h L14–19). Shapes are compared without leading 1s, so a reshape that only adds or drops them joins as a copy; an LSTM cell's four gate activations and its state update form one group whose outputs are the new cell and hidden states; and inputs that are row-strided views, such as the four gate slices of one matmul result, are read in place rather than copied.

Finding groups

• Fusible ops. The 13 unary ops, the six arithmetic binary ops, gelu in either form, and pow with an exponent of 2, 3, 0.5, −0.5 or −1 (or 1, a copy) (fusion.cc L35–85). Comparisons, where and clamp are not.
• Fusible nodes. One f32 output of the group's shape; unary inputs of exactly that shape; binary inputs laid out as the full output, a row repeated per row, one value per row, or a scalar, the same layouts the binary fast path takes (L102–126). Shapes are canonical, with leading 1s dropped (L90–96). A reshape, squeeze or unsqueeze joins only when it reads a value the group made, so a reshape of an outside value stays a free view. A group holds at most kMaxSteps = 48 steps, whose tile buffers stay within L2 (L87–88).
• Growing. Starting from a fusible node, FindFusionGroups takes every following node of the same shape while each is fusible, connected to the group or not (L181–189).
• Trimming. It drops trailing reshapes, which would only copy what the step before wrote, and groups shorter than two (L190–195). Each member's operands become references to an earlier step or to an outside input with its layout (L202–233).
• Outputs. A member's value becomes an output of the group when it is the last member, a graph output, unread, or read by a node after the group; every other value lives only in a tile buffer (L234–241).

Running a program

MakeFusedKernel views the output as rows of its last axis when any input is a row or column operand or a row-strided view, and otherwise as one long row (L255–263, L277–290). A full-size input is read in place when its last axis is contiguous and its other axes collapse into one evenly strided row axis; anything else is copied contiguous first (L128–147). It cuts rows into tiles of kTile = 512 floats, chosen so a program of a few dozen steps keeps every tile buffer in L1 (L25–26). Each task allocates steps × 512 floats of tile buffers once, then for each tile runs the program step by step with UnaryTile or BinaryTile. A step's operand is a tile buffer, the matching slice of a full input, the row slice of a bias, or one broadcast value; a step that makes one of the group's outputs writes straight into that output tensor instead of a tile buffer (L297–361). Tiles are split over the pool with a grain of about 64K element-steps per task (L302–306).

Fusion is on by default. rml-graph-run --no-fusion turns it off and --profile shows each fused program as its own row, which is the textual view of this layer the plan asked for (progress.md L158). There is no separate dump of a lowered IR: the layer is the executor's step list, and step_info() is how it is read (Executing a graph).

An LSTM step

FusesAnLstmCellAcrossItsGateSlices builds one LSTM step as PyTorch exports it and checks the steps the executor makes, for batch × hidden of 64 × 256, 3 × 5 and 1 × 700 (fusion_test.cc L145–205). Sixteen nodes become six steps:

Before this, every member had to read something the group made and a group had one output, so the gate activations, which do not feed each other, ran as separate steps, and a fused step copied any strided input into a contiguous tensor first. rml's checklist puts the change at about 39 to 26 ms for the PyTorch, JAX and TensorFlow lstm models at batch 64, and the weight packing below at 49 to 35 ms for torch_lstm_b64 on 4 threads (progress.md L165). The latest report has torch_lstm_b64 at 12.7 ms on 4 threads, against 6.9 for PyTorch eager timed the same day (frameworks.md L106); see Where it stands.

JAX and TensorFlow graphs used to reach this step with one more product each: their exporters leave the input projection x[:, t] @ W inside the unrolled loop. The export pass hoist_step_projections now makes it one product sliced per step, the shape PyTorch's export already has (Exporters and the model zoo). In interleaved A/B runs on 4 threads rml's checklist has jax/tf lstm_b64 about 1.15–1.2× faster and jax_rnn_b64 1.36×, with the batch-1 and batch-8 cells within noise or slightly faster (progress.md L169). The latest benchmark report was timed after both this and the im2col change below (progress.md L171).

Vectorized kernels

rml/graph/kernels.h holds the executor's hot non-GEMM ops. The design is in one comment block (kernels.cc L11–26):

• One build, three instruction sets. Each loop is marked target_clones("avx512f", "avx2", "default") on ELF x86-64, and the dynamic loader picks the best clone. Elsewhere, and in ASan and TSan builds (whose runtimes start after the ifunc resolvers run), it is one plain function.
• Vectorizable without -ffast-math. No calls and no branches inside the loops (selects instead), and reductions spread over kLanes = 16 independent accumulators, which keeps the rounding order fixed.
• Branch-free math. exp by range reduction and a degree-6 Cephes polynomial (about 1 ulp); tanh as an odd polynomial below 0.625 and 1 − 2/(e^2x + 1) above, both computed and one selected; erf as a Cephes polynomial below 1 and Abramowitz and Stegun 7.1.26 beyond (L36–108).
• Double where it matters. Row sums accumulate in double lanes, because a float sum over thousands of terms drifts past the 1e-5 the graph tests allow (L266–279), and layer norm splits the mean into a float and its rounding residual for rows with a tiny spread (L332–344).
• Chunks that pay for themselves. Work is split in chunks of kChunk = 2¹⁶ elements, or enough rows to cover that (L31–34, L112–130). Since ParallelFor runs a single chunk on the caller (thread_pool.h L49–56), a tensor of up to 65,536 floats never wakes the pool.

rml's own number for this work: PyTorch's batch-8 transformer block went from 23 ms to 6.7 ms on 4 threads once the executor used these kernels (progress.md L160).

SIMD GEMM

The GEMM that matmul, conv2d and sdpa reach is described in full on Linear algebra. In short (kernels.md L19–80): the micro-kernel is chosen at run time from AVX-512 (12 × 32 tile), AVX2 + FMA (6 × 16) or the portable vector-extension kernel; a single-row product skips packing and streams B once as axpys; products of at most 32 rows also skip packing and keep an 8 × 48 or 4 × 24 tile of C in registers. In the graph, a constant matmul weight is packed once when the executor is built rather than on every call (ops.cc L929–944). In two interleaved runs on four threads that took torch_lstm_b64 from 53 / 45 ms to 33 / 37 ms and torch_mlp_b64 from 1.05 / 1.39 ms to 0.82 / 0.82 ms (kernels.md L75–78). When a product has fewer than four (row block, column chunk) units per thread, the driver widens each task's column chunk, up to four times, to the width whose slowest thread computes the fewest panels; that took torch_rnn_b64 from 10.5 / 9.8 ms to 5.2 / 5.0 ms on four threads (kernels.md L67–74, details).

Folding the recurrence's adds into the step GEMM (prefilling C and running with β = 1, then also the activation as an epilogue, then all column slices in one dispatch) was tried three ways and dropped: in BM_RnnSteps a step took 54–62 µs fused against 58–62 µs unfused on 4 threads, because the per-step GEMM already runs at 245–290 GFLOP/s, about the machine's peak (kernels.md L258–273).

From kernels.md L99–110: the AVX-512 kernel is 6–7× the SSE-only kernel it replaced (16.5 and 60 GFLOP/s for 1024³). The few-row path took the batch-8 gan_generator graph from 0.50 to 0.15 ms and the RNN from 4.1 to 1.9 ms on one thread (L58–63).

Conv bias and im2col

JAX and TensorFlow export a convolution's bias as a separate add of a per-channel constant. With fusion on, the executor folds that add into the conv2d before it, which already takes an optional bias input, so the conv writes biased output in one pass: one conv2d+add step instead of two. It applies only when the conv has no bias of its own, its output is read only by the add and is not a graph output, and the constant is f32 with shape [O], [O, 1, 1] or [1, O, 1, 1] (executor.cc L23–52, Executing a graph). Together with the strided copy below, rml's checklist puts tf/jax unet_b64 at about 46 to 37 ms on 4 threads (progress.md L167).

The convolution's other half is building the patch columns. Im2ColPositions fills each row of cols (one channel and kernel tap) output row by output row, each a strided read of one input row with zeros where the tap lands in the padding. It now computes that valid range once per tap and fills the segments with plain loops, cloned for AVX-512 and AVX2 like the other kernels, not std::copy / std::fill: a segment is one output row, often 16 or 32 floats, where a library call per segment cost more than the copy (kernels.cc L446–493). rml's checklist has im2col about 2.3× faster on the U-Net shapes, and in interleaved A/B runs in its Linux container unet_b64 about 1.25–1.5× faster on 4 threads and 1.4× on 1, with unet_b8 and cnn_b64 1.3× (progress.md L170). It records only ratios, because that host ran about 40% slower than when the benchmark report was measured.

Buffers

The plan named "buffers planned at load". What landed is reuse rather than a static plan: each executor owns a BufferPool, every tensor a run allocates on the calling thread comes from it, and each freed buffer goes back keyed by its exact size, so a model with static shapes allocates nothing new after its first run (BufferPool, Executing a graph). Releases are planned at load: each step knows which values it is the last to read. Two other things are prepared once at load: the fusion programs, and packed constant weights.

Views that must be materialized (a reshape of a broadcast, a concat, a conv input) go through StridedCopy, which now merges axes, copies two-axis blocks with memcpy or fill_n rows, and threads copies of 64K elements or more (Tensors and dtypes). On the U-Net graphs at batch 64 and 4 threads, the upsample's reshape went from 8.2 to 1.7 ms (PyTorch) and 7.5 to 1.6 ms (TensorFlow) (kernels.md L230–256).

Fusion at export

Runtime fusion covers only runs of elementwise nodes. A spelled-out softmax or layer norm also needs row reductions, which split it into several steps, while the single op runs on one kernel from the table above. So the exporters' clean-up passes, run on both routes, rewrite the patterns back into single ops (passes.py L586–591):

All three are in the pass loop at this commit, and rml's checklist marks export-time layer norm and GELU fusion done (progress.md L163). The passes themselves are on Exporters and the model zoo.

Threads

A pool wake costs more than a few microseconds of work, so every parallel loop now has a grain sized to pay for one: 2¹⁶ elements in the graph kernels, about 64K element-steps in fused programs, and larger grains in the GEMM (Errors, memory and threads lists each). Pool workers also spin briefly before sleeping. The open item is the result: on the latest run, four threads is slower than one in 32 of 171 rml cells, almost all at batch 1 with sub-millisecond times (frameworks.md L375–376, progress.md L172).

Where it stands

Geometric mean over the 19 models of the faster framework mode's time divided by rml's (above 1, rml is faster), with the number of models where rml wins, on the Xeon container with torch 2.14.1, jax 0.12.0 and tensorflow 2.21.0. Every framework and rml cell was timed in one sweep on 2026-11-03 (frameworks.md L5–7, L343–362):

Against the fastest framework and mode per cell, rml wins 46 of 114 cells, with a geometric mean of 0.96: ahead on one thread at batches 1 and 8, behind at batch 64 and on four threads, where PyTorch leads (L364–367). The previous report put this at 80 cells and 1.31, but it timed only rml afresh against older framework cells, and rml's notes say it overstated rml (L344–347). The models furthest behind are named in the findings (L371–374): the U-Net (0.48×, against torch.compile's convolutions), the GCN (0.6×, against PyTorch eager on four threads), the LSTM (0.64×), the SSM, the CNN and the RNN, with the attention models at 0.9–0.96×. None of these models has a checklist item of its own yet. Cell by cell results and the method are on Benchmarks against the frameworks. Apple Silicon numbers (NEON, Accelerate) have not been produced (L349).

Tests