GOMA框架：几何抽象优化GEMM计算的全局最优映射
1. GOMA框架概述：几何抽象与全局最优映射
在深度学习加速领域，通用矩阵乘法（GEMM）作为Transformer等模型的核心计算算子，其执行效率直接影响整体系统性能。空间加速器通过硬件-算法的协同设计来优化GEMM计算，其中映射策略（mapping）的选择对性能和能耗具有决定性影响。传统映射空间探索方法面临组合爆炸问题，难以在可接受时间内获得全局最优解。
GOMA（Geometrically Optimal Mapping via Analytical Modeling）创新性地提出基于几何抽象的解析建模方法，将GEMM计算网格投影为三个正交平面，通过行走轴（walking axis）和层级分块（hierarchical tiling）精确建模数据复用模式。该框架首次实现了对任意（GEMM工作负载，目标硬件）组合的全局最优映射求解。
1.1 空间加速器中的映射挑战
在典型的空间加速器架构中（如图1所示），映射策略需要同时确定三个关键维度：
- 分块策略（Tiling）：决定各内存层级中暂存的计算/数据量
- 遍历顺序（Loop Permutation）：指定分块在时间维度上的推进方式
- 层级旁路（Level Bypass）：选择数据是否驻留特定内存层级
这些决策的组合形成了巨大的搜索空间。以典型卷积运算为例，其映射空间可超过10^20种可能，而GEMM的映射空间规模更为庞大。这种组合爆炸使得穷举搜索在计算上不可行。
关键观察：实验数据显示，同一GEMM工作负载在同一加速器上，仅因映射策略不同，能耗差异可达数个数量级（如图2所示）。这种敏感性凸显了映射优化的重要性。
1.2 现有方法的局限性
当前主流的映射空间探索方法可分为五类：
- 随机搜索：如Timeloop、Simba等，通用性强但采样效率低
- 黑盒启发式搜索：包括遗传算法、强化学习等，缺乏收敛保证
- 可微分近似：存在松弛-舍入误差，破坏最优性
- 剪枝枚举：理论上可收敛但搜索开销大
- 数学规划：如CoSA，但成本模型与真实能耗存在偏差
这些方法难以同时保证求解效率和解的质量，特别是在处理大规模GEMM工作负载时。GOMA通过几何抽象和解析建模，突破了这一瓶颈。
2. GOMA的核心方法论：计算网格的几何表示
2.1 三维计算网格与正交投影
GOMA将矩阵乘法P(x,y)=ΣA(x,z)B(y,z)抽象为三维计算网格：
G = {(x,y,z) | x∈[1,Lx], y∈[1,Ly], z∈[1,Lz]}其中每个点(x,y,z)对应一次乘加运算（MAC）。基于此表示，三个矩阵自然地与三个正交投影平面对齐：
- 矩阵A：x-z平面投影
- 矩阵B：y-z平面投影
- 矩阵P：x-y平面投影（部分和/输出）
这种几何对应关系形成了GOMA建模的基础。对于计算网格中的任何3D分块，其在三个平面上的投影面积分别表征了对应矩阵的数据需求规模。
2.2 层级分块的几何解释
忽略旁路情况下，多级内存层次可以理解为对计算网格G的渐进式分块覆盖（如图3右侧所示）。以典型的五级架构为例：
| 层级 | 硬件单元 | 计算粒度 | 几何对应 |
|---|---|---|---|
| DRAM | 主存 | 全局计算 | 整个网格 |
| SRAM | 片上缓存 | 大分块 | 蓝灰色立方体 |
| PE阵列 | 处理单元阵列 | 中分块 | 绿色立方体 |
| 寄存器文件 | PE寄存器 | 小分块 | 橙色立方体 |
| MACC | 乘加核心 | 单点计算 | 白色立方体 |
执行过程可描述为：全局计算集（DRAM）→SRAM分块→PE阵列分块→寄存器分块→逐点MAC运算。这种分层覆盖自然地反映了数据复用的机会。
2.3 行走轴与数据复用
分块策略仅确定了"每次计算多少"，而遍历顺序决定了"如何覆盖整个计算空间"。GOMA的关键洞见是：当3D分块沿某个轴推进时，三个投影中只有一个保持不变，其余两个需要更新。
具体规律为：
- 沿y轴推进：x-z投影（A）保持不变 → A可复用
- 沿x轴推进：y-z投影（B）保持不变 → B可复用
- 沿z轴推进：x-y投影（P）保持不变 → P可复用
这种机制解释了循环顺序如何影响能耗——在给定层级，哪个投影能保持更长时间不变，决定了哪种数据类型能在该层级获得更强的时序复用，从而减少跨层级数据传输。
3. GOMA的解析建模与优化
3.1 闭环能量目标函数
GOMA将映射选择表述为硬件和映射约束下的整数优化问题，其目标函数基于精确的解析能量模型。能量计算可分解为：
- 投影更新计数：根据行走轴确定各投影的更新频率
- 旁路门控：通过二进制开关矩阵B控制各项的激活
- 能量加权：按硬件提供的每字读写能耗加权求和
对于DRAM-SRAM链路（0-1），沿轴d的流量体积为：
N^(0-1)_d = 1[B^(1)_d=1] * V * (L^(0)_d*1[d=α0-1] + L^(1)_d*1[d≠α0-1])^-1其中V为总MAC数，α0-1为0-1阶段的行走轴。当d与行走轴一致时，投影更新可压缩为"每全局列头一次"。
3.2 归约轴的特殊处理
归约轴z具有特殊性，因为部分和P的更新遵循"读旧值→累加→写回"模式。GOMA引入边界系数ρ来统一处理：
ρ^(src-p)_z = 1 - 1/˜L^(src-p)_z其中˜L表示接收层级p的有效全局列计数。该系数将读旧流量表示为写回流量的比例，保持了两者计数基准的统一。
3.3 整数优化问题构建
综合以上组件，GOMA的优化问题可表述为：
最小化：总能量 = Σ(投影更新计数 × 单位成本) 约束条件： 1. 容量约束：各层级存储需求≤物理容量 2. PE数量约束：空间并行度≤硬件PE数 3. 整除约束：相邻层级分块尺寸需整数倍关系 决策变量： 1. 各层级分块尺寸{L^(p)_d} 2. 阶段行走轴{α} 3. 旁路配置{B}该问题可通过整数线性规划(ILP)或分支定界(B&B)等全局优化算法求解，并输出可验证的最优性证书。
4. 实现考量与优化效果
4.1 硬件成本模型集成
GOMA需要集成目标硬件的精确能耗参数，包括：
- 各存储层级的读写能耗（EDRAM_read, ESRAM_write等）
- 乘加核心的每操作能耗(eMACC)
- 空间归约能耗(Espa_reduct)
这些参数通常通过硬件性能建模或实际测量获得。在Timeloop兼容的实现中，默认设置Espa_reduct=0以保持一致性。
4.2 实际部署注意事项
- 分块尺寸选择：应确保寄存器文件分块能完全装入PE的寄存器组，同时考虑bank冲突避免
- 行走轴协调：相邻层级的行走轴选择会影响数据复用链，需全局协调
- 旁路权衡：旁路某层级虽节省该层读写，但可能增加上层供应压力
- 归约轴优化：对z轴的特殊处理能显著降低部分和通信开销
4.3 性能提升实测
在四种代表性加速器和LLM预填充工作负载上的实验表明：
- 能效比：相比现有最佳映射方法，EDP提升2.24-4.24倍
- 求解速度：求解时间减少3.83-73.6倍
- 可扩展性：可处理维度超过10^4的大规模GEMM
这些优势使GOMA特别适合大型语言模型(LLM)和扩散Transformer(DiT)等计算密集型应用的加速。
5. 扩展应用与未来方向
5.1 对卷积运算的推广
虽然GOMA本文聚焦GEMM，但其几何表示可扩展至高维计算网格。例如，卷积可视为更高维度的归约运算，其中：
- 输入特征图：多通道的3D网格
- 卷积核：4D张量（输出通道×输入通道×高×宽）
- 计算网格：5D空间（批×输出空间×核空间）
在这种扩展中，投影更新和行走轴的概念仍然适用，但需要调整维度计数和约束表达。
5.2 动态工作负载适配
当前GOMA针对静态GEMM工作负载。未来可探索：
- 动态形状适应：在线调整映射策略应对变化的问题规模
- 多算子联合优化：考虑GEMM与前后算子的数据流协调
- 温度感知映射：结合芯片热模型优化能耗分布
5.3 工具链集成
GOMA可与现有工具链深度集成：
- 编译器支持：作为MLIR或TVM的后端优化pass
- 设计空间探索：与架构参数协同优化
- 可视化调试：投影更新模式的可视化分析
这种几何抽象的映射优化方法，为空间加速器的自动代码生成和性能可移植性提供了新的理论基础。