值得收藏 | 深度剖析 TensorCore 卷積運算元實現原理

作者：章曉 | 曠視 MegEngine 架構師

一、前言

2020 年 5 月 Nvidia 釋出了新一代的 GPU 架構安培（Ampere）。其中和深度學習關係最密切的莫過於效能強勁的第三代的 TensorCore ，新一代的 TensorCore 支援了更為豐富的 DL（Deep Learning）資料型別，包括了新的 TesorFloat-32（TF32），Bfloat16（BF16）計算單元以及 INT8, INT4 和 INT1 的計算單元，這些計算單元為 DL 推理提供了全面的支援。為了發揮這些計算單元的能力，以往會由資深的 HPC 工程師手寫 GPU 彙編實現的卷積、矩陣乘運算元來挖掘硬體的能力。然而憑藉人力手工優化運算元的方式已經沒有辦法應對如此多的資料型別，因此對於 DL 應用的優化漸漸地越來越依賴一些自動化的工具，例如面向深度學習領域的編譯器。在這樣的趨勢下， Nvidia 開發了線性代數模板庫 CUTLASS ，抽象了一系列高效能的基本元件，可以用於生成各種資料型別，各種計算單元的卷積、矩陣乘運算元。 MegEngine 在 CUTLASS 的基礎上進行了二次開發，可以高效地開發新的高效能的運算元，快速地遷移到新的 GPU 架構。在上一篇 文章 中，我們已經簡單介紹了 MegEngine 的底層卷積運算元實現的使用方法，而本文將會深入介紹 MegEngine CUDA 平臺的底層卷積運算元的實現原理，並將會對 Nvidia CUTLASS 的 Implicit GEMM 卷積 文件 進行解讀和補充。

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

• 訪問全域性儲存（Global Memory）時，同一 Warp 中的相鄰執行緒訪問連續的地址，訪存請求會被合併，合併的訪存能夠最大化 Global Memory 的吞吐。
• 訪問 Global Memory 時，儘可能使用最寬的資料型別（float4）進行訪問，這樣可以最大化訪存指令的利用率。
• CUDA 的共享儲存（Shared Memory）按照每 4Bytes 劃分為一個 bank，共分為 32 個 bank。當同一 Warp 中的執行緒訪問同一 bank 的不同地址時會發生衝突（bank conflict）。無 bank conflict 的訪存模式才能最大化 Shared Memory 的吞吐。
• GPU 有視訊記憶體（Global Memory）、L2、L1（Shared Memory）、暫存器 4 個層次的儲存，直接訪問視訊記憶體的延遲很高，在優化 GEMM、Convolution 這樣的計算密集型運算元時，需要
  • 通過 L1 和暫存器的快取來減少 Global Memory 的訪存請求。
  • 通過大量的計算來隱藏不可避免的 Global Memory 訪存延遲。

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

• TileIterator : 用於訪問儲存中的一個Tile的資料。TileIterator 實現了advance()方法，支援在 Matrix , Tensor 等資料型別上進行遍歷。
• Fragment : 陣列型別，用於存放 TileIterator 讀取進來的資料。 Fragment 的資料通常存放在暫存器中。

然後我們簡單回顧一下 CUTLASS 設計的高效能的 GEMM 運算元的 Pipeline，按照 Pipeline 實現的運算元能夠在 CUDA 平臺上達到 cublas 的 90% 以上的效能。下圖演示了 CUTLASS 設計的 Pipeline 化的 GEMM 運算元：

1. 圖中第一行演示了由 PredicatedTileIterator 和 SmemTileIterator 配合完成從 Global Memory 到 Shared Memory 的資料搬運。
2. 第二行演示了 WarpTileIterator 負責從 Shared Memory 搬運資料到 Fragment 暫存器中。
3. 第三行展示了WarpMmaOperator 用 Fragment 暫存器中的矩陣資料執行矩陣乘加 (Matrix-Multiply-Add) 操作。

