前言

在這樣的趨勢下，Nvidia開發(fā)了線性代數(shù)模板庫CUTLASS，抽象了一系列高性能的基本組件，可以用于生成各種數(shù)據(jù)類型，各種計(jì)算單元的卷積、矩陣乘算子。MegEngine在CUTLASS的基礎(chǔ)上進(jìn)行了二次開發(fā)，可以高效地開發(fā)新的高性能的算子，快速地遷移到新的GPU架構(gòu)。

在上一篇文章中，我們已經(jīng)簡單介紹了MegEngine的底層卷積算子實(shí)現(xiàn)的使用方法，而本文將會(huì)深入介紹MegEngine CUDA平臺的底層卷積算子的實(shí)現(xiàn)原理，并將會(huì)對Nvidia CUTLASS的Implicit GEMM卷積文檔進(jìn)行解讀和補(bǔ)充。

因此，讀者在閱讀本文之前必須要了解的CUDA知識有：

?訪問全局存儲(chǔ)（Global Memory）時(shí)，同一Warp中的相鄰線程訪問連續(xù)的地址，訪存請求會(huì)被合并，合并的訪存能夠最大化Global Memory的吞吐。

?訪問Global Memory時(shí)，盡可能使用最寬的數(shù)據(jù)類型（float4）進(jìn)行訪問，這樣可以最大化訪存指令的利用率。

?CUDA的共享存儲(chǔ)（Shared Memory）按照每 4Bytes劃分為一個(gè)bank，共分為32個(gè)bank。當(dāng)同一 Warp中的線程訪問同一bank的不同地址時(shí)會(huì)發(fā)生沖突（bank conflict）。無bank conflict的訪存模式才能最大化 Shared Memory 的吞吐。

?GPU有顯存（Global Memory）、L2、L1（Shared Memory）、寄存器 4 個(gè)層次的存儲(chǔ)，直接訪問顯存的延遲很高，在優(yōu)化GEMM、Convolution這樣的計(jì)算密集型算子時(shí)，需要

–通過 L1 和寄存器的緩存來減少Global Memory的訪存請求。

–通過大量的計(jì)算來隱藏不可避免的Global Memory訪存延遲。

首先，我們需要了解CUTLASS引入的一些抽象概念

?TileIterator:用于訪問存儲(chǔ)中的一個(gè)Tile的數(shù)據(jù)。TileIterator實(shí)現(xiàn)了advance()方法，支持在Matrix、Tensor等數(shù)據(jù)類型上進(jìn)行遍歷。

?Fragment:數(shù)組類型，用于存放TileIterator讀取進(jìn)來的數(shù)據(jù)。Fragment的數(shù)據(jù)通常存放在寄存器中。

然后我們簡單回顧一下CUTLASS設(shè)計(jì)的高性能的GEMM算子的Pipeline，按照Pipeline實(shí)現(xiàn)的算子能夠在CUDA平臺上達(dá)到cublas的90%以上的性能。下圖演示了CUTLASS設(shè)計(jì)的Pipeline化的GEMM算子：

1.圖中第一行演示了由PredicatedTileIterator和SmemTileIterator配合完成從Global Memory到Shared Memory的數(shù)據(jù)搬運(yùn)。

2.第二行演示了WarpTileIterator負(fù)責(zé)從Shared Memory搬運(yùn)數(shù)據(jù)到Fragment寄存器中。

3.第三行展示了WarpMmaOperator用Fragment寄存器中的矩陣數(shù)據(jù)執(zhí)行矩陣乘加 (Matrix-Multiply-Add) 操作。

Implicit GEMM 算法

卷積映射為矩陣乘法

我們首先來看一下前向卷積算子的定義，假設(shè)輸入的feature map是x，卷積層的weight是w，輸出是y，其中x，y，w都是4維的Tensor，x的四個(gè)維度分別是NxICxIHxIW，w的四個(gè)維度分別是OCxICxFHxFW，y的四個(gè)維度分別是NxOCxOHxOW。那么輸出y和輸入x，w的數(shù)學(xué)關(guān)系式可以寫成

公式里的小寫字母代表了Tensor在每一維的坐標(biāo)，其中ih，iw和oh，ow，fh，fw的關(guān)系式可以寫為

