🤖 machine learning

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

Dit artikel stelt een door een wereldmodel geïnspireerde evaluator voor die de latente dynamiek van scheduling-acties leert om tensorprogramma-kandidaten efficiënt te rangschikken, waarbij het significante latentieverbeteringen bereikt ten opzichte van bestaande auto-schedulers zoals Ansor en tegelijkertijd het aantal vereiste metingen drastisch vermindert.

Oorspronkelijke auteurs: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Gepubliceerd 2026-06-09

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je de absoluut snelste route probeert te vinden om van je huis naar het huis van een vriend te rijden. Je hebt een kaart, maar de verkeersomstandigheden veranderen constant en er zijn miljoenen mogelijke routes die je zou kunnen nemen.

In de wereld van computerwetenschappen, specifiek voor Machine Learning, is dit precies wat een "compiler" doet. Het probeert de meest efficiënte manier te vinden voor een computer om complexe wiskundige taken uit te voeren (genaamd "tensor-programma's"). Het probleem is dat er zoveel manieren zijn om de code te schrijven dat het controleren van elke optie door het daadwerkelijk uit te voeren op de computer ontzettend traag en duur is. Het is alsoal elke mogelijke route te rijden om te zien welke het snelst is; je zou zonder brandstof komen te zitten voordat je de beste hebt gevonden.

De Oude Manier: Een Snapshot Maken

Voorheen gedroegen computerprogramma's die dit probeerden op te lossen (genaamd "auto-schedulers") zich als een fotograaf die een snapshot maakt van de eindbestemming. Ze keken naar de voltooide code, gokten hoe snel deze zou zijn, en besloten of het goed was.

Het artikel betoogt dat dit een slecht idee is omdat:

Het de reis negeert: Het begrijpt niet hoe de code daar terecht is gekomen. Twee verschillende routes kunnen op dezelfde plek eindigen, maar de ene kan een soepele snelweg zijn geweest terwijl de andere een hobbelige onverharde weg was. De snapshot ziet er hetzelfde uit, maar de ervaring (en snelheid) is anders.
Het in de war raakt door kleine details: Als je een paar woorden in de code verandert die de werking eigenlijk niet veranderen, kan het oude systeem denken dat het een compleet andere (en slechtere) route is.

Het Nieuwe Idee: Een "World Model" (De GPS-Simulator)

De auteurs stellen een nieuwe aanpak voor, geïnspireerd door World Models. Denk hierbij niet aan een fotograaf, maar aan een hoogwaardige GPS-simulator.

In plaats van alleen naar de eindbestemming te kijken, simuleert dit nieuwe systeem de hele reis stap voor stap in zijn "geest" (een wiskundige ruimte genaamd een "latent space").

Zo werkt het, met een kookanalogie:

De Ingrediënten (Begintoestand): Je begint met een rauw recept (de ongeoptimaliseerde code).
De Bewegingen van de Chef (Acties): De compiler maakt beslissingen, zoals "de uien snijden", "5 minuten fruiten" of "zout toevoegen".
De Oude Manier: De chef kijkt naar het voltooide gerecht en gokt: "Hm, dit smaakt oké."
De Nieuwe Manier (World Model): De chef heeft een mentale simulatie. Hij stelt zich voor: "Als ik de uien eerst snijd en dan fruit, dan is de textuur X. Als ik eerst fruit en dan snijd, dan is de textuur Y." Hij simuleert het proces van koken in zijn hoofd om de uiteindelijke smaak te voorspellen zonder het hele gerecht eerst daadwerkelijk te koken.

Hoe Ze Het Gebouwd Hebben

De onderzoekers hebben een systeem gebouwd met drie onderdelen:

De Vertaler (Encoder): Het zet de rommelige computercode om in een schone, wiskundige "gedachte" (een vector) die de computer gemakkelijk kan begrijpen.
De Simulator (Transition Model): Dit is de kerninnovatie. Het neemt de "gedachte" van de huidige code en past de "bewegingen van de chef" (scheduling acties) één voor één toe. Het voorspelt hoe de code er na elke stap uit zal zien, allemaal binnen het geheugen van de computer, zonder de code daadwerkelijk uit te voeren.
De Rechter (Ranking Model): Zododien de simulatie voltooid is, kijkt de rechter naar het voorspelde eindresultaat en zegt: "Deze route is waarschijnlijk de snelste," of "Die andere ziet er traag uit."

De Resultaten

Ze hebben dit getest op twee soorten computers: een krachtige CPU (Intel Xeon) en een hoogwaardige videokaart (NVIDIA RTX 4090).

Snellere Resultaten: Ze vonden betere code-schedules veel sneller dan de vorige beste methode (genaamd Ansor).
Minder Werk: Ze behaalden resultaten die net zo goed waren als de oude methode, maar ze hoefden slechts 10 keer minder "testritten" (metingen) uit te voeren.
Real-World Snelheid: Wanneer ze dit gebruikten om daadwerkelijke AI-modellen (zoals beeldherkenning of taalmodellen) te draaien, liepen de programma's 4 tot 5 keer sneller dan standaardversies, en in sommige gevallen zelfs tot 58 keer sneller.

