Axon: A Synthesizing Superoptimizer for Tensor Programs

Abstract: Writing high performance kernels for AI accelerators requires deep expertise in tiling, instruction selection, data layout, and operator fusion placing a significant burden on programmers. In this paper, we focus on tile based AI accelerator programs and present Axon, a synthesizing superoptimizer for tensor programs: it uses program synthesis to automatically generate target instructions from semantics specifications, and explores semantically equivalent program variants to select the best performing kernel empirically. Axon discovers algebraic transformations by propagating operators through computation graphs and uses SMT over unbounded tensors to guarantee that all transformations preserve semantics without requiring hand crafted rewrite rules. It then lowers tensor operations to target ISA instructions, explores tiling configurations constrained by hardware descriptions, and fuses operators and instructions to minimize memory traffic.

• The paper introduces Axon, a semantics-driven superoptimizer that uses SMT-verified transformations, symbolic tiling, ISA synthesis, and fusion to generate tensor kernels without hand-written rewrite rules.
• Axon evaluates all valid program variants in a persistent νGraph, then compiles and measures candidates on Trainium, achieving up to 19× speedups over the Neuron compiler and modest gains over Mirage.
• The results show that fusion and algebraic transformations are central to performance, while empirical search costs, backend compilation failures, limited kernel size, and target-specific semantics constrain scalability.

Overview

Axon is a synthesizing superoptimizer for tile-based AI accelerator programs. Given a tensor program expressed in a NumPy-like interface, it uses program synthesis to lower high-level tensor operations to target ISA instructions, explores all semantically equivalent program variants across algebraic transformations, tiling configurations, instruction selection, and fusion, and selects the best-performing kernel empirically on the target hardware (2606.26344). The system is evaluated on Amazon's Trainium against the production Neuron compiler, hand-optimized NKI kernels, and Mirage, a state-of-the-art search-based compiler.

The central design principle is that no optimization decision is made heuristically or prematurely. Instead of hand-crafted rewrite rules (as in TASO, Tensat, and Constable) or hand-written code-generation templates (as in Mirage), Axon derives every transformation from declarative semantics specifications and proves each one correct with an SMT solver over unbounded tensors. All valid variants are retained simultaneously in a persistent internal representation called the $\nu$Graph until concrete input shapes are known, at which point candidates are emitted, compiled, executed, and ranked by measured latency.

Motivation: algebraic transformation without rewrite rules

The motivating example is an RMSNorm followed by matrix multiplication, a pattern that appears before every linear projection in transformer layers such as those in LLaMA-3. Executed naively on Trainium, the normalization must complete before the matmul begins, underutilizing the chip's parallel compute engines. Axon discovers that because $\text{rms}(X)$ reduces along the last dimension of $X$—the same dimension contracted by matmul—the element-wise scaling commutes with the multiplication: $(X \odot \text{rms}(X)) \cdot W$ can be rewritten as $(X \cdot W) \odot \text{rms}(X)$. The transformed graph exposes two independent subgraphs that execute concurrently on different engines (Tensor Engine for matmul; Vector/Scalar Engines for RMS).

Notably, existing rewriting systems do not support this transformation. TASO, Tensat, and Constable rely on 91 hand-implemented rules, none covering this case; Mirage supports this particular rewrite but fails on the analogous Softmax+MatMul transformation. The authors attribute this to a structural limitation of graph rewriting systems generally: they match syntactic patterns rather than reasoning about operator semantics. Axon's operator propagation instead attempts swaps between adjacent operators and admits a swap only when an SMT solver proves semantic equivalence over all valid inputs, requiring no pre-enumerated rules.

The νGraph representation

A $\nu$Graph is a set of semantically equivalent dataflow graph variants together with per-operation metadata that accumulates as the graph is lowered through the pipeline. It extends equality saturation's e-graph idea from equivalence classes of sub-expressions within a term to equivalence classes of entire dataflow graphs. Each pipeline stage expands the variant set multiplicatively—for example, 2 algebraic variants × 4 tiling configurations × 3 instruction choices yield 24 candidate programs—and extraction instantiates symbolic parameters and emits concrete programs.

