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 Optimiseur de Code Puissant et Portable pour Accélérateurs Tensoriels

Informations Fondamentales

  • ID de l'article : 2505.18574
  • Titre : Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
  • Auteurs : Charles Hong, Sahil Bhatia, Alvin Cheung, Yakun Sophia Shao (UC Berkeley)
  • Classification : cs.PL cs.AI cs.AR cs.LG
  • Statut de publication : Préimpression. En révision.
  • Lien de l'article : https://arxiv.org/abs/2505.18574

Résumé

Les accélérateurs matériels, en particulier ceux conçus spécifiquement pour le traitement tensoriel, sont omniprésents dans les environnements informatiques actuels. Cependant, malgré les efforts considérables investis dans la construction de compilateurs, la programmation de ces accélérateurs tensoriels reste un défi, laissant la majeure partie de leur potentiel sous-exploitée. Cet article propose Autocomp, une méthode d'optimisation de code par recherche automatisée pilotée par LLM, permettant aux programmeurs d'accélérateurs de tirer parti des connaissances du domaine et des retours matériels. La méthode est réalisée par trois techniques clés : 1) la formulation de chaque processus d'optimisation comme une invite structurée en deux phases, divisée en phases de planification et de génération de code ; 2) l'insertion de connaissances du domaine lors de la planification par le biais d'un menu d'optimisation concis et adaptable ; 3) l'intégration des métriques de correction et de performance du matériel comme retour à chaque itération de recherche.

Contexte de Recherche et Motivation

Problèmes Fondamentaux

Les principaux défis de la programmation d'accélérateurs tensoriels incluent :

  1. Complexité de programmation : Contrairement à la programmation CPU généraliste, les accélérateurs tensoriels nécessitent une gestion explicite du mouvement des données, de la configuration d'état et de l'ordonnancement des opérations
  2. Coût d'adaptation des compilateurs : L'adaptation des compilateurs traditionnels aux nouvelles plates-formes matérielles nécessite un travail d'ingénierie considérable, les coûts de développement logiciel représentant 40-50% des coûts de développement du nouveau matériel
  3. Problème d'ordonnancement d'optimisation : Explosion combinatoire du problème de détermination des optimisations à appliquer et de leur ordre d'application
  4. Défi des langages à faibles ressources : Les architectures d'ensemble d'instructions (ISA) et les DSL des accélérateurs dédiés sont sous-représentés dans les corpus d'entraînement des LLM

Limitations des Approches Existantes

  1. Compilateurs traditionnels : XLA, TVM, Triton ne supportent que quelques backends matériels, principalement CPU et GPU
  2. Approches DSL : Halide, Exo fournissent des primitives pour exprimer les calculs tensoriels, mais le fardeau d'optimisation reste sur le programmeur
  3. Approches pilotées par les données : Nécessitent une grande quantité de données de performance pour l'entraînement, extrêmement rares pour les accélérateurs matériels spécialisés
  4. Application directe de LLM : Pour les langages d'accélérateurs à faibles ressources, la génération de code sans exemple est hautement peu fiable

Contributions Fondamentales

  1. Première méthode d'optimisation de code d'accélérateur tensoriel à faibles ressources pilotée par LLM : Proposition du cadre Autocomp, spécifiquement conçu pour les accélérateurs matériels dédiés
  2. Cadre d'optimisation hautement portable : Adaptation aux nouvelles plates-formes matérielles par modification des invites, réduisant considérablement les coûts d'ingénierie
  3. Performance exceptionnelle : Surpasse significativement les méthodes existantes sur trois plates-formes matérielles différentes
  4. Mécanisme de réutilisation d'ordonnancement : Démontre que les ordonnancements d'optimisation peuvent être réutilisés entre opérations tensorielles similaires, améliorant l'efficacité des échantillons

Détails de la Méthode

Définition de la Tâche

Entrée : Code d'accélérateur tensoriel non optimisé Sortie : Code fonctionnellement équivalent mais optimisé en performance Contraintes : Maintien de l'équivalence sémantique, vérification de la correction via le matériel

