GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators

This paper presents GOMA, a geometrically abstracted, globally optimal framework that uses analytical modeling to efficiently solve the combinatorial GEMM mapping problem for spatial accelerators, achieving significant improvements in energy-delay product and search speed over state-of-the-art methods.

The Gist

Imagine you are the manager of a massive, high-tech factory (a Spatial Accelerator) tasked with assembling millions of tiny Lego structures (performing Matrix Multiplications, the math behind AI).

Your factory has different storage rooms: a giant warehouse outside (DRAM), a large storage room inside (SRAM), workbenches (PE Arrays), and small toolboxes right next to the workers' hands (Registers).

The goal is to build the structures as fast and energy-efficiently as possible.

The Problem: The "Choice" Nightmare

To get the job done, you have to decide:

1. How to pack the Lego bricks: Do you bring in a huge box of bricks at once, or small handfuls? (This is Tiling).
2. The order of operations: Do you build the left side first, then the right? Or do you build the top, then the bottom? (This is Loop Permutation).
3. What to keep on the workbench: Do you keep the bricks in the toolbox, or just grab them from the warehouse every single time you need one? (This is Bypass).

The problem is that there are trillions of possible ways to arrange these choices. It's like trying to find the perfect outfit by trying on every single combination of shirts, pants, and shoes in the world. If you try them all one by one, you'll be old before you finish. If you just guess (which most current AI tools do), you might get a decent outfit, but you'll miss the perfect one that saves the most energy.

The Solution: GOMA (The "Geometric GPS")

The paper introduces GOMA, a new tool that acts like a GPS for your factory. Instead of guessing or trying every path, GOMA uses a clever trick called Geometric Abstraction.

1. The "Shadow" Analogy (Geometric Abstraction)

Imagine your Lego structure is a 3D cube. GOMA doesn't look at the cube itself; it looks at the shadows the cube casts on the walls.

• One shadow shows the "Left/Right" view.
• One shows the "Front/Back" view.
• One shows the "Top/Bottom" view.

GOMA realizes that every time you move a step in the factory, you only change two of these shadows, while one stays the same.

• If you move "Forward," the "Top" shadow doesn't change, so you don't need to fetch new "Top" bricks. You reuse what you have!
• If you move "Sideways," the "Front" shadow stays the same, so you reuse those bricks.

By tracking these shadows, GOMA can instantly calculate exactly how many bricks need to be moved between storage rooms. It turns a complex, messy puzzle into a simple math equation that can be solved in a split second.

2. The "Shortcut" (Analytical Modeling)

Most other tools try to simulate the factory running to see how much energy it uses. This is like actually building the Lego set to see how long it takes. It's slow.

GOMA, however, has a magic formula. Because it understands the geometry of the shadows, it can predict the energy cost of any arrangement instantly (in O(1) time, which means it takes the same tiny amount of time whether you have 10 bricks or 10 billion).

3. The "Perfect Solution" (Global Optimality)

Because GOMA has this instant formula, it can use a powerful math solver to look at all the possibilities at once and find the single best arrangement.

• Old tools: "I'll try 1,000 random outfits and pick the best one I found." (Might miss the perfect one).
• GOMA: "I have a map of the entire world. I can prove mathematically that this specific outfit is the absolute best possible one."

Why Does This Matter?

The researchers tested GOMA on modern AI models (like the ones powering chatbots and image generators).

• Energy Savings: GOMA found ways to run these models 2 to 4 times more efficiently than the current best tools. That means your phone battery lasts longer, or your data center uses less electricity.
• Speed: GOMA figured out the best plan 4 to 70 times faster than the competition. It didn't waste time guessing; it just calculated the answer.

The Takeaway

Think of GOMA as the difference between a tourist wandering around a city hoping to find the best restaurant, and a local who has a perfect map and knows exactly which restaurant is the best, how to get there, and how much it costs, without ever having to walk the whole city.

It takes a chaotic, overwhelming problem (choosing how to run AI on chips) and turns it into a clean, solvable geometry puzzle, guaranteeing the most efficient result every time.

Technical Summary

Here is a detailed technical summary of the paper "GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators."

1. Problem Statement

The paper addresses the critical challenge of Mapping Space Exploration (MSE) for General Matrix Multiplication (GEMM) on spatial accelerators (e.g., TPUs, Eyeriss, Gemmini).

• The Challenge: The mapping space (comprising tiling, loop permutation, and level bypass) grows combinatorially (often exceeding $10^{10}$ for GEMM). Finding the optimal mapping that minimizes energy and latency is computationally infeasible via exhaustive search.
• Limitations of Existing Methods:
  • Search-based (Random/Heuristic): Methods like genetic algorithms or reinforcement learning lack global optimality guarantees and often get stuck in local optima.
  • Differentiable Approximations: Relaxing integer constraints introduces approximation errors and rounding issues, breaking optimality guarantees.
  • Mathematical Programming: Existing approaches (e.g., CoSA) struggle to accurately model real hardware energy costs (often using proxies like resource utilization) and suffer from redundant search spaces, leading to slow solving times for large workloads.
• Goal: To develop a framework that can compute a provably globally optimal mapping for any (GEMM workload, hardware) pair within an acceptable time budget, while guaranteeing energy efficiency.

