江苏龙网

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

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


机器之心发布
作者：XiaoyuWang
九大章节 ， 一万余字 ， 这篇文章可能是目前为止Maxas汇编器工作原理最全面、最细致的解析 。
机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理
本文插图

在从事深度学习框架的实现工作时 ， 了解到 Nervana 有一个称为 Maxas 的汇编代码生成器项目 ， 可以生成性能超过 nVidia 官方版本的矩阵相乘的 GPU 机器码 ， 由此对其工作原理产生兴趣 。
项目地址：https://github.com/NervanaSystems/maxas
其作者 Scott Gray 在代码外提供了详细的文档（https://github.com/NervanaSystems/maxas/wiki/SGEMM） ， 但由于该算法的复杂性 ， 行文晦涩 ， 逻辑跳跃 ， 尤其是对一些方法的动机没有交待 ， 很容易迷失在细节中 。
本文可以看作按作者对该文档的理解进行的重写 ， 但求在细节上不厌其烦 ， 对其代码的前因后果作尽可能完整的交待 ， 不过大体结构还是按照 maxas 文档来安排 ， 文中图片也全部出自该文档 。
值得说明的是 Maxas 使用的算法完全依赖于 Maxwell 架构的一些特性 ，随着新一代 GPU 的架构的演进这个项目本身已经完全过时了 ， 但其解决问题的思路仍然值得借鉴 。
背景
单精度矩阵相乘（SGEMM）应该是每一个学习 CUDA 编程的程序员最熟悉的案例了 ， 从 nVidia 推出 CUDA 的第一版起这就是其官方教程里唯一的例子 ， 这不仅因 SGEMM 是几乎所有计算密集型软件最重要的底层模块 ， 更是因为通过这个例子可以很好地展示 GPU 中的各种优化技巧 ， 特别是对 GPU 内存层次的利用 。
对于两个 NxN 的矩阵 A 和 B 的相乘 ， 一个最简单的并行方法是对于其输出矩阵 C（大小同为）的每一个元素开一个线程 ， 该线程载入 A 的一行和 B 的一列 ， 然后对其做一次向量的内积 。 但问题是在 GPU 上访问显存的延时相当的大（~100 时钟周期） ， 如果 A 的一行因为在内存中是连续的还能够利用 GPU 的超大显存带宽一次载入多个元素平摊其载入时间以及缓存来降低延时 ， 对于 N 上千的大矩阵来说 B 的一个列中的元素的内存地址也要相隔上千 ， 意味着载入一次中除了需要的那个列元素外大部分数据都是无用的 ， 同时这种访问模式几乎不可能有缓存线命中 ， 总而言之这个并行方法的内存访问效率低到令人发指 。
对其的优化就要用到共享内存了 ， 共享内存是位于 GPU 上的片上缓存 ， 速度可与一级缓存相当 ， 而且同一个线程块中的线程可以通过共享内存交换数据 ， 唯一的缺点是容量有限 。 为了利用这一小块高速内存 ， 原来的并行算法必须有所改进：矩阵将在每个维度被划分成个小片 ， 输出矩阵 C 可以写为：
A 和 B 也做同样的划分 ， 其中不是元素而是小片矩阵 ， 当然小片大小为 1 时小片矩阵就退化为单个元素 。 显然矩阵乘法的定义依然在此适用： 。
如果把小片看作一个元素 ， 整个矩阵的规模相当于被缩小了 倍 。 对于每个小片的结果可以由一组线程负责 ， 其中每个线程对应小片中的一个元素 。 这个线程组将 A 的行小片和 B 的列小片一一载入共享内存 ， 在共享内存上对其做矩阵相乘 ， 然后叠加在原有结果上 。 这样对小片中的元素每载入一次可以被在高速的共享内存中被访问 K 次 ， 由此得到远高于原始方法的内存存取效率 。
机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理
本文插图

网上找的分片算法示意图 ， 图示比 CUDA 官方教程中的更加清楚一点 。
以上只是对这种算法的一个简单描述 ， 经过这样的优化整个算法已经可以达到相当高的效率了 ， 而且是进一步进行优化的基础 。 为了理解后文中的内容 ， 读者需要细读 CUDA 的官方教程以对这个分片算法非常熟悉 ， 并且对 nVidia GPU 的编程模型有比较深入的理解 。


推荐阅读


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

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