🤖 machine learning

Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search

Cet article propose un évaluateur inspiré par les modèles de monde qui apprend la dynamique latente des actions d'ordonnancement pour classer efficacement les candidats de programmes tensoriels, atteignant des améliorations de latence significatives par rapport aux auto-ordonnanceurs existants comme Ansor tout en réduisant considérablement le nombre de mesures requises.

Auteurs originaux : Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Publié 2026-06-09

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de trouver l'itinéraire le plus rapide pour conduire de votre domicile à la maison d'un ami. Vous avez une carte, mais les conditions de circulation changent constamment, et il existe des millions d'itinéraires possibles que vous pourriez emprunter.

Dans le monde de l'informatique, plus précisément pour l'Apprentissage Automatique (Machine Learning), c'est exactement ce qu'un « compilateur » fait. Il essaie de trouver la manière la plus efficace pour un ordinateur d'exécuter des tâches mathématiques complexes (appelées « programmes de tenseurs »). Le problème est qu'il existe tellement de façons différentes d'écrire le code que vérifier chaque option en l'exécutant réellement sur l'ordinateur est incroyablement lent et coûteux. C'est comme essayer de parcourir tous les itinéraires possibles juste pour voir lequel est le plus rapide ; vous tomberiez en panne d'essence avant d'avoir trouvé le meilleur.

L'ancienne méthode : Prendre un instantané

Auparavant, les programmes informatiques qui tentaient de résoudre cela (appelés « auto-ordonnanceurs » ou auto-schedulers) agissaient comme un photographe prenant un instantané de la destination finale. Ils regardaient le code fini, devinaient sa vitesse, et décidaient s'il était bon.

L'article soutient que c'est une mauvaise idée car :

Cela ignore le voyage : Cela ne comprend pas comment le code est arrivé là. Deux itinéraires différents peuvent mener au même endroit, mais l'un peut avoir été une autoroute fluide tandis que l'autre était une piste de terre cahoteuse. L'instantané se ressemble, mais l'expérience (et la vitesse) est différente.
Cela se laisse confondre par des détails mineurs : Si vous changez quelques mots dans le code qui ne modifient pas réellement son fonctionnement, l'ancien système pourrait penser qu'il s'agit d'un itinéraire complètement différent (et pire).

La nouvelle idée : Un « Modèle de Monde » (Le simulateur GPS)

Les auteurs proposent une nouvelle approche inspirée des Modèles de Monde (World Models). Voyez cela non pas comme un photographe, mais comme un simulateur GPS de haute technologie.

Au lieu de simplement regarder la destination finale, ce nouveau système simule l'intégralité du voyage étape par étape dans son « esprit » (un espace mathématique appelé « espace latent »).

Voici comment cela fonctionne, en utilisant une analogie culinaire :

Les ingrédients (État initial) : Vous commencez avec une recette brute (le code non optimisé).
Les mouvements du chef (Actions) : Le compilateur prend des décisions, comme « hacher les oignons », « faire revenir pendant 5 minutes » ou « ajouter du sel ».
L'ancienne méthode : Le chef regarde le plat final et devine : « Hmm, cela semble correct ».
La nouvelle méthode (Modèle de Monde) : Le chef possède une simulation mentale. Il imagine : « Si je hache les oignons puis que je les fais revenir, la texture sera X. Si je les fais revenir puis que je les hache, la texture sera Y. » Il simule le processus de cuisine dans sa tête pour prédire le goût final sans avoir à cuisiner tout le repas au préalable.

Comment ils l'ont construit

Les chercheurs ont construit un système composé de trois parties :

