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


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

图 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
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;


推荐阅读