博客 (下)MegEngine TensorCore 卷积算子实现原理

(下)MegEngine TensorCore 卷积算子实现原理
技术分享 发布时间:2021/05/18 16:17:30 浏览量:2141 标签:卷积 算子 CUTLASS

接上文MegEngine TensorCore 卷积算子实现原理(上)

显存 -> 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完成的

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

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

其中 bias 是一个PerChannel的 Tensor,代表了每个输出通道的偏置,z 是一个和卷积输出大小一致的 Tensor,用于ConvolutionElemwiseAdd的融合。 最后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 中进行数据交换。

参考资料

上一篇:分享

利用共享内存实现比 NCCL 更快的集合通信

下一篇: 分享

深度学习算子优化-FFT

相关推荐