De Kernboodschap

Het artikel beweert dat door de computer te leren het proces van optimalisatie (de reis) te begrijpen in plaats van alleen het resultaat (de snapshot), we de snelste code veel efficiënter kunnen vinden. Het is also kind van een GPS die het verkeer in zijn hoofd simuleert om de beste route te vinden, in plaats van alleen te gokken op basis van een foto van de bestemming.

Beperkingen genoemd in het artikel:

Dit systeem is een "rechter" die helpt om de beste route te kiezen; het bedenkt de routes niet zelf. Als de zoekmachine vanaf het begin geen goede routes voorstelt, kan de rechter dat niet oplossen.
Als de "reis" extreem lang en complex is, kunnen de kleine foutjes in de simulatie in de geest van de computer zich opstapelen, waardoor de voorspelling minder nauwkeurig wordt.
Het is ontworpen om opties tegen elkaar af te wegen (welke is sneller?), niet om de exacte tijd tot op de milliseconde te voorspellen.

Technische Samenvatting: Naar Compiler Wereldmodellen

Probleemstelling

Optimalisatie van tensorprogramma's is essentieel voor moderne machine learning-systemen, maar het identificeren van het meest efficiënte programma wordt gehinderd door een enorme combinatorische zoekruimte. Om de prohibitieve kosten van on-device metingen te vermijden, vertrouwen bestaande auto-schedulers (bijv. Ansor, AutoTVM) op geleerde kostenmodellen om kandidaat-programma's te rangschikken. Deze modellen behandelen kandidaten echter doorgaans als statische code-snapshots (bijv. de uiteindelijke TensorIR-tekst).

Dit statische perspectief introduceert twee kritieke beperkingen:

Verlies van Transformatiecontext: Het negeert het stapsgewijze transformatieproces dat het programma heeft gegenereerd. Bijgevolg kunnen evaluatoren niet begrijpen hoe een scheduling-beslissing in een latere stap afhankelijk is van de context die door eerdere beslissingen is gevestigd.
Gevoeligheid voor Syntactische Ruis: Statische encoders laten zich gemakkelijk misleiden door oppervlakkige variaties in de code. Bijvoorbeeld, twee structureel equivalente programma's (bijv. matrixvermenigvuldiging met lusgroottes $N=1024$ versus $N=1023$ ) kunnen door de handling van randvoorwaarden (loop peeling) aanzienlijk verschillende TensorIR-teksten produceren, wat ertoe kan leiden dat het model hun prestatiepotentieel verkeerd classificeert ondanks identiek hardwaregedrag.

Methodologie

De auteurs stellen een Compiler World Model-framework voor dat de evaluatie van een schedule herformuleert als actie-geconditioneerde latente dynamiek over programmastaten. In plaats van een definitieve statische code-artefact te evalueren, simuleert het framework de evolutie van de programmastaat door een continue latente ruimte.

Het framework bestaat uit drie geleerde componenten die binnen de TVM AutoScheduler zijn geïmplementeerd:

1. Tensor Program State Representation (Encoder)

Input: TensorIR-tekst (zowel pre-schedule als post-schedule staten).
Architectuur: Een trainbare CodeBERT-gebaseerde architectuur die is aangepast voor lange inputs. Het gebruikt een sliding window om overlappende chunks te tokeniseren en te encoderen, die vervolgens worden geaggregeerd via gemaskeerde mean pooling tot een dichte vectorrepresentatie $z = E(s)$ .
Training: Gebruikt contrastief leren om vergelijkbare TensorIR-staten (van hetzelfde workload-bron) dicht bij elkaar in de latente ruimte te brengen, terwijl ongelijksoortige staten van elkaar worden weggeduwd. Dit creëert een staatruimte waarin semantische equivalentie behouden blijft ondanks syntactische variaties.

2. Actie-geconditioneerde Multi-Step State Transition

Doel: De terminal-staat representatie voorspellen die wordt geïnduceerd door een sequentie van scheduling-acties zonder de tussenliggende AST's expliciet te materialiseren.
Mechanisme: Startend vanaf een initiële staat-embedding $z_0$ $z_{0}$ , simuleert het model de effecten van een sequentie van scheduling-acties $a_{1:T}$ $a_{1 : T}$ stap voor stap.
- Acties worden geserialiseerd en gecodeerd in dichte embeddings.
- Een transitie-model $F(\cdot)$ , geïmplementeerd met een TransH-gebaseerde architectuur (oorspronkelijk afkomstig uit kennisgraaf-representatieleer), voorspelt de volgende latente staat: $z_t = F(z_{t-1}, a_t)$ .
- Dit modelleert scheduling-acties als geometrische transformaties in een relatie-specifiek hypervlak, waardoor de herkomst van beslissingen (decision provenance) effectief wordt gevolgd en referentiële ambiguïteiten worden opgelost (bijv. wat een lusvariabele representeert na fusie).
Training: Gebruikt een multi-step rollout loss die een state-matching term (het minimaliseren van de afstand tot de grondwaarheid van toekomstige staten) combineert met een margin-gebaseerde ranking term (om ervoor te zorgen dat de voorspelde staat dichter bij de grondwaarheid ligt dan bij negatieve samples).

