黑科技：用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，输出