Le Traducteur (Encodeur) : Il transforme le code informatique désordonné en une « pensée » mathématique propre (un vecteur) que l'ordinateur peut comprendre facilement.
Le Simulateur (Modèle de transition) : C'est l'innovation centrale. Il prend la « pensée » du code actuel et applique les « mouvements du chef » (actions d'ordonnancement) un par un. Il prédit ce à quoi le code ressemblera après chaque étape, tout cela à l'intérieur de la mémoire de l'ordinateur, sans réellement exécuter le code.
Le Juge (Modèle de classement) : Une fois la simulation terminée, le juge regarde le résultat final prédit et dit : « Cet itinéraire est probablement le plus rapide » ou « Celui-là semble lent ».

Les résultats

Ils ont testé ce système sur deux types d'ordinateurs : un processeur puissant (Intel Xeon) et une carte graphique haut de gamme (NVIDIA RTX 4090).

Des résultats plus rapides : Ils ont trouvé de meilleurs ordonnancements de code beaucoup plus rapidement que la meilleure méthode précédente (appelée Ansor).
Moins de travail : Ils ont obtenu des résultats aussi bons que l'ancienne méthode, mais ils ont dû effectuer 10 fois moins de « tests de conduite » (mesures).
Vitesse en conditions réelles : Lorsqu'ils ont utilisé cela pour exécuter de véritables modèles d'IA (comme la reconnaissance d'images ou les modèles de langage), les programmes tournaient 4 à 5 fois plus vite que les versions standards, et dans certains cas, jusqu'à 58 fois plus vite.

L'essentiel

L'article affirme qu'en apprenant à l'ordinateur à comprendre le processus d'optimisation (le voyage) plutôt que seulement le résultat (l'instantané), nous pouvons trouver le code le plus rapide de manière beaucoup plus efficace. C'est comme avoir un GPS qui simule le trafic dans sa tête pour trouver le meilleur itinéraire, plutôt que de simplement deviner en se basant sur une photo de la destination.

Limites mentionnées dans l'article :

Ce système est un « juge » qui aide à choisir le meilleur itinéraire ; il n'invente pas les itinéraires lui-même. Si le moteur de recherche ne propose aucun bon itinéraire au départ, le juge ne peut pas le corriger.
Si le « voyage » est extrêmement long et complexe, la simulation dans l'esprit de l'ordinateur peut commettre de petites erreurs qui s'accumulent, rendant la prédiction moins précise.
Il est conçu pour classer les options entre elles (lequel est plus rapide ?), et non pour prédire le temps exact à la milliseconde près.

Résumé Technique : Vers des Modèles de Monde pour Compilateurs

Énoncé du Problème

L'optimisation des programmes de tenseurs est essentielle pour les systèmes d'apprentissage automatique modernes, pourtant l'identification du programme le plus efficace est entravée par un espace de recherche combinatoire massif. Pour éviter le coût prohibitif des mesures sur l'appareil, les auto-ordonnanceurs existants (par exemple, Ansor, AutoTVM) s'appuient sur des modèles de coût appris pour classer les candidats de programmes. Cependant, ces modèles traitent généralement les candidats comme des instantanés de code statiques (par exemple, le texte final de TensorIR).

Cette perspective statique introduit deux limitations critiques :

Perte de Contexte de Transformation : Elle ignore le processus de transformation étape par étape qui a généré le programme. Par conséquent, les évaluateurs ne peuvent pas comprendre comment une décision d'ordonnancement dans une étape ultérieure dépend du contexte établi par des décisions antérieures.
Sensibilité au Bruit Syntaxique : Les encodeurs statiques sont facilement induits en erreur par des variations de code superficielles. Par exemple, deux programmes structurellement équivalents (par exemple, une multiplication de matrices avec des tailles de boucle $N=1024$ vs $N=1023$ ) peuvent produire des textes TensorIR significativement différents en raison de la gestion des bordures (peeling de boucle), provoquant une erreur de classification par le modèle quant à leur potentiel de performance malgré un comportement matériel identique.

Méthodologie

Les auteurs proposent un cadre de Modèle de Monde de Compilateur qui reformule l'évaluation de l'ordonnancement comme une dynamique latente conditionnée par des actions sur des états de programmes. Au lieu d'évaluer un artefact de code statique final, le cadre simule l'évolution de l'état du programme à travers un espace latent continu.

Le cadre se compose de trois composants appris instanciés au sein de l'AutoScheduler de TVM :

1. Représentation de l'État du Programme de Tenseur (Encodeur)

Entrée : Texte TensorIR (états pré-ordonnancement et post-ordonnancement).
Architecture : Une architecture basée sur CodeBERT entraînable, adaptée aux entrées longues. Elle utilise une fenêtre glissante pour tokeniser et encoder des segments chevauchants, les agrégeant via un regroupement par moyenne masquée (masked mean pooling) en un vecteur de représentation dense $z = E(s)$ .
Entraînement : Utilise l'apprentissage contrastif pour rapprocher dans l'espace latent les états TensorIR similaires (provenant du même flux de travail) tout en éloignant les états dissemblables. Cela crée un espace d'états où l'équivalence sémantique est préservée malgré les variations syntaxiques.

2. Transition d'État Multi-étapes Conditionnée par l'Action

Objectif : Prédire la représentation de l'état terminal induite par une séquence d'actions d'ordonnancement sans matérialiser explicitement les AST intermédiaires.
Mécanisme : En partant d'un embedding d'état initial $z_0$ $z_{0}$ , le modèle déroule l'effet d'une séquence d'actions d'ordonnancement $a_{1:T}$ $a_{1 : T}$ étape par étape.
- Les actions sont sérialisées et encodées en embeddings denses.
- Un modèle de transition $F(\cdot)$ , implémenté à l'aide d'une architecture basée sur TransH (originaire de l'apprentissage de représentation de graphes de connaissances), prédit l'état latent suivant : $z_t = F(z_{t-1}, a_t)$ .
- Cela modélise les actions d'ordonnancement comme des transformations géométriques dans un hyperplan spécifique à la relation, permettant de suivre efficacement la provenance des décisions et de résoudre les ambiguïtés référentielles (par exemple, ce qu'une variable de boucle représente après une fusion).
Entraînement : Utilise une perte de déroulement multi-étapes combinant un terme de correspondance d'état (minimisant la distance par rapport aux états futurs de vérité terrain) et un terme de classement basé sur la marge (garantissant que l'état prédit est plus proche de la vérité terrain que des échantillons négatifs).

3. Évaluation de Candidat Basée sur le Classement

Entrée : La représentation de l'état terminal prédite $\hat{z}_T$ , les caractéristiques d'action structurées (par exemple, décompte des actions, longueur de séquence) et les caractéristiques matérielles optionnelles.
Sortie : Un score scalaire pour la hiérarchisation relative des candidats au sein du même flux de travail.
Architecture : Un modèle de classement XGBoost entraîné avec un objectif LambdaRank.
Stratégie d'Entraînement : Le modèle ne prédit pas le temps d'exécution absolu. Au lieu de cela, il apprend à classer les candidats au sein d'un groupe de flux de travail sur la base des temps d'exécution mesurés, traitant le problème comme une tâche d'apprentissage de classement (learning-to-rank).

Contributions Clés

Perspective de Modèle de Monde pour les Compilateurs : L'article adapte le paradigme du modèle de monde à l'optimisation de tenseurs, établissant un cadre où les schémas candidats sont modélisés à travers l'évolution d'états latents conditionnés par des actions plutôt que par des instantanés statiques.
Jeu de Données de Trajectoire Action-État : Construction d'un jeu de données basé sur TVM/TenSet qui organise les journaux de réglage (tuning logs) en trajectoires action-état alignées (état pré-ordonnancement, séquence d'actions, états intermédiaires, état post-ordonnancement) pour la prédiction multi-étapes.
Intégration de Recherche Efficace : Implémentation du cadre dans l'AutoScheduler de TVM, démontrant que la modélisation de la dynamique latente réduit considérablement les coûts d'auto-réglage tout en maintenant une haute qualité de recherche.

Résultats Expérimentaux

La méthode a été évaluée sur un processeur Intel Xeon Gold 6430 et un GPU NVIDIA RTX 4090 à travers sept modèles de réseaux neuronaux (incluant ResNet, BERT, GPT et OPT) et 103 sous-graphes extraits.

Performance vs Ansor : Avec un budget de 64 essais, la méthode améliore la latence des sous-graphes représentatifs par rapport à Ansor de 1,37× sur GPU et 1,54× sur CPU.
Efficacité d'Échantillonnage : La méthode atteint une qualité de recherche comparable à Ansor-10K (10 000 essais) en utilisant seulement 1 000 essais (10× moins de mesures). Plus précisément, elle surpasse légèrement Ansor-10K en moyenne arithmétique et égale sa moyenne géométrique à 2,2 % sur les sous-graphes représentatifs.
Accélération de Bout en Bout : Comparée à PyTorch/PyTorch-opt (cuDNN), la méthode délivre des accélérations de moyenne géométrique de 4,61× (CPU) et 3,67× (GPU), avec des pics d'accélération atteignant 58,61×.
Études d'Ablation : Le retrait du module de transition appris (revenant à un encodage d'état initial ou final statique) a entraîné des augmentations de latence significatives (médiane ~56–69 %), confirmant que la modélisation de l'évolution de l'état intermédiaire est la composante la plus critique.

