GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators

El marco \textsc{GOMA} propone un modelo analítico basado en abstracción geométrica para obtener automáticamente mapeos óptimos globalmente en aceleradores espaciales, mejorando significativamente la eficiencia energética y reduciendo el tiempo de búsqueda en comparación con los métodos actuales.

Lo esencial

¡Claro que sí! Imagina que tienes que organizar una gran fiesta de baile (que representa el cálculo matemático de una Inteligencia Artificial) en un gigantesco estadio (el chip de computadora o "acelerador espacial").

El problema es que hay millones de formas de organizar a los bailarines, mover las sillas y decidir quién trae la música. Si lo haces mal, la fiesta se vuelve un caos, la gente se cansa (gasta mucha energía) y tarda horas en terminar. Si lo haces bien, todo fluye, se gasta poca energía y termina en segundos.

Aquí está la explicación de GOMA (el nuevo método del paper) usando analogías sencillas:

1. El Problema: Un Laberinto Gigante

Antes de GOMA, los expertos intentaban encontrar la mejor forma de organizar la fiesta probando millones de combinaciones al azar o usando reglas de "intuición".

• El problema: El número de formas de organizar la fiesta es tan enorme (más que los granos de arena en la playa) que probar todas es imposible. Los métodos antiguos tardaban días en encontrar una solución "bueno, pero no perfecta", o a veces se quedaban atascados en una solución mala.

2. La Idea de GOMA: El Mapa Geométrico

Los autores de GOMA tuvieron una idea brillante: dejar de ver la fiesta como un caos y empezar a verla como una caja de cubos geométros.

• La Analogía de los Cubos: Imagina que toda la tarea de baile es un cubo gigante de 3D.
  • Los bailarines son los puntos dentro del cubo.
  • La música y los movimientos son las proyecciones de ese cubo en las paredes (arriba, abajo, a los lados).
• La Magia: En lugar de probar millones de formas de mover los cubos, GOMA usa matemáticas puras (modelado analítico) para dibujar un mapa exacto de cuánta energía se gasta en cada movimiento.
  • Es como tener una fórmula mágica que te dice: "Si mueves el cubo hacia la izquierda, gastas 2 monedas de energía. Si lo mueves hacia arriba, gastas 1".
  • Esta fórmula es tan rápida que puede calcular el costo de cualquier organización en una fracción de segundo (O(1), o sea, instantáneo).

3. La Solución: El Arquitecto Perfecto

Una vez que tienen esa fórmula mágica, GOMA no adivina ni prueba al azar. Actúa como un arquitecto matemático infalible:

1. Define las reglas: "El estadio tiene un tamaño fijo", "no podemos tener más bailarines que sillas", etc.
2. Busca el camino óptimo: Usa un algoritmo inteligente (optimización entera) para encontrar la única y mejor forma de organizar la fiesta que cumpla todas las reglas y gaste la mínima energía posible.
3. El Certificado de Garantía: A diferencia de los otros métodos que dicen "creo que esta es la mejor", GOMA entrega un certificado matemático que prueba: "Sí, esta es la mejor solución posible. No existe ninguna otra forma de hacerlo mejor".

4. ¿Por qué es tan importante? (Los Resultados)

El paper prueba GOMA con modelos de Inteligencia Artificial muy grandes (como los que usan Chatbots avanzados) en diferentes tipos de chips.

• Ahorro de Energía: GOMA encuentra formas de organizar la fiesta que hacen que el chip gaste 2 a 4 veces menos energía que los métodos actuales. ¡Es como si pudieras bailar toda la noche con la batería de un solo AA!
• Velocidad: Mientras que los otros métodos tardaban horas o días en decidir cómo organizar la fiesta, GOMA lo hace en segundos.
• Precisión: En lugar de adivinar, GOMA calcula exactamente qué pasa si saltas un nivel de memoria (una técnica llamada "bypass"), algo que otros métodos ignoraban o hacían mal.

En Resumen

Imagina que antes tenías que buscar la llave correcta en un montón de 10 millones de llaves a ciegas.
GOMA es como tener un detector de metales perfecto que te señala exactamente cuál es la llave correcta al instante, y además te da un certificado que dice: "Esta es la única llave que abre esta puerta".

Es una herramienta que permite a las computadoras ser más rápidas, más baratas (menos energía) y más inteligentes, simplemente organizando mejor sus tareas internas.

Resumen técnico

Resumen Técnico: GOMA

1. El Problema

La multiplicación de matrices general (GEMM) es el operador computacional dominante en modelos de aprendizaje profundo modernos, como los Transformers y los Grandes Modelos de Lenguaje (LLM). En los aceleradores espaciales (hardware dedicado para IA), la eficiencia de ejecución y el consumo energético dependen críticamente de la estrategia de mapeo (mapping).

Sin embargo, el espacio de búsqueda de mapeos presenta una explosión combinatoria. Para un GEMM típico, el espacio de mapeo supera $10^{10}$posibilidades, y para convoluciones puede exceder$10^{20}$.

• Desafío actual: Los métodos existentes (búsqueda aleatoria, heurística de caja negra, aproximaciones diferenciables o programación matemática) enfrentan un dilema: o bien son rápidos pero no garantizan la optimalidad global (quedándose en óptimos locales), o bien son precisos pero computacionalmente prohibitivos para espacios grandes.
• Objetivo: Encontrar un mapeo globalmente óptimo que minimice el costo (energía) dentro de un tiempo de solución aceptable, garantizando matemáticamente que no existe una solución mejor bajo las restricciones dadas.

2. Metodología: GOMA

GOMA propone un marco de trabajo basado en abstracción geométrica y modelado analítico para transformar el problema de búsqueda de mapeos en un problema de optimización entera resoluble de forma eficiente.