這里的stride_h, stride_w, pad_h, pad_w是卷積層的參數(shù)。根據(jù)im2col算法的原理，公式里定義的卷積運(yùn)算可以轉(zhuǎn)化為一個(gè)矩陣乘法，也即

其中

?矩陣A由weight轉(zhuǎn)化而來，是一個(gè)OC X IC·FH·FW的矩陣。

?矩陣B由feature map轉(zhuǎn)化而來，是一個(gè)IC·FH·FW X N·OH·OW的矩陣

?矩陣C代表了輸出的Tensor y，是一個(gè)OC X N·OH·OW的矩陣。

矩陣和Tensor在各個(gè)位置上的元素的對應(yīng)關(guān)系為

其中矩陣的下標(biāo)i，j，k和Tensor的坐標(biāo)之間的關(guān)系為

當(dāng)j已知時(shí)，可以用下面的關(guān)系式推算出feature map的坐標(biāo)

當(dāng)k已知時(shí)，可以推算出weight的坐標(biāo)

同時(shí)結(jié)合oh，ow，fh，fw,就可以計(jì)算出ih和iw。

根據(jù)上面的討論，我們可以把卷積的運(yùn)算過程，寫成一個(gè)隱式矩陣乘法(Implicit GEMM)的形式：

上面的Implicit GEMM算法仍然是串行的形式，接下來我們要把它改造成CUDA上的并行算法。首先我們對整個(gè)計(jì)算任務(wù)進(jìn)行分塊，讓每個(gè)線程塊負(fù)責(zé)計(jì)算并輸出大小為TILE_MxTILE_N的矩陣。于是算法變成了下面的形式：

為了提高訪存的效率，我們可以在GEMM_K這一維上也進(jìn)行分塊，每次將TILE_MxTILE_K的矩陣A和TILE_KxTILE_N的矩陣B緩存到Shared Memory里，避免重復(fù)的Global Memory訪存。于是，算法就變成了如下形式：

因?yàn)槲覀兛梢灾苯訌?fù)用CUTLASS里已經(jīng)實(shí)現(xiàn)好了高性能的WarpMmaOperator，所以實(shí)現(xiàn)基于Implicit GEMM的卷積算子只需要

? 適配DeviceConvolution、KernelConvolution

和ThreadblockConvolution，支持傳入Tensor類型和Convolution Layer的參數(shù)。

? 添加PredicateTileIterator支持讀取Tensor的一個(gè) Tile 的數(shù)據(jù)到Shared Memory中，并隱式地將讀入的數(shù)據(jù)組織成矩陣的形式。

? 算法的main loop中直接調(diào)用WarpTileIterator從Shared Memory讀取數(shù)據(jù)，然后由WarpGemmOperator完成Warp-level的GEMM運(yùn)算。

? EpilogueOperator適配卷積算子,將Accumulator的數(shù)據(jù)寫回Global Memory的Tensor中。

接下來我們會(huì)以INT8數(shù)據(jù)類型的TensorCore卷積算子來介紹MegEngine底層的卷積實(shí)現(xiàn)，本文會(huì)重點(diǎn)介紹 2、3、4 是如何實(shí)現(xiàn)的，關(guān)于如何使用已經(jīng)寫好的卷積算子，可以參考之前的文章。

Global Memory 數(shù)據(jù)布局（Layout）

為了最大化TensorCore類型的卷積算子的吞吐，MegEngine使用了128位的Global Memory訪存指令，因此在訪問Tensor的數(shù)據(jù)的時(shí)候要求地址滿足128位對齊。MegEngine使用了NCHW32的格式來存儲(chǔ)Tensor,NCHW32格式的特點(diǎn)為：

?Tensor的通道維度按照32個(gè)channel進(jìn)行分組，每 32個(gè)channel連續(xù)的存放在存儲(chǔ)中。

?Tensor的其余維度按照W、H、C、N的順序地址變化由快到慢的存放在存儲(chǔ)中。

由于采用了32個(gè)通道對齊的存儲(chǔ)格式，因此卷積layer要求輸入和輸出feature map的通道數(shù)都是32的倍數(shù)。

預(yù)處理訪存偏移量

