黑科技：用cutlass进行低成本、高性能卷积算子定制开发

图形处理器通用计算(GPGPU)是指利用 GPU 来计算原本由 CPU 处理的通用计算任务。由于现代 GPU 拥有强大的并行处理能力，通用 GPU 在面对矩阵乘法、卷积等大量并行的计算密集型算法时，性能远远超越了传统的 CPU。CUDA 是由 NVIDIA 推出的 GPGPU 高性能计算方案，目前大多数深度学习推理任务都可以通过 CUDA 来进行加速。

为了充分发挥 CUDA 平台的计算能力，NVIDIA 推出了高度优化的深度学习、线性代数算子库 cudnn、cublas、cutlass，以及 CUDA 平台上的深度学习推理框架 TensorRT。

cudnn、cublas 这样的基础算子原语库在常见的卷积层上性能表现很好，通常都能够满足用户的需求，但是在面对用户高度定制化的算法时，基础算子库往往并不能充分发挥硬件的性能。这是由于算子优化的长尾问题引起的，基础算子库引入了许多卷积优化的通用策略，但是这些优化的策略并不能覆盖所有的情况，实际算法中的卷积层有可能并不能从通用的优化策略中获得收益，从而无法充分发挥硬件的性能。

基础算子库的另一个问题是用户无法对这些基础算子进行定制化开发，当算法开发人员想为卷积算子添加一种新的激活函数，或者想添加一种特殊的卷积算子(比如：LocalConv)时，就会变得束手无策。
cutlass 是 NVIDIA 推出的一款线性代数模板库，它定义了一系列高度优化的算子组件，开发人员可以通过组合这些组件，开发出性能和 cudnn、cublas 相当的线性代数算子。但是 cutlass 仅支持矩阵乘法运算，不支持卷积算子，从而难以直接应用到计算机视觉领域的推理部署中。
TensorRT 是一款非常强大的深度学习推理部署框架，在 CUDA 平台上性能表现非常优秀，而且目前已经比较成熟，用户使用起来比较方便。然而 TensorRT 也存在着一些问题，对于开发人员来说，TensorRT 是一个黑盒，用户没有办法细粒度控制 TensorRT 内部的实现细节。

例如：在部署量化网络时，开发人员无法控制 TensorRT 底层的量化细节，有可能会出现部署和训练的精度对不齐的问题。再比如：TensorRT 在推理部署时，用户无法精细的控制算子的显存使用情况，有时 TensorRT 在运行网络时耗费了大量的显存，而用户却没有特别好的办法对此进行优化。

为了在 CUDA 平台上进行深度学习的推理部署，各大开源框架也都推出了各自的解决方案。

大部分开源训练框架在 CUDA 平台上的部署方案，都是基于模型转换工具，将网络转换成 TensorRT 支持的格式，然后交由 TensorRT 来执行推理任务。然而各大训练框架在算子的定义上会有细微的差别，这使得在模型转换的过程中会引入难以避免的性能、精度上的损失。
TVM 作为一款支持全平台的深度学习推理框架，对 CUDA 平台进行了比较好的支持。TVM 基于算子优化的原语定义了一系列矩阵乘法、卷积的模板，通过对模板进行运行时调优，来获得最优的性能。但是 TVM 采用的代码自动生成技术在 CUDA 平台上的效果和 cudnn、cublas 等手动调优的算子库还有不少差距，另外 TVM 在性能调优时需要耗费比较长的时间。上述两点原因阻碍了 TVM 在真实的推理部署场景中得到很好的应用。

由于官方库无法满足算法开发中的定制化需求，而开源界对 CUDA 平台的优化不够深入，无法满足算法部署中的性能需求，MegEngine 基于 cutlass 进行了二次开发，补充了 cutlass 对卷积算子的支持。用户通过自定义分块大小，可以很好的解决算子优化中的长尾问题。同时框架复用了 cutlass 里高度优化的算子组件，同时提炼了一套 CUDA 平台卷积算子的优化策略，让用户以较低的开发成本，完成定制化的卷积算子开发。

基于 CUTLASS 的卷积算子开发框架

算子优化的长尾问题

在实际的模型推理部署中，cudnn 这样的官方库的性能往往不够好。例如，cudnn 只对输出通道数多于 64 的情况进行了优化，而当通道数不足 64 的时候，cudnn 需要将通道数补齐 64，并且启动更多的线程数来进行计算，这不仅造成了计算资源的浪费，而且不能获得较好的算子性能。

