2025-11-11T09:13:09.652713

Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators

Hong, Bhatia, Cheung et al.
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.
academic

Autocomp: Un Potente e Portabile Ottimizzatore di Codice per Acceleratori Tensoriali

Informazioni Fondamentali

  • ID Articolo: 2505.18574
  • Titolo: Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
  • Autori: Charles Hong, Sahil Bhatia, Alvin Cheung, Yakun Sophia Shao (UC Berkeley)
  • Classificazione: cs.PL cs.AI cs.AR cs.LG
  • Stato di Pubblicazione: Preprint. In revisione.
  • Link Articolo: https://arxiv.org/abs/2505.18574

Riassunto

Gli acceleratori hardware, in particolare quelli progettati specificamente per l'elaborazione tensoriale, sono ormai ubiqui negli ambienti computazionali contemporanei. Tuttavia, nonostante i considerevoli sforzi investiti nella costruzione di compilatori, la programmazione di questi acceleratori tensoriali rimane una sfida, lasciando gran parte del loro potenziale sottoutilizzato. Questo articolo propone Autocomp, un metodo per ottimizzare il codice attraverso una ricerca automatizzata guidata da LLM, consentendo ai programmatori di acceleratori di sfruttare la conoscenza del dominio e il feedback hardware. Il metodo si realizza attraverso tre tecniche chiave: 1) formulazione di ogni processo di ottimizzazione come un prompt strutturato in due fasi, diviso in fase di pianificazione e generazione di codice; 2) inserimento della conoscenza del dominio durante la pianificazione attraverso un menu di ottimizzazione conciso e adattabile; 3) integrazione del feedback di correttezza e metriche di prestazione dall'hardware in ogni iterazione di ricerca.

Contesto di Ricerca e Motivazione

Problemi Fondamentali

Le sfide principali nella programmazione di acceleratori tensoriali includono:

  1. Complessità di Programmazione: A differenza della programmazione CPU generica, gli acceleratori tensoriali richiedono la gestione esplicita del movimento dei dati, della configurazione dello stato e della pianificazione delle operazioni
  2. Costo di Adattamento del Compilatore: L'adattamento dei compilatori tradizionali a nuove piattaforme hardware richiede un lavoro ingegneristico sostanziale, con i costi di sviluppo software che rappresentano il 40-50% dei costi totali di sviluppo del nuovo hardware
  3. Problema di Pianificazione dell'Ottimizzazione: Esplosione combinatoria nel determinare quali ottimizzazioni applicare e in quale ordine
  4. Sfida dei Linguaggi a Bassa Risorsa: I set di istruzioni architettoniche (ISA) e i DSL degli acceleratori dedicati sono sottorappresentati nei corpora di addestramento degli LLM

Limitazioni dei Metodi Esistenti

  1. Compilatori Tradizionali: XLA, TVM, Triton supportano solo pochi backend hardware, principalmente CPU e GPU
  2. Approcci DSL: Halide, Exo forniscono primitive per esprimere calcoli tensoriali, ma l'onere dell'ottimizzazione rimane sul programmatore
  3. Metodi Guidati dai Dati: Richiedono grandi quantità di dati di prestazione per l'addestramento, estremamente scarsi per gli acceleratori hardware specifici del dominio
  4. Applicazione Diretta di LLM: La generazione di codice zero-shot è altamente inaffidabile per i linguaggi di acceleratori a bassa risorsa

Contributi Fondamentali

  1. Primo Metodo di Ottimizzazione del Codice per Acceleratori Tensoriali a Bassa Risorsa Guidato da LLM: Propone il framework Autocomp, specificamente progettato per acceleratori hardware dedicati
  2. Framework di Ottimizzazione Altamente Portabile: Adattabile a nuove piattaforme hardware attraverso la modifica dei prompt, riducendo significativamente i costi ingegneristici
  3. Prestazioni Eccezionali: Supera significativamente i metodi esistenti su tre diverse piattaforme hardware
  4. Meccanismo di Riutilizzo della Pianificazione: Dimostra che le pianificazioni di ottimizzazione possono essere riutilizzate tra operazioni tensoriali simili, migliorando l'efficienza del campionamento

Dettagli del Metodo

Definizione del Compito

Input: Codice non ottimizzato per acceleratore tensoriale Output: Codice funzionalmente equivalente ma ottimizzato per le prestazioni Vincoli: Mantenere l'equivalenza semantica, garantire la correttezza attraverso la verifica hardware

