Source-linked AI summary
Stochastic Superoptimization
Eric Schkufza, Rahul Sharma, Alex Aiken
TL;DR
Superoptimization seeks highly performant code, but conventional decomposition and complete search approaches are limited by dependent decisions and immense program spaces. The paper formulates loop-free binary optimization as stochastic cost minimization and uses MCMC to explore candidate programs. STOKE expands the searchable program space and produces rewrites that can match or outperform production compilers and, in some cases, handwritten assembly.
Problem
Conventional compiler decomposition and complete superoptimization limit the ability to search broadly for highly performant loop-free binary programs.
Method
The paper encodes transformation correctness and performance improvement in a cost function and uses MCMC to sample low-cost rewrites from the loop-free assembly program space.
Results
STOKE produces rewrites that in many cases match or outperform production compilers, with some results also outperforming expert handwritten assembly.
Takeaways & Limitations
The approach substantially expands the program space STOKE can reason about while improving the quality of its generated rewrites.
Takeaways & Limitations
STOKE is limited to loop-free fragments and can fail when incorrect alternatives differ from the correct result by only a small amount.
Abstract
from arXiv · showhide
We formulate the loop-free, binary superoptimization task as a stochastic search problem. The competing constraints of transformation correctness and performance improvement are encoded as terms in a cost function, and a Markov Chain Monte Carlo sampler is used to rapidly explore the space of all possible programs to find one that is an optimization of a given target program. Although our method sacrifices com- pleteness, the scope of programs we are able to reason about, and the quality of the programs we produce, far exceed those of existing superoptimizers. Beginning from binaries com- piled by llvm -O0 for 64-bit X86, our prototype implemen- tation, STOKE, is able to produce programs which either match or outperform the code sequences produced by gcc with full optimizations enabled, and, in some cases, expert handwritten assembly.
1. INTRODUCTION
Stochastic superoptimization reformulates loop-free binary optimization as incomplete search over programs, jointly balancing correctness and performance. STOKE uses this approach to explore a broader program space and produce code that can outperform optimized compilers and expert assembly.
- Motivation: Traditional compiler optimization can miss the best code because instruction selection, register allocation, and target-dependent optimization are mutually dependent.Factoring optimization into independently solvable subproblems also creates the phase ordering problem.
- Limitations of prior approaches: Complete superoptimizers guarantee finding a satisfying program but are limited by enormous search spaces or dependence on expert-written rules.Explicit enumeration restricts program length, while implicit enumeration restricts the search space through expert knowledge.
- Approach: The proposed approach encodes correctness and speed as terms in a cost function over loop-free executable instruction sequences and minimizes that cost stochastically.The search sacrifices completeness in exchange for exploring a larger space of programs.
- Approach: MCMC sampling explores the irregular search space by producing low-cost program samples, rather than relying on exact optimization.The sampler maintains a current rewrite, proposes modifications, and accepts or rejects them until its computational budget is exhausted.
- Results: Starting from llvm -O0 code, STOKE found an OpenSSL Montgomery multiplication kernel 16 lines shorter and 1.6 times faster than gcc -O3.The discovered algorithm was slightly faster than expert handwritten assembly and was automatically verified equivalent to the original code.
2. RELATED WORK
Prior superoptimizers explore restricted program classes using enumeration or equality-preserving transformations, achieving completeness or scalability within limited domains. STOKE instead applies incomplete MCMC search directly to a high-dimensional space of loop-free assembly programs.
- Prior superoptimization: Earlier superoptimizers enumerate programs explicitly or implicitly within restricted classes and evaluate how well multiple programs satisfy performance constraints.These approaches differ from traditional compilers, which drive generation of a single program using performance constraints.
- Enumeration: Massalin’s system enumerated instruction sequences of increasing length but optimized sequences up to length 12 only after restricting enumerable opcodes to 10–15.STOKE’s 11-instruction kernel was found while considering a large subset of the hardware instruction set.
- Equality-based search: Equality-based systems such as Denali and Equality Saturation improve scalability by exploring only programs known to equal the input, but rely heavily on expert transformation knowledge.They apply equality-preserving transformation axioms to guide candidate exploration.
- Database search: Bansal’s database superoptimizer used register-renaming symmetries to support input sequences up to length 6 and output sequences up to length 3 without expert knowledge at query time.Its offline search hides the high cost of superoptimization and makes database queries inexpensive.
- STOKE: STOKE differs by using incomplete stochastic search and MCMC to explore the extremely high-dimensional, irregular space of loop-free assembly programs.The paper presents MCMC as a tractable general method for many optimization problems of this form.
3. COST MINIMIZATION
The paper formulates optimization as minimizing a cost that balances correctness preservation with performance, then uses MCMC to search irregular program spaces. Metropolis-Hastings proposals accept better rewrites and sometimes worse ones, helping avoid local minima.
- Cost function: Correctness and performance are encoded as separate terms in a cost function over candidate rewrites.The correctness term measures functional equivalence, while performance can represent runtime, code size, disk accesses, or power.
- Cost function: The correctness metric is zero exactly when target and rewrite functions produce identical live-output side effects for matching live inputs.The paper treats code sequences as functions of registers and memory contents and notes that undefined behavior need not be handled.
- MCMC sampling: MCMC is used because the resulting cost functions are highly irregular and unsuitable for exact optimization.Sampling increasingly concentrates on lower-cost regions and behaves practically like hill climbing that is robust to local minima.
- MCMC sampling: Metropolis-Hastings maintains a current rewrite, proposes a modified rewrite, and accepts or rejects it until the computational budget is exhausted.Ergodic proposals allow the chain to transform any point in the search space into any other point.
- MCMC sampling: Proposal distributions combine local modifications with global changes, while symmetric proposals simplify acceptance to a ratio computed directly from cost.This mixture supports both incremental exploration and major transformations of the current rewrite.
- MCMC sampling: Better proposals are always accepted, whereas worse proposals may be accepted with decreasing probability to prevent trapping in local minima.The acceptance probability becomes lower as the proposed rewrite is increasingly worse relative to the current one.
4. X86 BINARY OPTIMIZATION
STOKE makes binary optimization practical by replacing expensive exact checks with testcase-based correctness and runtime approximations inside an MCMC search. It further improves exploration through early termination, relaxed equality, and separate synthesis and optimization phases, while retaining important limits on validation throughput and partially correct computations.
- 4.1 Transformation Correctness: 64-bit X86 optimization is implemented with testcase-based correctness checks because symbolic validation throughput is too low for MCMC.The emulator performs just under 500,000 testcase executions per second, enabling useful approximate equality evaluation.
- 4.1 Transformation Correctness: The testcase approximation measures Hamming distance between live outputs and provides a smoother search landscape than a binary symbolic equality test.Register and memory discrepancies are counted over live outputs, while undefined behavior is penalized separately.
- 4.2 Performance: Runtime performance is approximated with static average instruction latency because compiling and repeatedly executing candidates is too expensive inside MCMC.The heuristic correlates well with actual runtimes, though outliers arise from unusually high or low micro-op-level instruction parallelism.
- 4.4 Separating Synthesis From Optimization: STOKE separates correctness-focused synthesis from subsequent optimization because local MCMC can rarely cross the low-probability path to a distinct expert-code region.The improved equality metric also permits correct values in arbitrary locations, implicitly exploring multiple computations in parallel.
- 4.5 Optimized Acceptance Computation: Early termination raises testcase throughput by stopping proposal evaluation once its running cost exceeds a rejection bound.During Montgomery multiplication synthesis, peak throughput exceeds 50,000 testcase evaluations per second.
- 4.7 Why and When Synthesis Works: The approach remains limited when useful partial correctness cannot be detected, such as computations reduced to a single boolean value.If synthesis fails, optimization can still proceed from the target’s region rather than discovering a separate equivalent-program region.
5. STOKE
STOKE optimizes loop-free 64-bit X86 binaries through parallel synthesis and optimization, using generated testcases, instrumentation, MCMC search, validation, and runtime-based reranking.
- System organization: STOKE accepts a target binary and annotated driver, generates random inputs and testcases, and distributes synthesis and optimization across small clusters.Candidates are returned after fixed computational budgets and subsequently reranked by actual runtime.
- Testcase generation: STOKE limits generated inputs to fixed-width bit strings and requires annotations for inputs used to form legal memory addresses.By default, it generates 32 testcases per target and sandboxes candidate memory accesses.
- Transformation validation: STOKE validates loop-free assembly equivalence by translating code sequences into SMT bit-vector formulas and querying whether live outputs can differ.The validator relates memory values at equal addresses and returns counterexamples when inequivalence is found.
- Validation assumptions: STOKE assumes stack addresses are constant offsets from the stack pointer and simplifies 64-bit multiplication and division to keep validation tractable.The stack-address assumption reduces expensive memory constraints, especially for llvm -O0 code with heavy stack traffic.
- Computational budget: Synthesis and optimization run in parallel on a cluster of 40 dual-core 1.8 GHz AMD Opterons, each with a 30-minute computational budget.The MCMC parameters for both phases are summarized in Figure 11.
6. EVALUATION
Across benchmarks from cryptography, linear algebra, and low-level programming, STOKE generally matches or exceeds production compiler performance and sometimes discovers algorithmically distinct rewrites, while exhibiting clear scope and synthesis limitations.
- Overall evaluation: STOKE consistently discovers rewrites matching gcc and icc with full optimizations, and in several cases exceeds both at performance comparable to expert handwritten assembly.Figure 10 summarizes performance improvements over llvm -O0, while Figure 12 reports the runtimes required to produce them.
- 6.1 Hacker’s Delight: For “Cycle Through 3 Values,” STOKE rediscovers the natural conditional-assignment algorithm that gcc and icc fail to recognize from the esoteric implementation.The five-line rewrite remains beyond brute-force enumeration-based superoptimizers.
- 6.1 Hacker’s Delight: STOKE replaces four-part multiplication in “Compute the Higher Order Half of a 64-bit Product” with a single multiplication using the appropriate bit-width intrinsic.The benchmark also includes typical superoptimizer rewrites such as using the popcnt intrinsic.
- 6.2 SAXPY: For SAXPY, STOKE discovers an SSE-vectorized rewrite that broadcasts the scalar and processes four elements at a time, unlike gcc -O3.The manually unrolled input and annotations expose aligned, non-aliasing memory accesses for this optimization.
- 6.3 Limitations: For linked-list traversal, STOKE removes stack traffic and replaces multiplication with bit shifting but fails to cache the head pointer across the loop.Because STOKE optimizes only the loop-free inner fragment, its rewrite is slower than gcc -O3.
- 6.3 Limitations: STOKE cannot synthesize rewrites for three Hacker’s Delight benchmarks whose correct results differ by one bit from simple incorrect alternatives.Its optimization phase nevertheless finds code comparable to production compiler output for those cases, motivating a more sophisticated cost function.
7. CONCLUSION AND FUTURE WORK
The paper reformulates loop-free binary superoptimization as stochastic search, using cost minimization and MCMC to balance correctness with performance. Although completeness is sacrificed, the approach expands the searchable program scope and improves rewrite quality, while STOKE still has substantial room for improvement.
- The framework reformulates program optimization as stochastic search over competing correctness and performance requirements.These requirements are represented as terms in a cost function.
- An MCMC sampler rapidly explores the cost function and produces low-cost samples corresponding to high-quality code sequences.
- Sacrificing completeness enables the system to reason about a broader program scope and produce higher-quality rewrites than existing superoptimizers.
- STOKE remains a prototype with substantial room for improvement despite often producing rewrites competitive with or better than production compiler output.Future work targets loop-containing code and synthesis costs robust to deceptively attractive but incorrect alternatives.