矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理

机器之心发布

作者：XiaoyuWang

九大章节，一万余字，这篇文章可能是目前为止Maxas汇编器工作原理最全面、最细致的解析。

在从事深度学习框架的实现工作时，了解到 Nervana 有一个称为 Maxas 的汇编代码生成器项目，可以生成性能超过 nVidia 官方版本的矩阵相乘的 GPU 机器码，由此对其工作原理产生兴趣。

项目地址： https://github.com/NervanaSystems/maxas

其作者 Scott Gray 在代码外提供了详细的文档（https://github.com/NervanaSystems/maxas/wiki/SGEMM），但由于该算法的复杂性，行文晦涩，逻辑跳跃，尤其是对一些方法的动机没有交待，很容易迷失在细节中。

本文可以看作按作者对该文档的理解进行的重写，但求在细节上不厌其烦，对其代码的前因后果作尽可能完整的交待，不过大体结构还是按照 maxas 文档来安排，文中图片也全部出自该文档。

值得说明的是 Maxas 使用的算法完全依赖于 Maxwell 架构的一些特性， 随着新一代 GPU 的架构的演进这个项目本身已经完全过时了，但其解决问题的思路仍然值得借鉴。

背景

单精度矩阵相乘（SGEMM）应该是每一个学习 CUDA 编程的程序员最熟悉的案例了，从 nVidia 推出 CUDA 的第一版起这就是其官方教程里唯一的例子，这不仅因 SGEMM 是几乎所有计算密集型软件最重要的底层模块，更是因为通过这个例子可以很好地展示 GPU 中的各种优化技巧，特别是对 GPU 内存层次的利用。

对于两个 NxN 的矩阵 A 和 B 的相乘，一个最简单的并行方法是对于其输出矩阵 C（大小同为）的每一个元素开一个线程，该线程载入 A 的一行和 B 的一列，然后对其做一次向量的内积。但问题是在 GPU 上访问显存的延时相当的大（~100 时钟周期），如果 A 的一行因为在内存中是连续的还能够利用 GPU 的超大显存带宽一次载入多个元素平摊其载入时间以及缓存来降低延时，对于 N 上千的大矩阵来说 B 的一个列中的元素的内存地址也要相隔上千，意味着载入一次中除了需要的那个列元素外大部分数据都是无用的，同时这种访问模式几乎不可能有缓存线命中，总而言之这个并行方法的内存访问效率低到令人发指。

对其的优化就要用到共享内存了，共享内存是位于 GPU 上的片上缓存，速度可与一级缓存相当，而且同一个线程块中的线程可以通过共享内存交换数据，唯一的缺点是容量有限。为了利用这一小块高速内存，原来的并行算法必须有所改进：矩阵将在每个维度被划分成  个小片，输出矩阵 C 可以写为：

A 和 B 也做同样的划分，其中不是元素而是小片矩阵，当然小片大小为 1 时小片矩阵就退化为单个元素。显然矩阵乘法的定义依然在此适用：。

如果把小片看作一个元素，整个矩阵的规模相当于被缩小了 倍。对于每个小片的结果可以由一组线程负责，其中每个线程对应小片中的一个元素。这个线程组将 A 的行小片和 B 的列小片一一载入共享内存，在共享内存上对其做矩阵相乘，然后叠加在原有结果上。这样对小片中的元素每载入一次可以被在高速的共享内存中被访问 K 次，由此得到远高于原始方法的内存存取效率。

网上找的分片算法示意图，图示比 CUDA 官方教程中的更加清楚一点。

以上只是对这种算法的一个简单描述，经过这样的优化整个算法已经可以达到相当高的效率了，而且是进一步进行优化的基础。为了理解后文中的内容，读者需要细读 CUDA 的官方教程以对这个分片算法非常熟悉，并且对 nVidia GPU 的编程模型有比较深入的理解。

基本思想

如上节所述，分片算法在利用了片上高速缓存之后，不但小片矩阵的乘法速度可以大大加快，还可以利用计算小片矩阵相乘的时间将下一个小片从主内存传送至片上共享内存，换句话说此时整个矩阵相乘的时间已经完全由小片矩阵相乘所决定，如果要进一步提高性能就要在小片矩阵相乘上做文章了。

在共享内存内部做矩阵相乘虽然已经很快了，但距离硬件性能的极限还是有距离，主要瓶颈是两个。首先共享内存的延时终究还是比不过寄存器，在 Maxwell/Pascal 上寄存器延迟时 6 个时钟周期，在共享内存上达到 23 个周期。此外，GPU 的运算单元无法直接操作共享内存的数据，需要有一个传输指令将其送到寄存器上，而这个 mov 指令会占用和实际计算指令几乎相当的时间，可谓相当大的开销。要达到硬件的峰值性能，我们希望 GPU 的每个运算单元的流水线都被运算指令占满，每个时钟周期都有结果被计算出来。为了在现有基础上达到这一目标，maxas 和之前的一些研究提出了如下方法：

