GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators

原文链接: arXiv:2603.07962

PDF

论文标题: GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators
作者: Wulve Yang, Hailong Zou, Rui Zhou, Jionghao Zhang, Qiang Li, Gang Li, Yi Zhan, Shushan Qiao
机构: 中国科学院微电子研究所，中国科学院大学
发表: IEEE Transactions on Computers, 2026

摘要

空间加速器上的通用矩阵乘法（GEMM）对映射选择高度敏感，执行效率和能耗都会受到显著影响。然而，映射空间呈现组合爆炸，使得在可接受的时间内获得最优映射极具挑战性。现有方法通常面临两大挑战：缺乏全局最优性保证，且随着映射空间增长变得极其缓慢。

本文提出 GOMA，一个基于几何抽象和解析建模的全局最优 GEMM 映射框架。GOMA 从第一性原理出发，引入 GEMM 映射的几何抽象，生成精确的解析能量目标函数，对任何给定映射的评估复杂度为 O(1)。GOMA 将映射选择表述为硬件和映射约束下的整数优化问题，使用解析能量模型作为目标函数自动化映射搜索。实验表明，在代表性加速器和大语言模型预填充工作负载上，GOMA 相比 SOTA 映射器将能效延迟积（EDP）提升 2.24–4.24 倍，同时加速求解时间 3.83–73.6 倍。

1. 问题定义

随着深度学习的广泛应用，核心算子如通用矩阵乘法（GEMM）和卷积在边缘设备和数据中心都产生了巨大的计算需求。其中，GEMM 在 Transformer 等模型中占据主导地位，已成为最关键的计算内核之一。

为了满足这一需求，用于 DNN 加速的空间加速器已成为学术界和工业界的重要研究方向。这一趋势催生了经典设计如 Eyeriss、Gemmini 和 Google TPU，以及大量近期工作。

“在此类加速器的硬件 - 算法协同设计中，映射是决定性能和能效的关键组件。”

映射描述了目标算法在特定硬件实例上的具体执行策略，包括计算分块（tiling）、循环置换（loop permutation）和层级旁路（level bypass）。映射既重要又具有挑战性：

• 影响巨大：即使在同一硬件上部署同一算法，不同映射也可能导致能效和性能相差几个数量级
• 空间巨大：典型卷积的映射空间可超过 10²⁰，而 GEMM 的映射空间远超 10¹⁰，使得穷举搜索最优映射在计算上不可行

近年来提出了许多映射空间探索（MSE）方法，包括随机搜索、黑盒启发式搜索、可微模型近似和数学规划。然而，这些方法在大规模离散空间中难以同时实现效率和解决方案质量：

• 搜索方法：通常难以在巨大离散空间中同时保证效率和解决方案质量
• 可微近似：通常松弛整数约束，引入近似误差并削弱结果保证
• 数学规划：现有工作（如 CoSA）仍难以精确表征真实硬件成本，整体求解效率有限

因此，在可接受时间内获得可证明的最优映射仍然是一个紧迫的开放问题。

2. 方法框架

GOMA 的核心创新在于将 GEMM 映射问题转化为一个几何遍历问题，通过解析建模实现快速全局最优求解。

图：GOMA 整体工作流程（来源：原文 Figure 5）

2.1 3D 计算网格几何抽象

GOMA 从第一性原理出发，将 GEMM 计算抽象为一个 3D 计算网格。对于矩阵乘法：

P(x, y) = Σ A(x, z) × B(y, z)  (z 从 1 到 L_z)

每个 MAC（乘累加）操作对应一个计算点 (x, y, z)，整个计算可表示为全局计算点集 G：

G = {(x, y, z) | x ∈ [1, L_x], y ∈ [1, L_y], z ∈ [1, L_z]}

图：GEMM 作为 3D 计算网格及其三个正交投影，分别对应 A(x,z)、B(y,z) 和部分和/输出 P(x,y)（来源：原文 Figure 3）

在这一几何表示下，三个矩阵自然地与三个正交投影平面对齐：

• A 矩阵 ↔ x-z 平面投影
• B 矩阵 ↔ y-z 平面投影
• 输出 P ↔ x-y 平面投影

对于任何 3D 计算瓦片（tile），它在三个正交平面上的覆盖面积分别给出了 A、B、P 的数据需求规模。

2.2 层级瓦片与内存层次结构