3. Ranking-gebaseerde Kandidaat Evaluatie

Input: De voorspelde terminal-staat representatie $\hat{z}_T$ , gestructureerde actie-features (bijv. aantal acties, sequentielengte) en optionele hardware-features.
Output: Een scalaire score voor relatieve prioritering van kandidaten binnen dezelfde workload.
Architectuur: Een XGBoost ranking-model getraind met een LambdaRank objectief.
Training Strategie: Het model voorspelt geen absolute runtime. In plaats daarvan leert het kandidaten binnen een workload-groep te rangschikken op basis van gemeten executietijden, waarbij het probleem wordt behandeld als een learning-to-rank taak.

Belangrijkste Bijdragen

World-Model Perspectief voor Compilers: Het artikel past het wereldmodel-paradigma aan voor tensoroptimalisatie en vestigt een framework waar kandidaat-schedules worden gemodelleerd via actie-geconditioneerde latente staatsevolutie in plaats van statische snapshots.
Action-State Trajectory Dataset: Constructie van een TVM/TenSet-gebaseerde dataset die tuning-logs organiseert in uitgelijnde actie-staat trajecten (pre-schedule staat, actie-sequentie, tussenliggende staten, post-schedule staat) voor multi-step voorspelling.
Efficiënte Zoekintegratie: Implementatie van het framework in de TVM AutoScheduler, waarbij wordt aangetoond dat het modelleren van latente dynamiek de autotuning-kosten aanzienlijk vermindert terwijl de hoge zoekkwaliteit behouden blijft.

Experimentele Resultaten

De methode werd geëvalueerd op een Intel Xeon Gold 6430 CPU en een NVIDIA RTX 4090 GPU over zeven neurale netwerkmodellen (waaronder ResNet, BERT, GPT en OPT) en 103 geëxtraheerde subgraphs.

Prestaties vs. Ansor: Onder een budget van 64 proeven (trials) verbetert de methode de latency van representatieve subgraphs ten opzichte van Ansor met 1.37× op GPU en 1.54× op CPU.
Sample Efficiëntie: De methode bereikt een vergelijkbare zoekkwaliteit als Ansor-10K (10.000 proeven) met slechts 1.000 proeven (10× minder metingen). Specifiek presteert het iets beter dan Ansor-10K in het rekenkundig gemiddelde en komt het binnen 2,2% overeen in het geometrisch gemiddelde op representatieve subgraphs.
End-to-End Speedup: Vergeleken met PyTorch/PyTorch-opt (cuDNN) levert de methode geometrische gemiddelde versnellingen van 4.61× (CPU) en 3.67× (GPU), met pieksnelheden die oplopen tot 58.61×.
Ablatie-studies: Het verwijderen van de geleerde transitie-module (terugkeer naar statische initiële of finale staat-encoding) resulteerde in significante latency-toenames (mediaan ~56–69%), wat bevestigt dat het modelleren van de tussenliggende staatsevolutie de meest kritieke component is.

Betekenis en Beperkingen

Betekenis:
Het artikel beargumenteert dat het behandelen van compileroptimalisatie als een dynamisch proces van staatstransities, in plaats van een statisch code-evaluatieprobleem, het mogelijk maakt om de semantische afhankelijkheden tussen scheduling-beslissingen te vangen. Deze aanpak onderdrukt effectief syntactische ruis en lost referentiële ambiguïteiten op, wat leidt tot een efficiëntere zoektocht met minder hardware-metingen.

Beperkingen (zoals vermeld door de auteurs):

Afhankelijkheid van de Zoekruimte: De methode verbetert de evaluatie van kandidaten, niet de generatie. De effectiviteit ervan wordt begrensd door het vermogen van de onderliggende scheduler om sterke schedules voor te stellen; een betere evaluator kan geen schedules herstellen die nooit worden voorgesteld.
Foutaccumulatie: Voorspellingsfouten kunnen accumuleren over lange of complexe scheduling-trajecten in de latente ruimte, wat de betrouwbaarheid in moeilijke gevallen kan verminderen.
Ranking vs. Predictie: Het model is getraind voor ranking binnen een workload, niet voor absolute latency-voorspelling. Dit maakt het ideaal voor online prioritering, maar minder geschikt voor scenario's die gekalibreerde prestatieschattingen over verschillende workloads vereisen.