如果我们利用 MegEngine 开源的 cutlass 算子开发框架，就可以很方便地对输出通道数较小的情况进行定制优化。

例如：当输入 feature map 的 4 维分别是 N=16, C=64, H=92, W=160 时，卷积核的大小为 3x3，输出的通道数为 32 时，我们可以通过如下的代码，添加一种新的分块大小，来处理输出通道数为 32 的情形：

// 定义输入 feature map tensor 的 layout
using LayoutSrc = cutlass::layout::TensorNCxHWx<32>;
// 定义输入 weight tensor 的 layout
using LayoutFilter = cutlass::layout::TensorCxRSKx<32>;
// 定义线程块的分块大小，M，N，K
using ThreadBlockShape = cutlass::gemm::GemmShape<32, 64, 64>;
// 定义 warp 的分块大小，M，N，K
using WarpShape = cutlass::gemm::GemmShape<32, 16, 64>;
// 定义 Matrix Multiply-Add 指令的矩阵分块大小，M，N，K
using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>;
// 定义卷积后处理 operator
using EpilogueOp = cutlass::epilogue::thread::
                      BiasAddLinearCombinationReluClamp<int8_t, 8, 
                          int32_t, int32_t, float>;
using Convolution = cutlass::convolution::device::Convolution<
  int8_t,       // 输入 feature map 的 data type
  LayoutSrc,    // 输入 feature map 的 layout
  int8_t,       // 输入 weight 的 data type
  LayoutFilter, // 输入 weight 的 layout
  int8_t,       // 输出 tensor 的 data type
  LayoutSrc,    // 输出 tensor 的 layout
  int32_t,      // 输入 bias 的 data type
  LayoutSrc,    // 输入 bias 的 layout
  int32_t,      // 矩阵乘法内部累加的 data type
  cutlass::convolution::ConvType::kConvolution, 
  cutlass::arch::OpClassTensorOp, 
  cutlass::arch::Sm75, 
  ThreadBlockShape, WarpShape, InstructionShape, 
  EpilogueOp, 
  cutlass::convolution::threadblock::
      ConvolutionNCxHWxThreadblockSwizzle<
          cutlass::Convolution::ConvType::kConvolution>, 
  2,           // 2 代表是否开启 shared memory ping-pong prefetch 优化
  16, 16>;     // tensor alignment, 代表 load/store 指令的位宽
               // 越宽指令吞吐量越高，有助于提升性能
Convolution conv_op;
typename Convoluition::Arguments args{...};
conv_op.initialize(args, workspace);
// 执行 convolution 算子
conv_op();

在 T4 卡上实测，我们通过 cutlass 自定义的算子实现比 cudnn 的性能快了 26%。

而在一些常见的卷积参数下，cutlass 定义的卷积算子的性能也是和 cudnn 的性能可比的，我们在 T4 卡上实测了 ResNet50 中一些常见卷积层的性能：

cutlass 在选取的 17 个卷积层下有 11 个卷积层的性能超过了 cudnn，余下的 6 个卷积层的性能也基本达到了 cudnn 的 80%以上。

算子融合

NVIDIA 的 Turing 架构显卡引入了 TensorCore int8 计算单元，GPU 的计算能力得到了极大的提升，然而 GPU 的访存能力并没有相应增长，这时候 GPU 的访存往往成为了推理性能的瓶颈。在这种场景下，我们就需要将访存密集型算子和计算密集型算子进行融合，减少访存密集型算子的开销。下面我们通过一个使用 TensorCore int8 推理加速的例子来介绍 MegEngine 和 cutlass 是如何进行算子融合的。

CUDA 平台上的 8-bit 量化卷积层采用的是 NCHW4 的数据布局(Layout)。不同于常见的 NCHW 的 Layout，这种 Layout 将 4 个通道打包在一起，连续的存放在内存中，然后按照 stride 从小到大依次存放 Tensor 的 W、H、C、N 四个维度的数据。为了使用 TensorCore 来进行加速，需要将 Tensor 的 Layout 转换为 NCHW32 的 Layout，这种 Layout 和 NCHW4 类似，只是将 32 个通道打包到一起存放到内存中。