忽略层级旁路，多级内存层次结构可以理解为通过层级瓦片逐步覆盖 3D 计算网格 G：

• MACC 层级：单个 MACC 点 (x, y, z) 是最小单位
• 寄存器文件层级：单个 PE 的寄存器文件保存小计算瓦片所需的 A/B/P
• PE 阵列层级：PE 阵列空间连接多个 PE 的小瓦片，形成更大的并行计算瓦片
• SRAM 层级：SRAM 保存更大规模的 A/B/P，跨多个阵列阶段重复供应数据
• DRAM 层级：对应全局计算点集，存储矩阵乘法所需的整个 A/B/P

整体执行可总结为：全局计算点集（DRAM）首先划分为大 SRAM 瓦片，然后划分为 PE 阵列瓦片，再划分为寄存器文件瓦片，最后在每个寄存器文件瓦片内逐点执行 MAC 操作。

2.3 遍历轴与数据复用

“循环顺序影响能量的原因可归结为：在给定层级，哪个投影可以保持更长时间不变，决定了哪种数据类型可以在该层级获得更强的时间复用。”

GOMA 引入行走轴（walking axis）概念来描述瓦片在高层级瓦片内的推进方向。当 3D 计算瓦片沿某个轴推进一步时，三个 2D 投影中恰好有一个保持不变，而另外两个发生变化：

• 沿 y 轴推进：x-z 投影保持不变 → A(x,z) 可实现时间复用
• 沿 x 轴推进：y-z 投影保持不变 → B(y,z) 可实现时间复用
• 沿 z 轴推进：x-y 投影保持不变 → P(x,y)（部分和）可实现时间复用

图：当瓦片沿一个轴（此处为 y）推进时，一个投影（x-z 平面）保持恒定并可复用，而另外两个投影更新（来源：原文 Figure 4）

因此，总流量自然分解为三部分：分别统计 y-z 平面（对应 B）、x-z 平面（对应 A）和 x-y 平面（对应 P）在遍历过程中的更新次数，然后乘以相应的投影面积得到流量体积。

2.4 层级旁路机制

层级旁路（level bypass）改变了数据实际采取的层级路径。如果某层级选择旁路某种数据类型，则该层级不再对该数据类型执行读/写或驻留操作。数据直接从上层供应到下层，直到到达第一个选择存储它的层级。

“旁路使得传输的能量归属不再局限于相邻层级。对于给定数据类型，其最近的上层驻留层级可能相隔数个层级。”

GOMA 采用接收者中心（receiver-centric）的方式组织能量核算：首先回答”该层级的遍历需要多少数据传输”，然后回答”每个数据项来自哪里”。不同的数据源导致每次传输的能量成本不同。

2.5 闭合形式能量目标函数

结合上述直觉，GOMA 的优化问题可总结为：在容量、并行性和可除性嵌套等硬约束下，选择：

1. 各层级的瓦片形状（沿 x/y/z 的尺寸）
2. 各阶段的行走轴（哪个内部维度推进）
3. 哪些数据类型在哪些层级驻留或旁路

以最小化总能量。

由于稳定的推导链，能量可以以闭合形式计算：

总能量 = Σ (投影更新次数 × 单位能量成本)

GOMA 的核心是将”数据流/映射”简化为层级几何遍历问题，并以投影更新作为统一对象组织计数。

2.6 优化公式

GOMA 将映射搜索形式化为约束优化问题：

minimize:   Ē_total = Ē_(src-1) + Ē_(src-3) + Ē_(src-4) + Ē_(4) (+ Ē_(leak))

subject to:
  - 容量约束（旁路数据除外）
  - PE 数量约束
  - 整数性和可除性约束

variables:
  - L^(1), L^(2), L^(3) ∈ ℕ³₊  (各层级瓦片尺寸)
  - α_(0-1), α_(1-2) ∈ {x, y, z}  (行走轴)
  - B^(1), B^(3) ∈ {0, 1}³  (旁路策略)

关键优势在于：对于任何固定的决策组合，能量评估简化为对 d ∈ {x, y, z} 的有限替换和求和。在固定层级数量和数据类型的情况下，其计算复杂度为常数 O(1)，不随问题规模或瓦片数量增长。

2.7 最优性保证

GOMA 提供两个层面的最优性保证：

1. 目标函数高保真度：使用 Timeloop 的能量模型作为参考实现评估闭合形式能量目标的保真度。在 8064 个映射配置中，8004 个（99.26%）与 Timeloop-model 完全一致（相对误差为 0）。所有样本的平均相对误差为 0.099%（对应能量一致性约 99.9%）。

