江苏龙网

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

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


载入的数据被暂时在寄存器上 ， 等待被储存到共享内存 。 共享内存中的数据排布如下图：
机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理
本文插图

图 4. 矩阵 A 被载入的数据在共享内存中的排布 。
因为共享内存的的偏移单位是 1 字节 ， 终于又回到了直观的表示 ， 由此可见以上两图的数据储存方式其实是完全一样的 ， 均为的列优先储存 ， 唯一的区别是在共享内存中 B 的数据会直接跟在 A 后面 。
【机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理】机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理
本文插图
图 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
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 loadedbar.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 。


推荐阅读


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

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