A. Abstracción Geométrica (El Núcleo)
En lugar de tratar el mapeo como una secuencia de bucles anidados, GOMA modela el GEMM como una cuadrícula de cómputo 3D ($x, y, z$), donde cada punto representa una operación MAC (multiplicación-acumulación).

• Proyecciones: Las matrices de entrada ($A, B$) y salida ($P$) se interpretan como proyecciones ortogonales de esta cuadrícula 3D sobre los planos $x-z$, $y-z$ y $x-y$.
• Eje de Travesía (Walking Axis): La ordenación de los bucles (permutación) determina en qué dirección avanza un "tile" (bloque de cálculo). GOMA observa que al avanzar en un eje, una proyección permanece constante (permitiendo reutilización temporal de datos) mientras las otras dos se actualizan.
• Conteo de Actualizaciones: El tráfico de datos se calcula contando cuántas veces se actualizan estas proyecciones durante la traversía. Esto permite derivar un modelo de tráfico exacto.

B. Modelado Analítico de Energía (Objetivo $O(1)$)
GOMA deriva una función objetivo de energía analítica y cerrada (closed-form) basada en principios primeros:

• Evaluación Constante: Para cualquier mapeo dado, la evaluación de la energía tiene una complejidad de $O(1)$, independiente del tamaño del problema o del número de tiles.
• Manejo de Bypass (Nivel de Memoria): El modelo incorpora explícitamente la decisión de "bypass" (saltar un nivel de memoria). Si un nivel se salta, el modelo recalcula la cadena de acceso, asignando el costo de energía al nivel de origen y destino reales, no solo a niveles adyacentes.
• Precisión: El modelo logra una consistencia del 99.9% con Timeloop (el estándar de la industria para modelado de energía), validado en miles de configuraciones.

C. Formulación como Problema de Optimización Entera
El problema se formula como un Programa Lineal Entero (ILP) con restricciones duras:

• Variables de decisión: Formas de los tiles en cada nivel jerárquico, ejes de travesía (permutación de bucles) y políticas de bypass por eje.
• Restricciones: Capacidad de memoria (SRAM, RegFile), número de unidades de procesamiento (PE) y divisibilidad anidada (los tamaños de los tiles deben ser divisibles entre niveles).
• Solución: Se utiliza un solucionador global (Gurobi con ramificación y acotación) para encontrar la solución óptima global. El proceso proporciona un certificado de optimalidad (gap de optimalidad = 0%), garantizando que la solución es la mejor posible bajo el modelo.

3. Contribuciones Clave

1. Abstracción Geométrica Unificada: Introducción de una nueva representación basada en proyecciones ortogonales de una cuadrícula de cómputo 3D, que unifica la tiling, la permutación de bucles y el bypass en un solo marco matemático.
2. Modelo de Energía Exacto y Rápido: Desarrollo de un objetivo de energía analítico de $O(1)$ que es extremadamente preciso (99.9% de consistencia con simuladores detallados) y permite la búsqueda automatizada sin aproximaciones heurísticas.
3. Optimalidad Global Verificable: Es la primera vez que se logra computar un mapeo globalmente óptimo para cualquier par (carga de trabajo GEMM, hardware) en tiempo razonable, con un certificado matemático de optimalidad.
4. Código de Código Abierto: El framework GOMA ha sido liberado en GitHub para la comunidad.

4. Resultados Experimentales

Los autores evaluaron GOMA en 24 casos que combinan 4 aceleradores espaciales representativos (desde Edge como Eyeriss/Gemmini hasta Data Center como A100/TPU) y cargas de trabajo de prellenado (prefill) de LLMs (Qwen y LLaMA de 0.6B a 70B).

• Eficiencia Energética (EDP): GOMA mejora el Producto Energía-Delay (EDP) entre 2.24x y 4.24x en comparación con los mapeadores del estado del arte (SOTA) como CoSA, FactorFlow, LOMA y SALSA.
• Velocidad de Solución: GOMA es significativamente más rápido, reduciendo el tiempo de búsqueda entre 3.83x y 73.6x. Mientras que otros métodos tardan minutos u horas en problemas grandes, GOMA encuentra soluciones óptimas en segundos (promedio de 0.65s por GEMM).
• Estabilidad: A diferencia de los métodos heurísticos que muestran fluctuaciones grandes dependiendo de la carga de trabajo, GOMA mantiene un rendimiento superior y estable en todos los escenarios, especialmente en matrices grandes donde la explosión combinatoria afecta más a los métodos de muestreo.
• Análisis por Capa: GOMA demuestra ventajas particularmente grandes en operaciones de matriz-matriz complejas (como las capas de atención y MLP), donde los métodos heurísticos fallan en encontrar soluciones cercanas al óptimo.

5. Significado e Impacto

El trabajo de GOMA representa un avance fundamental en el diseño de aceleradores espaciales:

• Cambio de Paradigma: Pasa de la búsqueda heurística (probabilística) a la optimización analítica determinista para la exploración de espacios de mapeo.
• Viabilidad Práctica: Demuestra que es posible obtener soluciones globalmente óptimas en tiempo real, lo cual es crucial para el despliegue eficiente de modelos de IA en hardware diverso.
• Fundamento para Futuras Investigaciones: Proporciona una base sólida para explorar escenarios más complejos, como la co-optimización software-hardware y la exploración de mapeos multi-capas, al ofrecer un modelo de costo preciso y resoluble matemáticamente.

En resumen, GOMA resuelve el problema de la "explosión combinatoria" en el mapeo de aceleradores espaciales mediante una elegante abstracción geométrica, logrando la combinación ideal de optimalidad garantizada y eficiencia computacional.