MegEngine的卷積實(shí)現(xiàn)在GEMM_K的維度上是按照(IC/32)·FH·FW·32的順序累加，寫成偽代碼的形式如下：

如果寫成一層循環(huán)，那么應(yīng)該寫成：

可以看到在迭代過程中，如果直接計(jì)算指針的偏移量的話，會(huì)引入很多除法和求余運(yùn)算。而在CUDA平臺上，整數(shù)的除法和求余的開銷是非常大的，因此我們將一些地址的偏移量在host端預(yù)先計(jì)算好，存到kernel param的buffer中，需要時(shí)從constant memory中直接讀取地址，避免除法和求余運(yùn)算。

對于每個(gè)線程來說，在主循環(huán)中指針移動(dòng)的offset如下圖所示：

如果地址的增量可以用delta來表示的話，那么delta是以FH*FW為周期的，即：

因此我們只需要大約O(FH·FW)的存儲(chǔ)空間。其中地址偏移量的計(jì)算邏輯可以參考代碼conv2d_tile_iterator_nt_src_fprop_precomp.h

由于kernel param buffer的大小為4KB，我們用了大約3KB來存儲(chǔ)地址的增量，所以MegEngine的卷積實(shí)現(xiàn)要求Convolution Layer的FH*FW的大小不能太大，但是一般情況下，3x3，5x5，7x7的卷積都可以處理。Nvidia官方實(shí)現(xiàn)的迭代順序與本文介紹的略有不同：

?官方實(shí)現(xiàn)需要將IC補(bǔ)齊為TILE_K的倍數(shù)，這樣在通道數(shù)較小時(shí)會(huì)浪費(fèi)一些計(jì)算量。

?官方實(shí)現(xiàn)的線程塊在訪問輸入feature map的時(shí)候地址的跨度比較大，降低了訪存的局部性，對cache不夠友好。

因此在性能方面，MegEngine的實(shí)現(xiàn)會(huì)更有優(yōu)勢，而官方實(shí)現(xiàn)的優(yōu)點(diǎn)是對Convolution Layer的參數(shù)沒有太多限制，通用性更好。

Warp-level Mma(Matrix-multiply-add)指令

cuda10.2引入了新的Warp-level的mma和ldmatrix指令，用戶可以通過mma指令使用 TensorCore 來進(jìn)行高速的矩陣乘加運(yùn)算，通過ldmatrix精細(xì)地控制Warp給TensorCore喂數(shù)據(jù)。其中mma指令的用法如下：

這條指令的語義是由一個(gè)Warp的32個(gè)線程同步地完成8x8x16的矩陣乘加運(yùn)算，它有三個(gè)輸入操作數(shù)，其中參與矩陣乘法運(yùn)算的分別是一個(gè)8x16的矩陣A 和一個(gè)16x8的矩陣B，這兩個(gè)輸入矩陣的數(shù)據(jù)分布在同一Warp的32個(gè)線程中。矩陣A的布局如下圖所示： 

? 同一Warp中的32個(gè)線程分為8組，每組四個(gè)線程，負(fù)責(zé)讀取8x16的矩陣中的一行。

? 每一組中的一個(gè)線程讀取每一行中相鄰的4個(gè)int8的數(shù)據(jù)，恰好填滿一個(gè)32位的寄存器。

類似的矩陣B的布局如下圖所示：

? 每4個(gè)線程一組，共分為8組，每組負(fù)責(zé)讀取16x8的矩陣中的一列。

? 每一組中的一個(gè)線程負(fù)責(zé)讀取一列中相鄰的4個(gè)數(shù)據(jù)。

參與累加運(yùn)算的矩陣C和輸出矩陣D的數(shù)據(jù)也同樣分布在32個(gè)線程中，它們的布局如下圖所示：

? 同樣每4個(gè)線程一組，每組負(fù)責(zé)讀入/輸出一行的數(shù)據(jù)。

? 每個(gè)線程負(fù)責(zé)輸出一行中的相鄰兩個(gè)int32類型的數(shù)據(jù)，恰好構(gòu)成一個(gè)64位的寄存器。

