Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators

TL;DR

Autocomp is an innovative LLM-driven approach that automates and enhances code optimization for tensor accelerators, significantly improving performance across multiple hardware platforms by leveraging domain knowledge and hardware feedback.

Contribution

The paper introduces Autocomp, a novel method combining structured prompts, domain knowledge, and hardware feedback to automate tensor accelerator code optimization, outperforming existing solutions.

Findings

Autocomp-optimized code is 5.6x faster than vendor libraries.

Outperforms expert hand-tuned code by 1.9x.

Achieves 3.8x higher performance than ML-based cost models.

Abstract

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.