江苏龙网

  1. 首页 > 资讯 > 科技 > >

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


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-7st.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 文档中涉及控制码的部分没有进行解读 。


推荐阅读


上一篇：机器之心开发者必看！行业数智专家传授OOTB生态化协作秘籍

下一篇：魅族天青色等魅族｜魅族17终于发布 黄章把魅族放在了未来