通過對mma指令的分析，如果Global Memory/Shared Memor中的數(shù)據(jù)是以行優(yōu)先 (RowMajor) 或者列優(yōu)先 (ColumnMajor) 的格式存儲(chǔ)的，那么當(dāng)同一Warp執(zhí)行空間上連續(xù)的兩個(gè)8x8x16的矩陣乘加運(yùn)算時(shí)，每個(gè)線程讀取的數(shù)據(jù)將會(huì)是跳躍的，執(zhí)行每次乘法都只能讀取32位寬的數(shù)據(jù)到寄存器中，而低位寬的Load指令通常沒有辦法最大化利用存儲(chǔ)的帶寬。因此Nvidia提供了ldmatrix的指令，可以讓同一Warp一次性讀取4個(gè)8x16的矩陣到寄存器中，這樣恰好可以讓W(xué)arp中的每個(gè)線程一次讀取128位的數(shù)據(jù)，最大化帶寬的利用率。 

ldmarix的用法如下所示：

上述這條指令恰好讀取了4個(gè)8x16的矩陣，每個(gè)線程恰好負(fù)責(zé)讀取矩陣的一行數(shù)據(jù)，讀取完成后，線程之間會(huì)進(jìn)行數(shù)據(jù)交換，將矩陣的數(shù)據(jù)重新分布到各個(gè)線程，讀取的過程如下圖所示：

這一節(jié)介紹了TensorCore相關(guān)的mma和ldmatrix指令，有了這兩條高性能的指令，我們還需要為數(shù)據(jù)設(shè)計(jì)巧妙的Shared Memory存儲(chǔ)格式，消除從Shared Memory讀取數(shù)據(jù)的bank conflict，從而提升Shared Memory的讀取效率。

Shared Memory的數(shù)據(jù)布局 

在介紹Shared Memory中的數(shù)據(jù)布局之前，我們需要了解Shared Memory的訪存特點(diǎn)。Shared Memory按照每4個(gè)字節(jié)組成一個(gè)bank，共劃分成了32個(gè)bank，同一Warp的線程訪問了相同bank的不同地址時(shí)會(huì)發(fā)生conflict，導(dǎo)致訪存的效率變慢。在同一Warp的線程訪問不同位寬的數(shù)據(jù)時(shí)，會(huì)有不同的行為：

? 每個(gè)線程訪問Shared Memory中32位的數(shù)據(jù)，訪存將在一個(gè)階段內(nèi)完成。

? 每個(gè)線程訪問Shared Memory中64位的數(shù)據(jù)，訪存會(huì)在兩個(gè)階段內(nèi)完成：

– 第一個(gè)階段：前16個(gè)線程訪存128字節(jié)的數(shù)據(jù)。

– 第二個(gè)階段：后16個(gè)線程訪存128字節(jié)的數(shù)據(jù)。

? 每個(gè)線程訪問Shared Memory中的128位的數(shù)據(jù)，訪存會(huì)在四個(gè)階段內(nèi)完成：

– 每個(gè)階段由8個(gè)線程完成128字節(jié)的數(shù)據(jù)的訪存。

如果上述過程中每個(gè)階段都沒有bank conflict，則能夠達(dá)到最大的Shared Memory訪存效率。 

通常為了避免Shared Memory的bank conflict，我們會(huì)對Shared Memory的數(shù)據(jù)進(jìn)行padding，讓線程訪問的數(shù)據(jù)錯(cuò)開，避免落在同一bank中。但是這樣做的問題是會(huì)使得kernel需要Shared Memory的Size變大，但是SM上的L1 cache(Shared Memory)又是有限的，所以padding會(huì)降低kernel的occupancy，進(jìn)而就會(huì)降低kernel的性能。 

因此CUTLASS設(shè)計(jì)了一種Shared Memory的交錯(cuò)布局方式，它能夠在不進(jìn)行padding的前提下，使得線程訪存的地址沒有bank conflict。接下來，我們以64x64的矩陣為例來詳細(xì)介紹數(shù)據(jù)在Shared Memory中的布局。

首先，線程讀取數(shù)據(jù)的粒度都是128位，也即16個(gè)INT8類型的數(shù)據(jù)，因此我們在演示數(shù)據(jù)的布局時(shí)總是以16個(gè)數(shù)據(jù)為一組。如果矩陣是以行優(yōu)先(RowMajor)的格式來組織的，那么在邏輯上的布局如下圖所示：

從圖中可以看到

