值得收藏 | 深度剖析 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