在并行处理器设计中,我们希望最大化访存吞吐,让更多的数据分布在不同的 bank,而非在一个 bank 中产生堵塞。一种场景是面对多应用并行,这往往可以通过划分上下文基地址隔离;而另一种场景则是高并行同一个数据共用基地址,本文针对该场景下常见情形 Tensor Data Layout 进行讨论。本文旨在针对硬件设计者介绍从数据逻辑到物理实现完整映射关系,其中很多推导过程比较繁琐,如果是编程实践中大致记得 row-major, col-major, swizzle data layout 几种常见 layout 即可。
高维数据组织:从逻辑表示到物理表示
假设有 B 个深度为 D 宽度为 W 的 bank,其最小单位可用一个三维向量 \(P\) 表示,比如 \((10, 2,4)\) 表示访问 bank=2, depth=10 那一地址中的第 width=4 个 scalar,scalar 粒度可根据需要设置,常见寻址用 byte。引入转置表示以统一向量用列向量表示,便于后续公式推导。
其中 depth , bank , width 三个维度地位并不对称,depth 需要多个时间周期来访问不同 depth,而 bank 和 width 存在空间上的并行性,因此定义 \(P_{s}=[bank,width]^T\) 以便后续分析。RAM 是一种通过牺牲在 depth 上的并行读取从而节约读写外围电路面积开销进而实现高密度存储数据的结构,以时间换空间,恰恰适应程序数据生存周期很长,但调用次数相对稀疏的特点。
定义三维向量到地址的映射关系是 \(f\),即 :
常见映射关系为 \(Addr=f(P)=[S_{d},S_{b},S_{w}]\cdot P=[W \times B,W,1]\cdot P\),单从表达式看 \(f\) 并非存在唯一逆映射,考虑 \(P\) 每个元素的取值范围以及所有都是整数,存在唯一逆映射 \(f^{-1}\)。这里的地址是相对前文定义 scalar 粒度而言,并非 SRAM 的地址控制信号, SRAM 的地址控制信号为 \(P_{b}=[depth, bank]^T\)。为了设计方便起见,\(D,B,W\) 都会设计为 2 的倍数,因此可以直接从 Addr 截取相应的 bit 表示对应的选择关系。按照 bit 排列的顺序,分为 high-order interleaving 和 low-order interleaving[1]。
一个 N 维个张量数据 \(X\in R^N\) 中每个元素存在一个 N 维向量 \(I\in R^1\) 用于索引,定义其与地址的映射关系为:
比如二维矩阵索引 \(I=[row, col]^T\) Row-Major 映射关系定义为 \(Addr=g(I)=Base + Offset = Base + [COL, 1]\cdot I\),而 Col-Major 映射关系定义为 \(Addr=Base + Offset = Base + [1,ROW]\cdot I\)。由于同一个数据往往共用一个基地址,普遍讨论的 data layout 一般指 offset 和 \(T\) 之间关系。至此,我们可定义数据逻辑表示到存储物理表示的函数关系:
计算单元-Layout 耦合策略
SIMD 计算单元往往存在一个特定的逻辑并行顺序,比如每个 PE 是一个乘法累加器的 Output Stationary 的脉动阵列,计算 (T,I) (O,I) -> (T,O)
的 GEMM 运算,每个周期 input feature 需要在 T 维度并行数据,weight 需要在 O 维度并行的数据。
在 NVIDIA GPU 上,矩阵乘法 \(D=AB+C\) 可调用 wmma
或者 mma
指令完成,以 wmma
为例,分为三个阶段 PTX 指令组合完成
wmma.load
从 memory load 源操作数 \(A,B,C\) 每个操作数都要写一条 ptx 指令,可指定原操作数来自 global memory 或 shared memory [2]wmma.mma
进行 GEMM 计算wmma.store
将 D 保存到 memory 中
虽然 load 支持显式指定 A、B 、C 不同的 layout,但特定的 data layout 下似乎会发生 bank serialization[3][4]。猜测 GPU 是先将任意的 data-layout 以特点的 layout 顺序加载到 register file 中,如果发生 bank serialization 是在 load 阶段而非 mma 阶段。
计算单元对特定的 data layout 需求来自于 SIMD 并行 PE 之间的物理互联-逻辑运算映射关系,反过来说,如果能够在计算单元实现某种 reconfigurable 控制流&互联重构,就可以对数据存储 layout 具有一定鲁棒。进一步讨论存在两种策略,一种是则是前文所述的计算单元 reconfigurable,一种是在存储上实现 layout 鲁棒,具体表现是程序员-编译器的 layout 设计和 bank 之间的 NoC 网络。由于处理器中往往存在多个计算单元而通过统一的存储共享上下文,一般往往使得计算单元和 layout 耦合,将复杂度转移到共用的存储通路设计上。 感性计算举个例子,假设一个处理器中有 N 个计算单元和一个共用 buffer,实现重构的复杂度是 k>1 而不实现的复杂度是 1,计算单元 reconfigurable 的开发复杂度是 \(O(kN+1)\),而存储单元 NoC 的开发复杂度是 \(O(N+k)\)。
实现 NV 风格的较高鲁棒访问 register file,一是需要生成多个独立的 bank 控制流,即等于 bank 数量的可编程的地址生成单元,据说 NV register file 是 4-bank dual port rf 设计[5],则一共是 8 个地址生成单元;二是需要给出 NoC 的控制信号,对于 \(N\) 输入 \(M\) 输出的 cross bar,一共需要 \(M\) 个 \(log_2(N)\) 比特控制信号,这部分最好也用一个可编程单元完成。对于 NoC 引发的 overhead ,之前的 blog [6]已经有过讨论不再叙述。
Bank Conflict 的数学表达
假设计算单元单位周期对数据 layout 需求是 \(\{T_{0},T_{1},..., T_{N-1}\}\),通过 \(f^{-1} \circ g\) 可得到对应物理存储 \(\{P_{0},P_{1},...,P_{N-1}\}\) 进而得到 \(\{P_{s,0},P_{s,1},...,P_{s,N-1}\}\)。不产生 bank conflict 的定义是,这 N 个物理存储表示在 \((bank, width)\) 维度上相同的个数小于等于相应读写端口数量(Dual-port, Two-port, Single-port),即集合 \(\{P_{s,0},P_{s,1},...,P_{s,N-1}\}\) 的“众数”频数小于等于端口数量。
举例:传统线性 Layout
对于 \(f\) 和 \(f^{-1}\),其定义如下:
线性 layout 即 row-major 或 col-major, \(g\) 和 \(g^{-1}\) 定义如下:
则有:
一般来说,计算单元对于矩阵的访存需求都是沿着某一个维度(而非对角线),即 \(I=c\times e_{1}+i \times e_{2}\) ,其中 \(e_{1}, e_{2}\) 是 row 或者 col 的单位向量,\(c\) 是一个常数,\(i\) 为各不相同的多个取值。附计算 bank conflict 示例代码 [7]。
而 Swizzle Layout 则是跳过了中间地址转换过程,直接构造 \(P\) 和 \(I\) 之间的数学关系,并基于 \(I=c\times e_{1}+i \times e_{2}\) 的假设,在 \(\{e_{1},e_{2}\}=\{e_{row}, e_{col}\}\) 或 \(\{e_{1},e_{2}\}=\{e_{col}, e_{row}\}\) 都有 \(P\) 各不相同,即对任意方向读取都满足 bank-free。
https://www.geeksforgeeks.org/computer-organization-architecture/types-of-memory-interleaving/ ↩︎
https://www.cnblogs.com/devil-sx/p/19091444 ↩︎
https://arxiv.org/abs/2410.20399 ↩︎
https://leimao.github.io/blog/Row-Major-VS-Column-Major/ ↩︎
https://www.zhihu.com/question/608936006/answer/1947296899364295591 ↩︎
https://www.cnblogs.com/devil-sx/p/18692062 ↩︎
https://github.com/Devil-SX/Bank-Conflict-Calculation ↩︎