從卷積到矩陣乘

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

im2col 算法

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

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

Implicit Batched GEMM

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

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

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

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

Direct Conv

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

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

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

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

運(yùn)行時(shí)間

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

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

? ?