二、Implicit GEMM 演算法

卷積對映為矩陣乘法

我們首先來看一下前向卷積運算元的定義，假設輸入的 feature map 是 x，卷積層的 weight 是 w，輸出是 y，其中 x,y,w 都是 4 維的 Tensor，x 的四個維度分別是 NxICxIHxIW，w 的四個維度分別是 OCxICxFHxFW，y 的四個維度分別是 NxOCxOHxOW。那麼輸出 y 和輸入 x, w 的數學關係式可以寫成

\[\text{y}( \text{n}, \text{oc}, \text{oh}, \text{ow} ) = \sum_{\text{ic}} \sum_{\text{fh}} \sum_{\text{fw}} \text{x} (\text{n}, \text{ic}, \text{ih}, \text{iw}) \cdot \text{w} ( \text{oc}, \text{ic}, \text{fh}, \text{fw} ) \]

公式裡的小寫字母代表了 Tensor 在每一維的座標，其中 ih，iw 和 oh，ow，fh，fw 的關係式可以寫為

ih = oh * stride_h - pad_h + fh
iw = ow * stride_w - pad_w + fw

這裡的stride_h, stride_w, pad_h, pad_w是卷積層的引數。
根據 im2col 演算法的原理，公式裡定義的卷積運算可以轉化為一個矩陣乘法，也即

C = Matmul(A, B)

其中

• 矩陣 A 由 weight 轉化而來，是一個\(\text{OC}\times\text{IC}\cdot\text{FH}\cdot\text{FW}\)的矩陣。
• 矩陣 B 由 feature map 轉化而來，是一個\(\text{IC}\cdot\text{FH}\cdot\text{FW}\times\text{N}\cdot\text{OH}\cdot\text{OW}\)的矩陣
• 矩陣 C 代表了輸出的 Tensor y，是一個\(\text{OC}\times\text{N}\cdot\text{OH}\cdot\text{OW}\)的矩陣。

矩陣和 Tensor 在各個位置上的元素的對應關係為

\[\begin{equation} \begin{aligned} A_{ik} &= \text{w}\left(\text{oc}, \text{ic}, \text{fh}, \text{fw}\right) \\ \end{aligned} \end{equation} \]

\[\begin{equation} \begin{aligned} B_{kj} &= \text{x}\left(\text{n}, \text{ic}, \text{ih}, \text{iw}\right) \\ \end{aligned} \end{equation} \]

\[\begin{equation} \begin{aligned} C_{ij} &= \text{y}\left(\text{n}, \text{oc}, \text{oh}, \text{ow}\right) \end{aligned} \end{equation} \]

其中矩陣的下標\(i, j, k\)和 Tensor 的座標之間的關係為

i = oc
j = n * OH * OW + oh * OW + ow
k = ic * FH * FW + fh * FW + fw

當 \(j\) 已知時，可以用下面的關係式推算出 feature map 的座標

n = j / (OH * OW)
j_res = j % (OH * OW)
oh = j_res / OW
ow = j_res % OW

當 \(k\) 已知時，可以推算出 weight 的座標

ic = k / (FH * FW)
k_res = k % (FH * FW)
fh = k_res / FW
fw = k_res % FW

同時結合 oh, ow, fh, fw，就可以計算出 ih 和 iw。
根據上面的討論，我們可以把卷積的運算過程，寫成一個隱式矩陣乘法 (Implicit GEMM) 的形式：

GEMM_M = OC
GEMM_N = N * OH * OW
GEMM_K = IC * FH * FW
for i in range(GEMM_M):
    oc = i
    for j in range(GEMM_N):
        accumulator = 0
        n = j / (OH * OW)
        j_res = j % (OH * OW)
        oh = j_res / OW
        ow = j_res % OW
        for k in range(GEMM_K):
            ic = k / (FH * FW)
            k_res = k % (FH * FW)
            fh = k_res / FW
            fw = k_res % FW
            ih = oh * stride_h - pad_h + fh
            iw = ow * stride_w - pad_w + fw
            accumulator = accumulator + x(n, ic, ih, iw) * w(oc, ic, fh, fw)
        y(n, oc, oh, ow) = accumulator