1. 利用新加入的向量指令，一个指令可以传输四个连续的浮点数，大大减少传输指令的数量，并且有利于用计算指令隐藏传输指令消耗的时间；

2. 通过交叉布置计算指令和传输指令，实现数据预读取和计算的流水线；

3. 分片算法利用高速的共享内存缓存主显存上需要多次存取的数据，那么把这个思路发展下去，在小片矩阵内部作进一步分片，利用寄存器去缓存共享内存的数据，得到进一步的加速。但是这个新的分片算法和之前的有所不同，也带来了额外的困难。

为了实现这些方法需要对 GPU 指令和寄存器的精确控制，已经不在 CUDA 语言表达能力的范围之内，所以其实现必须由 GPU 原生汇编语言完成（并非 PTX 这样的伪汇编语言），但不妨碍用表达能力更强的类似 C 的伪代码来说明这个实现。从伪代码到实际的汇编代码有相当直接的转换方法，在 maxas 中用 perl 实现了这一转换。

maxas 算法概要

只考虑两个矩阵相乘，在之前的直观算法中，计算一个 C 矩阵的元素是按照矩阵乘法的定义，取 A 中的一行和 B 中的一列做内积。A 中的一行和 B 中的一列都要被用到 64 次。如果要充分利用寄存器的优势三个的矩阵（每个矩阵占 16KB）都要放在寄存器中对寄存器文件（每个 SM 64K）是巨大的压力，更严重的问题是和共享内存不同，寄存器在线程间是不能共享的，如果每个线程要在自己的寄存器中保存所负责的 C 矩阵的那部分，它也必须在寄存器中保存所用到的 A 的行和 B 的列。结果是大量寄存器用于保存重复的数据，而且对共享内存的访问很可能造成 bank 冲突，进一步恶化延时。

如果换一个思路，不从输出矩阵 C 的角度，而从输入矩阵的角度，不难发现 A 的第 k 列仅被用于和 B 的第 k 行的元素相乘，也就是说如果取 A 的第 k 列和 B 的第 k 行，将其中所有元素对两两相乘并加到其所贡献的输出矩阵元素上：

这些行和列就完成了其在矩阵相乘中的使命，可以被扔掉了。这种算法可以大大减少输入矩阵对寄存器的占用，而且载入个数据就可以进次加乘运算，完全符合利用寄存器进一步缓存共享内存数据的要求。不难看出该方法在 A 的列和 B 的行大小不一样时依然可以适用，只要它们的列指标和行指标相同。

maxas 对于小片矩阵乘法是用 64 个线程来并行实现的，其中每个线程负责计算个矩阵的乘积，64 个线程按照布局，这样就确定了小片的大小为一个边长个元素的矩阵（每线程 8 元素 x8 线程）。这一点区别于原始分片算法中每个线程计算矩阵中的一个元素，也是充分利用寄存器的超低延迟的关键。

图2. maxas 计算两个 64x64 矩阵相乘的示意图，绿色的 4x4 小片是线程 0 负责的那部分元素，黄色是其他线程负责那部分的左上角元素。图中只标出了左上角 4x4 矩阵的线程号，其他 15 个只是其重复。每个黑框中包含 32 线程，代表一个 warp。这块 32x32 矩阵计算完成后，线程 0 以及其他线程保持相对位置，依次平移到另外三个绿色小片标出的地方，完成这个 64x64 矩阵的计算。左边的向量是 A 矩阵的一个列，上方的向量是 B 矩阵中与之对应的行，其中标为绿色的数据（各 8 个浮点数）是线程 0 所需要用到的，其他线程需要的不难类推。

maxas 的整个实现就是为了这个算法服务的，后文中所要解决的问题也是该算法在实现过程中所遇到的。以上的那些参数选择，比如为什么选择 64 个线程，都是根据 GPU 硬件资源决定的，以便在满足每个线程所需的寄存器资源基础上，创建尽可能多的线程 warp，以便调度器在某些 warp 等待数据时将别的 warp 切换进来进行计算。

将输入矩阵载入共享内存

为了实现以上算法，这 64 个线程首先被用来将两个输入矩阵所需要的部分从主显存载入共享内存。这里需要指出一个原文中没有提到的假设，即在 maxas 中默认矩阵是按照行优先储存的，即矩阵的每一列在内存中是连续的，否则无法解释后面的一系列算法。由于算法的不同载入的方法也有所不同，并且在原方法基础上增加了一些优化：

