Axon: A Synthesizing Superoptimizer for Tensor Programs

2026-06-24

AI Summary

This work addresses the high barrier to entry in manually developing high-performance tensor kernel functions for AI accelerators, a process that traditionally demands deep expertise in tiling strategies, instruction selection, data layout optimization, and operator fusion. To alleviate this burden, the paper proposes an automated synthesis approach that eliminates the need for handcrafted rules by integrating program synthesis, SMT solving, and algebraic transformations of computation graphs. The method formally verifies semantic equivalence over unbounded tensors and systematically explores the space of tiling strategies and instruction/operator fusion under hardware constraints. This enables the automatic generation of kernels that are not only semantically correct but also highly efficient, significantly reducing memory traffic and improving execution performance while substantially lowering the development effort required.

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.

Problem

• tensor programs
• AI accelerators
• high performance kernels
• program synthesis
• operator fusion

Innovation

• program synthesis
• superoptimization
• tensor programs
• SMT solving
• operator fusion