實錄 | MegEngine 大 Kernel 卷積工程優化實踐
本篇整理自 3.19 日 “Large Kernel Makes CNN Great Again” 專題 Meetup 中《MegEngine 大 Kernel 卷積工程優化實踐》分享。作者：王彪 | 曠視 MegEngine 異構計算組負責人。
視訊回顧 01:10:55 直達
從卷積到矩陣乘
矩陣乘（GEMM）具有計算密度大和易並行等優良特點，傳統領域例如晶片產業和高效能運算領域往往將 GEMM 作為 benchmark 並已經將其優化到接近硬體理論峰值。 為了獲得更好的效能收益，im2col 演算法將 GEMM 帶進了卷積神經網路的工程優化領域。Implicit GEMM 演算法進一步解決了 im2col 固有的多餘視訊記憶體佔用和冗餘的前後處理問題，這在儲存受限的硬體例如 GPU 上尤為重要，使得 GEMM 在卷積優化中的重要性進一步提升。加之硬體廠商也開始越來越多的對矩陣乘提供硬體支援，例如各種 MMA 指令和 TensorCore 的加入，這些原因共同促使現有很多優化演算法庫已經將 im2col/Implicit GEMM 作為其預設的卷積優化方案。
im2col 演算法
本文中假設卷積的輸入 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 上做滑窗，每次將一個小立方體的資料按照 chw 的順序展開成一列。從上到下，從左到右滑完整個 input 之後將會得到一個行 K = ic*kh*kw，列 N = n*oh*ow 的矩陣記作矩陣 B。此時我們計算 GEMM(A, B) 就可以得到卷積的結果矩陣 C，其行 M = oc，列 N = n*oh*ow。
在 GEMM 的計算過程中，根據 m, n 和 k 三個維度下標可以推匯出資料在卷積的輸入輸出以及 kernel 中的下標。所以可以將 im2col 和 GEMM 兩個過程融合在一起從而達到降低視訊記憶體佔用和效能加速的效果，這其實就是 Implicit GEMM 的原理。本文不過多介紹，感興趣的可以閱讀之前的技術文章。
Implicit Batched GEMM
上一篇文章主要介紹了 MegEngine 大 kernel depthwise 卷積優化的背景和動機，本篇文章將介紹具體的優化思路和工程實踐。藉助 im2col/Implicit GEMM 演算法，GEMM 在傳統的針對 dense 卷積的優化中表現出來了優良的效能。所以針對大 kernel depthwise 卷積也應該嘗試使用 GEMM 實現。如
回想一下 dense 卷積用 kernel 去卷 input，對 input 做 im2col 變換。現在大 kernel depthwise 卷積用 input 去卷 kernel，所以此時應該對 kernel 做 im2col 變換。演算法過程沒有本質區別，只需要在 im2col 中將 kernel 看做 input，將 input 看做 kernel 即可。由於 depthwise 卷積是逐通道做卷積的，因此 im2col 變換也需要逐通道做。如下圖所示，每個通道變換之後都會產生一個 M = n, N = oh*ow, K = ih*iw 的 GEMM。根據上一篇文章的分析，Batched GEMM 相比於 Batched GEMV 更容易打滿硬體裝置的浮點計算峰值。
CUTLASS 是 NVIDIA 的開源模板庫，它旨在提供一種用較小的成本寫出一個性能不是那麼差的 GEMM 的能力。CUTLASS 內建了針對 GEMM 的 meta schedule，能夠讓計算儘量掩蓋訪存延遲從而達到不錯的效能。曠視早在 CUTLASS 官方開源其卷積實現之前就基於 CUTLASS 做出了自己的卷積實現，時至今日已經打磨出了一個更適合內部業務的曠視版 CUTLASS。此處的 Implicit Batched GEMM 也是基於曠視版 CUTLASS 實現的，程式碼已經隨 MegEngine v1.8.2 開源出來了，實現細節就不過多介紹了。如下圖的實驗資料顯示隨著 kernel size 的增加，Implicit Batched GEMM 的效能大致是呈線性增長的，部分情況下可以逼近理論峰值。
Implicit Batched GEMM 的優點一方面是可以複用成熟的 GEMM 優化思路和基礎設施，還可以方便使用 TensorCore 進行加速；另一方面如果在推理的時候不要求可變 shape 的話，對 kernel 的 im2col 變換可以提前算好進一步加速。當然它的缺點也很明顯，比如小 batch 情況下依然會退化成 Bacthed GEMV。如果用 M*N*K*2 來近似 GEMM 的計算量的話，不難發現 Implicit Batched GEMM 的計算量相比 dense 卷積轉成的 GEMM 增大了 \(\frac{ih*iw}{kh*kw}\) 倍，這意味著 Implicit Batched GEMM 在 input 顯著大於 kernel size 時效能不佳。如下圖所示的實驗結果也顯示著當 input 大於 kernel size 時，隨著 input 的增加 Implicit Batched GEMM 的效能有明顯下滑。需要一種新的優化方法來迎合下游如檢測、分割等業務裡的大 input size 的需求，這種方法在小 batch size 或者大輸入下的效能表現也要足夠好。
Direct Conv
由於大 kernel depthwise 卷積的計算密度比較高，所以其實簡單實現一版效能基本都能達到峰值效能的 70%-80%。Driect Conv 的寫法其實有很多，這裡只提供一種寫法思路供參考。如下圖所示，為了更好的利用 CUDA 的多級儲存以最大利用頻寬，Direct Conv 採用多級分塊策略。每個 Thread Block 負責計算 output 的一個分塊，然後每個 Warp 對 Thread Block Tile 按行進一步做分塊。為了適應更大的 kernel size，我們在 Thread level 上不僅針對 output 做了分塊，還對 kernel 做了分塊。
簡單舉個例子介紹 Thread level 的分塊策略。假設 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 個執行緒各自的結果規約到一起才是最終結果。此處藉助了 Warp Shuffle API__shuffle_xor_sync，它實現了一種蝶形規約，其原理如下圖所示。由於只需要將每 4 個執行緒的結果規約到一起就行了，所以只需要進行 2 次 __shuffle_xor_sync 即可，最後將 outupt 寫回。
實驗資料顯示在 input 大小為 48 時 Direct Conv 的效能已經略高於 Implicit Batched GEMM 了，intput 為 64 時 Direct Conv 的效能會顯著高於 Implicit Batched GEMM。得益於 MegEngine 的運算元自動選擇機制，使用者使用的時候不用指定具體的實現方式，MegEngine 會自動選擇最佳實現。
執行時間
為了衡量運算元的優劣，前面的實驗都是從運算元絕對效能和硬體理論峰值相比的角度設計的。為了讓使用者有更直觀的感受，我們同樣測試了大 kernel depthwise 卷積的執行時間。實驗環境為 2080Ti @ cuda10.1 + cudnn7.6.3，所用的資料型別為 fp32， batch size 為 64，channel 為 384，用 24 個 layer 進行前向和反向計算。從下圖可見 MegEngine 比 PyTorch(with cudnn) 最高快 10 倍以上，優化後的 MegEngine 在 31×31 的 kernel size 上和 PyTorch 9×9 的訓練時間相當。
如下圖所示，只測試一個 layer 的前向推理，其他的配置和訓練保持一致。經過優化後的 MegEngine 比 cudnn 最高快 8 倍多，並且 fp16 相比 fp32 也有 2 倍多加速，歡迎嘗試一下混合精度訓練。程式碼已經隨著 MegEngine v1.8.2 開源，使用 v1.9 （即將釋出）效果更佳~
附
GitHub：MegEngine 天元 （歡迎 star~
歡迎加入 MegEngine 技術交流 QQ 群：1029741705