Axon: A Synthesizing Superoptimizer for Tensor Programs
Axon is een synthetiserende superoptimizer die automatisch hoogperformante tensorprogramma's voor AI-acceleratoren genereert door middel van programma-synthese en SMT-gebaseerde semantische verificatie om algebraïsche transformaties te ontdekken, tiling en instructie-selectie te optimaliseren, en operatoren te fuseren zonder afhankelijk te zijn van handmatig vervaardigde herschrijfregels.
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 het perfecte gebak wilt bakken, maar je werkt in een hightech keuken met zeer specifieke, veeleisende ovens en mengers. Je weet wat je wilt bakken (het recept), maar om de machine zo snel mogelijk te laten doen wat hij moet doen, moet je precies weten hoe je de ingrediënten laadt, wanneer je moet roeren en welke specifieke knoppen je op de oven moet indrukken.
Lamaag was het schrijven van code voor AI-chips (de "ovens") als het met de hand bakken van dat gebak, één kruimeltje per keer. Je moest een meesterkok zijn die de exacte mechanica van de oven begreep om een goed resultaat te krijgen. Als je een kleine fout maakte in hoe je de ingrediënten laadde, kwam het gebak traag en droog uit de oven.
Axon is een nieuwe tool die werkt als een geautomatiseerde, superintelligente sous-chef die niet alleen je recept volgt, maar ook betere manieren bedenkt om het te bakken.
Hier is hoe Axon werkt, opgedeeld in eenvoudige stappen:
1. Het "Wat" versus het "Hoe"
Normaal gesproken schrijft een programmeur een recept met de tekst: "Meng bloem en eieren, en bak dan." Ze maken zich geen zorgen over hoe de oven ze mengt.
- Het Probleem: De oven (de AI-chip) is veeleisend. Hij heeft verschillende compartimenten (zoals een grote mengkom, een kleine garde en een verwarmingselement). Als je de opdracht geeft om alles in de grote kom te mengen terwijl het ook met de garde had gekund, verspil je tijd.
- Axon's Taak: Axon neemt jouw eenvoudige recept en zegt: "Wacht, ik kan dit sneller doen." Het volgt niet alleen het recept; het herschrijft de instructies om aan te sluiten bij de specifieke sterktes van de oven.
2. Het "Magische Herordenen" (Algebraïsche Transformaties)
Stel je voor dat je recept zegt: "Hak eerst alle wortels, meng ze dan in de soep, en voeg dan zout toe."
- De Oude Manier: Je hakt, dan meng je, dan voeg je zout toe. Je moet de kom telkens heen en weer dragen tussen de snijplank en het fornuis.
- Axon's Manier: Axon kijkt naar de wiskunde en realiseert zich: "Hé, ik kan het zout toevoegen terwijl de wortels worden gehakt, of ik kan de wortels in de soep mengen voordat ik ze hak, en het zal precies hetzelfde smaken!"
- Het Resultaat: Door de volgorde van operaties te herordenen, laat Axon de oven twee dingen tegelijk doen (zoals hakken en verhitten gelijktijdig) in plaats van één na het ander. Het bewijst wiskundig dat het gebak nog steeds hetzelfde zal smaken, ook al zijn de stappen anders.
3. De "Verpakkingspuzzel" (Tiling)
Stel je voor dat je een enorme berg ingrediënten hebt, maar je mengkom is klein. Je moet ingrediënten verplaatsen van de voorraadkast (traag) naar de kom (snel), een beetje mengen, weer terugverplaatsen, en herhalen.
- Het Probleem: Als je te veel tegelijk verplaatst, loopt de kom over. Als je te weinig verplaatst, ben je de hele tijd heen en weer aan het lopen naar de voorraadkast.
- Axon's Manier: Axon is een meester in verpakken. Het probeert miljoenen verschillende manieren om de ingrediënten in "batches" (tegels/tiles) te groeperen. Het vraagt: "Wat als ik 100 wortels tegelijk verplaats? Wat als ik er 50 verplaats?" Het test deze verschillende groeperingen om de perfecte grootte te vinden die de kom vol houdt zonder dat deze overloopt, waardoor de reisjes naar de voorraadkast tot een minimum worden beperkt.
4. De "Instructiehandleiding" (ISA Synthese)
De oven heeft een specifiek bedieningspaneel met knoppen zoals "Hoge Hitte", "Garde" en "Vouwen".
- Het Probleken: Soms zegt het recept "Mixen", maar de oven heeft geen "Mix"-knop. Hij heeft alleen "Garde" en "Vouwen".
- Axon's Manier: Axon vertaalt "Mixen" naar de exacte sequentie van "Garde en dan Vouwen" die de oven begrijpt. Het vindt zelfs afkortingen, zoals het besef dat "Vouwen terwijl je naar het fornuis beweegt" een enkele knopdruk is op deze specifieke oven, wat tijd bespaart.
5. De "Proeverij" (Empirische Selectie)
Dit is het belangrijkste deel. Axon gokt niet welke manier het beste is.
- Het Proces: Het genereert duizenden verschillende versies van het recept (sommigen met een andere volgorde, anderen met andere batchgroottes, anderen met verschillende knopdrukken).
- De Test: Het voert al deze versies daadwerkelijk uit op de echte oven (de AI-chip) en meet de tijd. Het kiest de versie die het snelst klaar is. Het is alsoort bij het bakken van 1.000 kleine taartjes om te zien welke er in 5 minuten uit de oven komt in plaats van 10.
De Resultaten: Hoeveel Sneller?
Het paper testte Axon op Amazon's Trainium-chips (een type AI-oven) met 20 verschillende soorten "gebak" (AI-taken).
- Vs. De Standaard Oven: Vergeleken met de standaard manier van bakken (Amazon's Neuron compiler), maakte Axon enkelvoudige taken tot wel 3,7 keer sneller en complexe taken tot wel 19 keer sneller.
- Vs. Expert Chefs: Zelfs in vergelijking met recepten geschreven door menselijke experts die de oven van binnen en van buiten kennen, was Axon soms 1,35 keer sneller.
- Vs. Andere Robots: Het versloeg een andere geautomatiseerde tool genaamd "Mirage" met gemiddeld ongeveer 10%, en het kon 16 soorten gebak bakken die de andere tool helemaal niet kon afhandelen.
Samenvatting
Axon is een tool die een eenvoudige beschrijving neemt van wat een AI moet doen en automatisch de absoluut snelste manier vindt om dit op een specifieke computerchip te doen. Dit doet het door wiskundig te bewijzen dat het stappen kan herordenen, door miljoenen manieren te testen om data te verpakken, en door de code daadwerkelijk uit te voeren om de winnaar te kiezen. Het neemt de noodzaak weg voor mensen om experts te zijn in de minuscule, saaie details van de hardware, waardoor zij gewoon het recept kunnen schrijven terwijl Axon het bakken afhandelt.
Verdrinkt u in papers in uw vakgebied?
Ontvang dagelijkse digests van de nieuwste papers die bij uw onderzoekswoorden passen — met technische samenvattingen, in uw taal.