MegEngine TensorCore 卷积算子实现原理（下）-博客-MegEngine

显存 -> Shared Memory 的数据搬运

这一节我们会介绍从显存 (Global Memory) 到 Shared Memory 的数据搬运。显存到 Shared Memory 的数据搬运是由 Conv2dTileSrcIteratorFpropPrecomp 来完成的，本文并不会详细地解读代码的实现，而是描述线程搬运数据的过程，帮助大家建立直观的印象，更好地理解代码。
如果以上一节中 Shared Memory 的逻辑布局为例，同一 Warp 中每个线程读取的数据的逻辑布局如下图所示，每个线程读取 16 个 INT8 类型的数据，恰好构成一个 Vector。

而在实际的物理显存中，线程访问的数据分布如下图所示：

我们可以看到每个线程读取了 128 位的数据。
相邻的线程读取的数据在物理上是连续的。

因此线程从 Global Memory 读取数据的 pattern 可以满足合并访存的要求，同时以最大的数据位宽进行访存，最大化了显存带宽的利用率。
然后如果将线程读取的数据映射到 Shared Memory 的物理地址，我们可以看到

每 8 个线程向 Shared Memory 写入 128 字节的数据，恰好落在 Shared Memory 的 32 个不同的 bank 中。
同一 Warp 的访存分为四个阶段完成，每个阶段都没有 bank conflict。

下图演示了一个 Warp 写入 Shared Memory 的过程：

Shared Memory -> 寄存器的数据搬运

Shared Memory 到寄存器的数据搬运是由 MmaTensorOpMultiplicandTileIterator 完成的。同一 Warp 在每一轮迭代过程会读取 4 个 8x16 的矩阵到寄存器中，每个线程会读取一行的数据。例如第一轮迭代时，线程读取的数据在逻辑上的布局如下图所示：

而实际上数据在 Shared Memory 里的物理布局如下图：

可以看到：

每个线程读取了 128 位的数据，因此访存分为四个阶段来进行。
每一阶段的 8 个线程读取的数据恰好落在了 Shared Memory 的 32 个 bank 中，并且线程访存的数据之间不存在冲突。

当进行到第二轮迭代时，每个线程访问的数据的物理布局如下图：

同样的访存的每一个阶段都不存在 bank conflict。

Accumulator 写回全局存储

在 int8 的情况下，同一 Warp 负责输出 64x64 的结果，kernel 会分成 8 次写回 Global Memory，每次写回 32x8 的矩阵。这样保证了每次将 Tensor 按照 NCHW32 格式写回显存时，同一 Warp 的 32 个线程恰好写了物理上连续的 256 字节的数据，而每个线程写回 8 个字节，保证了可以使用64位宽的数据类型进行显存的写操作，尽可能提高带宽的利用率。
由于mma指令的特点，输出矩阵的数据分布在各个线程上，而为了能够合并访存，即：让相邻线程写回的地址是连续的，我们利用 Shared Memory 对同一 Warp 中 32 个线程的数据进行了交换。数据交换后，每个线程拥有连续的 8 个通道的数据，且线程写的地址是连续的，保证了写回 Global Memory 满足合并访存的要求。
线程交换数据的过程如下图所示：

每一轮迭代，Warp 中的 32 个线程将 32x16 的矩阵数据写入到 Shared Memory 中。接着如下图所示，每个线程会把连续的 8 个 channel 的数据读到寄存器中。

Shared Memory 的数据交换是由以下两个Iterator完成的

InterleavedTileIteratorTensorOp 完成了每一轮迭代将 32x8 的数据写入到 Shared Memory 中。
InterleavedSharedLoadIteratorTensorOp 负责将连续的 8 个 channel 的数据读到Fragment寄存器中。

当线程将交换后的数据读到Fragment寄存器之后，会由EpilogueOp，在卷积的基础上完成BiasAdd的运算。以 BiasAddLinearCombinationRelu 为例，它实际上完成了下面的运算：

accumulator = conv(x, w)
y = alpha * accumulator + beta * bias + gamma * z

其中 bias 是一个PerChannel的 Tensor，代表了每个输出通道的偏置，z 是一个和卷积输出大小一致的 Tensor，用于Convolution和ElemwiseAdd的融合。
最后EpilogueOp的输出会由 TensorPredicatedTileIteratorTensorOp 真正地写回到 Global Memory 中。每个线程写回的数据如下图所示：

可以看到线程写回的 pattern 满足合并访存的要求，因此能最大化 Global Memory 写的效率。

总结

本文介绍了 MegEngine 底层的卷积算子实现原理，算子性能可以达到cudnn的80%以上，测速结果可以参见文章。

MegEngine 会对卷积实现进行持续优化，进一步提升算子的性能，目前来看有以下两点可做的优化：

借鉴 Nvidia 官方 CUTLASS ImplicitGEMM Convolution 实现对 mask 的处理，提高TileIterator对于 mask 判断的效率。
现在的卷积实现在写回显存时利用 Shared Memory 进行数据交换是存在 bank conflict 的。后续会考虑两点优化
- 对 Shared Memory 的数据布局进行探索，消除 bank conflict，优化 Shared Memory 数据交换的效率。
- 对 Global Memory 中的 Weight Tensor 的布局进行探索，提高每个 Thread 上 accumulator 的局部性，避免在 Shared Memory 中进行数据交换。

显存 -> Shared Memory 的数据搬运

Shared Memory -> 寄存器的数据搬运

Accumulator 写回全局存储

总结

参考资料