机器之心矩阵相乘在GPU上的终极优化:深度解析Maxas汇编器工作原理( 六 )


C 矩阵的计算:寄存器分配和计算顺序
现在所需要的数据已经被尽可能高效地被送到寄存器了 , 似乎可以直接使用 FFMA 加乘指令对它们直接进行操作了 , 毕竟这才是矩阵相乘内核应该做的事 。 不幸的是在此之前还要解决一个可能是整个项目中最大的麻烦 , 就是寄存器访问的 bank 冲突 。
为了在一个流计算单位内容纳大量线程 , GPU 准备了多达 32K 的寄存器文件 , 显然其访问无法和 CPU 一样直接 , 而是和共享内存一样要通过 bank 进行 , 因此 bank 冲突也是难免的 , 而一旦出现会对性能造成很大影响 。 Maxwell 上的寄存器文件有 4 个 32 位的 bank , 每个寄存器通过对其编号的 mod 4 操作被映射到对应的 bank 上 。 如果在 C 矩阵的计算中出现以下指令就会出现寄存器 bank 冲突:
FFMA R0, R4, R5, R0; # R0, R4 在 bank 0 , R5 在bank 1 , R0和R4产生bank冲突
后来每一代 GPU 架构的寄存器 bank 都会有变动 , 比如 Volta 架构就是分为 2 个 64 位的 bank , 这也是 maxas 无法在现在的主流 GPU 上发挥性能的主要原因 。
直接使用汇编语言的一大优势就是可以通过手动分配寄存器尽可能减少 bank 冲突:
0-63 号分配给 C 矩阵;
64-71 和 80-87 分配给 A 矩阵 , 72-79 和 88-95 分配给 B 矩阵(分配两倍于实际大小用于流水线预读取) 。
因为是用向量指令载入 , 分配给 A 和 B 的每四位寄存器编号必须是连续的 , 也就是所有四个 bank 都会连续出现 , 因此在 A 和 B 的寄存器选择上并没有优化空间 , maxas 能做到的是尽量调整分配给 C 的 63 个寄存器的顺序 , 其所采用的编号方案如下图:
机器之心矩阵相乘在GPU上的终极优化:深度解析Maxas汇编器工作原理
本文插图
图 6. 每个线程内部所用数据的寄存器编号 , 每个寄存器所在 bank 用不同颜色标出 , 如果某个 C 元素的颜色和其对应的 A 或 B 元素相同就会发生 bank 冲突 , 这种元素在图中用黑框标出 。
显然这是调整寄存器编号能得到的最好结果 , 图中黑框标出的 bank 冲突不管如何调整 C 矩阵的编号是无法避免的 , 因为其来源是 A 和 B 用到了同一个 bank , 而 A 和 B 中的操作数既需要占据所有四个 bank(每个 bank 两个数) , 又需要与另一个矩阵中的其他所有操作数配对 , A 的每一寄存器必然会和 B 中的两个寄存器产生 bank 冲突 。 事实上如果 C 使用最简单的寄存器编号方式 , 比如在第一行占据 0~7 , 那么其中每一个寄存器都会和对应的 B 操作数发生 bank 冲突 , 就是非常坏的一种编号方法 。
要进一步消除通过寄存器分配所不能消除的那部分 bank 冲突 , 需要用到 Maxwell 提供的一种操作数重用特性 。 在发出一条指令时 , 可以将其的一些操作数设为重用 , 硬件将把该操作数送往一个重用缓存 。 送如果后一条指令在同一位置要用到同一个操作数 , 则该指令可以直接从缓存而不用通过 bank 去取这个操作数 , 从而避免 bank 冲突 。
FFMA R2, R4.reuse, R5, R2; # 此处指定R4将会被重用 , 将其放入缓存FFMA R0, R4.reuse, R5, R0; # R0和R4本来产生bank冲突 , 但因为上一条指令缓存了第二个操作数R4 , 只要到bank中取R0即可 , 从而避免了bank冲突FFMA R0, R5, R4, R0; # R0和R4的bank冲突依然会发生 , 因为所缓存的R4在第2个操作数 , 但本指令中R4落在第3个操作数上 。
如果在线程中简单地一行行或一列列遍历图 6 中 C 矩阵的 64 个寄存器 , 并且将 A 的寄存器设为重用 , 因为就可以解决 16 个 A 和 B 的寄存器 bank 冲突中的 14 个 , 不能解决的是寄存器 R3 和 R35 , 因为它们是该行的第一个用到该 A 操作数的指令 , 之前没有指令将其送入重用缓存 。 知道原因后这两个 bank 冲突也可以被轻松化解 , 只要在遍历每行时偶数行从右到左(0 行从 26 到 3)奇数行从左到右(1 行从 7 到 30) 。 但是 maxas 即使对此还是不满足 , 它在前述的来回遍历基础上又加上一个漩涡提出了一个更诡异的遍历方法:


推荐阅读