在使用 MegEngine 进行推理部署时，只要用户打开了 TensorCore 的优化选项，MegEngine 就会在图优化阶段插入合适的 Tensor Reformat 算子来完成 Layout 的转换，如图 2 中的第一个阶段图变换所示。接下来 MegEngine 会将消去冗余的 Tensor Reformat 算子，得到图 2 中的第二个阶段的计算序列。

结合 cutlass，MegEngine 还可以进一步对计算图进行优化。首先，我们发现池化(Pooling)算子和它后面相连的 Reformat 算子是可以交换的。交换两个算子的顺序之后，计算图最前面的 Elemwise、Convolution、Reformat 这三个算子可以通过 cutlass 融合成一个超级卷积算子(Super Conv)，这样就得到了图 2 中最后的计算图。在优化后的计算图中，TensorCore 引入的访存密集型算子已经全部融合进卷积算子中了，这样优化后的推理网络可以完全享受到 TensorCore 的加速效果，而没有额外的 Tensor Reformat 的开销。

那么如何使用 cutlass 的算子融合功能呢？cutlass 已经提供了 NCHW4 和 NCHW32 这两种 Layout 相互转换的高性能读写组件，只需要将卷积的 operator 和相应的后处理(Epilogue)的 operator 组合起来就可以定义 Convolution+Reformat 的融合算子了。图 3 中示例代码展示了如何用 cutlass 定义一个输入 Tensor 为 NCHW4 Layout，输出 Tensor 为 NCHW32 Layout 的卷积算子。

// 定义输入 feature map tensor 的 layout
using LayoutSrc = cutlass::layout::TensorNCxHWx<4>;
// 定义输入 weight tensor 的 layout
using LayoutFilter = cutlass::layout::TensorCxRSKx<4>;
// 定义输出 tensor 的 layout
using LayoutDst = cutlass::layout::TensorNCxHWx<32>;
// 定义线程块的分块大小，M，N，K
using ThreadBlockShape = cutlass::gemm::GemmShape<64, 128, 32>;
// 定义 warp 的分块大小，M，N，K
using WarpShape = cutlass::gemm::GemmShape<64, 32, 32>;
// 定义 Matrix Multiply-Add 指令的矩阵分块大小，M，N，K
using InstructionShape = cutlass::gemm::GemmShape<1, 1, 4>;
// 定义卷积后处理 operator
using EpilogueOp = cutlass::epilogue::thread::
                      BiasAddLinearCombinationReluClamp<int8_t, 4, 
                          int32_t, int32_t, float>;
using Convolution = cutlass::convolution::device::Convolution<
  int8_t,       // 输入 feature map 的 data type
  LayoutSrc,    // 输入 feature map 的 layout
  int8_t,       // 输入 weight 的 data type
  LayoutFilter, // 输入 weight 的 layout
  int8_t,       // 输出 tensor 的 data type
  LayoutDst,    // 输出 tensor 的 layout
  int32_t,      // 输入 bias 的 data type
  LayoutDst,    // 输入 bias 的 layout
  int32_t,      // 矩阵乘法内部累加的 data type
  cutlass::convolution::ConvType::kConvolution, 
  cutlass::arch::OpClassSimt, 
  cutlass::arch::Sm61, 
  ThreadBlockShape, WarpShape, InstructionShape, 
  EpilogueOp, 
  cutlass::convolution::threadblock::
      ConvolutionNCxHWxThreadblockSwizzle<
          cutlass::Convolution::ConvType::kConvolution>, 
  2,           // 2 代表是否开启 shared memory ping-pong prefetch 优化
  4, 16>;      // tensor alignment, 代表 load/store 指令的位宽
               // 越宽指令吞吐量越高，有助于提升性能
Convolution conv_op;
typename Convoluition::Arguments args{...};
conv_op.initialize(args, workspace);
// 执行 convolution 算子
conv_op();

我们在 T4 卡上对 ResNet50 中的第一个卷积层进行了测试，当输出 Tensor 是 NCHW4 Layout 时的耗时是 3.03ms，Tensor Reformat 算子的耗时是 0.309ms，Convolution+Reformat 算子的耗时也是 3.03ms，但是融合后减少了 Tensor Reformat 算子开销，性能提升了约 10%。

定制卷积算子

在一些高度定制化的场景，算法工程师会提出一些新的卷积算子来提升网络的性能。例如，识别任务中 Local 算子、Google Brain 提出的 CondConv 算子等等，这些算子引入了更多的参数量，来提升模型的推理精度。