2. Methodology: GOMA Framework

GOMA (Geometrically Optimal Mapping) introduces a novel approach based on geometric abstraction and analytical modeling to transform the mapping problem into a solvable integer optimization problem.

A. Geometric Abstraction (The Core Intuition)

Instead of treating loops and dataflow as abstract code, GOMA models GEMM as a 3D Compute Grid ($x, y, z$):

• Compute Grid: The entire computation is a set of points $(x, y, z)$.
• Projections: The three input/output matrices ($A, B, P$) correspond to the orthogonal projections of this grid onto the $x-z$, $y-z$, and $x-y$ planes.
• Traversal & Reuse:
  • Walking Axis: As a 3D tile advances along a specific axis (e.g., $y$), the projection orthogonal to that axis (e.g., $x-z$) remains unchanged, enabling temporal reuse of that data.
  • Update Counting: Data movement is quantified by counting how many times each projection plane must be updated during traversal. This allows for an exact calculation of data traffic.
• Level Bypass: The framework models "bypass" decisions (skipping a memory level) by rewriting the data access chain. If a level is bypassed, data is supplied directly from a higher level, altering the energy attribution.

B. Analytical Energy Modeling

GOMA derives a closed-form, $O(1)$ energy objective function based on the geometric model:

• Traffic Calculation: It calculates the number of word transfers for each data type ($A, B, P$) between memory levels (DRAM $\to$ SRAM $\to$ PE-array $\to$ Regfile $\to$ MACC) based on tile sizes and walking axes.
• Energy Weighting: These traffic counts are multiplied by hardware-specific unit energy costs (read/write/leakage).
• Boundary Handling: Special handling is applied to the reduction axis ($z$) for partial sums, distinguishing between the first accumulation (write-only) and subsequent steps (read-accumulate-write).
• Accuracy: The model achieves 99.9% consistency with the widely used Timeloop simulator, validating its fidelity without requiring simulation.

C. Global Optimization Formulation

The mapping selection is formulated as a constrained Integer Linear Programming (ILP) problem:

• Variables: Tiling dimensions ($L^{(p)}_d$), walking axes ($\alpha$), and bypass decisions ($B$).
• Constraints: Hardware capacity (SRAM/Regfile), PE count, and divisibility/nesting constraints.
• Solver: The problem is solved using a branch-and-bound solver (Gurobi). Because the objective function is exact and the search space is reduced via geometric folding, the solver can find the global optimum and provide a verifiable optimality certificate (gap = 0%).

3. Key Contributions

1. Geometric Abstraction: A first-principles derivation of GEMM mapping as a hierarchical geometric traversal problem, yielding an exact analytical energy objective with $O(1)$ evaluation complexity.
2. Global Optimality: The first mapping framework to guarantee a globally optimal solution for GEMM on spatial accelerators within seconds, outputting a verifiable certificate of optimality.
3. Unified Optimization: A unified integer optimization formulation that jointly optimizes tiling, loop permutation, and level bypass under hard hardware constraints.
4. High Fidelity: An energy model that matches the Timeloop simulator with 99.9% consistency across diverse workloads and mapping configurations.

4. Experimental Results

The authors evaluated GOMA on four representative accelerator templates (Eyeriss-like, Gemmini-like, A100-like, TPU v1-like) and 12 LLM prefill workloads (ranging from 0.6B to 70B parameters).

• Energy-Delay Product (EDP):
  • GOMA outperforms State-of-the-Art (SOTA) mappers (including CoSA, FactorFlow, LOMA, SALSA, and Timeloop-Hybrid) by 2.24× to 4.24× in EDP.
  • It consistently finds the lowest EDP across all hardware scales and model sizes.
• Solving Speed:
  • GOMA accelerates the time-to-solution by 3.83× to 73.6× compared to baselines.
  • It solves complex cases in seconds (average 0.65s per GEMM), whereas baselines often take minutes or fail to converge.
• Scalability:
  • While heuristic methods struggle with large mapping spaces (showing instability and suboptimal results), GOMA's analytical approach maintains stable performance and determinism even for ultra-large workloads (e.g., 70B models on A100-like hardware).
• Bypass Impact: The study highlights that explicit modeling of level bypass is a critical degree of freedom that significantly impacts energy efficiency, a feature often ignored or poorly handled by other mappers.

5. Significance

• Theoretical Breakthrough: GOMA bridges the gap between the need for exact hardware cost modeling and the need for fast, global optimization. It proves that with the right geometric abstraction, the "combinatorial explosion" of mapping can be tamed analytically.
• Practical Impact: For the rapidly growing field of Large Language Models (LLMs) and Diffusion Transformers, where GEMM dominates compute, GOMA provides a tool to automatically generate the most energy-efficient execution strategies for specific hardware targets.
• Foundation for Future Work: By providing a reusable, globally optimal foundation, GOMA enables more complex future research, such as multi-layer mapping exploration and full software-hardware co-design, without being bottlenecked by slow or suboptimal mapping search.

In summary, GOMA represents a paradigm shift from heuristic search to analytical global optimization for spatial accelerator mapping, delivering provably optimal solutions with unprecedented speed and accuracy.