? 每16個(gè)元素分為一組，被稱為一個(gè)Vector，被染上了不同的顏色。

? 每行相鄰的32個(gè)元素被稱為一個(gè)Crosswise，恰好是NCHW32格式中的一組channel的數(shù)據(jù)。

在Shared Memory的物理存儲(chǔ)中，矩陣的數(shù)據(jù)進(jìn)行了重新排列，如下圖所示：

我們可以看到Shared Memory的物理布局有以下特點(diǎn)：

? 每4行的一個(gè)Crosswise的數(shù)據(jù)作為一組，連續(xù)存放在Shared Memory中，緊接著會(huì)存放這4行的下一個(gè)Crosswise的數(shù)據(jù)。

? 每組數(shù)據(jù)包含了8個(gè)Vector，占據(jù)了128個(gè)字節(jié)，恰好是Shared Memory中的32個(gè)不同的bank。

?每組數(shù)據(jù)在排列是進(jìn)行了交錯(cuò)，保證了ldmatrix時(shí)不會(huì)發(fā)生bank conflict。

顯存 -> Shared Memory 的數(shù)據(jù)搬運(yùn)

這一節(jié)我們會(huì)介紹從顯存(Global Memory)到Shared Memory的數(shù)據(jù)搬運(yùn)。顯存到Shared Memory的數(shù)據(jù)搬運(yùn)是由Conv2dTileSrcIteratorFpropPrecomp來完成的，本文并不會(huì)詳細(xì)地解讀代碼的實(shí)現(xiàn)，而是描述線程搬運(yùn)數(shù)據(jù)的過程，幫助大家建立直觀的印象，更好地理解代碼。

如果以上一節(jié)中Shared Memory的邏輯布局為例，同一Warp中每個(gè)線程讀取的數(shù)據(jù)的邏輯布局如下圖所示，每個(gè)線程讀取16個(gè)INT8類型的數(shù)據(jù)，恰好構(gòu)成一個(gè)Vector。

而在實(shí)際的物理顯存中，線程訪問的數(shù)據(jù)分布如下圖所示：

? 我們可以看到每個(gè)線程讀取了128位的數(shù)據(jù)。

? 相鄰的線程讀取的數(shù)據(jù)在物理上是連續(xù)的。

因此線程從Global Memory讀取數(shù)據(jù)的pattern可以滿足合并訪存的要求，同時(shí)以最大的數(shù)據(jù)位寬進(jìn)行訪存，最大化了顯存帶寬的利用率。 

然后如果將線程讀取的數(shù)據(jù)映射到Shared Memory的物理地址，我們可以看到 

? 每8個(gè)線程向Shared Memory寫入128字節(jié)的數(shù)據(jù)，恰好落在Shared Memory的32個(gè)不同的bank中。

? 同一Warp的訪存分為四個(gè)階段完成，每個(gè)階段都沒有bank conflict。

下圖演示了一個(gè)Warp寫入Shared Memory的過程：

Shared Memory -> 寄存器的數(shù)據(jù)搬運(yùn)

Shared Memory到寄存器的數(shù)據(jù)搬運(yùn)是由MmaTensorOpMultiplicandTileIterator完成的。同一 Warp 在每一輪迭代過程會(huì)讀取4個(gè)8x16的矩陣到寄存器中，每個(gè)線程會(huì)讀取一行的數(shù)據(jù)。例如第一輪迭代時(shí)，線程讀取的數(shù)據(jù)在邏輯上的布局如下圖所示： 

而實(shí)際上數(shù)據(jù)在Shared Memory里的物理布局如下圖：

可以看到：

? 每個(gè)線程讀取了128位的數(shù)據(jù)，因此訪存分為四個(gè)階段來進(jìn)行。

? 每一階段的8個(gè)線程讀取的數(shù)據(jù)恰好落在了Shared Memory的32個(gè)bank中，并且線程訪存的數(shù)據(jù)之間不存在沖突。

當(dāng)進(jìn)行到第二輪迭代時(shí)，每個(gè)線程訪問的數(shù)據(jù)的物理布局如下圖：

同樣的訪存的每一個(gè)階段都不存在bank conflict。

Accumulator 寫回全局存儲(chǔ)