Signification et Limites

Signification :
L'article soutient que traiter l'optimisation du compilateur comme un processus dynamique de transitions d'état, plutôt que comme un problème d'évaluation de code statique, permet aux évaluateurs de capturer les dépendances sémantiques entre les décisions d'ordonnancement. Cette approche supprime efficacement le bruit syntaxique et résout les ambiguïtés référentielles, menant à une recherche plus efficace avec moins de mesures matérielles.

Limites (telles que citées par les auteurs) :

Dépendance à l'Espace de Recherche : La méthode améliore l'évaluation des candidats, pas leur génération. Son efficacité est limitée par la capacité du planificateur sous-jacent à proposer des ordres forts ; un meilleur évaluateur ne peut pas récupérer des schémas qui ne sont jamais proposés.
Accumulation d'Erreurs : Les erreurs de prédiction peuvent s'accumuler sur des trajectoires d'ordonnancement longues ou complexes dans l'espace latent, réduisant potentiellement la fiabilité pour les cas difficiles.
Classement vs Prédiction : Le modèle est entraîné pour le classement au sein d'un flux de travail, et non pour la prédiction de latence absolue. Cela le rend idéal pour la hiérarchisation en ligne, mais moins adapté aux contextes nécessitant des estimations de performance calibrées entre différents flux de travail.