Architettura Fondamentale: Framework di Ottimizzazione in Due Fasi

Fase 1: Generazione del Piano (Plan Generation)

La struttura del prompt contiene:

  1. Descrizione ISA dell'Acceleratore: Semantica delle istruzioni, specifiche di indirizzamento della memoria, descrizione della struttura hardware
  2. Codice Attuale: Codice da ottimizzare
  3. Feedback di Prestazione: Metriche di latenza (cicli), utilizzo della memoria, ecc.
  4. Menu di Ottimizzazione: Opzioni di ottimizzazione di alto livello predefinite (come tiling dei cicli, riordinamento, fusione, ecc.)
  5. Informazioni di Iterazione di Ricerca: Numero di iterazione attuale, per guidare le scelte di ottimizzazione

Fase 2: Implementazione del Codice (Implementation)

La struttura del prompt contiene:

  1. Descrizione ISA dell'Acceleratore: Come nella Fase 1
  2. Codice Attuale: Come nella Fase 1
  3. Piano Generato: Piano di ottimizzazione specifico dall'output della Fase 1
  4. Esempi di Apprendimento in Contesto: Esempi di codice per ottimizzazioni complesse (come tiling)
  5. Istruzioni di Implementazione: Istruzioni in linguaggio naturale per applicare il piano e generare codice ottimizzato

Strategia di Ricerca a Fascio

  • Ricerca a fascio con larghezza B=6, esplorando in parallelo molteplici traiettorie di ottimizzazione
  • Filtraggio di Correttezza: Verifica dei candidati di codice attraverso suite di test funzionali
  • Selezione di Prestazione: Conservazione solo dei candidati con prestazioni superiori al nodo genitore
  • Ottimizzazione Iterativa: Processo di ricerca con budget fisso di T iterazioni

Innovazioni Tecniche Chiave

1. Tecniche di Aumento della Diversità

  • Dropout del Menu di Ottimizzazione: Rimozione casuale di opzioni di menu parziali durante la pianificazione (probabilità 70%)
  • Integrazione LLM: Distribuzione delle richieste tra più LLM per aumentare la diversità delle risposte

2. Integrazione del Feedback Hardware

  • Metriche di prestazione in tempo reale (latenza, utilizzo della memoria) guidano le scelte di ottimizzazione successive
  • Simulazione precisa ai cicli o misurazione delle prestazioni a livello di chip

3. Meccanismo di Riutilizzo della Pianificazione

  • Registrazione di sequenze di pianificazione di alta qualità
  • Riutilizzo di pianificazioni note per operazioni tensoriali simili (stesso rapporto di aspetto o dimensioni condivise)
  • Ulteriore ottimizzazione dopo ricerca leggera

Configurazione Sperimentale

Piattaforme Hardware

  1. Gemmini: Generatore di acceleratori open-source, supporta array sistolici e acceleratori tensoriali di stile vettoriale
  2. AWS Trainium: Acceleratore tensoriale commerciale ad alte prestazioni, utilizza Neuron Kernel Interface (NKI)
  3. NVIDIA L40S GPU: GPU moderna per data center, include Tensor Core dedicati

Benchmark

  • Gemmini: GEMM e convoluzione di ResNet-50, controllo predittivo del modello TinyMPC
  • Trainium: Operatori di deep learning a livello tutorial e avanzato (RMSNorm, LayerNorm, GEMM, Mamba, ecc.)
  • GPU: Benchmark KernelBench Level 1

Metodi di Confronto

  1. Librerie Software di Alto Livello: Libreria software Gemmini, PyTorch NeuronX, PyTorch
  2. Codice di Basso Livello Non Ottimizzato: Exo non ottimizzato, codice tutorial nki-samples
  3. Codice Ottimizzato Manualmente: Implementazioni di ottimizzazione manuale a livello di esperti
  4. Modello di Costo ML: TVM MetaSchedule (GPU)
  5. FSM Hardware: Macchina a stati finiti hardware Gemmini (limite di riferimento)

Risultati Sperimentali

Prestazioni Principali