上面的 Implicit GEMM 演算法仍然是序列的形式，接下來我們要把它改造成 CUDA 上的並行演算法。首先我們對整個計算任務進行分塊，讓每個執行緒塊負責計算並輸出大小為TILE_MxTILE_N的矩陣。於是演算法變成了下面的形式：

for i_out in range(GEMM_M / TILE_M):
    for j_out in range(GEMM_N / TILE_N):
        ThreadblockConvolution(x, w, y)
        
def ThreadblockConvolution(x, w, y):
    accumulate[TILE_M, TILE_N] = 0
    for i_in in range(TILE_M):
        oc = i_out * TILE_M + i_in
        for j_in in range(TILE_N):
            j = j_out * TILE_N + j_in
            n = j / (OH * OW)
            j_res = j % (OH * OW)
            oh = j_res / OW
            ow = j_res % OW
            for k in range(GEMM_K):
                ic = k / (FH * FW)
                k_res = k % (FH * FW)
                fh = k_res / FW
                fw = k_res % FW
                ih = oh * stride_h - pad_h + fh
                iw = ow * stride_w - pad_w + fw
                accumulator(i_in, j_in) = accumulator(i_in, j_in) 
                                        + x(n, ic, ih, iw) * w(oc, ic, fh, fw)
            y(n, oc, oh, ow) = accumulator(i_in, j_in)

為了提高訪存的效率，我們可以在GEMM_K這一維上也進行分塊，每次將TILE_MxTILE_K的矩陣 A 和TILE_KxTILE_N的矩陣 B 快取到 Shared Memory 裡，避免重複的 Global Memory 訪存。於是，演算法就變成了如下形式：

for i_out in range(GEMM_M / TILE_M):
    for j_out in range(GEMM_N / TILE_N):
        ThreadblockConvolution(x, w, y)

def ThreadblockConvolution(x, w, y):
    accumulator[TILE_M, TILE_N] = 0
    smem_A[TILE_M, TILE_K] = 0
    smem_B[TILE_K, TILE_N] = 0
    for i_in in range(TILE_M):
        oc = i_out * TILE_M + i_in
        for j_in in range(TILE_N):
            j = j_out * TILE_N + j_in
            n = j / (OH * OW)
            j_res = j % (OH * OW)
            oh = j_res / OW
            ow = j_res % OW
            for k_out in range(GEMM_K / TILE_K):
                load_tile_to_smem(x, A_smem)
                load_tile_to_smem(w, B_smem)
                WarpGemm(A_smem, B_smem, accumulator)
            y(n, oc, oh, ow) = accumulator(i_in, j_in)

def WarpGemm(A_smem, B_smem, accumulator):
    for k_in in range(TILE_K):
        accumulator(i_in, j_in) = accumulator(i_in, j_in) 
                                + A_smem(i_in, k_in) * B_smem(k_in, j_in)

因為我們可以直接複用 CUTLASS 裡已經實現好了高效能的WarpMmaOperator，所以實現基於 Implicit GEMM 的卷積運算元只需要

• 適配DeviceConvolution、KernelConvolution和ThreadblockConvolution，支援傳入 Tensor 型別和 Convolution Layer 的引數。
• 新增PredicateTileIterator支援讀取 Tensor 的一個 Tile 的資料到 Shared Memory 中，並隱式地將讀入的資料組織成矩陣的形式。
• 演算法的 main loop 中直接呼叫WarpTileIterator從 Shared Memory 讀取資料，然後由WarpGemmOperator完成 Warp-level 的 GEMM 運算。
• EpilogueOperator適配卷積運算元，將 Accumulator 的資料寫回 Global Memory 的 Tensor 中。

