从卷积到矩阵乘法矩阵乘法（GEMM）具有计算密度高、易并行等优良特性。芯片行业、高性能计算领域等传统领域往往以GEMM为基准，并且已经将其优化到接近硬件理论峰值。为了获得更好的性能增益，im2col算法将GEMM带入了卷积神经网络的工程优化领域。ImplicitGEMM算法进一步解决了im2col固有的冗余显存使用和冗余前后处理问题，这些问题在GPU等存储受限的硬件上尤为重要，进一步提升了GEMM在卷积优化中的重要性。此外，硬件厂商也开始为矩阵乘法提供越来越多的硬件支持，比如增加了各种MMA指令和TensorCore。这些原因共同促使许多现有的优化算法库使用im2col/ImplicitGEMM作为其默认的卷积优化。计划。im2col算法假设卷积的输入shape为(n,ic,ih,iw)，kernel为(oc,ic,kh,kw)，输出为(n,oc,oh,ow)。im2col算法的流程如下图所示。简单reshapekernel得到一个行M=oc，列K=ic*kh*kw的矩阵，称为矩阵A。用一个kernel大小的立方体作为输入的滑动窗口，展开每次按chw的顺序将一个小立方体变成一列。将整个输入从上到下，从左到右滑动后，将得到一个行K=ic*kh*kw列N=n*oh*ow的矩阵作为矩阵B。此时，我们可以计算GEMM(A,B)得到卷积结果矩阵C，其行M=oc，列N=n*oh*ow。在GEMM的计算过程中，可以根据m、n、k的三维下标推导出卷积中数据的输入输出和核中的下标。因此，可以将im2col和GEMM这两个进程整合起来，达到降低内存占用，加速性能的效果。这其实就是ImplicitGEMM的原理。本文不做过多介绍，有兴趣的可以阅读之前的技术文章。ImplicitBatchedGEMM上一篇文章主要介绍了MegEngine大核深度卷积优化的背景和动机。本文将介绍具体的优化思路和工程实践。借助im2col/ImplicitGEMM算法，GEMM在传统的密集卷积优化中表现出了优异的性能。因此，GEMM也应该用于大核深度卷积。上面分析过，直接使用与denseconvolution相同的方法将im2col算法应用于大核depthwiseconvolution，会产生BatchedGEMV，很难达到硬件浮点计算的峰值。显然，为了将GEMM应用于大核深度卷积，我们应该改变我们的想法。如下图，denseconvolution一般是用kernel对输入进行卷积计算输出，但是当kernelsize很大甚至大于输入时，实际上应该是用输入对kernel进行devolved计算输出.回想一下，稠密卷积使用内核对输入进行反演，并对输入进行im2col变换。现在大核depthwiseconvolution是用输入对kernel进行deconvolute，所以此时要对kernel进行im2col变换。算法过程没有本质区别，只需要在im2col中以kernel为输入，input为kernel即可。由于depthwiseconvolution是逐通道进行的，im2col变换也需要逐通道进行。如下图所示，每次通道变换后，都会生成一个M=n，N=oh*ow，K=ih*iw的GEMM。根据上一篇文章的分析，BatchedGEMM比BatchedGEMV更容易填充硬件设备浮点运算的峰值。CUTLASS是NVIDIA的开源模板库，旨在提供以较小的成本编写性能不差的GEMM的能力。CUTLASS内置了GEMM的元调度，可以让计算尽可能覆盖内存访问延迟，从而获得良好的性能。旷视早在CUTLASS正式开源其卷积实现之前，就已经基于CUTLASS做了自己的卷积实现，打磨了一个更适合内部业务的旷视版本的CUTLASS。这里的ImplicitBatchedGEMM也是基于旷视版的CUTLASS实现的。代码已经用MegEnginev1.8.2开源，实现细节不再过多介绍。下图所示的实验数据表明，随着kernelsize的增大，ImplicitBatchedGEMM的性能大致呈线性增长，在某些情况下可以逼近理论峰值。ImplicitBatchedGEMM的优势在于，一方面可以复用成熟的GEMM优化思想和基础设施，可以使用TensorCore进行加速；另一方面，如果在推理过程中不需要可变形状，则可以将内核的im2col变换进一步加速。当然，它的缺点也很明显。比如在小批量的情况下还是会退化成BacthedGEMV。如果用M*N*K*2来逼近GEMM的计算量，不难发现ImplicitBatchedGEMM的计算量比之前增加了$\frac{ih*iw}{kh*kwGEMM由密集卷积转换而来}$次，这意味着当输入明显大于内核大小时，ImplicitBatchedGEMM表现不佳。下图所示的实验结果也表明，当输入大于核大小时，ImplicitBatchedGEMM的性能会随着输入的增加而显着下降。需要一种新的优化方法来满足检测和分割等下游服务中大输入尺寸的需求，并且该方法的性能在小批量或大输入下应该足够好。由于largekerneldepthwiseconvolution的计算密度比较高，DirectConv简单实现第一个版本基本可以达到70%-80%的peakperformance。DriectConv的写法其实有很多种，这里只提供一种写法供参考。如下图所示，为了更好的利用CUDA的多级存储来最大化带宽利用率，DirectConv采用了多级分块策略。每个ThreadBlock负责计算一个块的输出，然后每个Warp进一步逐行阻塞ThreadBlockTile。为了适应更大的kernelsize，我们不仅在Thread层面将输出分块，还将kernel分块。举一个简单的例子来介绍线程级别的阻塞策略。假设ThreadBlock大小为128，Thread组织成32×4的形式，每行4个线程负责计算一行输出。内核也分为四列，每行四个线程负责读取内核的一列。如下图，Thread0读取内核的第0列和输入的第0-3列，计算出4个输出；线程1读取内核的第2列和输入的第1-4列，计算出4个outputs输出。线程2和线程3等等。由于内核是分块的，每行4个线程计算完成后，每个Thread都持有输出的部分求和，最终的结果需要归约为4个线程各自的结果。这里借助WarpShuffleAPI__shuffle_xor_sync实现了一个蝴蝶协议，其原理如下图所示。由于只需要将每4个线程的结果一起reduce，所以只需要执行两次__shuffle_xor_sync，最后把outupt写回去。实验数据表明，当输入大小为48时，DirectConv的性能略高于ImplicitBatchedGEMM，而当输入大小为64时，DirectConv的性能明显高于ImplicitBatchedGEMM。得益于MegEngine的算子自动选择机制，用户在使用时无需指定具体的实现方式，MegEngine会自动选择最佳实现方式。运行时间为了衡量算子的优劣，前面的实验都是从比较算子的绝对性能和硬件的理论峰值的角度来设计的。为了让用户有更直观的体验，我们还测试了大核depthwiseconvolution的运行时间。实验环境为2080Ti@cuda10.1+cudnn7.6.3，使用的数据类型为fp32，batchsize为64，channel为384，正反计算分别使用24层。从下图可以看出，MegEngine比PyTorch（带cudnn）快了10倍，优化后的MegEngine在31×31的内核大小上与PyTorch9×9的训练时间相同。如下图所示，只测试了一层的正向推理，其他配置与训练一致。优化后的MegEngine比cudnn快8倍，fp16也比fp32快2倍以上。欢迎尝试混合精度训练。代码已经用MegEnginev1.8.2开源，使用v1.9会更好（即将推出）~