Piattaforma Gemmini

  • Benchmark GEMM: Miglioramento di 5,6× rispetto alla libreria software Gemmini, superamento del codice ottimizzato manualmente da esperti di 1,4×
  • Benchmark di Convoluzione: Miglioramento di 2,6× rispetto alla libreria software, superamento dell'ottimizzazione manuale di 1,1×
  • Algebra Lineare a Grana Fine: Superamento del codice non ottimizzato di 2,7×, persino superamento dell'implementazione FSM hardware ottimizzata da esperti di 1,6× (propagazione in avanti)

Piattaforma AWS Trainium

  • Carichi di Lavoro Tutorial: Superamento del codice ottimizzato manualmente di 1,36× (media geometrica), superamento del codice compilato PyTorch NeuronX di 13,52×
  • Carichi di Lavoro Avanzati: Superamento del codice ottimizzato a livello di esperti di 1,9× (media geometrica), convoluzione profonda 1D con miglioramento fino a 17,37×

NVIDIA L40S GPU

  • Benchmark KernelBench: Superamento di PyTorch di 2,05× (media geometrica), superamento di TVM MetaSchedule di 3,8×
  • Prestazioni superiori a PyTorch in tutti i benchmark, mentre TVM supera PyTorch solo in 2 benchmark

Analisi degli Esperimenti di Ablazione

Esperimenti di ablazione dettagliati verificano l'importanza di ogni componente:

  • ISA dell'Acceleratore: La rimozione causa un calo significativo delle prestazioni, ma consente comunque miglioramenti
  • Menu di Ottimizzazione: Completamente necessario, la rimozione causa il completo deterioramento delle prestazioni di ottimizzazione
  • Dropout del Menu: Impatto significativo sulle prestazioni, prevenendo il bias del modello verso opzioni di menu limitate
  • Integrazione LLM: Fornisce importante diversità, le prestazioni con un singolo modello sono inferiori
  • Feedback di Prestazione Hardware: Utile ma effetto limitato, poiché il menu di ottimizzazione contiene già metriche rilevanti

Effetto del Riutilizzo della Pianificazione

  • Con budget di 100 campioni: pianificazione riutilizzata raggiunge accelerazione di 4,6×, senza riutilizzo solo 3,7×
  • Con budget di 200 campioni: pianificazione riutilizzata raggiunge accelerazione di 5,0×, senza riutilizzo solo 4,2×
  • Dimostra la generalizzabilità della pianificazione, riducendo efficacemente i costi di ricerca per benchmark simili

Lavori Correlati

Ottimizzazione del Codice per Acceleratori Tensoriali

  • Modelli di Prestazione: Timeloop, MAESTRO utilizzano modelli di architettura hardware di alto livello
  • Metodi Automatizzati: Machine learning, programmazione lineare, ottimizzazione black-box, apprendimento per rinforzo
  • Limitazioni: Le astrazioni esistenti trascurano le ottimizzazioni specifiche dell'implementazione e a livello di istruzione

Ottimizzazione del Codice Guidata da LLM

  • Ambito di Applicazione: Ricerca evolutiva, generazione aumentata da recupero, ottimizzazione iterativa, post-addestramento del modello
  • Ottimizzazione a Livello di Sistema: CUDA, funzioni intrinseche SIMD
  • Lacune di Ricerca: Mancanza di lavori di ottimizzazione del codice guidati da LLM per hardware dedicato (non CPU/GPU)

Conclusioni e Discussione

Conclusioni Principali

  1. Efficacia dell'Ottimizzazione Guidata da LLM: Autocomp supera significativamente i metodi tradizionali su molteplici piattaforme hardware
  2. Portabilità Estremamente Elevata: Adattamento a nuovo hardware attraverso la sola modifica dei prompt, costo ingegneristico minimo
  3. Valore del Riutilizzo della Pianificazione: Le pianificazioni di ottimizzazione mostrano buona generalizzabilità, migliorando significativamente l'efficienza del campionamento

Intuizioni Tecniche

  1. Necessità del Design in Due Fasi: La separazione delle fasi di pianificazione e implementazione aumenta il tasso di successo per compiti di ottimizzazione complessi
  2. Importanza della Conoscenza del Dominio: La conoscenza specialistica del dominio fornita dal menu di ottimizzazione è critica per le prestazioni
  3. Valore del Feedback Hardware: Le metriche di prestazione in tempo reale guidano efficacemente la selezione della direzione di ottimizzazione

