GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators

GOMA is een op analytische modellering gebaseerd framework dat via een geometrische abstractie binnen een acceptabele tijd een globaal optimale mapping voor GEMM-werklasten op ruimtelijke acceleratoren berekent, wat resulteert in aanzienlijke verbeteringen in energie-efficiëntie en doorlooptijd ten opzichte van de huidige stand van de techniek.

De kern

Stel je voor dat je een gigantische, complexe puzzel moet oplossen: het uitvoeren van wiskundige berekeningen (specifiek matrixvermenigvuldiging, de ruggengraat van moderne AI) op een speciale computerchip. Deze chip heeft verschillende lagen van geheugen, zoals een grote opslagruimte (DRAM), een snellere werkbank (SRAM), en heel veel kleine rekenmachientjes (PE's) die de feitelijke werk doen.

Het probleem? Er zijn ontelbaar veel manieren om deze puzzelstukjes op deze chip te leggen. Je kunt kiezen welke data waar wordt opgeslagen, in welke volgorde de berekeningen plaatsvinden, en welke stukjes data je direct overslaat om tijd te besparen.

Als je dit allemaal willekeurig probeert (zoals veel huidige systemen doen), duurt het eeuwen om de beste oplossing te vinden, en vind je vaak toch niet de allerbeste.

Hier komt GOMA (Geometrically Optimal Mapping via Analytical Modeling) om de hoek kijken. Het is een slimme nieuwe methode die deze puzzel oplost in een flits, terwijl het garandeert dat je de perfecte oplossing hebt gevonden.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "3D-kaart" in plaats van een lijstje

Stel je de berekening voor als een enorme, driedimensionale kubus van blokjes (zoals een Rubik's kubus, maar dan met miljarden blokjes).

• Huidige methoden: Kijken naar deze kubus als een lange, saaie lijst met regels en proberen die te husselen.
• GOMA's aanpak: Kijkt naar de kubus als een geometrisch object. Het ziet de kubus van drie kanten (voorkant, zijkant, bovenkant). In plaats van te husselen, berekent GOMA precies hoeveel "verkeer" er nodig is om deze kubus te vervoeren tussen de verschillende geheugenniveaus. Het is alsof je in plaats van te raden welke route de snelste is, direct de afstand en het brandstofverbruik van elke mogelijke route op een kaart kunt aflezen.

2. De "Loopbaan" van de data

Stel je voor dat je een pakketje (data) moet bezorgen.

• De "Loopbaan" (Walking Axis): GOMA beslist welke kant op je het pakketje het beste kunt vervoeren. Als je in de richting van de "z-as" loopt, blijft het pakketje op de "x-as" en "y-as" rustig liggen. Dat betekent dat je die data niet steeds opnieuw hoeft te halen; je kunt het hergebruiken. GOMA berekent wiskundig precies welke loopbaan het minst energie kost, net als een GPS die de route met de minste stoplichten en de kortste afstand kiest.

3. Het "Slaan" van de brug (Bypass)

Soms is het slimmer om een tussenstation over te slaan.

• De Analogie: Stel je hebt een postkantoor (SRAM) en een brievenbus bij de deur (Register). Soms is het sneller en goedkoper om de post direct van het centrale depot (DRAM) naar de brievenbus te sturen, zonder dat het eerst naar het postkantoor gaat.
• GOMA beslist automatisch welke data deze "brug" mag overslaan en welke data wel even moet rusten. Dit bespaart enorm veel energie, omdat het vermijden van onnodige stops de batterij van je apparaat (of datacenter) spaart.

4. De "Perfecte" Oplossing in Seconden

De meeste andere systemen zijn als een zoektocht in een donker bos: ze lopen wat rond, hopen op geluk, en hopen dat ze de schat vinden. Als ze het bos te groot is, raken ze verdwaald of geven ze op.

GOMA is als een luchtfoto van dat bos.

• Het maakt een wiskundig model (een formule) dat precies zegt hoeveel energie elke route kost.
• Omdat het een formule is, kan het de oplossing direct berekenen in plaats van te gissen.
• Het geeft niet alleen de oplossing, maar ook een certificaat dat zegt: "Dit is wiskundig bewezen de beste oplossing die er bestaat." Je hoeft niet te twijfelen of er misschien nog een betere is.

Wat levert dit op?

In tests met grote AI-modellen (zoals die gebruikt worden voor chatbots) op verschillende soorten chips:

• Energiebesparing: GOMA gebruikt 2 tot 4 keer minder energie dan de beste bestaande systemen.
• Snelheid: Het vindt de oplossing 4 tot 73 keer sneller. Waar andere systemen uren nodig hebben om te denken, doet GOMA dit in seconden.

Kortom: GOMA is de slimme architect die, in plaats van te bouwen en te slopen om te zien wat werkt, direct de perfecte blauwdruk tekent die de minste energie kost en het snelst gebouwd kan worden. Het maakt het gebruik van dure AI-chips veel efficiënter en goedkoper.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators" in het Nederlands.

Probleemstelling

General Matrix Multiplication (GEMM) is een van de meest kritische berekeningskernen in moderne deep learning-modellen, zoals Transformers en Large Language Models (LLMs). Op ruimtelijke versnellers (spatial accelerators) is de prestatie en energie-efficiëntie van GEMM extreem gevoelig voor de gekozen "mapping" (toewijzing van de berekening op de hardware). Een mapping bepaalt hoe data wordt getiled (opgedeeld), de volgorde van loops (permutatie) en of data bepaalde geheugenniveaus overslaat (bypass).

Het centrale probleem is dat de zoekruimte voor optimale mappings combinatorisch explodeert. Voor een typische GEMM-operatie kan deze ruimte groter zijn dan $10^{10}$ mogelijkheden. Bestaande methoden voor het verkennen van deze ruimte (Mapping Space Exploration - MSE) hebben twee grote tekortkomingen:

1. Geen garantie op optimaliteit: Methoden zoals random search, heuristieken (genetische algoritmen, RL) en differentieerbare benaderingen vinden vaak suboptimale oplossingen omdat ze de discrete aard van het probleem verwaarlozen of lokale optima vastlopen.
2. Schaalbaarheid: Methoden die exacte oplossingen proberen te vinden (zoals wiskundige programmering) zijn vaak te traag of modelleren de hardwarekosten onnauwkeurig, waardoor ze niet binnen een acceptabele tijd een globaal optimum kunnen vinden voor grote werklasten.

Methodologie: GOMA

GOMA (Geometrically Optimal Mapping via Analytical Modeling) is een framework dat voor het eerst een globaal optimale mapping berekent binnen een korte tijd, met een verifieerbaar bewijs van optimaliteit. De kern van de methode bestaat uit drie stappen:

1. Geometrische Abstractie (3D Compute Grid)
In plaats van te werken met complexe loop-nesting, modelleert GOMA de GEMM-berekening als een 3D compute-grid ($x, y, z$), waarbij elke punt een MAC-operatie (Multiply-Accumulate) vertegenwoordigt.

• De drie invoer/uitvoer-matrices (A, B, P) worden gezien als orthogonale projecties van dit 3D-grid op de vlakken $x-z$, $y-z$ en $x-y$.
• Tiling wordt vertaald naar het hiërarchisch bedekken van dit grid.
• Loop-permutatie wordt vertaald naar de "wandelas" (walking axis): de richting waarin een tile beweegt. Als een tile beweegt langs een as, blijven de projecties loodrecht op die as tijdelijk onveranderd, wat temporale hergebruik (reuse) van data mogelijk maakt.
• Level Bypass wordt gemodelleerd als een beslissing of een data-type een geheugenniveau overslaat, wat de energiekosten van data-overdracht direct beïnvloedt.

2. Analytisch Energie-model (O(1) Evaluatie)
GOMA leidt een exacte, gesloten-formule energie-objectief af.

• Het berekent het dataverkeer niet door simulatie, maar door het tellen van het aantal updates van de drie projecties tijdens de traversie van het grid.
• Het model houdt rekening met de specifieke semantiek van de reductie-as ($z$), waarbij de eerste stap in een accumulatieketen geen "read-old" vereist.
• Het resultaat is een energieberekening die O(1) (constante tijd) duurt, ongeacht de grootte van de matrix of het aantal tiles. Dit model toont 99,9% consistentie met de gevestigde Timeloop-simulator.

3. Integer Optimalisatie
Het zoekprobleem wordt omgezet in een geconstrueerd integer optimalisatieprobleem.

• Doel: Minimaliseren van de totale energie (of EDP).
• Variabelen: Tile-groottes op elk niveau, keuze van de wandelassen voor elke stap, en bypass-configuraties.
• Constraints: Hardware-capaciteit (SRAM, Register File), aantal Processing Elements (PEs), en deelbaarheid (divisibility) tussen niveaus.
• Oplossing: Het probleem wordt opgelost met een globale solver (zoals Gurobi met Branch-and-Bound). Omdat de energie-functie exact en analytisch is, kan de solver een verifieerbaar bewijs van optimaliteit leveren (een gap van 0% tussen de onder- en bovengrens).

Belangrijkste Bijdragen

1. Nieuwe Geometrische Abstractie: Een eerste-principes benadering die GEMM-mapping reduceert tot een hiërarchisch geometrisch traversieprobleem, resulterend in een exacte analytische energieformule.
2. GOMA Framework: Een uniform formulering van mapping als een integer optimalisatieprobleem dat tiling, loop-permutatie en bypass simultaan optimaliseert.
3. Globale Optimaliteit: Voor het eerst in mapping-exploratie kan GOMA binnen seconden een globaal optimale mapping vinden voor elke (werklast, hardware)-combinatie, inclusief een verifieerbaar optimaliteitscertificaat.
4. Open Source: De code is beschikbaar gesteld op GitHub.

Resultaten

De auteurs hebben GOMA getest op vier representatieve versnellers (Eyeriss-achtig, Gemmini-achtig, A100-achtig, TPU v1-achtig) en werklasten van LLMs (Qwen3 en LLaMA-3 series) in de prefill-fase.

• Energie-Delay Product (EDP): GOMA verbetert de EDP met 2,24x tot 4,24x ten opzichte van state-of-the-art (SOTA) mappers (zoals CoSA, FactorFlow, LOMA, SALSA en Timeloop Hybrid).
• Snelheid: De tijd om tot een oplossing te komen (time-to-solution) is 3,83x tot 73,6x sneller dan bestaande methoden. GOMA lost complexe gevallen vaak op in seconden, terwijl andere methoden minuten of uren nodig hebben of niet convergeren.
• Stabiliteit: Heuristische methoden vertonen grote fluctuaties in kwaliteit afhankelijk van de werklast, terwijl GOMA consistent het beste resultaat levert.
• Scalability: Waar andere methoden (zoals CoSA) langzamer worden naarmate de werklast groter wordt (door combinatorische explosie), behoudt GOMA een stabiele en snelle oplostijd dankzij de analytische O(1) evaluatie.

Betekenis en Impact

GOMA lost een langdurig open probleem op in het co-design van hardware en algoritmen: het vinden van een verifieerbaar optimale mapping voor ruimtelijke versnellers binnen een praktische tijdslimiet.

• Het bewijst dat analytische modellering superieur is aan simulatie-gebaseerde zoektochten voor dit type probleem.
• Het biedt een fundamentele basis voor toekomstige onderzoekrichtingen, zoals multi-layer mapping exploratie en software-hardware co-optimisatie.
• Voor de industrie betekent dit dat versnellers voor LLMs en andere AI-werklasten aanzienlijk energie-efficiënter kunnen worden ingesteld zonder handmatige tuning of riskante heuristieken.