This design defers all choices to empirical selection, avoiding greedy sequential decisions that foreclose later opportunities. The cost is combinatorial growth: complex kernels such as GQA and Gated MLP produce roughly 30K candidates, evaluable in about two hours as a one-time cost. The approach is tractable only because ML kernels typically contain fewer than ten operators per layer; the paper does not claim scalability beyond kernel-sized programs.

Pipeline stages

Operator propagation. A worklist algorithm attempts to propagate each operator downward through the computation graph by swapping it with successors, using progressively increasing "clone degree" (how many of the operator's inputs pass through its successor). Each candidate swap is checked for semantic equivalence; saturation occurs when no further swaps are possible. Downward propagation suffices because pushing operators down necessarily moves others up.

Symbolic tiling. Each operation is annotated with symbolic strip ($[n_0,n_1]$), block ($[b_0,b_1]$), and tile ($[t_0,t_1]$) dimensions, reflecting the two-level HBM/on-chip-SRAM hierarchy. Block sizes range from 1 to 32 tiles (bounded by SBUF capacity); strip dimensions along non-reduction axes are shared across operators to enable later fusion. After algebraic transformation, parallel branches receive independent tiling parameters optimized for their respective engines.

ISA synthesis. A two-phase sketch-driven synthesizer lowers each tiled 2-D operation to target instructions. Phase 1 builds an instruction pool filtered by shared constituent operations; Phase 2 recursively fills sketch holes and checks complete candidates for provable equivalence. On Trainium, a depth limit of 3 suffices because the ISA is coarse-grained. For example, the synthesizer discovers that matmul(x, y) lowers to nc_matmul(nc_transpose(x), y), since nc_matmul consumes its first operand transposed. Semantics are specified at two levels: hardware-agnostic tensor operator semantics (shared across targets) and per-target ISA semantics that additionally encode hardware tile constraints (e.g., $K \leq 128$, $\text{rms}(X)$0, $\text{rms}(X)$1 for nc_matmul).

ISA fusion. A second round of algebraic transformations operates at the instruction level, necessary because synthesis introduces new operations (e.g., transposes) invisible to the original graph. Instruction fusion replaces instruction sequences with fused equivalents via depth-1 synthesis—for instance, discovering that $\text{rms}(X)$2, eliminating a layout transformation that Mirage cannot perform. Operator/loop fusion merges adjacent nodes sharing strip dimensions when no reduction intervenes, including nodes from independent subgraphs for deep loop-nest fusion.

Code emission. Concrete NKI programs are extracted by instantiating symbolic parameters, compiled by the Neuron compiler (which handles SBUF/PSUM allocation), executed on hardware, and validated numerically against a reference implementation on random inputs ($\text{rms}(X)$3 on FP32). Because Axon does not control on-chip buffer allocation directly, some symbolically feasible candidates are rejected by the backend compiler.

Correctness guarantees

Equivalence checking adapts TensorRight's methodology, originally built to verify XLA rewrite rules, to Axon's transformations. Tensors are encoded as uninterpreted functions over symbolic indices, so proofs hold for tensors of arbitrary size with known rank—dynamic shapes are supported naturally. Element-wise and layout operations are checked directly; reductions and matmuls use arithmetic axioms such as distributivity of scalar multiplication over summation; nonlinear functions (exp, sigmoid) are treated as uninterpreted, so swaps involving them are conservatively rejected. Solver timeouts also result in conservative rejection.

Two limitations are stated plainly. First, checking uses real-valued arithmetic rather than floating-point theory: swapping multiply past matmul takes 0.15 s with reals versus 247.75 s with floats—a 1650× slowdown—so floating-point rounding differences are not covered by the proof and are instead handled by random-input testing. Second, conservative treatment of nonlinear functions may reject valid transformations. Relative to Mirage's probabilistic equivalence checking, which provides statistical guarantees over finite inputs, Axon's proofs hold over all inputs, though only in the real-arithmetic model.

Evaluation results

The evaluation covers 20 benchmarks (11 individual operators, 9 multi-operator kernels including GQA and Gated MLP) on a single NeuronCore of a trn1.32xlarge instance, comparing against Neuron compiler 2.21, hand-optimized NKI samples, and Mirage's NKI extension.

Against the Neuron compiler, Axon achieves geomean speedups of 1.23×–1.90× on individual operators, peaking at 3.7× (SiLU at 16384×16384). Gains grow with tensor size, attributed to better temporal locality and DMA utilization than the compiler's fixed tiling heuristics. On multi-operator kernels, speedups reach 19× on Transpose+MatMul at 8192×8192×8192 (geomean 1.32×), 2.39× on QKV Projection (geomean 1.24×), and 1.64× on GQA (geomean 1.14×); RMSNorm+MatMul achieves up to 1.47× (geomean 1.11×).

Against hand-optimized NKI kernels (available only for Cumsum and RoPE), Axon matches expert code on most configurations and outperforms it by up to 1.35×, with geomean speedups of 1.07× and 1.06× respectively.

Against Mirage, Axon achieves geomean speedups of 10% on RMSNorm+MatMul, 6% on Softmax+MatMul, and 10% on Matmul+Matmul+SiLU+Mul—the four benchmarks Mirage's NKI extension supports—and synthesizes kernels for 16 additional benchmarks Mirage cannot handle. The paper also observes that Mirage's hand-crafted templates sometimes produce code slower than the Neuron compiler baseline, whereas Axon consistently matches or exceeds it.

The ablation study isolates contributions: disabling fusion causes geomean slowdowns of 35–50% across four multi-operator benchmarks, making it the most impactful optimization; disabling algebraic transformations causes 15–35% slowdowns. The two are complementary, targeting memory traffic and exposed parallelism respectively.

Compilation overhead is dominated by Phase 2 (backend compilation and execution, >97% of total time): totals range from ~5 minutes for simple operators to ~1.7 hours for SwiGLU Gated MLP, with 60–76% of emitted candidates compilable. This is a one-time cost amortized by caching, comparable to Mirage's reported search times of up to 4 hours and to schedule-based compilers like Ansor.

Limitations and open questions

The paper concedes several boundaries explicitly. Porting to a new target requires new ISA semantics, hardware constraints, and a code emitter; accelerators with fundamentally different execution models—many-core designs like Meta's MTIA and Microsoft's Maia, or GPU thread-block models as exposed by Triton—are not addressed by the current design. Evaluation is limited to first-generation Trainium due to limited availability of newer generations; transfer to Trainium2/3 is expected based on architectural similarity but not demonstrated. Benchmarks exclude convolutions, sparse operators, and training-specific backward passes, and all measurements are kernel-level rather than end-to-end inference latency. The multiplicative variant explosion is tractable only for small kernels, and the reliance on the backend compiler for buffer allocation means a substantial fraction of candidates (24–40%) waste compilation effort. Whether the real-arithmetic verification plus random testing methodology suffices for kernels where numerical drift compounds across fused operations remains an open question the paper does not resolve.

Conclusion

Axon demonstrates that a semantics-driven synthesis pipeline—operator propagation verified by SMT over unbounded tensors, sketch-driven ISA lowering, symbolic tiling, and two-level fusion, unified in a $\text{rms}(X)$4Graph—can automatically produce Trainium kernels that exceed a mature production compiler by up to 19×, surpass hand-tuned NKI code, and modestly outperform Mirage while covering far more benchmarks. Its principal contribution is replacing hand-crafted rewrite rules and templates with provable, specification-derived transformations, at the cost of a one-time empirical search whose tractability is currently bounded to kernel-sized programs on engine-parallel accelerator architectures.