接下來我們會以 INT8 資料型別的 TensorCore 卷積運算元來介紹 MegEngine 底層的卷積實現，本文會重點介紹 2、3、4 是如何實現的，關於如何使用已經寫好的卷積運算元，可以參考之前的 文章。

Global Memory 資料佈局（Layout）

為了最大化 TensorCore 型別的卷積運算元的吞吐，MegEngine 使用了 128 位的 Global
Memory 訪存指令，因此在訪問 Tensor 的資料的時候要求地址滿足 128 位對齊。MegEngine 使用了 NCHW32 的格式來儲存 Tensor，NCHW32 格式的特點為：

• Tensor 的通道維度按照 32 個 channel 進行分組，每 32 個 channel 連續的存放在儲存中。
• Tensor 的其餘維度按照 W、H、C、N 的順序地址變化由快到慢的存放在儲存中。

由於採用了 32 個通道對齊的儲存格式，因此卷積 layer 要求輸入和輸出 feature map 的通道數都是 32 的倍數。

預處理訪存偏移量

MegEngine 的卷積實現在GEMM_K的維度上是按照\((\text{IC}/32)\cdot \text{FH}\cdot \text{FW}\cdot32\)的順序累加，寫成虛擬碼的形式如下：

kInterleaved = 32
for ic_out in range(IC//kInterleaved):
    for fh in range(FH):
        for fw in range(FW):
            for ic_in in range(kInterleaved):
                # do mma
                ......

如果寫成一層迴圈，那麼應該寫成：

kInterleaved = 32
for k in range(GEMM_K):
    chw = k // kInterleaved
    ic_in = k % kInterleaved
    ic_out = chw // (FH * FW)
    chw_res = chw % (FH * FW)
    fh = chw_res // FW
    fw = chw_res % FW
    pointer += ic_out * C_STRIDE + fh * H_STRIDE + fw * W_STRIDE
    # do mma
    ......

可以看到在迭代過程中，如果直接計算指標的偏移量的話，會引入很多除法和求餘運算。而在 CUDA 平臺上，整數的除法和求餘的開銷是非常大的，因此我們將一些地址的偏移量在 host 端預先計算好，存到 kernel param 的 buffer 中，需要時從 constant memory 中直接讀取地址，避免除法和求餘運算。
對於每個執行緒來說，在主迴圈中指標移動的 offset 如下圖所示：

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

delta(step, TILE_K) = delta(step + (FH * FW), TILE_K)

因此我們只需要大約\(\text{O}\left(\text{FH}\cdot\text{FW}\right)\)的儲存空間。其中地址偏移量的計算邏輯可以參考程式碼 conv2d_tile_iterator_nt_src_fprop_precomp.h。由於 kernel param buffer 的大小為 4KB，我們用了大約 3KB 來儲存地址的增量，所以 MegEngine 的卷積實現要求 Convolution Layer 的FH*FW的大小不能太大，但是一般情況下，3x3, 5x5, 7x7 的卷積都可以處理。Nvidia 官方實現的迭代順序與本文介紹的略有不同：

• 官方實現需要將IC補齊為TILE_K的倍數，這樣在通道數較小時會浪費一些計算量。
• 官方實現的執行緒塊在訪問輸入 feature map 的時候地址的跨度比較大，降低了訪存的區域性性，對 cache 不夠友好。

因此在效能方面，MegEngine 的實現會更有優勢，而官方實現的優點是對 Convolution Layer 的引數沒有太多限制，通用性更好。

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

cuda10.2 引入了新的 Warp-level 的mma和ldmatrix指令，使用者可以通過mma指令使用 TensorCore 來進行高速的矩陣乘加運算，通過ldmatrix精細地控制 Warp 給 TensorCore 喂資料。其中mma指令的用法如下：

unsigned A, B;  // input matrix fragment data
int C[2], D[2]; // accumulators
asm volatile(
    "mma.sync.aligned.m8n8k16.rol.col.satfinite.s32.s8.s8.s32 {%0,$1}, {%2}, {%3}, {%4,%5};\n"
    : "=r"(D[0]), "=r"(D[1])
    : "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));

這條指令的語義是由一個 Warp 的 32 個執行緒同步地完成 8x8x16 的矩陣乘加運算，它有三個輸入運算元，其中參與矩陣乘法運算的分別是一個 8x16 的矩陣 A 和一個 16x8 的矩陣 B，這兩個輸入矩陣的資料分佈在同一 Warp 的 32 個執行緒中。
矩陣 A 的佈局如下圖所示：

• 同一 Warp 中的 32 個執行緒分為 8 組，每組四個執行緒，負責讀取 8x16 的矩陣中的一行。
• 每一組中的一個執行緒讀取每一行中相鄰的 4 個 int8 的資料，恰好填滿一個 32 位的暫存器。

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

• 每 4 個執行緒一組，共分為 8 組，每組負責讀取 16x8 的矩陣中的一列。
• 每一組中的一個執行緒負責讀取一列中相鄰的 4 個資料。

參與累加運算的矩陣 C 和輸出矩陣 D 的資料也同樣分佈在 32 個執行緒中，它們的佈局如下圖所示：

• 同樣每 4 個執行緒一組，每組負責讀入/輸出一行的資料。
• 每個執行緒負責輸出一行中的相鄰兩個 int32 型別的資料，恰好構成一個 64 位的暫存器。

通過對mma指令的分析，如果 Global Memory/Shared Memory 中的資料是以行優先 (RowMajor) 或者列優先 (ColumnMajor) 的格式儲存的，那麼當同一 Warp 執行空間上連續的兩個 8x8x16 的矩陣乘加運算時，每個執行緒讀取的資料將會是跳躍的，執行每次乘法都只能讀取 32 位寬的資料到暫存器中，而低位寬的 Load 指令通常沒有辦法最大化利用儲存的頻寬。因此 Nvidia 提供了ldmatrix的指令，可以讓同一 Warp 一次性讀取 4 個 8x16 的矩陣到暫存器中，這樣恰好可以讓 Warp 中的每個執行緒一次讀取 128 位的資料，最大化頻寬的利用率。
ldmarix的用法如下所示：

unsigned addr;  // shared memory pointer
int x, y, z, w; // loaded data
int4 data;      // loaded fragment
asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];"
    : "=r"(x), "=r"(y), "=r"(z), "=r"(w)
    : "r"(addr));
