江苏龙网

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

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


本文插图

图2. maxas 计算两个 64x64 矩阵相乘的示意图 ， 绿色的 4x4 小片是线程 0 负责的那部分元素 ， 黄色是其他线程负责那部分的左上角元素 。 图中只标出了左上角 4x4 矩阵的线程号 ， 其他 15 个只是其重复 。 每个黑框中包含 32 线程 ， 代表一个 warp 。 这块 32x32 矩阵计算完成后 ， 线程 0 以及其他线程保持相对位置 ， 依次平移到另外三个绿色小片标出的地方 ， 完成这个 64x64 矩阵的计算 。 左边的向量是 A 矩阵的一个列 ， 上方的向量是 B 矩阵中与之对应的行 ， 其中标为绿色的数据（各 8 个浮点数）是线程 0 所需要用到的 ， 其他线程需要的不难类推 。
maxas 的整个实现就是为了这个算法服务的 ， 后文中所要解决的问题也是该算法在实现过程中所遇到的 。 以上的那些参数选择 ， 比如为什么选择 64 个线程 ， 都是根据 GPU 硬件资源决定的 ， 以便在满足每个线程所需的寄存器资源基础上 ， 创建尽可能多的线程 warp ， 以便调度器在某些 warp 等待数据时将别的 warp 切换进来进行计算 。
将输入矩阵载入共享内存
为了实现以上算法 ， 这 64 个线程首先被用来将两个输入矩阵所需要的部分从主显存载入共享内存 。 这里需要指出一个原文中没有提到的假设 ， 即在 maxas 中默认矩阵是按照行优先储存的 ， 即矩阵的每一列在内存中是连续的 ， 否则无法解释后面的一系列算法 。 由于算法的不同载入的方法也有所不同 ， 并且在原方法基础上增加了一些优化：
1. B 矩阵用到的是行数据 ， 而默认的行优先矩阵储存中行数据是不连续的 ， 需要做转
置 ， 行变为列 ， 这样 A 和 B 的载入方法可以完全相同 ， 以降低代码复杂度；
2. A 和 B 矩阵被作为一维纹理载入 ， 这样不但可以利用纹理数据的缓存 ， 而且不用耗费时间进行数组越界检查 ， 因为纹理载入会自动将越界的数据置为 0 ， 对于矩阵相乘不会有任何影响 。
了解以上预处理可以更方便地理解后面的伪代码 。 创建纹理和转置的工作应该是在 GPU 内核执行前完成的 ， 不影响内核执行的性能 。
纹理内存中的数据也是分段被载入共享内存的 ， 不过按照原方法每段载入的应该是一个个的小片 ， 但为了充分利用寄存器资源 ， maxas 采用了完全不同的计算方法 。 如果线程块计算的是 ， 首先将矩阵 A 按每 64 行一条划分为的条 ， 所需的输入数据全部在第条上 ， 当然这一条数据还还是很大 ， 需要进一步分次载入 ， maxas 中每次载入并消耗 ， 分次完成 。 对于矩阵 B 方法类似 ， 只不过是按列划分为的条 ， 转置后载入方法和 A 完全相同 。 其内存布局如下图所示：
机器之心矩阵相乘在GPU上的终极优化：深度解析Maxas汇编器工作原理
本文插图

图 3. 每次循环中被一个 warp 载入共享内存的一段纹理 ， 可以看作 Bj 或转置后的 Ai ， 这样这个矩阵其实又回到了常规的列优先储存 。 这个图转置后看其实更直观 。
图中每一个格子代表一个线程负责载入的数据单元 ， 绿色格子是线程 0 要先后载入的四个单元 ， 黄色格子是其余 31 个线程第一次载入的那部分数据 。 整个 warp 每次载入两行 ， 如此重复四次完成 。 在 GPU 上执行单位是 32 个线程组成的 warp ， 所以 64 个线程是分为两个 warp 执行 。 其中一个 warp（线程 0-31）载入 A 另一个（线程 32-63）载入 B 。 此图有一个容易造成困惑的地方是图中的矩阵形状为而不是 ，这是因为后面每个线程会用到向量指令一次载入 4 个浮点数 ， 即每个格子本身就是四个浮点数 。 在后面的代码中会看到在纹理内存上使用向量指令时偏移量会相对实际元素的数量除以 4 。
这个加载方法显然不是唯一的 ， 我的理解是因为 A 和 B 的载入方法完全一样 ， 只是所用的纹理不同 ， 所以相比一个线程同时加载 A 和 B 可以减少与计算无关的指令的代码量 。


推荐阅读


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

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