Limitazioni

  1. Dipendenza dalle Capacità LLM: Le prestazioni del metodo sono limitate dalle capacità di generazione e ragionamento del codice dell'LLM sottostante
  2. Costo di Ricerca: Richiede molteplici chiamate LLM e simulazione hardware, con costo computazionale relativamente elevato
  3. Specificità del Dominio: Il menu di ottimizzazione richiede progettazione manuale per diverse piattaforme hardware
  4. Ambito di Valutazione: Principalmente focalizzato su carichi di lavoro di calcolo tensoriale, l'applicabilità ad altri tipi di calcolo rimane sconosciuta

Direzioni Future

  1. Generazione Automatica del Menu: Ricerca di metodi per costruire automaticamente menu di ottimizzazione
  2. Migrazione della Pianificazione tra Piattaforme: Esplorazione della migrazione della conoscenza di pianificazione tra diverse piattaforme hardware
  3. Ottimizzazione dell'Efficienza dei Costi: Riduzione del numero di chiamate LLM e simulazioni hardware nel processo di ricerca
  4. Applicazioni più Ampie: Estensione ad altri acceleratori dedicati non tensoriali

Valutazione Approfondita

Punti di Forza

  1. Forte Innovatività: Prima applicazione di LLM all'ottimizzazione del codice per acceleratori tensoriali a bassa risorsa, percorso tecnico innovativo
  2. Alto Valore Pratico: Affronta effettivi punti critici ingegneristici, riducendo significativamente i costi di sviluppo software per nuovo hardware
  3. Valutazione Completa: Valutazione complessiva su tre diverse piattaforme hardware, risultati convincenti
  4. Metodo Generico: Il design del framework mostra buona scalabilità e portabilità
  5. Prestazioni Eccezionali: Supera significativamente i migliori metodi esistenti su molteplici benchmark

Insufficienze

  1. Costo Computazionale: Richiede numerose chiamate LLM e simulazioni hardware, potenzialmente limitando l'applicazione pratica
  2. Dipendenza dalla Progettazione Manuale: Il menu di ottimizzazione richiede ancora progettazione esperta, grado di automazione limitato
  3. Limitazioni di Valutazione: Principalmente focalizzato su tipi specifici di calcolo tensoriale, generalizzabilità da verificare
  4. Analisi Teorica Insufficiente: Mancanza di garanzie teoriche sulla convergenza e ottimalità del metodo

Valutazione dell'Impatto

  1. Valore Accademico: Inaugura l'applicazione di LLM nell'ottimizzazione della compilazione per hardware dedicato, significato accademico importante
  2. Impatto Industriale: Promette di ridurre significativamente i costi di sviluppo dello stack software per nuovo hardware, valore industriale importante
  3. Riproducibilità: Gli autori si impegnano a open-source l'implementazione e i prompt, favorendo la ricerca successiva
  4. Ispirazione: Fornisce un nuovo percorso tecnico per l'ottimizzazione della compilazione di altri hardware dedicati

Scenari Applicabili

  1. Sviluppo di Prototipi di Nuovo Hardware: Generazione rapida di codice ottimizzato per acceleratori tensoriali appena progettati
  2. Costruzione di Compilatori DSL: Come complemento o alternativa ai compilatori tradizionali
  3. Strumenti di Ottimizzazione delle Prestazioni: Assistenza agli sviluppatori nell'ottimizzazione del codice acceleratore esistente
  4. Ricerca e Insegnamento: Fornire strumenti automatizzati per la programmazione e l'ottimizzazione degli acceleratori

Bibliografia

L'articolo cita numerosi lavori correlati, principalmente includenti:

  • Progettazione di acceleratori hardware (Gemmini, TPU, Trainium, ecc.)
  • Compilatori e DSL (XLA, TVM, Halide, Exo, ecc.)
  • Generazione di codice LLM (CodeGen, Codex, ecc.)
  • Metodi di ottimizzazione automatizzata (apprendimento per rinforzo, algoritmi evolutivi, ecc.)

Valutazione Complessiva: Questo è un articolo di ricerca di alta qualità che fornisce importanti contributi nel nuovo campo interdisciplinare dell'applicazione di LLM all'ottimizzazione della compilazione per hardware dedicato. Il metodo è fortemente innovativo, la valutazione sperimentale è completa e il valore pratico è significativo. Sebbene rimanga spazio per miglioramenti in termini di costo computazionale e grado di automazione, apre nuove direzioni per lo sviluppo del campo, con importante valore sia accademico che industriale.