從卷積到矩陣乘

矩陣乘（GEMM）具有計算密度大和易并行等優(yōu)良特點，傳統(tǒng)領(lǐng)域例如芯片產(chǎn)業(yè)和高性能計算領(lǐng)域往往將 GEMM 作為 benchmark 并已經(jīng)將其優(yōu)化到接近硬件理論峰值。 為了獲得更好的性能收益，im2col 算法將 GEMM 帶進了卷積神經(jīng)網(wǎng)絡(luò)的工程優(yōu)化領(lǐng)域。Implicit GEMM 算法進一步解決了 im2col 固有的多余顯存占用和冗余的前后處理問題，這在存儲受限的硬件例如 GPU 上尤為重要，使得 GEMM 在卷積優(yōu)化中的重要性進一步提升。加之硬件廠商也開始越來越多的對矩陣乘提供硬件支持，例如各種 MMA 指令和 TensorCore 的加入，這些原因共同促使現(xiàn)有很多優(yōu)化算法庫已經(jīng)將 im2col/Implicit GEMM 作為其默認的卷積優(yōu)化方案。

im2col 算法

本文中假設(shè)卷積的輸入 shape 為 (n, ic, ih, iw)，kernel 為 (oc, ic, kh, kw)，output 為 (n, oc, oh, ow)。im2col 算法的過程如下圖所示，簡單的將 kernel reshape 一下就得到了一個行 M = oc，列 K = ic*kh*kw 的矩陣記作矩陣 A。用 kernel 大小的立方體在 input 上做滑窗，每次將一個小立方體的數(shù)據(jù)按照 chw 的順序展開成一列。從上到下，從左到右滑完整個 input 之后將會得到一個行 K = ic*kh*kw，列 N = n*oh*ow 的矩陣記作矩陣 B。此時我們計算 GEMM(A, B) 就可以得到卷積的結(jié)果矩陣 C，其行 M = oc，列 N = n*oh*ow。

在 GEMM 的計算過程中，根據(jù) m, n 和 k 三個維度下標可以推導(dǎo)出數(shù)據(jù)在卷積的輸入輸出以及 kernel 中的下標。所以可以將 im2col 和 GEMM 兩個過程融合在一起從而達到降低顯存占用和性能加速的效果，這其實就是 Implicit GEMM 的原理。本文不過多介紹，感興趣的可以閱讀之前的技術(shù)文章。

Implicit Batched GEMM

上一篇文章主要介紹了 MegEngine 大 kernel depthwise 卷積優(yōu)化的背景和動機，本篇文章將介紹具體的優(yōu)化思路和工程實踐。借助 im2col/Implicit GEMM 算法，GEMM 在傳統(tǒng)的針對 dense 卷積的優(yōu)化中表現(xiàn)出來了優(yōu)良的性能。所以針對大 kernel depthwise 卷積也應(yīng)該嘗試使用 GEMM 實現(xiàn)。如前文分析，直接使用和 dense 卷積一樣的方法將 im2col 算法應(yīng)用到大 kernel depthwise 卷積將會產(chǎn)生一個 Batched GEMV，這很難達到硬件浮點計算峰值。顯然為了將 GEMM 應(yīng)用在大 kernel depthwise 卷積上，我們應(yīng)該轉(zhuǎn)變一下思維。如下圖所示，dense 卷積一般用 kernel 去卷 input 并算得 output，但當 kernel size 很大甚至可能比 input 還大時，此時其實應(yīng)該用 input 去卷 kernel 并算得 output。

回想一下 dense 卷積用 kernel 去卷 input，對 input 做 im2col 變換。現(xiàn)在大 kernel depthwise 卷積用 input 去卷 kernel，所以此時應(yīng)該對 kernel 做 im2col 變換。算法過程沒有本質(zhì)區(qū)別，只需要在 im2col 中將 kernel 看做 input，將 input 看做 kernel 即可。由于 depthwise 卷積是逐通道做卷積的，因此 im2col 變換也需要逐通道做。如下圖所示，每個通道變換之后都會產(chǎn)生一個 M = n, N = oh*ow, K = ih*iw 的 GEMM。根據(jù)上一篇文章的分析，Batched GEMM 相比于 Batched GEMV 更容易打滿硬件設(shè)備的浮點計算峰值。

CUTLASS 是 NVIDIA 的開源模板庫，它旨在提供一種用較小的成本寫出一個性能不是那么差的 GEMM 的能力。CUTLASS 內(nèi)置了針對 GEMM 的 meta schedule，能夠讓計算盡量掩蓋訪存延遲從而達到不錯的性能。曠視早在 CUTLASS 官方開源其卷積實現(xiàn)之前就基于 CUTLASS 做出了自己的卷積實現(xiàn)，時至今日已經(jīng)打磨出了一個更適合內(nèi)部業(yè)務(wù)的曠視版 CUTLASS。此處的 Implicit Batched GEMM 也是基于曠視版 CUTLASS 實現(xiàn)的，代碼已經(jīng)隨 MegEngine v1.8.2 開源出來了，實現(xiàn)細節(jié)就不過多介紹了。如下圖的實驗數(shù)據(jù)顯示隨著 kernel size 的增加，Implicit Batched GEMM 的性能大致是呈線性增長的，部分情況下可以逼近理論峰值。

