1. GOMA框架概述几何抽象与全局最优映射在深度学习加速领域通用矩阵乘法GEMM作为Transformer等模型的核心计算算子其执行效率直接影响整体系统性能。空间加速器通过硬件-算法的协同设计来优化GEMM计算其中映射策略mapping的选择对性能和能耗具有决定性影响。传统映射空间探索方法面临组合爆炸问题难以在可接受时间内获得全局最优解。GOMAGeometrically 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。基于此表示三个矩阵自然地与三个正交投影平面对齐矩阵Ax-z平面投影矩阵By-z平面投影矩阵Px-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)_d1] * 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)或分支定界(BB)等全局优化算法求解并输出可验证的最优性证书。4. 实现考量与优化效果4.1 硬件成本模型集成GOMA需要集成目标硬件的精确能耗参数包括各存储层级的读写能耗EDRAM_read, ESRAM_write等乘加核心的每操作能耗(eMACC)空间归约能耗(Espa_reduct)这些参数通常通过硬件性能建模或实际测量获得。在Timeloop兼容的实现中默认设置Espa_reduct0以保持一致性。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设计空间探索与架构参数协同优化可视化调试投影更新模式的可视化分析这种几何抽象的映射优化方法为空间加速器的自动代码生成和性能可移植性提供了新的理论基础。