data = make_int4(x, y, z, w);

上述這條指令恰好讀取了 4 個 8x16 的矩陣，每個執行緒恰好負責讀取矩陣的一行資料，讀取完成後，執行緒之間會進行資料交換，將矩陣的資料重新分佈到各個執行緒，讀取的過程如下圖所示：

這一節介紹了 TensorCore 相關的mma和ldmatrix指令，有了這兩條高效能的指令，我們還需要為資料設計巧妙的 Shared Memory 儲存格式，消除從 Shared Memory 讀取資料的 bank conflict，從而提升 Shared Memory 的讀取效率。

Shared Memory 的資料佈局

在介紹 Shared Memory 中的資料佈局之前，我們需要了解 Shared Memory 的訪存特點。Shared Memory 按照每 4 個位元組組成一個 bank，共劃分成了 32 個 bank，同一 Warp 的執行緒訪問了相同 bank 的不同地址時會發生 conflict，導致訪存的效率變慢。在同一 Warp 的執行緒訪問不同位寬的資料時，會有不同的行為：

• 每個執行緒訪問 Shared Memory 中 32 位的資料，訪存將在一個階段內完成。
• 每個執行緒訪問 Shared Memory 中 64 位的資料，訪存會在兩個階段內完成：
  • 第一個階段：前 16 個執行緒訪存 128 位元組的資料。
  • 第二個階段：後 16 個執行緒訪存 128 位元組的資料。