但是在 CUDA 平台上，这些算子往往没有比较好的优化实现，这就阻碍了这些算子在实际的推理任务中落地。我们发现这些算子的计算过程和普通的卷积算子大体相同，只是访问卷积核的方式略有不同。

我们可以在 cutlass 的卷积算子定义前处理(Prologue)的 operator 来改变卷积算子访问卷积核的方式，同时复用 cutlass 中的高性能卷积组件，来实现性能较优的 Local 算子和 CondConv 算子。在旷视的人脸识别业务中，我们基于 cutlass 实现了高性能的量化 CondConv 算子已经得到了落地，在不影响推理性能的情况下，获得了免费的涨点。

自定义激活函数

目前 NVIDIA 提供的 cudnn 算子库中卷积算子支持的激活函数只有 ReLU，如果算法工程师在模型中想开一下脑洞，使用一些新颖的激活函数(例如：HSwish)，那么这样的激活函数是不能被融合进卷积算子中的，这样会造成模型推理耗时增加，在一些对推理延时要求高的场景下，新型激活函数就不能真正得到落地。

如果借助 cutlass，就可以比较轻松地解决自定义激活函数的问题，我们只需要添加一种新的后处理(Epilogue)operator 就可以实现新的激活函数了。例如，下面的代码定义了 HSwish 的激活函数：

template <typename ElementOutput_,  
          int Count,  
          typename ElementAccumulator_ = ElementOutput_,
          typename ElementBias_ = ElementOutput_,    
          typename ElementCompute_ = ElementOutput_,
          FloatRoundStyle Round = FloatRoundStyle::round_to_nearest,
          typename Policy = NumericArrayConverterPolicy<
                  ElementOutput_, Count, 
                  ElementAccumulator_, ElementBias_,
                  ElementCompute_, Round>>
class BiasAddLinearCombinationHSwishClamp {
    /// 定义 Param、构造函数等，这里省略部分代码
    /// ...
public:
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(FragmentAccumulator const& accumulator,
                              FragmentBias const& bias,
                              FragmentOutput const& source) const {
        SourceConverter source_converter;
        AccumulatorConverter accumulator_converter;
        BiasConverter bias_converter;


        ComputeFragment converted_source = source_converter(source);
        ComputeFragment converted_accumulator =
                accumulator_converter(accumulator);
        ComputeFragmentBias converted_bias = bias_converter(bias);


        ComputeFragment intermediate;


        multiplies<ComputeFragment> mul_add_source;
        multiply_add<ComputeFragment> mul_add_accumulator;
        multiply_add<ComputeFragmentBias> mul_add_bias;
        HSwish<ComputeFragment> hswish;


        minimum<ComputeFragment> min_accumulator;
        maximum<ComputeFragment> max_accumulator;


        /// 计算+bias
        intermediate =
                mul_add_source(gamma_, converted_source);
        intermediate =
                mul_add_accumulator(alpha_, converted_accumulator,
                                    intermediate);
        intermediate = mul_add_bias(beta_, converted_bias,
                                    intermediate);
        /// 计算 HSwish 激活                     
        intermediate = hswish(scale_, inv_scale_, intermediate);


        ElementCompute const kClamp = ElementCompute(
                (1U << (sizeof_bits<ElementOutput>::value - 1)) - 1);


        intermediate =
                max_accumulator(intermediate, -kClamp - ElementCompute(1));
        intermediate = min_accumulator(intermediate, kClamp);


        /// 转换成输出的 data type
        OutputConverter destination_converter;
        return destination_converter(intermediate);
    }
};

只需要要将新定义的 Epilogue operator 传入 Convolution operator 的模板，就可以得到一个融合了新的激活函数的卷积算子了。

CUDA 平台的推理部署

到目前为止，最新版本的 MegEngine 已经集成了由 cutlass 实现的卷积算子。

按照[文档]介绍的方法 dump 量化好的模型，就可以使用 MegEngine 来完成推理的部署了。

[文档地址]

https://megengine.org.cn/doc/advanced/inference_in_nvidia_gpu.html#inference-in-nvidia-gpu

我们可以用 load_and_run 工具来对模型测速。

[如何使用 load_and_run]

https://megengine.org.cn/doc/advanced/how_to_use_load_and_run.html#how-to-use-load-and-run