在int8的情況下，同一Warp負(fù)責(zé)輸出64x64的結(jié)果，kernel會(huì)分成8次寫回Global Memory，每次寫回32x8的矩陣。這樣保證了每次將Tensor按照 NCHW32格式寫回顯存時(shí)，同一Warp的32個(gè)線程恰好寫了物理上連續(xù)的256字節(jié)的數(shù)據(jù)，而每個(gè)線程寫回8個(gè)字節(jié)，保證了可以使用64位寬的數(shù)據(jù)類型進(jìn)行顯存的寫操作，盡可能提高帶寬的利用率。 

由于mma指令的特點(diǎn)，輸出矩陣的數(shù)據(jù)分布在各個(gè)線程上，而為了能夠合并訪存，即：讓相鄰線程寫回的地址是連續(xù)的，我們利用Shared Memory對同一Warp中32個(gè)線程的數(shù)據(jù)進(jìn)行了交換。數(shù)據(jù)交換后，每個(gè)線程擁有連續(xù)的8個(gè)通道的數(shù)據(jù)，且線程寫的地址是連續(xù)的，保證了寫回Global Memory滿足合并訪存的要求。線程交換數(shù)據(jù)的過程如下圖所示： 

每一輪迭代，Warp中的32個(gè)線程將32x16的矩陣數(shù)據(jù)寫入到Shared Memory中。接著如下圖所示，每個(gè)線程會(huì)把連續(xù)的8個(gè)channel的數(shù)據(jù)讀到寄存器中。

Shared Memory的數(shù)據(jù)交換是由以下兩個(gè)Iterator完成的

?InterleavedTileIteratorTensorOp完成了每一輪迭代將32x8的數(shù)據(jù)寫入到Shared Memory中。

?InterleavedSharedLoadIteratorTensorOp負(fù)責(zé)將連續(xù)的8個(gè)channel的數(shù)據(jù)讀到Fragment寄存器中。

當(dāng)線程將交換后的數(shù)據(jù)讀到Fragment寄存器之后，會(huì)由EpilogueOp，在卷積的基礎(chǔ)上完成BiasAdd的運(yùn)算。

以BiasAddLinearCombinationRelu為例，它實(shí)際上完成了下面的運(yùn)算：

其中bias是一個(gè)PerChannel的Tensor，代表了每個(gè)輸出通道的偏置，z是一個(gè)和卷積輸出大小一致的Tensor，用于Convolution和ElemwiseAdd的融合。

最后EpilogueOp的輸出會(huì)由TensorPredicatedTileIteratorTensorOp真正地寫回到 Global Memory中。每個(gè)線程寫回的數(shù)據(jù)如下圖所示：

可以看到線程寫回的pattern滿足合并訪存的要求，因此能最大化Global Memory寫的效率。

總結(jié)

本文介紹了MegEngine底層的卷積算子實(shí)現(xiàn)原理，算子性能可以達(dá)到cudnn的80%以上，測速結(jié)果可以參見文章

MegEngine會(huì)對卷積實(shí)現(xiàn)進(jìn)行持續(xù)優(yōu)化，進(jìn)一步提升算子的性能，目前來看有以下兩點(diǎn)可做的優(yōu)化：

? 借鑒Nvidia官方CUTLASS ImplicitGEMM Convolution實(shí)現(xiàn)對mask的處理，提高TileIterator對于mask判斷的效率。

? 現(xiàn)在的卷積實(shí)現(xiàn)在寫回顯存時(shí)利用Shared Memory進(jìn)行數(shù)據(jù)交換是存在bank conflict的。后續(xù)會(huì)考慮兩點(diǎn)優(yōu)化：

–對Shared Memory的數(shù)據(jù)布局進(jìn)行探索，消除 bank conflict，優(yōu)化Shared Memory數(shù)據(jù)交換的效率。

–對Global Memory中的Weight Tensor的布局進(jìn)行探索，提高每個(gè)Thread上accumulator的局部性，避免在Shared Memory中進(jìn)行數(shù)據(jù)交換。

參考資料

? Warp-level Matrix Fragment Mma PTX文檔

? CUTLASS Implicit GEMM Convolution官方文檔

? Volta architecture and performance optimization

? Developing CUDA kernels to push Tensor Cores to the absolute limit on Nvidia A100