• 每個執行緒訪問 Shared Memory 中的 128 位的資料，訪存會在四個階段內完成：
  • 每個階段由 8 個執行緒完成 128 位元組的資料的訪存。

如果上述過程中每個階段都沒有 bank conflict，則能夠達到最大的 Shared Memory 訪存效率。
通常為了避免 Shared Memory 的 bank conflict，我們會對 Shared Memory 的資料進行 padding，讓執行緒訪問的資料錯開，避免落在同一 bank 中。但是這樣做的問題是會使得 kernel 需要 Shared Memory 的 Size 變大，但是 SM 上的 L1 cache(Shared Memory) 又是有限的，所以 padding 會降低 kernel 的 occupancy，進而就會降低 kernel 的效能。
因此 CUTLASS 設計了一種 Shared Memory 的交錯佈局方式，它能夠在不進行 padding 的前提下，使得執行緒訪存的地址沒有 bank conflict。接下來，我們以 64x64 的矩陣為例來詳細介紹資料在 Shared Memory 中的佈局。首先，執行緒讀取資料的粒度都是 128 位，也即 16 個 INT8 型別的資料，因此我們在演示資料的佈局時總是以 16 個資料為一組。如果矩陣是以行優先 (RowMajor) 的格式來組織的，那麼在邏輯上的佈局如下圖所示：

從圖中可以看到

• 每 16 個元素分為一組，被稱為一個 Vector，被染上了不同的顏色。
• 每行相鄰的 32 個元素被稱為一個 Crosswise，恰好是 NCHW32 格式中的一組 channel 的資料。

在 Shared Memory 的物理儲存中，矩陣的資料進行了重新排列，如下圖所示：

我們可以看到 Shared Memory 的物理佈局有以下特點：

• 每 4 行的一個 Crosswise 的資料作為一組，連續存放在 Shared Memory 中，緊接著會存放這 4 行的下一個 Crosswise 的資料。
• 每組資料包含了 8 個 Vector，佔據了 128 個位元組，恰好是 Shared Memory 中的 32 個不同的 bank。
• 每組資料在排列是進行了交錯，保證了ldmatrix時不會發生 bank conflict。

視訊記憶體 -> Shared Memory 的資料搬運

這一節我們會介紹從視訊記憶體 (Global Memory) 到 Shared Memory 的資料搬運。視訊記憶體到 Shared Memory 的資料搬運是由 Conv2dTileSrcIteratorFpropPrecomp 來完成的，本文並不會詳細地解讀程式碼的實現，而是描述執行緒搬運資料的過程，幫助大家建立直觀的印象，更好地理解程式碼。
如果以上一節中 Shared Memory 的邏輯佈局為例，同一 Warp 中每個執行緒讀取的資料的邏輯佈局如下圖所示，每個執行緒讀取 16 個 INT8 型別的資料，恰好構成一個 Vector。

而在實際的物理視訊記憶體中，執行緒訪問的資料分佈如下圖所示：

• 我們可以看到每個執行緒讀取了 128 位的資料。
• 相鄰的執行緒讀取的資料在物理上是連續的。

因此執行緒從 Global Memory 讀取資料的 pattern 可以滿足合併訪存的要求，同時以最大的資料位寬進行訪存，最大化了視訊記憶體頻寬的利用率。
然後如果將執行緒讀取的資料對映到 Shared Memory 的實體地址，我們可以看到

• 每 8 個執行緒向 Shared Memory 寫入 128 位元組的資料，恰好落在 Shared Memory 的 32 個不同的 bank 中。
• 同一 Warp 的訪存分為四個階段完成，每個階段都沒有 bank conflict。

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

Shared Memory -> 暫存器的資料搬運

Shared Memory 到暫存器的資料搬運是由 MmaTensorOpMultiplicandTileIterator 完成的。同一 Warp 在每一輪迭代過程會讀取 4 個 8x16 的矩陣到暫存器中，每個執行緒會讀取一行的資料。例如第一輪迭代時，執行緒讀取的資料在邏輯上的佈局如下圖所示：