2. 求解器级全局最优性：使用 Gurobi 与分支定界法求解约束优化问题。对于每个实例，求解器维护并输出当前最佳可行解对应的目标上界、可证明的目标下界以及两者之间的最优性间隙。仅当间隙收敛到 0 时终止，从而提供可直接验证的证明信息（UB/LB 和间隙），确保报告的优化映射在建模目标和约束下是全局最优的。

3. 实验设置

3.1 工作负载

GOMA 针对 GEMM 作为核心算子。研究选择代表性大模型，涵盖边缘和中心/云部署，使用预填充阶段作为评估场景：

• 边缘 LLM：Qwen3-0.6B 和 LLaMA-3.2-1B
• 中心 LLM：Qwen3-32B 和 LLaMA-3.3-70B

对于每个模型类别，评估三种输入长度设置（短/中/长）：

• 边缘场景：{1k, 8k, 32k}
• 中心场景：{2k, 32k, 128k}

共 12 个工作负载。对每个工作负载，枚举预填充阶段涉及的所有矩阵乘法算子，分为 8 种类型：attn_q_proj、attn_kv_proj、attn_score、attn_context、attn_output、mlp_gate_up、mlp_down 和 lm_head。

3.2 目标加速器

选择四个代表性空间加速器模板，在统一的 timeloop/accelergy 框架下建模：

构建 24 个案例：6 个边缘工作负载仅在两个边缘模板上评估（6×2=12），6 个中心工作负载仅在两个中心模板上评估（6×2=12）。

3.3 基线方法

与以下代表性映射算法/框架比较：

• Timeloop-mapper (Hybrid)
• LOMA [12]
• SALSA [14]
• CoSA [17]
• FactorFlow [23]

3.4 评估指标

使用能量 - 延迟积（EDP）作为整体指标：

EDP = E × T

虽然优化目标是总能量 E，但报告 EDP 以与先前工作保持一致。在 PE 数量约束下，GOMA 找到的映射实现 100% PE 利用率，因此 T 达到计算下界。在此设置下，最小化 E 等价于最小化 EDP。

4. 实验结果

4.1 整体 EDP 比较

下表总结了 24 个评估案例的归一化 EDP（归一化到 GOMA，越低越好）：

从整体结果可得出三个关键结论：

1. 全场景一致领先：GOMA 在所有工作负载和所有加速器模板上实现最低 EDP，在模型规模（0.6B→70B）、序列长度（1k→128k）和硬件资源（256→65k PE）上均表现出稳健优势。

2. 启发式方法显著不稳定：LOMA、SALSA 和 FactorFlow 等启发式和多阶段搜索方法在某些案例中可以接近最优，但在其他案例中偏差很大，表现出工作负载依赖性波动。这是因为目标函数受多级分块、循环置换和离散可除性/容量约束共同影响，形成高度非凸和不连续的组合景观。

3. 旁路是影响 EDP 的关键自由度：Timeloop Hybrid 在 Eyeriss-like 和 Gemmini-like 模板上除 GOMA 外取得最佳整体性能，很大程度上得益于其能够在搜索过程中探索跨多个内存层级的旁路组合，减少不必要的驻留，产生额外的能量和性能增益。

4.2 每层 EDP 分解

每个案例包含多个 GEMM 实例。按层（每 GEMM）的 EDP 分解显示：

• lm_head 具有”矩阵 - 向量”形状，使得多个映射器更容易接近最优。在 Gemmini-like 示例中，不同方法与 GOMA 在 lm_head 上的差距很小。Timeloop Hybrid 甚至可以达到最优。

• 矩阵 - 矩阵 GEMM 是差距的主要来源，且优势随规模放大。在 Gemmini-like 示例中，GOMA 在 attn_q_proj、attn_kv_proj、mlp_gate_up 和 mlp_down 等多个矩阵 - 矩阵工作负载上已显示出相对于基线的稳定领先。在更大规模设置（如 A100-like + 70B/128k）中，这些领先在主要大型 GEMM 上进一步放大。

4.3 映射器运行时间

除了映射质量，求解/搜索速度也是部署自动映射的核心指标。下表总结了 24 个评估案例的归一化映射器运行时间（归一化到 GOMA，越低越快）：