Implicit Batched GEMM 的優(yōu)點一方面是可以復(fù)用成熟的 GEMM 優(yōu)化思路和基礎(chǔ)設(shè)施，還可以方便使用 TensorCore 進行加速；另一方面如果在推理的時候不要求可變 shape 的話，對 kernel 的 im2col 變換可以提前算好進一步加速。當然它的缺點也很明顯，比如小 batch 情況下依然會退化成 Bacthed GEMV。如果用 M*N*K*2 來近似 GEMM 的計算量的話，不難發(fā)現(xiàn) Implicit Batched GEMM 的計算量相比 dense 卷積轉(zhuǎn)成的 GEMM 增大了 $\frac{ih*iw}{kh*kw}$ 倍，這意味著 Implicit Batched GEMM 在 input 顯著大于 kernel size 時性能不佳。如下圖所示的實驗結(jié)果也顯示著當 input 大于 kernel size 時，隨著 input 的增加 Implicit Batched GEMM 的性能有明顯下滑。需要一種新的優(yōu)化方法來迎合下游如檢測、分割等業(yè)務(wù)里的大 input size 的需求，這種方法在小 batch size 或者大輸入下的性能表現(xiàn)也要足夠好。

Direct Conv

由于大 kernel depthwise 卷積的計算密度比較高，所以其實簡單實現(xiàn)一版性能基本都能達到峰值性能的 70%-80%。Driect Conv 的寫法其實有很多，這里只提供一種寫法思路供參考。如下圖所示，為了更好的利用 CUDA 的多級存儲以最大利用帶寬，Direct Conv 采用多級分塊策略。每個 Thread Block 負責計算 output 的一個分塊，然后每個 Warp 對 Thread Block Tile 按行進一步做分塊。為了適應(yīng)更大的 kernel size，我們在 Thread level 上不僅針對 output 做了分塊，還對 kernel 做了分塊。

簡單舉個例子介紹 Thread level 的分塊策略。假設(shè) Thread Block size 是 128，Thread 被組織成 32×4 的形式，每一行的 4 個線程負責計算 output 的一行。將 kernel 也切分成四列，每一行的 4 個線程分別負責讀取 kernel 的一列。如下圖所示，Thread 0 讀取 kernel 的第 0 列和 input 的第 0-3 列，計算得到 4 個 output；Thread 1 讀取 kernel 的第 2 列和 input 的第 1-4 列，計算得到 4 個 output。Thread 2 和 Thread 3 以此類推。

由于對 kernel 做了分塊，所以每行的 4 個線程計算完畢之后每個 Thread 持有的是 output 的部分和，需要將 4 個線程各自的結(jié)果規(guī)約到一起才是最終結(jié)果。此處借助了 Warp Shuffle API__shuffle_xor_sync，它實現(xiàn)了一種蝶形規(guī)約，其原理如下圖所示。由于只需要將每 4 個線程的結(jié)果規(guī)約到一起就行了，所以只需要進行 2 次 __shuffle_xor_sync 即可，最后將 outupt 寫回。

實驗數(shù)據(jù)顯示在 input 大小為 48 時 Direct Conv 的性能已經(jīng)略高于 Implicit Batched GEMM 了，intput 為 64 時 Direct Conv 的性能會顯著高于 Implicit Batched GEMM。得益于 MegEngine 的算子自動選擇機制，用戶使用的時候不用指定具體的實現(xiàn)方式，MegEngine 會自動選擇最佳實現(xiàn)。

運行時間

為了衡量算子的優(yōu)劣，前面的實驗都是從算子絕對性能和硬件理論峰值相比的角度設(shè)計的。為了讓用戶有更直觀的感受，我們同樣測試了大 kernel depthwise 卷積的運行時間。實驗環(huán)境為 2080Ti @ cuda10.1 + cudnn7.6.3，所用的數(shù)據(jù)類型為 fp32， batch size 為 64，channel 為 384，用 24 個 layer 進行前向和反向計算。從下圖可見 MegEngine 比 PyTorch(with cudnn) 最高快 10 倍以上，優(yōu)化后的 MegEngine 在 31×31 的 kernel size 上和 PyTorch 9×9 的訓(xùn)練時間相當。

如下圖所示，只測試一個 layer 的前向推理，其他的配置和訓(xùn)練保持一致。經(jīng)過優(yōu)化后的 MegEngine 比 cudnn 最高快 8 倍多，并且 fp16 相比 fp32 也有 2 倍多加速，歡迎嘗試一下混合精度訓(xùn)練。代碼已經(jīng)隨著 MegEngine v1.8.2 開源，使用 v1.9 （即將發(fā)布）效果更佳~

? ?