arrow
Return

Tensor Program Superoptimization through Cost-Guided Symbolic Program Synthesis

delete2026-01-01
delete0
PRE
AI
A
Alexander Brauckmann *
A
Aarsh Chaube
J
José Wesley de Souza Magalhães
E
Elizabeth Polgreen
M
Michael F. P. O’Boyle
DOI:10.1109/CGO68049.2026.11395198delete
deleteOriginal
deleteOriginal request for help
deleteShare
deleteSave
Abstract

Abstract

En 中文
Modern tensor compiler frameworks like JAX and PyTorch accelerate numerical programs by compiling or mapping Domain-Specific Language (DSL) code to efficient executables. However, they rely on a fixed set of transformation rules and heuristics, which means they can miss profitable optimization opportunities. This leaves significant optimization potential unused for programs that fall outside these fixed patterns. This paper presents STENSO, a tensor DSL program superoptimizer that discovers such missing rewrites. STENSO's core is a symbolic program synthesis based search algorithm that systematically explores the space of equivalent programs. By combining symbolic execution and sketch-based program synthesis, it generates equivalent candidate implementations. To make the search computationally tractable, STENSO further integrates a cost model with a branch-and-bound algorithm scheme. This effectively prunes the search space, arriving at optimal solutions in a reasonable time. We evaluate STENSO on over 30 benchmarks. The discovered programs achieve geometric mean speedups of 3.8x over NumPy and 1.6x over state-of-the-art compilers like JAX and PyTorch-Inductor. These results underscore the limitations of heuristic-based compilation and demonstrate STENSO's effectiveness in finding such optimizations automatically.
Keywords:
Superoptimization
Tensor Programs
Program Synthesis
Compiler Optimization

Journal

2
2026 IEEE/ACM INTERNATIONAL SYMPOSIUM ON CODE GENERATION AND OPTIMIZATION, CGO
IF:
0
Papers:
48
Citations:
0

Organization

U
University of Edinburgh
Scholars:
5.2W
Papers: 4.6W
Citations: 71