GOMA 的案例级运行时间几何平均为 5.22 秒，平均每个 GEMM 仅 0.65 秒，每层最大运行时间仅 3.6 秒，满足实时映射需求。

在小型边缘案例上，CoSA 的运行时间有时可以接近 GOMA，甚至在少数点上略优。然而，随着工作负载规模和资源规模增加，CoSA 和其他基线通常表现出 substantial 运行时间爆炸。相比之下，GOMA 的解析目标具有常数时间评估，并使用显式折叠的低维整数决策变量，使得分支定界剪枝更高效，即使在大规模问题上也能保持稳定和可控的求解时间。

5. 优点与局限

优点

1. 全局最优性保证：GOMA 是首个能够在短时间内为任何给定（GEMM 工作负载，目标硬件实例）对计算全局最优映射的方法，并输出可验证的最优性证书。

2. O(1) 能量评估：解析闭合形式能量目标函数对任何给定映射的评估复杂度为常数 O(1)，不随问题规模增长。

3. 高保真度模型：与 Timeloop-model 的能量一致性达到 99.9%，在 8064 个映射配置中 99.26% 完全一致。

4. 显著性能提升：在 24 个评估案例上，相比 SOTA 映射器实现 2.24–4.24 倍 EDP 提升和 3.83–73.6 倍求解时间加速。

5. 统一框架：将分块、循环置换和旁路联合表述为整数优化问题，在容量、并行性和可除性嵌套等硬约束下求解。

局限

1. 当前仅支持 GEMM：虽然论文提到可将”计算网格”从 3D 推广到更高维度以支持卷积等算子，但当前实现主要聚焦于 GEMM。

2. 依赖准确的硬件能量参数：能量模型的准确性依赖于硬件规格提供的每访问读/写能量常数。如果这些参数不准确，可能影响最优映射的质量。

3. 整数优化求解器依赖：使用 Gurobi 等商业求解器进行全局优化，可能需要许可证。对于极大规模问题，求解时间可能增加（尽管实验中最大每层运行时间仅 3.6 秒）。

4. 固定硬件模板：当前工作针对特定类别的主流加速器模板，对于非常规硬件架构可能需要调整建模。

6. 总结

GOMA 提出了一种统一的 GEMM 映射优化框架，涵盖映射规范、能量评估和全局最优求解。具体而言：

• GOMA 使用几何抽象将层级执行简化为投影更新数的闭合形式计数
• 通过每轴驻留/旁路（层级旁路）重写跨层级访问链
• 实现了对任何给定映射的 O(1) 解析能量评估
• 将分块、循环置换和旁路联合表述为整数优化问题，在容量、并行性和可除性嵌套等硬约束下求解
• 通过全局求解输出可验证的最优性证书，在建模目标和约束下严格保证全局最优性

实验结果表明，GOMA 的解析模型与 Timeloop-model 高度一致。在多个代表性硬件模板和 LLM 预填充工作负载上，GOMA 一致实现最低 EDP，相比现有 SOTA 映射器提供 2.24–4.24 倍 EDP 改进和 3.83–73.6 倍求解时间加速。

随着空间加速器部署需求的持续增长，GOMA 的解析建模和全局最优求解能力为更复杂场景（如多层映射探索和软件 - 硬件协同优化搜索）提供了可重用的基础。

参考文献

1. W. Yang et al., “GOMA: Geometrically Optimal Mapping via Analytical Modeling for Spatial Accelerators,” IEEE Transactions on Computers, 2026. arXiv:2603.07962 [cs.AR]

2. A. Parashar et al., “Timeloop: A systematic approach to DNN accelerator evaluation,” in IEEE ISPASS, 2019.

3. A. Symons et al., “LOMA: Fast auto-scheduling on DNN accelerators through loop-order-based memory allocation,” in IEEE AICAS, 2021.

4. V. J. B. Jung et al., “SALSA: Simulated annealing based loop-ordering scheduler for DNN accelerators,” in IEEE AICAS, 2023.

5. Q. Huang et al., “CoSA: Scheduling by constrained optimization for spatial accelerators,” in IEEE/ACM ISCA, 2021.

6. M. Ronzani and C. Silvano, “FactorFlow: Mapping GEMMs on spatial architectures through adaptive programming and greedy optimization,” in ASPDAC, 2025.

本文基于论文 arXiv:2603.07962 生成，旨在提供技术总结和解读。如需引用，请参考原论文。