Architecture Fondamentale : Cadre d'Optimisation en Deux Phases

Phase 1 : Génération de Plan (Plan Generation)

La structure d'invite comprend :

  1. Description de l'ISA de l'accélérateur : Sémantique des instructions, normes d'adressage mémoire, description de la structure matérielle
  2. Code actuel : Code à optimiser
  3. Retour de performance : Métriques telles que la latence (nombre de cycles) et l'utilisation mémoire
  4. Menu d'optimisation : Options d'optimisation de haut niveau prédéfinies (comme le pavage de boucles, la réorganisation, la fusion, etc.)
  5. Informations d'itération de recherche : Numéro d'itération actuel, guidant le choix d'optimisation

Phase 2 : Implémentation de Code (Implementation)

La structure d'invite comprend :

  1. Description de l'ISA de l'accélérateur : Identique à la phase 1
  2. Code actuel : Identique à la phase 1
  3. Plan généré : Plan d'optimisation spécifique issu de la phase 1
  4. Exemples d'apprentissage contextuel : Exemples de code pour les optimisations complexes (comme le pavage)
  5. Instructions d'implémentation : Instructions en langage naturel pour appliquer le plan et générer le code optimisé

Stratégie de Recherche par Faisceau

  • Recherche par faisceau de largeur B=6, explorant en parallèle plusieurs trajectoires d'optimisation
  • Filtrage de correction : Vérification des candidats via suite de tests fonctionnels
  • Filtrage de performance : Conservation uniquement des candidats surpassant le nœud parent
  • Optimisation itérative : Processus de recherche avec budget T itérations fixe

Innovations Techniques Clés

1. Techniques d'Amélioration de la Diversité

  • Suppression du Menu d'Optimisation : Suppression aléatoire partielle des options de menu lors de chaque planification (probabilité 70%)
  • Intégration de LLM : Distribution des requêtes entre plusieurs LLM pour augmenter la diversité des réponses

2. Intégration du Retour Matériel

  • Métriques de performance en temps réel (latence, utilisation mémoire) guidant le choix d'optimisation suivant
  • Simulation au cycle précis ou mesure de performance au niveau de la puce