1. B 矩阵用到的是行数据，而默认的行优先矩阵储存中行数据是不连续的，需要做转

置，行变为列，这样 A 和 B 的载入方法可以完全相同，以降低代码复杂度；

2. A 和 B 矩阵被作为一维纹理载入，这样不但可以利用纹理数据的缓存，而且不用耗费时间进行数组越界检查，因为纹理载入会自动将越界的数据置为 0，对于矩阵相乘不会有任何影响。

了解以上预处理可以更方便地理解后面的伪代码。创建纹理和转置的工作应该是在 GPU 内核执行前完成的，不影响内核执行的性能。

纹理内存中的数据也是分段被载入共享内存的，不过按照原方法每段载入的应该是一个个的小片，但为了充分利用寄存器资源，maxas 采用了完全不同的计算方法。如果线程块计算的是，首先将矩阵 A 按每 64 行一条划分为的条，所需的输入数据全部在第条上，当然这一条数据还还是很大，需要进一步分次载入，maxas 中每次载入并消耗，分次完成。对于矩阵 B 方法类似，只不过是按列划分为的条，转置后载入方法和 A 完全相同。其内存布局如下图所示：

图 3. 每次循环中被一个 warp 载入共享内存的一段纹理，可以看作 Bj 或转置后的 Ai，这样这个矩阵其实又回到了常规的列优先储存。这个图转置后看其实更直观。

图中每一个格子代表一个线程负责载入的数据单元，绿色格子是线程 0 要先后载入的四个单元，黄色格子是其余 31 个线程第一次载入的那部分数据。整个 warp 每次载入两行，如此重复四次完成。在 GPU 上执行单位是 32 个线程组成的 warp，所以 64 个线程是分为两个 warp 执行。其中一个 warp（线程 0-31）载入 A 另一个（线程 32-63）载入 B。此图有一个容易造成困惑的地方是图中的矩阵形状为而不是， 这是因为后面每个线程会用到向量指令一次载入 4 个浮点数，即每个格子本身就是四个浮点数。在后面的代码中会看到在纹理内存上使用向量指令时偏移量会相对实际元素的数量除以 4。

这个加载方法显然不是唯一的，我的理解是因为 A 和 B 的载入方法完全一样，只是所用的纹理不同，所以相比一个线程同时加载 A 和 B 可以减少与计算无关的指令的代码量。

载入的数据被暂时在寄存器上，等待被储存到共享内存。共享内存中的数据排布如下图：

图 4. 矩阵 A 被载入的数据在共享内存中的排布。

因为共享内存的的偏移单位是 1 字节，终于又回到了直观的表示，由此可见以上两图的数据储存方式其实是完全一样的，均为的列优先储存，唯一的区别是在共享内存中 B 的数据会直接跟在 A 后面。

图 5. 载入输入矩阵 A 和 B 的示意图，注意图中 lda, ldb, ldc, bx, by, k 的意义，这些变量在后面的代码中都会被用到。

以下伪代码是整个过程的描述，英文注释为 maxas 文档中所带，额外的注释为中文。

