OpenPPL 中的卷积优化技巧
目前 OpenPPL 在 FP16 精度 T4 GPU 上取得了不错的性能，单 batch 模型的推理性能全面超越 TensorRT 8.0，多 batch 模型的推理性能达到 TensorRT 8.0 的 90%+。作者：刘宇玺本文将介绍 OpenPPL 目前开源的卷积算子实现方法，从算法层面介绍 OpenPPL 在 Tensor Core 上的设计方法和优化经验。不涉及 mma 指令、ldmatrix 指令等的使用技巧，比如 conflict-free 的数据搬运、shared memory 的数据排布等，相关知识可以参考 ptx 手册及英伟达官方的教程。
文章共分为五个部分：
[*]第一部分介绍 Tensor Core 上运行卷积算子的背景知识；
[*]第二部分介绍基于 Tensor Core 计算模式的特点，将卷积共分为三种类型；
[*]第三部分介绍针对各种卷积类型 OpenPPL 设计的卷积算法，以及生成卷积 kernel 的代码生成技术；
[*]第四部分将给出目前 OpenPPL 卷积算子/模型与 TensorRT 的性能对比；
[*]最后一部分介绍 OpenPPL CUDA 明年的一些规划。
一、Tensor Core 上的卷积运算
TensorCore 的编程方法
Tensor Core是从Volta架构开始英伟达引入的特殊计算硬件单元，它的输入端是两个固定尺寸的矩阵，矩阵相乘后快速得到输出矩阵。使用Tensor Core有以下4种方式。
[*]直接调用cuDNN API。cuDNN库中有专门针对Tensor Core提供的卷积算子API，可以直接调用cuDNN API来使用Tensor Core。但是它的缺点是Tensor Core类似于黑盒，无法进行定制化的算法开发。
[*]使用CUDA wmma API。wmma API对Tensor Core提供了更易用的编程抽象，方便了程序员调用，缺点是它封装了若干条底层的mma指令，指令粒度较粗，相比ptx/sass指令的实现有性能差距。很多基于wmma API深度学习编译框架（如TVM, MLIR，AKG等），无论是template-based还是polyhedral-based，由于底层都基于wmma API，自动生成的kernel与cuDNN/cuBLAS库提供的kernel相比性能有所差距。
[*]使用内联ptx汇编语言。ptx是架构兼容的中间级汇编语言，是英伟达支持的最底层语言，有完整的官方文档和完善的工具链。使用mma ptx指令可以较细粒度地操控Tensor Core，cutlass库基于mma ptx汇编生成的kernel性能已经接近cuDNN/cuBLAS的性能。
[*]使用sass汇编。虽然英伟达没有开放sass指令集，但github上已有开源的图灵汇编器（turingas）、安培汇编器（CuAssembler）来操作sass汇编。采用sass汇编编写卷积kernel不依赖nvcc编译器，可以进行细粒度的代码控制，如指令重排、 寄存器分配等，能够获得最佳的性能。缺点是需要有扎实的汇编知识（参见CloudCore知乎博文），代码调试困难。
OpenPPL采用内联ptx汇编的方式进行Tensor Core的编程。
Tensor Core上的卷积算法
由于Tensor Core的运算模式是固定大小的小矩阵乘运算，因此可以将卷积运算映射为矩阵运算。有两种映射方式，一种是显式矩阵乘卷积（explicit GEMM convolution），另一种是隐式矩阵乘卷积（implicit GEMM convolution）。二者的计算方式相同，区别在于是否需要额外的存储空间来放置临时数据。显式矩阵乘是在临时空间内将卷积操作的特征图和卷积核转换为两个矩阵，即Im2Col操作。隐式矩阵乘通过索引计算省去了Im2Col操作，从而避免了开辟临时空间。TensorRT和cutlass在Tensor Core计算单元上主要采用的是隐式矩阵乘算法。
整个隐式矩阵乘卷积的算法流程如下图所示，其中，IH/IW/IC代表输入特征图的长，宽，通道，FH/FW代表卷积核的长和宽，OH/OW/OC代表输出特征图的长，宽，通道，Batch大小为N。X代表通道的数目，其含义可以参考TensorRT中提出的五维数据排布格式NCxHWX，将X个通道放在inner-most维度（X一般是mma指令k维度大小的倍数，如8/16/32/64等）。
图 1 隐式矩阵乘卷积的算法流程
隐式矩阵乘卷积的算法流程如上图所示，其主要的区别在于滑窗的顺序。滑窗顺序（sliding order）指的是卷积核滑过特征图时元素之间的计算顺序，主要有三种滑窗顺序，如下表所示。
表 1 滑窗顺序对比滑窗类型滑窗顺序排布类型适合的计算核Channel-firstFH → FW → ICNHWCTensor CoreChannel-lastIC→ FH → FWNCHWCUDA CoreChannel-interleavedICx → FH → FW → xNHWC/NCxHWXTensor Core
[*]Channel-first滑窗顺序。它的优点是有更好的访存带宽。由于采用NHWC的排布格式，IC维度是连续存储的，channel-first的滑窗顺序可以保证通道维度在连续读取的。然而这种滑窗顺序的缺点是数据局部性较差，滑窗在IC维度的移动无法复用相邻OH/OW点的读取数据。
[*]Channel-last滑窗顺序。这种滑窗顺序的优点是有更好的数据复用，滑窗在FH/FW维度的移动可以复用相邻OH/OW点的读取数据。缺点是访存不连续，会造成访存带宽的降低。例如当滑窗到下一个FH或IC时，特征图的数据会在较远的地址存放，造成访存的不连续。特别的，当孔洞卷积大于1时，滑窗在FW维度的访存也是不连续的。因此，这种滑窗顺序不适合在Tensor Core上执行，无法高效利用Tensor Core的计算能力。
[*]Channel-interleaved滑窗顺序是结合了Channel-first和Channel-last两种滑窗顺序的优点，既可以有效利用窗口滑动时的数据复用，又可以保证访存数据是连续读取的。但是缺点如图所示，需要更复杂的索引控制以及相应的排布支持。OpenPPL同cutlass一样，都采用了这种滑窗顺序。
为了减少读取特征图数据时的索引计算，特征图索引的偏移地址（delta address）会提前计算好，保存在查找表中（Lookup Table，即LUT）,以参数的形式传递给卷积kernel。当采用时Channel-first滑窗顺序时，LUT的大小为FH*FW；而当采用Channel-interleaved滑窗顺序时，LUT的大小是FH*FW*IC/X。为了防止LUT表过大，OpenPPL采用了增量式偏移地址（incremental delta address）的方式，将LUT表大小压缩到FH*FW。
二、卷积的分类
虽然卷积映射为GEMM有统一的数学表示，但是由于卷积参数千变万化，不同的卷积参数运行在Tensor Core上呈现出不同的程序特征。根据卷积参数的变化特征和Tensor Core的计算特点，OpenPPL将卷积大致分为三类，并根据这三种分类设计了相应的卷积算法。
下图展示了典型的Resnet50模型的卷积参数变化情况，其中包含了大图少通道和小图多通道两种类型。特征图的三个尺寸“长x宽x通道”标注于特征图的上方，如224x224x3、56x56x64等。卷积核尺寸“长x宽”标注于卷积核的下方，如7x7、1x1等。
图 2Resnet50模型的卷积参数变化
[*]大图少通道类型：这种类型的卷积一般存在于神经网络的开始阶段，它的特点是特征图的长和宽较大（即大图），比如224，但是通道数目比较小（即多通道），比如3。这种卷积类型属于访存密集型运算，Tensor Core的计算能力对于这种类型得不到充分的发挥。
[*]小图多通道类型：这种类型的卷积存在于神经网络的中间和末尾阶段，它的特点是特征图的长和宽较小（即小图），比如56/28/14等，但是通道数目比较大（即多通道），比如64/128/256等。这种卷积类型也属于访存密集型计算，由于特征图的长和宽较小，无法充分利用GPU的并行计算资源。可以从挖掘规约维度（K维度）的并行性来充分利用GPU的硬件资源。
[*]多图多通道类型: 这种类型的卷积是小图多通道类型的多batch情况。当batch增多时并行性也得到增加。这种卷积类型属于计算密集型运算，比较适合发挥Tensor Core的计算能力。
大图少通道类型
图 3 大图少通道卷积转化为矩阵乘的尺寸图
把大图少通道的卷积类型转换成矩阵乘后的尺寸如上图所示，m维度的尺寸很大，可以生成足够多的线程块来使用GPU的硬件资源；而k维度的尺寸却很小，甚至小于mma指令中MxNxK的K的大小。这里以Turing GPU上FP16精度的mma指令为例，其MxNxK的大小为16x8x8；当卷积核的通道数小于8时，就要在通道末尾填充0（即padding 0），使其达到8的倍数后才能发射到Tensor Core上执行。通道末尾填充的0越多，无效的冗余操作就越多，Tensor Core的利用率就越低。对于常见的3通道RGB输入卷积层，Tensor Core的利用率只有3/8。
同时从上图也可以看出，batch的增加只会继续增加m维度的大小，并不会改变卷积的类型。即，多batch情况下的大图少通道类型依旧是大图少通道类型。
小图多通道类型
把小图多通道的卷积类型转换成矩阵乘后，矩阵乘的尺寸如下图所示，m/n维度的尺寸很小，而累加K维度的尺寸却很大。例如，当m/n维度的大小为256时，如果使用常见的Mtile=Ntile=128进行分块，这样只能产生4个线程块，大部分GPU的SM将处于空闲状态。
图 4 小图多通道卷积转化为矩阵乘的示意图
矩阵乘通常采用splitK算法来解决并行性不足的问题。SplitK算法就是在规约维度（K维度）上进行进一步的分块，来提高GPU的资源利用率。例如，当通道数IC大小为1024时，采用Ktile=128的分块后，K维度可以划分出1024/128=8个线程块，这样可以有效地提高GPU资源的利用率。cuBLAS和cutlass库的矩阵乘算法中都包含splitK支持。
多图多通道类型
图 5 多图多通道卷积转化为矩阵乘的尺寸图
把多图多通道的卷积类型转换成矩阵乘后，矩阵乘的尺寸如上图所示。小图多通道类型的m维度（OH*OW）尺寸很小，但随着batch的增加，m维度的尺寸变得很大，可以产生足够多的线程块来使用GPU的SM核；同时，规约K维度的尺寸也很大，每个warp有足够多的计算任务。因此，与前两种卷积类型不同（访存密集型），多图多通道的卷积类型属于计算密集型的任务。
可以从优化访存效率的角度来优化卷积性能。多图多通道类型的另一个特点是相比m/k维度而言，n维度的尺寸较小，卷积核矩阵的尺寸比特征图矩阵小很多。因此可以将卷积核矩阵尽量常驻在L2 cache中，通过复用卷积核数据来优化L2的性能。
三、三种卷积算法
根据三种卷积类型的特征，OpenPPL提出了相应的卷积算法，即，针对大图少通道卷积类型的IDXN算法，针对小图多通道卷积类型的2SPK算法，针对多图多通道卷积类型的SWZL算法。
IDXN算法
针对大图少通道的卷积类型，OpenPPL提出了IDXN算法，该算法的流程如下图所示。
图 6IDXN算法流程图
IDXN算法的第一个特点是与传统的增量式偏移（incremental delta）查找表不同，采用了累积式偏移（accumulative delta）查找表。这种查找表利用了IC通道数目少的特点，将FH*FW个卷积核的偏移地址累积式地计算出来，在寻址时通过一次计算便可得到相应的地址，不需要多次累加计算。例如对于大小为3x3、通道数为3的卷积核，增量式偏移需要3x3xCeilDiv(3, 8)=9次才能遍历完卷积核,而累积式偏移仅需要CeilDiv(3x3x3,8)=4次便可遍历完卷积核。
IDXN算法的第二个特点是不使用共享内存（SMem）来暂存数据。在使用Tensor Core时，mma指令经常与ldmatrix指令搭配使用，即，mma指令计算所需的矩阵A、矩阵B的数据从GMem读入后会放到SMem中，然后ldmatrix指令按照特定的规则将数据从SMem中读入各个线程的寄存器中。然而这种使用方式适合于计算密集型的任务，大图少通道类型属于访存密集型的任务，程序的瓶颈在访存带宽上，因而使用SMem暂存数据的收益很小，相反还会增加数据读取的延迟。
IDXN算法的第三个特点是不对输出结果在SMem中重排，而是直接输出到GMem中。这是由于n维度很小，向量化输出带来的收益很小。把数据放到SMem中进行重排反而会增加一次读写SMem的延迟，因此IDXN算法直接将结果从寄存器输出的GMem中。
2SPK算法
2SPK算法是针对小图多通道的卷积类型，提出的双层splitK的卷积算法。它借鉴矩阵乘中的splitK算法，在规约维度进行双层的任务划分。即除了在通道层IC进行划分外，也对卷积核FW*FH进行划分，这样就可以产出更多的split，更充分地利用GPU的并行硬件资源。
例如，当对3x3卷积核、通道=1024进行卷积操作时，第一层对通道按照split-Channel=8进行划分，每个split的channel数为1024/8=256；第二层对3x3卷积核按照split-Filter=9进行划分，每个split的Filter大小为3x3/9=1；通过这两层划分后产生的split总数为8x9=72。相比通道层划分只能产生8的split数量，双层的任务划分可以产生72的split数量，进一步提高了并行性。
图 7双层规约方式的示意图
2SPK算法采用双层规约的方式，即采用SMem规约和GMem规约相结合的方式，对split的中间结果进行累加。如上图所示，第一级规约将线程块TB内的结果在SMem累加，输出中间结果到GMem上；第二级规约再执行单独的kernel，将GMem的中间结果累加，得到最后的结果。双层的规约方式结合了SMem规约与GMem规约的特点，提高了规约的效率。
SWZL算法
针对多图多通道的卷积类型，OpenPPL提出了SWZL算法。由于该卷积类型属于计算密集型的任务，可以从优化L2 cache访存效率的角度来提高Tensor Core的使用率。SWZL算法的计算流程与cutlass的卷积算法比较类似，如下图所示。
图 8SWZL算法流程图
SWZL算法最主要的特点是采用了牛耕式（Boustrophedon）的线程块调度策略。GPU的线程块调度采用轮询的方式来调度，即先调度X维度的线程块，再调度Y维度的线程块，最后调度Z维度的线程块。这种轮询的调度策略对矩阵数据的访问也是轮询的，如果发生了抖动缓存的效率会降低。SWZL算法利用卷积核矩阵尺寸较小的特点，牛耕式的发射线程块，提高了L2 cache的效率。
此外，SWZL算法使用了ldmatrix指令来从SMem中给Tensor Core加载数据；同时使用向量化输出的方式，提高了输出数据的效率。
自动代码生成
针对Tensor Core，这里以FP16精度的MxNxK=16x8x8的mma指令为例，介绍2SPK算法在Tensor Core上的实现方式，IDXN算法和SWZL算法以此类推。
由于卷积尺寸千变万化，固定分块大小的kernel并不能适用于所有的卷积尺寸上。cutlass以及TensorRT都采用自动代码生成的方式，遍历尽可能多的分块大小来生成卷积kernel。OpenPPL采用基于宏的代码生成方式，由宏来定义TB，Warp,K维度及stage的分块大小，这些宏传递给函数模板，最终生成相应的卷积kernel。
[*]Warp级分块：Warp级的分块大小是mma指令MxN大小的倍数，上限是硬件的物理容量。Warp级的分块大小如下表所示（由于寄存器容量的限制，无法生成128x64的warp分块）：
表 2Warp级的分块大小M \ N81632641616x816x1616x3216x643232x832x1632x3232x646464x864x1664x3264x64128128x8128x16128x32\2. TB级分块：TB级的分块大小是Warp分块大小的倍数，上限是SMem的大小以及TB内最大的线程数目。TB级的分块大小如下表所示：
表 3TB级的分块大小。表中的数字表示相应Warp级分块大小的倍数M \ N12411x11x21x422x12x22x444x14x2\3. K维度的分块：K维度的分块表示每个TB在一次循环中处理的K维度元素的个数，它是mma指令中K维度大小的倍数。由于2SPK算法中使用了SMem来进行规约，因此还需对TB内的规约次数进行遍历。K维度的分块大小如所示：
表 4K维度的分块大小。以k32_s16为例，s16表示每个split的大小为16，k32表示K维度的大小是32，此时TB内规约大小为2。 K / S816321xk8_s8k16_s16k32_s322xk16_s8k32_s16k64_s324xk32_s8k64_s16k128_s324. Pipeline Stage的分块。为了更好地掩盖访存的延迟，可以预取GMem中的数据到SMem中，Tensor Core可以及时地得到数据，从而使计算和访存更好地相互重叠。目前OpenPPL支持的stage的分块大小是1和2。
上述四种TB，Warp，K维度和stage的分块组合最终生成了2SPK卷积算法的全部kernel。
四、OpenPPL卷积算子/模型性能
本部分给出目前OpenPPL框架在T4 GPU上FP16精度的推理性能。目前OpenPPL支持两种方式来生成卷积kernel：Static模式和JIT模式。两种模式的区别请参看《OpenPPL CUDA运行时编译机制》。此处给出的性能是Static模式下的卷积kernel性能。
图 9Resnet34模型的各层卷积配置信息
这里选择Resnet34模型为例（卷积层数较少且具有典型性），分析各种卷积算法的性能。其中C1层属于大图少通道类型，C2~C11层属于小图多通道类型。
图 10单batch模式Resnet34模型的各层卷积性能对比
单batch模式下，Resnet34模型各卷积层的性能如上图所示，OpenPPL相比TensorRT有20%的性能提升。其中，C1层采用IDXN算法，相比TensorRT性能优势明显，有超过60%的加速比；C2~C11层采用2SPK算法，C8~C11层的性能优势明显，有40%+的性能提升。与此同时，cutlass和cuDNN库在单batch的卷积层上性能不佳，性能低于TensorRT的50%，它们在大图少通道和小图多通道类型上没有进行深入的优化。
图 11多batch模式Resnet34模型的各层卷积性能对比
Batch=32模式下，Resnet34模型各卷积层的性能如上图所示，OpenPPL与cuDNN、TensorRT的性能相差无几，在各个卷积层上性能互有胜负。这表明cuDNN和TensorRT在多Batch模式下的大图少通道和多图多通道类型上进行了深入的优化。OpenPPL在C1层依然采用IDXN算法；C2~C11层却采用了SWZL算法。与此同时，cutlass库在多Batch的卷积层上依然性能不佳，不到TensorRT的60%。其中C1、C3、C4、C6、C7、C9、C10层的性能远低于其他三者。
图 12单batch模式下的模型的性能对比
得益于卷积层的性能优势，OpenPPL在单batch的CV类公开网络上取得了全面超越TensorRT的性能，如上图所示。其中除VGG19模型稍慢于TensorRT外（4%），其他模型的性能都超过了TensorRT，如DenseNet模型有40%的性能提高，ShuffleNet模型有48%的性能提高，Resnet34模型有18%的性能提高。
图 13多batch模式下的模型的性能对比
OpenPPL在batch=32的CV类公开网络上的模型推理性能如上图所示，达到了TensorRT的95%。通过nvprof进行性能对比，发现多batch的卷积性能略低于TensorRT，这与上图的模型性能结果相符。OpenPPL今后将继续提高多图多通道卷积类型的性能。
总结
本文介绍了OpenPPL在卷积算子方面的优化经验，取得了不错的效果。接下来OpenPPL将支持T4 GPU上int8精度的推理，并且明年将开始全面支持Ampere架构的A30 GPU。
最后，欢迎大家使用OpenPPL!
参考文献
[*]Chetlur, Sharan, et al. "cudnn: Efficient primitives for deep learning." arXiv preprint arXiv:1410.0759 (2014).
[*]Bhaskaracharya, Somashekaracharya G., Julien Demouth, and Vinod Grover. "Automatic kernel generation for volta tensor cores." arXiv preprint arXiv:2006.12645 (2020).
[*]Jorda, Marc, Pedro Valero-Lara, and Antonio J. Pea. "Performance evaluation of cudnn convolution algorithms on nvidia volta gpus." IEEE Access 7 (2019): 70461-70473.
[*]https://docs.nvidia.com/cuda/cublas/index.html
[*]https://github.com/NVIDIA/cutlass
[*]https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html
[*]https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-1688
[*]https://github.com/NVIDIA/cutlass/blob/master/media/docs/implicit_gemm_convolution.md
[*]https://github.com/NVIDIA/cutlass/blob/master/media/docs/efficient_gemm.md
页:
[1]