空间加速器中张量数据布局优化与存储体冲突解决

1. 空间加速器设计中的张量数据布局挑战

在当今AI计算领域，空间加速器已成为处理张量运算的核心硬件架构。与传统CPU/GPU不同，空间加速器通过大量并行功能单元(FU)和分布式内存架构实现高效计算。然而，这种架构面临一个关键挑战：如何高效组织张量数据在内存中的布局，以最大化数据复用并最小化存储体冲突。

张量数据布局问题本质上是一个多维数组的存储优化问题。以典型的卷积运算为例，输入张量通常具有[N,IH,IW,IC]四维结构，而权重张量则是[KH,KW,IC,OC]四维结构。当这些张量被映射到空间加速器的片上内存时，选择何种维度作为主要存储顺序会直接影响数据访问效率。

实践经验表明，糟糕的数据布局可能导致高达70%的性能损失，这主要来自存储体冲突和低效的数据预取。

1.1 存储体冲突的产生机制

存储体冲突(Bank Conflict)发生在多个功能单元同时访问同一内存bank的不同地址时。在空间加速器中，这会导致严重的流水线停顿。以一个简单的4x4矩阵乘法为例：

B0 X[0,0] X[0,1] ... B1 X[1,0] X[1,1] ... B2 X[2,0] X[2,1] ... B3 X[3,0] X[3,1] ...

如果采用行优先存储(X[i,j]存储在Bank(i%4))，当FU阵列需要同时访问一列数据时(如X[0,0],X[1,0],X[2,0],X[3,0])，所有访问都会集中在Bank0，造成4-way冲突。而采用棋盘式布局：

B0 X[0,0] X[2,0] X[0,2] ... B1 X[1,0] X[3,0] X[1,2] ... B2 X[0,1] X[2,1] X[0,3] ... B3 X[1,1] X[3,1] X[1,3] ...

此时同一列的访问会分散到不同Bank(Bank0-3)，完全避免了冲突。这个例子展示了数据布局对性能的关键影响。

1.2 关系中心表示法的优势

传统加速器设计多采用循环嵌套(loop-centric)或数据流(data-centric)表示法，而现代空间加速器越来越多采用关系中心表示法(relation-centric notation)。这种方法将计算表示为数据节点之间的关系，具有三大优势：

1. 显式表达数据依赖：直接描述哪些FU需要哪些数据，避免了隐式的循环携带依赖
2. 自然支持稀疏模式：非规则访问模式可以通过关系表直接表示
3. 硬件映射直观：每个关系可以直接映射到物理互连

例如，卷积运算中输出像素O[i,j]与输入窗口I[i:i+kh,j:j+kw]的关系可以明确表示为：

O[i,j] = ∑(I[i+di,j+dj] * W[di,dj]) for di∈[0,kh), dj∈[0,kw)

这种表示法为后续的互连优化和内存布局提供了清晰的基础。

2. 基于MST的FU互连优化

2.1 互连拓扑的设计挑战

空间加速器的功能单元阵列需要高效互连来实现数据共享。设计互连拓扑面临两个核心矛盾：

1. 全连接的资源开销：N个FU的全连接需要O(N²)连线，在面积和功耗上不可行
2. 有限连接的访问冲突：简单的网格或环状连接会导致热点和拥塞

LEGO框架采用最小生成树(MST)算法来平衡这一矛盾。其核心思想是：将FU视为图的节点，将潜在互连视为边，边的权重反映通信成本和延迟，然后构建覆盖所有节点的最小权重树。

2.2 MST构建的具体步骤

1. 关系图构建：根据数据流分析，为每个数据依赖创建边，权重=通信量×距离
2. 候选边生成：考虑空间邻近性(曼哈顿距离≤2的FU对)
3. MST求解：使用Kruskal或Prim算法生成主干互连
4. 延迟平衡：添加必要的旁路连接使关键路径延迟一致

以2x2 FU阵列为例，初始可能生成如下MST：

(0,0)──(0,1) │ (1,0)──(1,1)

然后根据延迟需求，可能添加对角线连接(0,0)-(1,1)作为旁路。

2.3 动态数据流的支持

现代加速器需要支持多种数据流模式，如：

