向量化访存是指将多个内存访问操作合并为一个内存访问操作。这样可以减少内存访问的次数,提高内存访问的效率。在高性能领域,这又可以叫做 SIMD(Single Instruction Multiple Data)。在 CPU 侧,有适用 intel 平台的 SSE 指令集,适用 Arm 端的 Neon 指令集;在 GPU 侧除 cuda 外还有 opencl 框架,这些工具都支持向量化读取和向量化计算。

而在本节中,我们将介绍如何通过向量化访存来提高矩阵乘法的性能。

1. 优化思路

上一个 Kernel 中加载矩阵 A 共享内存的代码如下:

for (uint load_offset = 0; load_offset < BM; load_offset += stride_A){    smem_A[(inner_row_A + load_offset) * BK + inner_col_A] = A[(inner_row_A + load_offset) * K + inner_col_A];}

可以看到每次从 A 中读取一个元素, 且每次读取的元素不是连续的。我们可以使用向量读取指令 LDS.128 优化 Shared Memory 访问(对应 float4 数据类型)可以提高访存效率。GPU 是以 4 维向量为基本单位进行计算的,4 个浮点数组成的 float4 向量是 GPU 最基本的类型,使用 GPU 对两个 float4 进行向量计算与对两个整数或两个浮点数进行计算一样,只需要一个指令即可完成。

若每个线程每次取 1 个浮点数,每个线程需要消耗 4 次内存指令,才能将全局内存搬运至共享内存,若采用 float4 向量内存指令,每个线程每次可以搬运 4 个浮点数,则每个线程仅需要执行一次内存指令即可完成搬运。

LDS.128 指令可以一次性读取 4 个 float 类型的数据。

在将共享内存的值读取到寄存器当中的时候也可以使用用向量化加速。但是需要做一定的处理, 上一个 Kernel 中, 寄存器 A 中的值并不是顺序读取的。因此我们在将数据写入到共享内存的时候需要做一个转置。

同样在将结果写回到全局内存时,也可以一次型写入 4 个 float 类型的数据。这样可以减少全局内存访问的次数。

算法整体流程如下:

picture 0

picture 0

本 Kerne 和上一个 Kernel 的主要区别就在于如何加载数据到共享内存中。A 矩阵加载过程如下图所示:

picture 1

picture 1

共享内存的大小是 BM * BK, 我们每次读取 4 个元素。因为后续再把共享内存的数据加载到寄存器的时候我们需要连续的读取。因此我们需要把数据转置一下。B 矩阵由于本身后续加载到寄存器的时候就是连续的读取, 因此不需要转置。

后续的流程和上一个 Kernel 是一样的。区别就是读取和写入的时候使用了向量化的指令。下面让我们结合代码来看一下具体的实现。

2. 代码实现

在代码实现上, 我认为最难理解的就是共享内存的加载逻辑以及使用向量化访存后坐标的计算。