tid = threadId.x;bx  = blockId.x; // 可以看作C_ij的iby  = blockId.y; // 可以看作C_ij的j
blk = tid >= 32 ? by    : bx; ldx = tid >= 32 ? ldb/4 : lda/4; // lda和ldb为A的行数或B的列数，ldx为其在向量载入中的表示（每次4个浮点数故除以4），可以看成每一行的跨度tex = tid >= 32 ? texB  : texA; // 64个线程分为2个warp，第一个载入纹理A，第二个载入纹理B
tid2  = (tid >> 4) & 1;  // 此处将32个线程分为两组，tid2=0为第一组载入图三中的第一行，tid2=1载入第二行tid15 = tid & 15; // 本线程在每一行中列的位置
// 这些track变量表示本线程需要载入的数据在tex中的偏移，乘以4即在$A_i$或$B_j_T$中的偏移track0 = blk*64/4 + tid15 + (ldx * tid2);  // 本线程载入数据的起始位置，解释见后文track2 = track0 + ldx*2; // 本线程载入的第二部分数据，在第一部分两行后，以此类推track4 = track0 + ldx*4;track6 = track0 + ldx*6;
end = track0 + (k-8)*ldx; // 载入结束的标志，其中k为A的列数和B的行数，即两个相乘矩阵的公共维度，对于NxN的矩阵， k=N。因为每次载入8行所以减去8作为最后一次的载入的标记
writeS = tid15*4*4 + tid2*64*4; // 本线程载入数据在共享内存中的偏移，注意其相当于track0的第二项的16倍，因为向量载入的偏移1代表（4个数x每个浮点数4字节）writeS += tid >= 32 ? 2048 : 0; // 如果载入的是矩阵B的数据，放在矩阵A的数据之后，而矩阵A占用（64列x8行x每元素4字节）=2048字节
while (track0 < end){    tex.1d.v4.f32.s32 loadX0, [tex, track0];    tex.1d.v4.f32.s32 loadX2, [tex, track2];    tex.1d.v4.f32.s32 loadX4, [tex, track4];    tex.1d.v4.f32.s32 loadX6, [tex, track6];
    st.shared.v4.f32 [writeS + 4*0*64], loadX0;    st.shared.v4.f32 [writeS + 4*2*64], loadX2;    st.shared.v4.f32 [writeS + 4*4*64], loadX4;    st.shared.v4.f32 [writeS + 4*6*64], loadX6;
    // our loop needs one bar sync after share is loaded    bar.sync 0;
    // Increment the track variables and swap shared buffers after the sync.    // We know at this point that these registers are not tied up with any in flight memory op.    track0 += ldx*8;    track2 += ldx*8;    track4 += ldx*8;    track6 += ldx*8;    writeS ^= 4*16*64; // 又见magic number！其意义是A和B在共享内存中一共占4*16*64=4096字节，但是共享内存一共分配了8192字节的两组，每次载入后用这个操作切换到另外一组，其目的是实现一个流水线，在一个warp载入数据进一组时另一个warp就可以操作另一组已经完成载入的数据了
    // Additional loop code omitted for clarity.}

以上代码中还有两个问题需要用一定篇幅来澄清：

1. track0 = blk*64/4 + tid15 + (ldx * tid2); 中的第一项 blk*64/4 是用来从纹理中选择输入矩阵 A 或转置矩阵 B 中第 blk 条左上角相对整个输入矩阵左上角的一维偏移。由于所有条的左上角都在输入矩阵的第一列中，而行优先储存中第一列中任一点的偏移就是其行数，对于第 blk 条左上角就是 blk*64 ，而 / 4 来自向量因子。 tid15 + (ldx * tid2) 的意义比较清楚，即本线程在图 3 中对应的黄色格点相对绿色格点的位置， tid15 是列坐标， tid2 是行坐标，在一维表示时需要乘以该方向的跨度 ldx 。

2. 代码中用了四个 track 变量来记录四次载入的偏移，很容易想到只用一个 track 变量，载入一次后对偏移加两行再做一次载入来完成同样的工作:

tex.1d.v4.f32.s32 loadX0, [tex, track];track += ldx*2;tex.1d.v4.f32.s32 loadX2, [tex, track];track += ldx*2;tex.1d.v4.f32.s32 loadX4, [tex, track];track += ldx*2;tex.1d.v4.f32.s32 loadX6, [tex, track];track += ldx*2;

这样做的问题是 tex.1d.v4.f32.s32 指令发出后其 track 操作数是不会被保存的，为了保证其不被下一个增量指令修改，必须要插入控制代码等待前一条载入指令完成，而最坏情况下该指令可能花去上百个时钟周期。用四个偏移变量就可以不用等待传输完成就可以将四条载入指令一一发射出去，也是起到了流水线的作用。其代价是每个线程需要额外占用三个寄存器，所幸 GPU 上有足够的寄存器可以提供。

将共享内存中的数据载入寄存器

上节的工作完成后，共享内存中就有 A 和 B 的数据各 8 行，每行 64 个浮点数。将其各取出一行就可以将其中的元素进行前述的加乘操作，完成后各再取出一行直到共享内存中的 8 行数据被用完，此时其他 warp 应该已经在共享内存的另一组完成了从纹理内存的传输，计算线程只需切换到另一组进行计算即可。

如图 2 所示，对于每个线程，其实只需要 64 个浮点数中的 8 个，其在 A 和 B 向量中位置可以根据图上的计算出，具体计算过程在代码中是通过一顿骚位操作实现的，在此可以提前做一说明：

readAs = ((tid >> 1) & 7) << 4;

图中线程 2*i 和 2*i+1 会用到同一段 A，可以写作 (i/2) % 8。这段位操作就是这个表达式的位操作实现。<<4 实现了每段 A 的列向量间隔 16 字节，这是向量载入 4 个浮点数的大小。

readBs = (((tid & 0x30) >> 3) | (tid & 1)) << 4 + 2048;

B 的行向量中的选择更复杂一点，首先观察到对于偶数线程每隔 16 个线程 B 方向就要差 2 段（8 个浮点数），所以对于 6 个比特位表示的 64 线程， tid & 0x30 表示其中代表 tid mod 16 的后四位可以被遮盖掉，只有前两位对选择 B 有意义。其后的 >>3 其实是先 >>4 将前两位拉到个位数，再 *2 来表达相差的两段。 | (tid & 1) 等价于 +(tid & 1) ，表达的线程 2*i+1 永远会选择线程 2*i 后的那段数据，这也补上了之前相差两段中缺失的部分。

在图 2 中可能很早就有人注意到其中的线程排布顺序非常别扭，并没有按照线程号按行或列一个个排下来，其原因是为了避免共享内存访问的 bank 冲突。bank 冲突的定义和发生的条件在 CUDA 官方文档中有详细说明，简单地说就是共享内存的访问按地址被分成若干个 bank（最简单的方法是做求余数的操作），如果两个线程要访问位于同一 bank 的共享内存，其访问无法同时完成，访存时间成倍增加，增加的倍数由一个 bank 最多被多少个线程同时访问决定。当然这是最一般的情况，GPU 中提供了一些机制，比如广播，尽量减轻 bank 冲突对访问时间的影响。图 2 所示的线程顺序就是为了消除 bank 冲突所作的调整后的最佳排序。另一个别扭的地方是每个线程计算 4 个而不是直接计算，这也是为了避免 bank 冲突的技巧，在每个线程的实际计算中 4 个完全等价于 1 个矩阵。

在实现代码中还用到了一个技巧，虽然每线程只需要输入 16 个输入数据，实际分配的寄存器是这个数字的两倍，目的和前述的类似，是为了用两组寄存器实现流水线，即每个线程在用一行数据作计算时预先读取读取下一行的数据。

readAs = ((tid >> 1) & 7) << 4;readBs = (((tid & 0x30) >> 3) | (tid & 1)) << 4 + 2048;
while (track0 < end){    // Process each of our 8 lines from shared    for (j = 0; j < 8; j++)    {        // We fetch one line ahead while calculating the current line.        // Wrap the last line around to the first.        prefetch = (j + 1) % 8; // +1代表预读取                // Use even/odd rows to implement our double buffer.        if (j & 1)  // 在两组寄存器j0和j1直接切换        {            // 共享内存到寄存器的传输还是使用向量指令            // 在maxas的代码中可见j0Ax<00-03>是4个连续的寄存器，一个指令就可以完成到这4个寄存器的传输而无需一一写出            ld.shared.v4.f32 j0Ax00, [readAs + 4*(prefetch*64 + 0)];            ld.shared.v4.f32 j0By00, [readBs + 4*(prefetch*64 + 0)];            ld.shared.v4.f32 j0Ax32, [readAs + 4*(prefetch*64 + 32)];            ld.shared.v4.f32 j0By32, [readBs + 4*(prefetch*64 + 32)];        }        else        {            ld.shared.v4.f32 j1Ax00, [readAs + 4*(prefetch*64 + 0)];            ld.shared.v4.f32 j1By00, [readBs + 4*(prefetch*64 + 0)];            ld.shared.v4.f32 j1Ax32, [readAs + 4*(prefetch*64 + 32)];            ld.shared.v4.f32 j1By32, [readBs + 4*(prefetch*64 + 32)];        }    }    // swap our shared memory buffers after reading out 8 lines    readAs ^= 4*16*64;    readBs ^= 4*16*64;
    // Additional loop code omitted for clarity.}

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 个寄存器的顺序，其所采用的编号方案如下图：

图 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 即使对此还是不满足，它在前述的来回遍历基础上又加上一个漩涡提出了一个更诡异的遍历方法：

 1,  0,  2,  3,  5,  4,  6,  7, 33, 32, 34, 35, 37, 36, 38, 39, 45, 44, 46, 47, 41, 40, 42, 43, 13, 12, 14, 15,  9,  8, 10, 11,17, 16, 18, 19, 21, 20, 22, 23, 49, 48, 50, 51, 53, 52, 54, 55, 61, 60, 62, 63, 57, 56, 58, 59, 29, 28, 30, 31, 25, 24, 26, 27

按照文档中的说法，每个操作数有 8 字节的重用缓存可以缓存两个 4 字节寄存器，而逐行遍历只用了其中一个用于缓存 A 寄存器的数据，所以缓存使用率偏低。我推测考虑到有 B 操作数也有重用缓存但没有利用，逐行遍历的重用缓存利用率为 4/8/2=25%。对于来回遍历的利用率估算不是那么直观，文档直接给出了其利用率为 39%，而漩涡遍历的利用率可以达到 49%。从 maxas 最后生成的汇编代码来看其中有的指令确实对 A 和 B 同时使用了重用缓存，同时在对每个操作数缓存了两个寄存器：

--:-:-:-:1      FFMA R37, R71.reuse, R72.reuse, R37;--:-:-:-:1      FFMA R36, R71.reuse, R73.reuse, R36;#前2条指令对第3个操作数缓存了R72和R73，它们被接下来的2条指令用到--:-:-:-:1      FFMA R38, R69.reuse, R73, R38; --:-:-:-:1      FFMA R39, R69.reuse, R72, R39;#前4条指令对第2个操作数缓存了R69和R71，它们被接下来的4条指令用到--:-:-:-:1      FFMA R45, R71.reuse, R74.reuse, R45;--:-:-:-:1      FFMA R44, R71, R75.reuse, R44;--:-:-:-:1      FFMA R46, R69.reuse, R75.reuse, R46;--:-:-:-:1      FFMA R47, R69, R74.reuse, R47;

不过，在来回遍历可以完全解决 bank 冲突的情况下依然试图提高重用缓存的使用率的目的并不在于提高重用率，而且因为 FFMA 指令之间插入的从共享内存到寄存器的载入指令。这样做的目的是为了载入指令的延迟可以被不依赖于其数据的计算指令所填充。遍历 C 矩阵寄存器的前八条指令和其间插入的载入指令如下：

01:-:-:-:0      FFMA R1, R66.reuse, R72.reuse, R1;--:-:-:-:1      LDS.U.128 R80, [R106+0x200];--:-:-:-:1      FFMA R0, R66, R73.reuse, R0;--:-:-:-:0      FFMA R2, R64.reuse, R73.reuse, R2;--:-:-:-:1      LDS.U.128 R88, [R107+0x200];--:-:-:-:1      FFMA R3, R64, R72.reuse, R3;--:-:-:-:0      FFMA R5, R67.reuse, R72.reuse, R5;--:-:-:-:1      LDS.U.128 R84, [R106+0x300];--:-:-:-:1      FFMA R4, R67, R73.reuse, R4;--:-:-:-:0      FFMA R6, R65.reuse, R73.reuse, R6;--:-:1:-:1      LDS.U.128 R92, [R107+0x300];--:-:-:-:1      FFMA R7, R65, R72.reuse, R7;

由于载入指令的运行时间需要 20 多个时钟周期（对应于共享内存的延迟），在这段时间内其第一个的操作数都有可能被通过 bank 访问到，此时有可能 其后的计算指令也被发射并需要访问同一个 bank，这就造成了一个延迟的 bank 冲突。 不过这只是一个基本原理，maxas 的遍历顺序是如何具体避免这样的 bank 冲突目前还没有完全搞清楚。

 从本节可以看出即使所有计算全部在寄存器中进行，还是要用到两个技巧来得到最佳性能：

1. 最优化的寄存器编号

2. 最优化的遍历顺序

至于计算本身已经显得如此简单以至于 maxas 文档都懒得提了。

传输 C 矩阵到主显存

每个线程块完成所负责的那个小片矩阵的计算后，最后一个任务就是将其从寄存器传输到主显存。由于寄存器无法在线程间共享（其实有_shfl_sync() 指令可以但是此处不适用），所以每个线程必须将所计算的 4 个矩阵先传输至共享内存。之所以不从寄存器直接传输到主显存是因为按照现有的线程布局无法利用 GPU 的超大带宽。要充分利用 GPU 的带宽我们希望每个 warp 的 32 个线程所同时传输的数据是连续的，这样一个时钟周期里就可以一下子传输 128 字节的数据，如果这些数据离得太远，在最坏可能下需要分 32 次才能传输出去。

根据图 2 的线程布局，每一列连续的 64 字节数据分布在 8 个线程中，比如第 1 列前 4 行的结果都保存在线程 0,2,4,6,8,10,12,14 所控制的寄存器中，每个线程在该行有 8 个寄存器，而且为了避免 bank 冲突这 8 个寄存器都不是连续的，因此不能使用向量传输指令，所以需要分 8 次才能完成一个 warp 的传输，而此时 warp 中的其他 24 个线程将无所事事，因为其数据都不在这一列上。为了解决这个问题可以首先利用共享内存暂存所有线程的数据，然后用令一种线程布局将共享内存中的连续的数据同时传输出去。

首先还是要先将寄存器中的数据写入共享内存。每个线程的寄存器中所保存的四个矩阵是分成在列上对齐的两对。按图 6 所示的 maxas 的寄存器分配，每一列上有八个寄存器。比如第一列有寄存器 3,7,1,5，属于同一个矩阵的一列，以及 35,39,33,37，属于令一个矩阵的一列。由于结果矩阵 C 也是按照行优先储存的，如果将寄存器 3,7,1,5 拷贝到 4 个连续的寄存器（maxas 中命名为 cs<0-3>），35,39,33,37 拷贝到 cs<4-7>，就可以用向量储存指令在两个指令内将 8 个数拷贝到共享内存中对应的位置。图 7 中的左图是这个过程的示意图，可以看作将图 2. 的矩阵每隔四列抽出一条来拼在一起。完成后在共享内存中得到一个的矩阵，其中每一列都是连续的且对应于 C 矩阵中的一列。这时候改变一下线程次序，令 warp 中一个线程传输该列上一个字节的数据，就可以完成一次传输 32 个连续的浮点数。这个共享内存中的缓冲区可以利用之前为载入 A 和 B 所分配的空间，在完成 C 的计算后 A 和 B 的数据已经没有用了。

图 7. 左图为寄存器写入共享内存的线程布局，右图为此后从同一块共享内存读取的线程布局。本图中每一列是图 2 中矩阵 C 的一列，相邻的 2 列在矩阵 C 中间隔 4 列。

该方法的实现代码如下。虽然这个方法需要反复在寄存器和共享内存之间搬运数据，共享内存的延迟可以通过在 2 个 warp 间切换任务而得到掩盖，毕竟比多次访问主显存要快多了。其中值得注意的是虽然这个方法明显是分步完成的，但是代码中没有任何同步指令，这是因为所有的数据交换都是在同一个 warp 的 32 个线程直接完成的，warp 之间的数据保持独立。GPU 的执行模型可以严格保证同一 warp 内的同一指令可以同时完成。能够省去同步指令也是图 2 中并行方法的一个优势。

tid31 = tid & 31;tid32 = tid & 32; // 只可能取两个值，0 为第一个warp，32为第二个warp
// Remove the high bits if present from the last loop's xor. // Also remove the 2048 added onto readBs.// 之前对A和B在共享内存中分配了两倍于所需的容量（4096字节），一块用于已载入数据的计算，一块用于载入下一段A和B，每一块的前2048字节存放A，后2048字节存放B//这个AND操作等价于取第一块存放A的内存来存放CreadAs &= 0x7ff;  // 本线程左上角数据在64x64矩阵中的行坐标readBs &= 0x7ff;  // 本线程左上角数据在64x64矩阵中的列坐标
// Write to shared using almost the same shared mapping as before but collapse readBs down to stride one.writeCs = (readBs / 4) * 64 + readAs; // 本线程左上角数据在64x64矩阵中的相对左上角的一维偏移，根据行坐标和列坐标计算出。64/4是行优先储存矩阵进行向量传输的跨度。对于线程0这就是图7左图中的绿色格子中最上面的一个。
// Read out with a mapping amenable to coalesced global writesreadCs = ((tid32 << 3) + tid31) << 2; // 本线程在读取时开始的一维偏移，不考虑用于将元素个数转换为字节数的<<2操作，可以看到每个warp中的31个线程的编号tid31对应连续的readCs，也就是图7右图中的一整列黄色格子，第一个warp对应左边一列，第二个warp对应右边一列，中间相隔4行，也就是64*4==32<<3（第二个warp中tid32就是32）个元素。由图7右图可见很明显一个warp可以一次传输32个连续的浮点数。
ldc4 = ldc * 4; // ldc是矩阵C在行优先储存中列方向的跨度，因子4代表其单位为字节而不是浮点数。
cx = bx*64 + tid31; // cx可看作所要写入主显存的那一整列的在整个矩阵C中所对应的行数cy = by*64 + (tid32 >> 1); // cy可看作所要写入主显存的那一整列的在整个矩阵C中所对应的列数，显然对于同一个warp列数是一样的
// Cy00, Cy04, Cy08, Cy12 是图7右图中上面那四个绿色格点在整个矩阵C中的偏移// 虽然它们在共享内存中相隔1列，在矩阵C中它们之间的间隔是4列，所有它们之间的偏移是ldc4*4Cy00 = (cy*ldc + cx) * 4 + C;Cy04 = Cy00 + ldc4 * 4;Cy08 = Cy00 + ldc4 * 8;Cy12 = Cy00 + ldc4 * 12;
foreach copy vertical line of 8 registers from C into .v4.f32 cs0 and cs4{    // 步骤1. 从不连续的寄存器到共享内存    // Feed the 8 registers through the warp shuffle before storing to global    // 这里有一个缺失的步骤，每个线程首先将其上下对其的两个4x4矩阵中取出同一列的各四个元素，此时它们为了避免bank冲突    // 不得不位于不连续的寄存器上，这个步骤将其复制到8个连续的额外的寄存器cs0-cs7，上面的矩阵使用cs0-3，下面的使用cs4-7    st.shared.v4.f32 [writeCs + 4*00], cs0; // 将连续的寄存器cs0,cs1,cs2,cs3用向量指令传输到共享内存中该4个数对应的位置    st.shared.v4.f32 [writeCs + 4*32], cs4; // 和上一行同样的操作，因为上下两个4x4矩阵间隔32个数，需要对写入位置增加4*32字节的偏移
    // 步骤2. 从共享内存读取到寄存器，重用cs寄存器，不过此时并不用到其连续的特性    // cs0, cs2, cs4, cs6位于同一行上，在行优先储存中相差64个元素，4*64字节    ld.shared.f32 cs0, [readCs + 4*(0*64 + 00)];    // cs1,cs3,cs5,cs7位于另一行上，由图7右图可见与上一行相差32个元素    ld.shared.f32 cs1, [readCs + 4*(0*64 + 32)];    ld.shared.f32 cs2, [readCs + 4*(1*64 + 00)];    ld.shared.f32 cs3, [readCs + 4*(1*64 + 32)];    ld.shared.f32 cs4, [readCs + 4*(2*64 + 00)];    ld.shared.f32 cs5, [readCs + 4*(2*64 + 32)];    ld.shared.f32 cs6, [readCs + 4*(3*64 + 00)];    ld.shared.f32 cs7, [readCs + 4*(3*64 + 32)];       // 步骤3. 将cs寄存器的数写入主显存，对于整个warp相当于将一列连续的32个浮点数写入主显存。逻辑上可以看作是步骤2的反过程，除了改列的位置在共享内存和主显存中有所不同。    st.global.f32 [Cy00 + 4*00], cs0;    st.global.f32 [Cy00 + 4*32], cs1;    st.global.f32 [Cy04 + 4*00], cs2;    st.global.f32 [Cy04 + 4*32], cs3;    st.global.f32 [Cy08 + 4*00], cs4;    st.global.f32 [Cy08 + 4*32], cs5;    st.global.f32 [Cy12 + 4*00], cs6;    st.global.f32 [Cy12 + 4*32], cs7;        // 在下一次循环中输出本线程所计算的4个4x4矩阵的下一列，对应于C矩阵中的下一列，注意不要和图7中共享内存中的下一列混淆。     Cy00 += ldc4;    Cy04 += ldc4;    Cy08 += ldc4;    Cy12 += ldc4;
    // After processing forth set shift over to the stride 32 registers    // 补充说明，4次循环后4个4x4矩阵的左边两个已经传输到主显存了，接下去要传输右边的两个。    // 左边和右边两对4x4矩阵在C矩阵中对应的位置可以通过平移32列而重合，考虑到矩阵本身宽度有4列（在之前4次循环中已经通过 += ldc4 4次得到实现）    // 实际需要额外平移的是左右两对4x4矩阵的间距32-4=28列，这是这就是28这个magic number的来由    if (4th iteration)    {        Cy00 += ldc4 * 28;        Cy04 += ldc4 * 28;        Cy08 += ldc4 * 28;        Cy12 += ldc4 * 28;    }}

maxas 文档中另有 一张图表达这个过程，但可能由于未能对该充分理解，感觉其意义不大反而容易造成混淆故没有在此引用。代码本身已经足够描述这一过程了。

完成到主显存的传输后，maxas 所生成的 GEMM 内核的任务就完成了

256 线程实现

在以上所描写的每线程块 64 线程的基础上，可以将其扩展 4 倍到 256 线程，每个线程所做的工作不变。这样每个线程块所计算的小片矩阵的两个维度各扩大 2 倍，达到，此时输入矩阵的载入和结果的输出会有相应的变化，但理解了 64 线程实现后这些变化就非常简单，在此无需赘述。对于比较大的矩阵 256 线程实现有一些性能上的优势，详细测试结果参见 maxas 文档。

结语

本文虽然尽可能详尽地对原文档中的伪代码进行了注释，但这还是相对高层的实现，具体到 GPU 机器码还有一个重要的课题，即控制码没有在本文中涉及。考虑到本文的目的仅是介绍一些 GPU 优化的思路和实现方法，对此 maxas 文档中涉及控制码的部分没有进行解读。

总的来说，maxas 所用的优化思路还是比较清晰的，按其说法之前已经有文献提出了，其最困难的地方在于 nVidia 不愿意透露其硬件的实现细节，以至于都需要其作者经过艰苦的反向工程猜测出来的才能达到硬件性能的极限。可能作者自己搭建了一个测试平台来快速验证某些指令的细微差别所带来的性能的影响。无论如何这是一个伟大的工作，值得任何一位有志于冲击硬件性能极限的工程师深入研究。

原文链接：https://www.jianshu.com/p/e01024892afb

本文为机器之心发布，转载请联系本公众号获得授权。

✄------------------------------------------------

加入机器之心（全职记者 / 实习生）：hr@jiqizhixin.com

投稿或寻求报道：content @jiqizhixin.com

广告 & 商务合作：bd@jiqizhixin.com