江苏龙网

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

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


基本思想
如上节所述 ， 分片算法在利用了片上高速缓存之后 ， 不但小片矩阵的乘法速度可以大大加快 ， 还可以利用计算小片矩阵相乘的时间将下一个小片从主内存传送至片上共享内存 ， 换句话说此时整个矩阵相乘的时间已经完全由小片矩阵相乘所决定 ， 如果要进一步提高性能就要在小片矩阵相乘上做文章了 。
在共享内存内部做矩阵相乘虽然已经很快了 ， 但距离硬件性能的极限还是有距离 ， 主要瓶颈是两个 。 首先共享内存的延时终究还是比不过寄存器 ， 在 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 线程） 。 这一点区别于原始分片算法中每个线程计算矩阵中的一个元素 ， 也是充分利用寄存器的超低延迟的关键 。
机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理


推荐阅读


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

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