而實際上資料在 Shared Memory 裡的物理佈局如下圖：

可以看到：

• 每個執行緒讀取了 128 位的資料，因此訪存分為四個階段來進行。
• 每一階段的 8 個執行緒讀取的資料恰好落在了 Shared Memory 的 32 個 bank 中，並且執行緒訪存的資料之間不存在衝突。

當進行到第二輪迭代時，每個執行緒訪問的資料的物理佈局如下圖：

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

Accumulator 寫回全域性儲存

在 int8 的情況下，同一 Warp 負責輸出 64x64 的結果，kernel 會分成 8 次寫回 Global Memory，每次寫回 32x8 的矩陣。這樣保證了每次將 Tensor 按照 NCHW32 格式寫回視訊記憶體時，同一 Warp 的 32 個執行緒恰好寫了物理上連續的 256 位元組的資料，而每個執行緒寫回 8 個位元組，保證了可以使用64位寬的資料型別進行視訊記憶體的寫操作，儘可能提高頻寬的利用率。
由於mma指令的特點，輸出矩陣的資料分佈在各個執行緒上，而為了能夠合併訪存，即：讓相鄰執行緒寫回的地址是連續的，我們利用 Shared Memory 對同一 Warp 中 32 個執行緒的資料進行了交換。資料交換後，每個執行緒擁有連續的 8 個通道的資料，且執行緒寫的地址是連續的，保證了寫回 Global Memory 滿足合併訪存的要求。
執行緒交換資料的過程如下圖所示：

每一輪迭代，Warp 中的 32 個執行緒將 32x16 的矩陣資料寫入到 Shared Memory 中。接著如下圖所示，每個執行緒會把連續的 8 個 channel 的資料讀到暫存器中。

Shared Memory 的資料交換是由以下兩個Iterator完成的

• InterleavedTileIteratorTensorOp 完成了每一輪迭代將 32x8 的資料寫入到 Shared Memory 中。
• InterleavedSharedLoadIteratorTensorOp 負責將連續的 8 個 channel 的資料讀到Fragment暫存器中。

當執行緒將交換後的資料讀到Fragment暫存器之後，會由EpilogueOp，在卷積的基礎上完成BiasAdd的運算。以 BiasAddLinearCombinationRelu 為例，它實際上完成了下面的運算：

accumulator = conv(x, w)
y = alpha * accumulator + beta * bias + gamma * z

其中 bias 是一個PerChannel的 Tensor，代表了每個輸出通道的偏置，z 是一個和卷積輸出大小一致的 Tensor，用於Convolution和ElemwiseAdd的融合。
最後EpilogueOp的輸出會由 TensorPredicatedTileIteratorTensorOp 真正地寫回到 Global Memory 中。每個執行緒寫回的資料如下圖所示：

可以看到執行緒寫回的 pattern 滿足合併訪存的要求，因此能最大化 Global Memory 寫的效率。

三、總結

本文介紹了 MegEngine 底層的卷積運算元實現原理，運算元效能可以達到cudnn的80%以上，測速結果可以參見文章。

MegEngine 會對卷積實現進行持續優化，進一步提升運算元的效能，目前來看有以下兩點可做的優化：

• 借鑑 Nvidia 官方 CUTLASS ImplicitGEMM Convolution 實現對 mask 的處理，提高TileIterator對於 mask 判斷的效率。
• 現在的卷積實現在寫回視訊記憶體時利用 Shared Memory 進行資料交換是存在 bank conflict 的。後續會考慮兩點優化
  • 對 Shared Memory 的資料佈局進行探索，消除 bank conflict，優化 Shared Memory 資料交換的效率。
  • 對 Global Memory 中的 Weight Tensor 的佈局進行探索，提高每個 Thread 上 accumulator 的區域性性，避免在 Shared Memory 中進行資料交換。

參考資料

• 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