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

在从事深度学习框架的实现工作时，了解到Nervana有一个称为Maxas的汇编代码生成器项目https://github.com/NervanaSystems/maxas，可以生成性能超过nVidia官方版本的矩阵相乘的GPU机器码，由此对其工作原理产生兴趣。其作者Scott Gray在代码外提供了详细的文档https://github.com/NervanaSystems/maxas/wiki/SGEMM，但由于该算法的复杂性，行文晦涩，逻辑跳跃，尤其是对一些方法的动机没有交待，很容易迷失在细节中。本文可以看作按作者对该文档的理解进行的重写，但求在细节上不厌其烦，对其代码的前因后果作尽可能完整的交待，不过大体结构还是按照maxas文档来安排，文中图片也全部出自该文档。

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

背景

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

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

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

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

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

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

 

maxas对于小片矩阵乘法是用64个线程来并行实现的，其中每个线程负责计算2x2个矩阵的乘积，64个线程按照8X8布局，这样就确定了小片的大小为一个边长N=2x4x8=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内核执行前完成的，不影响内核执行的性能。
纹理内存中的数据也是分段被载入共享内存的，不过按照原方法每段载入的应该是一个个64x64的小片，但为了充分利用寄存器资源，maxas采用了完全不同的计算方法。如果线程块计算的是，首先将矩阵A按每64行一条划分为64*N的条，所需的输入数据全部在第条上，当然这一条数据还还是很大，需要进一步分次载入，maxas中每次载入并消耗64*8，分N/8次完成。对于矩阵B方法类似，只不过是按列划分为N*64的条，转置后载入方法和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的i
by  = 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个4x4而不是直接计算8x8，这也是为了避免bank冲突的技巧，在每个线程的实际计算中4个4x4完全等价于1个8x8矩阵。

在实现代码中还用到了一个技巧，虽然每线程只需要输入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个4x4矩阵先传输至共享内存。之所以不从寄存器直接传输到主显存是因为按照现有的线程布局无法利用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个线程将无所事事，因为其数据都不在这一列上。为了解决这个问题可以首先利用共享内存暂存所有线程的数据，然后用令一种线程布局将共享内存中的连续的数据同时传输出去。

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

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

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

maxas文档中另有一张图表达这个过程，但可能由于未能对该充分理解，感觉其意义不大反而容易造成混淆故没有在此引用。代码本身已经足够描述这一过程了。
完成到主显存的传输后，maxas所生成的GEMM内核的任务就完成了

256线程实现

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

结语

本文虽然尽可能详尽地对原文档中的伪代码进行了注释，但这还是相对高层的实现，具体到GPU机器码还有一个重要的课题，即控制码没有在本文中涉及。考虑到本文的目的仅是介绍一些GPU优化的思路和实现方法，对此maxas文档中涉及控制码的部分没有进行解读。
总的来说，maxas所用的优化思路还是比较清晰的，按其说法之前已经有文献提出了，其最困难的地方在于nVidia不愿意透露其硬件的实现细节，以至于都需要其作者经过艰苦的反向工程猜测出来的才能达到硬件性能的极限。可能作者自己搭建了一个测试平台来快速验证某些指令的细微差别所带来的性能的影响。无论如何这是一个伟大的工作，值得任何一位有志于冲击硬件性能极限的工程师深入研究。