例如ResNet-18 测试结果如下图所示：

./load_and_run resnet18.mge --input ./cat.npy --enable-nchw32 --fast-run
mgb load-and-run: using MegBrain 8.9999.0(0) and MegDNN 9.3.0
[09 14:14:14 from_argv@mgblar.cpp:1169][WARN] enable nchw32 optimization
load model: 3018.428ms
=== prepare: 182.441ms; going to warmup
[09 14:11:11 invoke@system.cpp:492][ERR] timeout is set, but no fork_exec_impl not given; timeout would be ignored
[09 14:11:11 invoke@system.cpp:492][ERR] timeout is set, but no fork_exec_impl not given; timeout would be ignored
[09 14:11:11 invoke@system.cpp:492][ERR] timeout is set, but no fork_exec_impl not given; timeout would be ignored
warmup 0: 481.411ms
=== going to run input for 10 times
iter 0/10: 19.432ms (exec=0.754,device=19.307)
iter 1/10: 18.537ms (exec=0.899,device=18.497)
iter 2/10: 18.802ms (exec=0.727,device=18.762)
iter 3/10: 18.791ms (exec=0.653,device=18.759)
iter 4/10: 18.614ms (exec=0.761,device=18.585)
iter 5/10: 18.529ms (exec=0.708,device=18.499)
iter 6/10: 18.660ms (exec=0.706,device=18.634)
iter 7/10: 18.917ms (exec=0.667,device=18.894)
iter 8/10: 19.093ms (exec=0.655,device=19.070)
iter 9/10: 19.211ms (exec=0.630,device=19.187)
=== finished test #0: time=188.586ms avg_time=18.859ms sd=0.304ms minmax=18.529,19.432

可以看到，在 T4 卡上，ResNet18 的 end-to-end 时间大概是 18.86ms，如果使用 TensorRT 来部署的话，end-to-end 时间大概是 16.85ms。MegEngine 在 CUDA 平台上的推理性能能达到 TensorRT 的 90%左右，总的来说还是可以比较的。在一些推理延时要求不高，但是高度定制化，推理精度要求高的部署场景下，直接使用 MegEngine 的 CUDA 平台推理部署方案还是能满足需求的。

总结

本文介绍了最新版的 MegEngine 中基于 cutlass 开发的卷积算子优化的框架。在接下来几篇文章，我们会继续介绍 cutlass 优化卷积算子的原理，以及如何使用 cutlass 在 MegEngine 里添加一个高性能的自定义卷积算子。

借助 cutlass 框架，开发人员可以开发自定义分块大小的卷积算子，解决推理优化中的长尾问题，可以支持自定义激活函数，可以完成卷积算子和访存密集型算子的融合，还可以定制性能还不错的变种卷积算子。

我们非常欢迎大家来使用 MegEngine 在 CUDA 平台的推理部署功能，以及基于 cutlass 的卷积算子的定制化功能，也非常期待开发者们能在使用过程中提出宝贵的意见，使得 MegEngine 和 cutlass 卷积框架能够在高度定制化的推理部署场景下帮助到广大的深度学习开发者。

参考文献

[1] Kerr, A., (2020). Developing CUDA kernels to push tensor cores to the absolute limit on NVIDIA A100. In: GPU Technology Conference.

[2] Chetlur, S., Woolley, C., Vandermersch, P., Cohen, J., Tran, J., Cantanzaro, B., & Shelhamer, E. (2014). cudnn: Efficient primitives for deep learning. arXiv preprint arXiv: 1410.0759.

[3] Vanholder, H. (2016). Efficient Inference with TensorRT. In: GPU Technology Conference.

[4] Chen, T., Moreau, T., Jiang, Z., Zheng, L., Yan, E., Shen, H., ... & Guestrin, C. (2018). TVM: An automated end-to-end optimizing compiler for deep learning. In: Proceedings of the 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI). (pp. 578-594).

[5] Yang, B., Bender, G., Le, Q.V., & Ngiam, J. (2019). CondConv: Conditionally parameterized convolutions for efficient inference. In: Advances in Neural Information Processing Systems. (pp. 1305-1316).

[6] Ma, N., Zhang, X., Huang, J., & Sun, J. (2020). WeightNet: Revisiting the design space of weight network. In: Proceedings of the European Conference on Computer Vision (ECCV).

欢迎访问