3. Mécanisme de Réutilisation d'Ordonnancement

  • Enregistrement des séquences d'ordonnancement de haute qualité
  • Réutilisation des ordonnancements connus pour les opérations tensorielles similaires (même rapport d'aspect ou dimensions partagées)
  • Optimisation supplémentaire après recherche légère

Configuration Expérimentale

Plates-formes Matérielles

  1. Gemmini : Générateur d'accélérateur open-source, supportant les tableaux systoliques et les accélérateurs tensoriels de style vectoriel
  2. AWS Trainium : Accélérateur tensoriel commercial haute performance, utilisant l'Interface Noyau Neuron (NKI)
  3. GPU NVIDIA L40S : GPU moderne de centre de données, incluant des Tensor Cores dédiés

Benchmarks

  • Gemmini : GEMM et convolutions de ResNet-50, contrôle prédictif de modèle TinyMPC
  • Trainium : Opérateurs d'apprentissage profond de niveau tutoriel et avancé (RMSNorm, LayerNorm, GEMM, Mamba, etc.)
  • GPU : Benchmarks KernelBench Niveau 1

Méthodes de Comparaison

  1. Bibliothèques logicielles avancées : Bibliothèque Gemmini, PyTorch NeuronX, PyTorch
  2. Code bas niveau non optimisé : Exo non optimisé, code tutoriel nki-samples
  3. Code optimisé manuellement : Implémentations d'optimisation manuelle au niveau expert
  4. Modèles de coût ML : TVM MetaSchedule (GPU)
  5. FSM matériel : Machine à états finis matérielle Gemmini (limite de référence)

Résultats Expérimentaux

Performance Principale

Plate-forme Gemmini

  • Benchmark GEMM : Amélioration de 5,6× par rapport à la bibliothèque Gemmini, surpassant le code optimisé manuellement par experts de 1,4×
  • Benchmark de convolution : Amélioration de 2,6× par rapport à la bibliothèque, surpassant l'optimisation manuelle de 1,1×
  • Algèbre linéaire granulaire : Surpasse le code non optimisé de 2,7×, surpassant même l'implémentation FSM matériel optimisée par experts de 1,6× (propagation avant)

Plate-forme AWS Trainium

  • Charges de travail tutoriel : Surpasse le code optimisé manuellement de 1,36× (moyenne géométrique), surpasse le code compilé PyTorch NeuronX de 13,52×
  • Charges de travail avancées : Surpasse le code optimisé au niveau expert de 1,9× (moyenne géométrique), convolution profonde 1D avec amélioration jusqu'à 17,37×

GPU NVIDIA L40S

  • Benchmark KernelBench : Surpasse PyTorch de 2,05× (moyenne géométrique), surpasse TVM MetaSchedule de 3,8×
  • Surpasse PyTorch dans tous les benchmarks, tandis que TVM surpasse PyTorch dans seulement 2 benchmarks

Analyse des Expériences d'Ablation

Les expériences d'ablation détaillées valident l'importance de chaque composant :

  • ISA de l'accélérateur : La suppression entraîne une baisse significative de performance, mais amélioration toujours réalisable
  • Menu d'optimisation : Complètement nécessaire, la suppression entraîne une dégradation complète de la performance d'optimisation
  • Suppression du menu : Impact significatif sur la performance, prévenant la polarisation du modèle vers des options de menu limitées
  • Intégration de LLM : Fournit une diversité importante, la performance d'un modèle unique étant inférieure
  • Retour de performance matériel : Utile mais effet limité, car le menu d'optimisation inclut déjà les métriques pertinentes

Effet de Réutilisation d'Ordonnancement

  • Avec budget de 100 échantillons : ordonnancement réutilisé atteint 4,6× accélération, sans réutilisation seulement 3,7×
  • Avec budget de 200 échantillons : ordonnancement réutilisé atteint 5,0× accélération, sans réutilisation seulement 4,2×
  • Démontre la généralisabilité de l'ordonnancement, réduisant efficacement le coût de recherche pour les benchmarks similaires

Travaux Connexes

Optimisation de Code d'Accélérateur Tensoriel

  • Modèles de performance : Timeloop, MAESTRO utilisant des modèles d'architecture matérielle de haut niveau
  • Méthodes automatisées : Apprentissage automatique, programmation linéaire, optimisation boîte noire, apprentissage par renforcement
  • Limitations : Les abstractions existantes ignorent les optimisations spécifiques à l'implémentation et au niveau des instructions

Optimisation de Code par LLM

  • Portée d'application : Recherche évolutive, génération augmentée par récupération, optimisation itérative, post-entraînement de modèle
  • Optimisation au niveau système : CUDA, fonctions intrinsèques SIMD
  • Lacune de recherche : Absence de travaux d'optimisation de code par LLM pour matériel spécialisé (non-CPU/GPU)

Conclusion et Discussion

Conclusions Principales

  1. Efficacité de l'optimisation pilotée par LLM : Autocomp surpasse significativement les méthodes traditionnelles sur plusieurs plates-formes matérielles
  2. Portabilité extrêmement élevée : Adaptation aux nouvelles plates-formes par simple modification des invites, coût d'ingénierie minimal
  3. Valeur de la réutilisation d'ordonnancement : Les ordonnancements d'optimisation présentent une bonne généralisabilité, améliorant significativement l'efficacité des échantillons

Perspectives Techniques

  1. Nécessité de la conception en deux phases : La séparation des phases de planification et d'implémentation améliore le taux de réussite des tâches d'optimisation complexes
  2. Importance des connaissances du domaine : Les connaissances spécialisées du domaine fournies par le menu d'optimisation sont essentielles à la performance
  3. Valeur du retour matériel : Les métriques de performance en temps réel guident efficacement le choix de direction d'optimisation

Limitations

  1. Dépendance aux capacités du LLM : La performance de la méthode est limitée par les capacités de génération et de raisonnement de code du LLM sous-jacent
  2. Coût de recherche : Nécessite plusieurs appels LLM et simulations matérielles, coût computationnel élevé
  3. Spécificité au domaine : Le menu d'optimisation doit être conçu manuellement pour différentes plates-formes matérielles
  4. Portée d'évaluation : Principalement concentrée sur les charges de travail de calcul tensoriel, applicabilité à d'autres types de calcul inconnue

Directions Futures

  1. Génération automatique de menu : Recherche de méthodes pour construire automatiquement les menus d'optimisation
  2. Migration d'ordonnancement inter-plates-formes : Exploration de la migration des connaissances d'ordonnancement entre différentes plates-formes matérielles
  3. Optimisation de l'efficacité des coûts : Réduction du nombre d'appels LLM et de simulations matérielles dans le processus de recherche
  4. Applications plus larges : Extension à d'autres accélérateurs spécialisés au-delà du calcul tensoriel

Évaluation Approfondie

Points Forts

  1. Innovation forte : Première application de LLM à l'optimisation de code d'accélérateur tensoriel à faibles ressources, approche technique novatrice
  2. Valeur pratique élevée : Résout les points douloureux d'ingénierie réels, réduisant considérablement les coûts de développement logiciel pour nouveau matériel
  3. Évaluation complète : Évaluation exhaustive sur trois plates-formes matérielles différentes, résultats convaincants
  4. Méthode générique : Conception du cadre avec bonne extensibilité et portabilité
  5. Performance exceptionnelle : Surpasse significativement les meilleures méthodes existantes sur plusieurs benchmarks

Insuffisances

  1. Coût computationnel : Nécessite un grand nombre d'appels LLM et simulations matérielles, pouvant limiter l'application pratique
  2. Dépendance à la conception manuelle : Le menu d'optimisation nécessite toujours une conception manuelle avec connaissances d'expert, degré d'automatisation limité
  3. Limitations d'évaluation : Principalement concentrée sur des types spécifiques de calcul tensoriel, généralisabilité à vérifier
  4. Analyse théorique insuffisante : Absence de garanties théoriques sur la convergence et l'optimalité de la méthode

Évaluation de l'Impact

  1. Valeur académique : Inaugure l'application de LLM à l'optimisation de compilation pour matériel spécialisé, importance académique significative
  2. Impact industriel : Susceptible de réduire considérablement les coûts de développement de pile logicielle pour nouveau matériel, valeur industrielle importante
  3. Reproductibilité : Les auteurs s'engagent à open-sourcer l'implémentation et les invites, favorisant la recherche ultérieure
  4. Caractère inspirant : Fournit une nouvelle voie technique pour l'optimisation de compilation d'autres matériels spécialisés

Scénarios d'Application

  1. Développement de prototypes matériels : Génération rapide de code optimisé pour accélérateurs tensoriels nouvellement conçus
  2. Construction de compilateurs DSL : Complément ou alternative aux compilateurs traditionnels
  3. Outil d'optimisation de performance : Aide les développeurs à optimiser le code d'accélérateur existant
  4. Recherche et enseignement : Fournit des outils automatisés pour la programmation et l'optimisation d'accélérateurs

Références

L'article cite un grand nombre de travaux connexes, incluant principalement :

  • Conception d'accélérateurs matériels (Gemmini, TPU, Trainium, etc.)
  • Compilateurs et DSL (XLA, TVM, Halide, Exo, etc.)
  • Génération de code par LLM (CodeGen, Codex, etc.)
  • Méthodes d'optimisation automatisée (apprentissage par renforcement, algorithmes évolutifs, etc.)

Évaluation Globale : Ceci est un article de recherche de haute qualité apportant des contributions importantes dans le domaine émergent et interdisciplinaire de l'application de LLM à l'optimisation de compilation pour matériel spécialisé. La méthode est fortement innovante, l'évaluation expérimentale est complète, et la valeur pratique est significative. Bien qu'il y ait encore place à l'amélioration en termes de coût computationnel et de degré d'automatisation, cet article ouvre une nouvelle direction pour le développement du domaine, possédant une valeur académique et industrielle importante.