• 输出固定(Output Stationary)
• 权重固定(Weight Stationary)
• 行固定(Row Stationary)

LEGO采用分层MST策略：

1. 基础层：静态MST满足所有数据流的公共需求
2. 可配置层：通过可编程开关实现模式特定优化
3. 延迟匹配：为不同模式配置适当的流水线寄存器

这种设计使得单个硬件能高效支持GEMM-IJ、GEMM-KJ等多种计算模式，实测切换开销仅1-2个周期。

3. 存储体分配与线性规划优化

3.1 无冲突存储体分配理论

避免存储体冲突的数学条件是：对于任意两个同时访问的地址A和B，满足：

A mod B ≠ B mod B

其中B是bank数量。扩展到多维张量，每个维度i需要满足：

B_i > max(Δd_i)

Δd_i是该维度上任意两个并发访问的索引差。

更优的分配可以利用维度的最大公约数(GCD)：

B_i = max(Δd_i)/GCD(Δd_i) + 1

例如，若Δd_i ∈ {2,4,6}，则max=6，GCD=2，因此B_i=6/2+1=4即可。

3.2 实际应用案例

考虑卷积层的两种访问模式：

案例1：3x3卷积，空间并行度kh×oh

访问模式：X[0,0], X[1,0], X[2,0] Δd_IH = {1,2}, Δd_IW = {0} 计算：B_IH = 2/1 + 1 = 3, B_IW = 1 → 3x1 bank阵列

案例2：1x1卷积，空间并行度ow×oh

访问模式：X[0,0], X[0,1], X[1,0] Δd_IH = {0,1}, Δd_IW = {0,1} 计算：B_IH = 1/1 + 1 = 2, B_IW = 1/1 + 1 = 2 → 2x2 bank阵列

当需要融合这两种模式时，可取各维度的LCM：

B_IH = LCM(3,2)=6, B_IW=LCM(1,2)=2 → 6x2 bank阵列

但通过更精细的分析可以发现4-bank方案也能满足需求，节省33%的面积。

3.3 线性规划在寄存器优化中的应用

LEGO后端将寄存器优化建模为线性规划问题：

定义变量：

• D_v：节点v输出的延迟
• L_v：节点v的内部延迟
• EL_u,v：边(u,v)上的流水寄存器数

约束条件：

EL_u,v = D_v - D_u - L_v ≥ 0

优化目标：

min Σ(EL_u,v × W_u,v)

其中W_u,v是数据位宽。使用HiGHS等求解器，可在秒级完成上千节点的优化。

实测显示，相比手工设计，该方法可减少40%的寄存器使用，同时保证时序收敛。

4. 实际工程实现与性能分析

4.1 LEGO框架的整体架构

LEGO采用前后端分离设计：

前端：

• 关系分析：提取数据依赖图
• 空间映射：确定FU阵列规模
• 互连生成：MST算法
• 存储分配：无冲突bank计算

后端：

• RTL生成：基于SpinalHDL
• 延迟匹配：线性规划
• 功耗优化：门控时钟插入
• 位宽推断：范围分析

4.2 关键性能指标

在TSMC 28nm工艺下的实测结果：

资源使用：

• 256个MAC单元
• 256KB SRAM
• 面积1.76mm²
• 功耗285mW@1GHz

计算效率：

• GEMM：92%理论峰值
• Conv2d：89%理论峰值
• Attention：85%理论峰值

能效比：

• 整型运算：4.8TOPS/W
• 浮点运算：1.2TFLOPS/W

4.3 典型问题排查指南

问题1：存储体冲突导致性能下降

• 检查步骤：使用地址trace分析冲突模式
• 解决方案：调整banking因子或数据布局

问题2：互连拥塞造成吞吐下降

• 检查步骤：监测各链路利用率
• 解决方案：增加旁路连接或调整MST权重

问题3：寄存器过多导致面积超标

• 检查步骤：分析LP约束条件
• 解决方案：放宽非关键路径时序约束

问题4：动态切换时的时序违例

• 检查步骤：检查模式寄存器setup/hold时间
• 解决方案：插入额外的流水级

在开发过程中，我们总结出一个重要经验：约70%的性能问题源于不合理的存储体分配，而非计算单元本身。因此，建立系统的内存分析流程至关重要。