Source-linked AI summary
Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
Charles Hong, Sahil Bhatia, Alvin Cheung, Yakun Sophia Shao
TL;DR
Programming tensor accelerators remains difficult because high-performance code requires hardware-specific expertise, while existing compilers and data-driven methods have limited backend coverage or scarce training data. Autocomp uses portable, iterative LLM-driven search with structured planning, domain-specific optimization menus, and hardware correctness and performance feedback. Across three platforms, it outperforms vendor, expert, and machine-learning-based baselines, while its schedules can be reused for similar operations.
Problem
High-performance tensor accelerator code is difficult to write, and existing compilers, DSLs, and data-driven methods have limited portability or rely on scarce accelerator performance data.
Source: §1 · p. 1Method
Autocomp performs iterative LLM-driven optimization using structured prompts, an adaptable optimization menu, and hardware correctness and performance feedback.
Source: §1 · p. 2Results
Autocomp outperforms prior approaches across three hardware platforms, including 5.6× faster Gemmini GEMMs than the vendor library, 1.9× faster Trainium operators than hand-optimized code, and 3.8× faster L40S kernels than TVM MetaSchedule.
Source: §1 · p. 2Takeaways & Limitations
Autocomp-generated schedules can be reused for similar tensor operations, delivering up to 24% greater speedups under a fixed sample budget.
Source: §1 · p. 2
Abstract
from arXiv · showhide
Hardware accelerators, especially those designed for tensor processing, have become ubiquitous in today's computing landscape. However, even with significant efforts in building compilers, programming these tensor accelerators remains challenging, leaving much of their potential underutilized. Recently, large language models (LLMs), trained on large amounts of code, have shown significant promise in code generation and optimization tasks, but generating low-resource languages, such as specialized tensor accelerator code still poses a significant challenge. We tackle this challenge with Autocomp, an approach that empowers accelerator programmers to leverage domain knowledge and hardware feedback to optimize code via an automated LLM-driven search. We accomplish this by: 1) formulating each optimization pass as a structured two-phase prompt, divided into planning and code generation phases, 2) inserting domain knowledge during planning via a concise and adaptable optimization menu, and 3) integrating correctness and performance metrics from hardware as feedback at each search iteration. Across three distinct hardware platforms, we demonstrate that Autocomp-optimized code runs 5.6x faster than the vendor-provided library (Gemmini), outperforms expert-level hand-tuned code by 1.9x (AWS Trainium), and achieves 3.8x higher performance than a machine learning-based cost model for GPUs (NVIDIA L40S). Additionally, we demonstrate that optimization schedules generated from Autocomp can be reused across similar tensor operations, improving speedups by up to 24% under a fixed sample budget.
1 INTRODUCTION
Autocomp addresses the difficulty of optimizing low-resource tensor accelerator code by combining LLM-driven search with domain knowledge and hardware feedback. Across three hardware platforms, it produces faster code than vendor libraries, expert hand-tuned code, and machine-learning-based scheduling methods.
1 INTRODUCTION
- Autocomp targets a gap left by compilers and DSLs that support few hardware backends and require substantial engineering or manual tuning for new accelerators.Data-driven approaches also depend on performance data that is scarce for domain-specific accelerators.Source: §1 · pp. 1–2
- Autocomp uses an iterative LLM-driven search that selects optimizations, generates DSL code, validates correctness, and benchmarks candidates for feedback.Its prompts encode DSL syntax, optimization rules, and performance feedback.Source: §1 · p. 2
- Structured prompts make Autocomp portable: adapting to new hardware requires changing the optimization menu and ISA description rather than writing a complex compiler backend.The framework was applied to three distinct hardware backends.Source: §1 · p. 2
- Autocomp generates code that outperforms baselines across three hardware platforms, including vendor libraries, expert hand-tuned code, and ML-based scheduling.It also reuses generated schedules for similar tensor operations, improving speedups by up to 24% under a fixed sample budget.Source: §1 · p. 2
- 5.6× faster Gemmini GEMMs than the vendor-provided library and 1.4× faster than expert-level hand-tuned code demonstrate Autocomp’s performance on Gemmini.Source: §1 · p. 2
- On AWS Trainium, Autocomp reduces the runtime of a wide range of hand-optimized tensor operators by an average of 1.9×.Source: §1 · p. 2
- On an NVIDIA L40S GPU, Autocomp’s KernelBench kernels run 2× faster than PyTorch and 3.8× faster than TVM MetaSchedule.Source: §1 · p. 2
2 BACKGROUND
Tensor accelerator optimization centers on explicit data movement, configuration state, and overlap between computation and data movement. Existing abstractions and LLM optimization methods leave important specialized-hardware and instruction-level challenges unresolved.
2.1 Code Optimization for Tensor Accelerators
- Tensor accelerator optimization minimizes data movement, manages configuration state, and schedules operations to overlap computation with data movement.Because fixed-size matrix instructions are often fixed, optimization focuses on surrounding software behavior.Source: §2.1 · pp. 2–3
- Hoisting and software pipelining can improve accelerator execution but require coordinated changes to loop nests, pointers, and indices in low-resource DSLs.Source: §2.1 · p. 3, Figure 3 · p. 3, Figure 4 · p. 3
- Existing tensor-accelerator abstractions capture aspects such as data movement but neglect implementation-specific and instruction-level optimizations.LLM-based code generation can directly rewrite accelerator code across a broader optimization space.Source: §2.1 · p. 3
2.2 LLM-Based Code Optimization
- Prior LLM code-optimization work uses evolutionary search, retrieval augmentation, compiler feedback, and model post-training across code-related tasks.Source: §2.2 · p. 3
- Specialized-hardware code optimization remains underexplored: zero-shot generation for such languages is highly unreliable, while Autocomp combines techniques to optimize accelerator code.The cited prior work targets system-level programming such as CUDA and SIMD intrinsics rather than specialized hardware.Source: §2.2 · p. 3
3 THE AUTOCOMP APPROACH
Autocomp addresses low-resource tensor accelerator code generation through structured planning, implementation, and hardware-guided search. It further improves practicality by increasing search diversity and reusing schedules across similar tensor operations.
3.1 Rationale
- Direct LLM rewriting is inadequate because tensor accelerator DSLs are low-resource languages and provide little guidance about which optimizations to apply or their order.Autocomp addresses these difficulties with decomposed planning and implementation rather than a single rewrite request.Source: §3.1 · p. 3
3.2 Plan-then-Implement: Two-Phase Optimization
- Autocomp separates optimization into plan generation and code implementation, using an iterative beam search to explore multiple optimization trajectories.Each plan selects one menu optimization and describes its transformation before the LLM generates functionally equivalent code.Source: §3.1 · p. 4, §3.2 · p. 4
- Hardware feedback combines accelerator instructions, current-code performance, utilization metrics, optimization options, rules, and search-iteration context.The planning prompt uses latency and, for Gemmini, scratchpad and accumulator utilization to guide subsequent optimization choices.Source: §3.2 · p. 4
3.3 Beam Search
- Beam search retains only functionally correct candidates that improve parent latency, then keeps the best B candidates for the next iteration.Candidates are tested against a reference implementation and measured through cycle-accurate simulation or on-chip performance measurement.Source: §3.3 · p. 5
3.4 Increasing Plan and Code Diversity
- Plan and code diversity is increased through optimization-menu dropout and LLM ensembling, while preventing repeated selection of the same optimization.The authors report that these techniques, along with other components, are ablated in Appendix B.Source: §3.4 · p. 5
3.5 Schedule Reuse
- Schedule reuse reduces new-workload search by applying optimization choices from a recorded schedule before a smaller full-search refinement.Reuse targets GEMMs with the same aspect ratios or two shared dimensions, reducing LLM calls and search time through smaller beam widths and sample counts.Source: §3.5 · p. 5
4 EVALUATION
Autocomp is evaluated across Gemmini, AWS Trainium, and NVIDIA L40S, with portability achieved by changing prompts rather than implementing compiler backends. Across workloads, it outperforms software libraries, hand-optimized implementations, and TVM-based scheduling while using hardware feedback and domain-specific optimization menus.
4.2 Discussion on Autocomp’s Portability
- Autocomp was ported to Trainium and GPUs by updating the optimization menu and ISA description, requiring less than a day of effort per backend.No complex compiler backend code was needed for the new platforms.Source: §4.2 · p. 6
4.3 Baselines
- The evaluation compares high-level libraries, unoptimized and hand-optimized low-level code, TVM MetaSchedule, and Gemmini hardware FSMs where available.Baseline availability differs across hardware platforms.Source: §4.3 · pp. 6–7, Table 1 · p. 7
4.4.1 Matrix multiplication (GEMM)
- 5.6×, 2.9×, and 1.4× are Gemmini GEMM speedups over the software library, unoptimized Exo code, and hand-optimized Exo code, respectively.Autocomp reaches at least 85% hardware-FSM utilization and 91% on average, enabled by software pipelining, double buffering, tiling, and loop-order exploration.Source: §4.4.1 · p. 7, Figure 7 · p. 8
4.4.2 Convolution
- 1.1× geomean and up to 1.2× are Autocomp’s convolution speedups over hand-optimized Exo code, while it reaches 2.6× over Gemmini’s software library.The smaller opportunity than GEMM reflects higher arithmetic intensity and stronger baseline utilization for convolution.Source: §4.4.2 · p. 8, Figure 8 · p. 8
4.4.3 Fine-Grained Linear Algebra
- 1.6× is Autocomp’s forward-pass speedup over an expert hardware FSM for fine-grained linear algebra, with 2.7× geomean speedup over unoptimized code.Shared-load hoisting, fine-grained software pipelining, and removal of blocking operations exploit software flexibility.Source: §4.4.3 · p. 8, Figure 9 · p. 8
4.5 Trainium Evaluation
- 1.9× is Autocomp’s geomean improvement over expert-optimized Trainium advanced workloads, while tutorial workloads improve 1.36× over hand-optimized code and 13.52× over PyTorch-compiled code.A 17.37× speedup is reported for 1D depthwise convolution through tile-size, loop-order, and reuse transformations.Source: §4.5.2 · p. 9, §4.5.1 · p. 9, Figure 10 · p. 9
4.6 NVIDIA L40S GPU Evaluation
- 2.05× is Autocomp’s geomean speedup over PyTorch on every selected L40S benchmark, whereas TVM averages 0.54× PyTorch and beats Autocomp only once.Autocomp explores Tensor Core implementations through PyTorch APIs and CUDA libraries while satisfying accuracy checks.Source: §4.6 · pp. 9–10, Figure 12 · p. 10
5 IMPROVING SAMPLE EFFICIENCY WITH SCHEDULE REUSE
Autocomp reuses schedules from structurally similar GEMMs to improve optimization under fixed sample budgets. Across reuse targets, schedule reuse consistently improves speedups over full search and reduces search cost.
Reuse setup
- Schedules optimized for a base GEMM are reused on structurally similar square, column-dominant, and row-dominant GEMMs.The base benchmark uses full search, while reuse targets share an aspect ratio or two dimensions.Source: §5 · p. 10, Table 2 · p. 10
Schedule reuse results
- For a 512×512×512 square GEMM reused from 1024×1024×1024, reuse consistently delivers better speedups with the same number of samples.Figure 13 shows the corresponding optimization trace comparison.Source: §5 · p. 10, Figure 13
- With 100 samples, schedule reuse achieves 4.6× geomean speedup over Gemmini’s software library, versus 3.7× without reuse.At 200 calls, reuse reaches 5.0× versus 4.2× without reuse.Source: §5 · p. 10
6 CONCLUSION
The paper presents Autocomp as an LLM-based flow for automatically optimizing low-resource accelerator code. Across three hardware platforms, it outperforms baselines on real-world benchmarks while requiring less manual effort than traditional compiler development or hand optimization.
Conclusion
- Across three hardware platforms, Autocomp generates code that significantly outperforms baselines on real-world benchmarks with less manual effort.The paper positions this flow as a potential component of the accelerator design process.Source: §6 · p. 10
- Autocomp automatically optimizes low-resource accelerator code at superhuman levels.Source: §6 · p. 10
- The approach is presented as a potential key component in the accelerator design process.Source: §6 · p. 10
A.1 Gemmini Experiments
The Gemmini appendix documents the benchmark sources used for the evaluation. It identifies the vendor library and Exo schedules used as comparison implementations.
Benchmark sources
- Gemmini benchmarks were uploaded to Google Drive for the evaluation.Source: §A.1 · p. 15
- Gemmini’s software-library baseline comes from the gemmini-rocc-tests repository.Source: §A.1 · p. 15
- Exo comparison code was generated using published schedules from Exo’s GitHub repository.Source: §A.1 · p. 15
A.2 Trainium Experiments
The Trainium evaluation uses tutorial and advanced operators with configurations drawn from reference materials. Implementations were adapted to run in isolation on a trn1.2xlarge instance, using original shapes where they could expose meaningful improvements.
Evaluation setup
- Trainium evaluation configurations are detailed in Tables 3 and 4 for tutorial and advanced operators.Source: Table 3 · p. 15, Table 4 · p. 15
- Reference code was copied from the nki-samples repository and modified slightly for isolated execution on a trn1.2xlarge instance.Source: §A.2 · p. 15
- The evaluation generally preserves original nki-samples shapes when examples exist and are large enough to show meaningful performance improvements.For causal self-attention with seq=16384, latency is measured between a query tile of length seqQ = 128 and key/value arrays of length seqKV = 16384.Source: §A.2 · p. 15
A.3 NVIDIA L40S GPU Experiments
The GPU evaluation uses KernelBench Level 1 operators, with benchmark shapes drawn from KernelBench v0.1 and exact PyTorch operators identified in its repository.
A.3 NVIDIA L40S GPU Experiments
- KernelBench Level 1 operators provide the GPU evaluation workload set.The selected benchmark details are listed in Table 5.Source: Table 5 · p. 15
- The evaluation uses shapes from KernelBench v0.1 and operators corresponding to the exact PyTorch implementations in the KernelBench repository.Source: §A.3 · p. 15
B ABLATION STUDIES
Ablations show that Autocomp depends on domain knowledge, diverse optimization proposals, hardware feedback, and model diversity, while remaining portable across prompts and robust across seeds.
B.1 Accelerator ISA
- Removing accelerator ISA information lowers correctness and prevents some instruction-level optimizations, although other features still improve the original code.Examples include first-compute handling and negative-scaled bias loading, which require ISA knowledge.Source: §B.1 · pp. 15–16, Table 6 · p. 16
B.2 Optimization Menu
- Removing the optimization menu completely deteriorates optimization performance and produces repetitive, less diverse plans.The menu supplies domain knowledge and optimization diversity during planning.Source: §B.2 · p. 16
B.3 Optimization Menu Dropout
- Menu dropout increases optimization-plan diversity by preventing models from concentrating on a limited set of options.Without dropout, models are biased toward a smaller subset of menu choices, and Table 6 reports a significant performance effect.Source: §B.3 · p. 16, Table 6 · p. 16
B.4 Hardware Performance Feedback
- Hardware performance feedback helps, but its effect can be limited because menu options already encode some information reflected by measured utilization metrics.Latency, scratchpad utilization, and accumulator utilization augment rather than replace the optimization menu.Source: §B.4 · p. 16, Table 6 · p. 16
B.5 LLM Ensembling
- Ensembling LLMs improves plan and code diversity, while individual models yield significantly lower performance than the ensemble.Alternative models still produce substantial speedups, demonstrating flexibility across LLM choices.Source: §B.5 · p. 16, §B.6 · pp. 16–17, Table 6 · p. 16
B.7 Variance Checking
- Across three seeds, the lowest-performing run remained within 11.8% of peak GEMM performance and 8.5% of peak convolution performance.The small spread indicates robust search behavior across repeated runs.Source: Table 6 · p. 16, §B.7 · p. 17
C PROMPTS
- The appendix presents hardware-specific ISA specifications, optimization menus, rules, and in-context examples for Gemmini, Trainium, and GPUs.These materials document how prompts are adapted across accelerator backends and optimization stages.Source: Figure 14 · p. 18, Figure 15 · p. 19, Figure 16 · p. 19
D.1 16x16 INT8 Gemmini—12544x64x256 GEMM
On a 12544x64x256 GEMM, Autocomp transforms unoptimized Gemmini code through diverse scheduling changes, achieving higher utilization than both unoptimized and hand-optimized implementations.
D.1 16x16 INT8 Gemmini—12544x64x256 GEMM
- 5.21× and 4.87× speedups follow accumulator and scratchpad double buffering, respectively, to overlap computation with data movement.The two transformations use alternating buffers for accumulators and A tiles.Source: §D.1 · p. 26
- 5.53× speedup follows loop fusion that removes redundant accumulator-zeroing loads by using preload overwrite behavior.The optimization eliminates the inner loop that loads zeros into the accumulator.Source: §D.1 · p. 26
- Autocomp distributes gains across many optimization steps rather than concentrating them in one or two transformations.The search applies arithmetic simplification, tiling, double buffering, software pipelining, first-compute handling, and loop transformations.Source: §D.1 · pp. 26–27
- Smaller tiles, explicit double buffering, and software pipelining improve overlap between data movement and computation while reducing CPU overhead.Autocomp uses 32×256 A tiles divided into two 16×256 tiles, unlike the 128×256 tiles in Exo Opt.Source: §D.1 · pp. 26–27
- 93% compute utilization is achieved by the Autocomp-optimized GEMM, versus 43% for hand-optimized code and 28% for unoptimized code.Figures 28–32 show the unoptimized, hand-optimized, and Autocomp-generated implementations.Source: Figure 28 · p. 28, Figure 29 · p. 29, Figure 30 · p. 30
D.2 4x4 FP32 Gemmini—TinyMPC Primal Update Forward Pass (Fine-Grained Linear Algebra)
For the TinyMPC primal-update forward pass, Autocomp improves low-level Gemmini code through loop and data-movement transformations while approaching the performance of, and surpassing, a hand-optimized hardware FSM implementation.
D.2 4x4 FP32 Gemmini—TinyMPC Primal Update Forward Pass (Fine-Grained Linear Algebra)
- 2.00× speedup follows moving CPU-based negation and addition onto the accelerator while keeping intermediate results resident in the accumulator.Negative scaling handles subtraction, and a dummy zero vector enables accumulation of the transformed d_i term.Source: §D.2 · p. 33
- 2.95× speedup follows loop fusion that removes unnecessary x[NHORIZON] computation while retaining reduced fence usage.The preceding 2.66× step unrolled the loop and aggressively removed fences; fusion keeps the fence reduction but undoes unrolling.Source: §D.2 · p. 33
- Autocomp’s main advantage over the hardware FSM is hoisting shared Kinf, Adyn, and Bdyn loads out of the loop.The coarse-grained FSM handles data movement in hardware, preventing this software-level reuse transformation.Source: §D.2 · p. 33
- The hardware FSM retains advantages in operator fusion, while Autocomp removes CPU fences and more completely hoists configuration overhead.Autocomp separately computes the d[i] addition, whereas the FSM fuses it into the bias; the FSM also requires fences after each matrix multiplication.Source: §D.2 · pp. 33–34
- Autocomp identifies the major optimization opportunities but handles subtraction of d[i] slightly suboptimally.The generated code remains well-commented, making its optimization sequence easier to inspect and reuse.Source: §D.2 · p. 34
- 15.7% theoretical maximum compute utilization is achieved by Autocomp, versus 10.9% for hand-optimized hardware FSM code and 5.7% for unoptimized code.Figures 33–37 compare the starting, hand-optimized, and Autocomp-generated implementations.Source: Figure 33 · p. 35, Figure 36 · p. 37, Figure 35 · p. 37