0
雷鋒網 AI 研習社按，日前，阿里機器翻譯團隊和 PAI 團隊發表博文，闡述將 TVM 引入 TensorFlow，可以帶來至少 13 倍的 batch 矩陣相乘（matmul）加速。雷鋒網 AI 研習社將原文編譯整理如下：
背景
神經機器翻譯（NMT）是一種端到端的自動翻譯方法，可能克服傳統的基于短語的翻譯系統的缺點。最近，阿里巴巴集團正致力于在全球電子商務中部署 NMT 服務。
目前，我們將 Transformer（https://arxiv.org/pdf/1706.03762.pdf） 作為 NMT 系統的核心組成。相較于傳統基于 RNN/LSTM 的方法，它更適合于高效的離線訓練，有著相同或更高的精度。
Transformer 在時間步長中打破了相關性，對離線訓練更友好，但在在線推理上，它并沒有那么高效。我們在生產環境中發現初版 Transformer 的推理速度大約比 LSTM 版本慢 1.5 倍到 2 倍。為了提高推理性能，我們已經進行了一些優化，包括圖級別的 op 融合、循環不變節點外提（loop invariant node motion）。我們觀察到一個特殊問題：batch 矩陣相乘是 Transformer 中的一個關鍵問題，目前它在 cuBLAS 中的實現并未得到很好的優化。
圖1：Transformer 模型架構
下圖表明，通過 TVM 生成的內核可以帶來至少 13 倍的 batch 矩陣相乘加速，伴隨算子融合，速度將更快。
batch 矩陣相乘
為什么選擇利用 batch 矩陣相乘
在 Transformer 中，batch 矩陣相乘被廣泛應用于 multi-head attention 的計算。利用 batch 矩陣相乘，可以并行運行 attention 層中的 multiple heads，這有助于提高硬件的計算效率。
圖2：左圖為 Scaled Dot-Product Attention，右圖為并行運行若干 attention 層的 Multi-Head Attention
我們在推理階段對 Transformer 模型進行了全面分析，結果表明，batch 矩陣相乘計算的開銷達到 GPU 內核執行時間的 30%。當使用 nvprof 對 cuBLAS batch 矩陣相乘內核做一些第一原理（first-principle）分析，很明顯，這種方法的表現并不好，同時我們還發現幾個有趣的現象。
什么是 batch 矩陣相乘
通常，batch 矩陣相乘計算會在一批矩陣上執行矩陣-矩陣乘法。batch 被認為是「統一的」，即所有實例都具有相同的維度（M，N，K）、leading 維度 (lda，ldb，ldc) 和它們各自的 A、B、C 矩陣的轉置。
batch 矩陣相乘計算具體可以描述如下：
void BatchedGemm(input A, input B, output C, M, N, K, batch_dimension) {
for (int i = 0; i < batch_dimension; ++i) {
DoGemm(A[i],B[i],C[i],M,K,N)
}
}
batch 矩陣相乘形狀
在語言翻譯任務中，batch 矩陣相乘的形狀比在其他工作負載下的常規矩陣相乘計算要小得多。Transformer 的形狀與輸入語句的長度和解碼器步長有關。一般來說小于 30。
至于 batch 維度，當給定推理 batch 大小時，它是固定數字。例如，如果 batch size 是 16，beam size 是 4，batch 維度是 16 * 4 * #head (在 multi-headattention 中 head 的數目，通常為 8)。矩陣 M、K、N 的范圍在 [1, max decode length] 或 [1, max encode length] 內。
batch 矩陣相乘的性能問題
首先，我們在理論上對 batch 矩陣相乘內核進行了 FLOP 分析。結果非常有趣：所有 batch 矩陣相乘的計算強度都是受限的（TFLOP 數少于 1）。
然后，我們通過 nvprof 描述了多形狀 batch 矩陣相乘的 cuBLAS 性能。下面的表格中是使用 NVIDIA M40 GPU（CUDA 8.0）得到的一些指標。
即使形狀不同（在 M、N、K 間變化），所有 maxwell_sgemmBatched_128x128_raggedMn_tn 調用執行的都是相同的 FLOP 數，這比理論值大得多。從中可以推斷，所有這些不同的形狀最終都會被填充成確定的形狀。在所有的形狀中，即使在最好的情況下，理論 FLOP 只占實際執行 FLOP 的 2.74%，因此大多數計算都是多余的。類似地，調用另一個 cuBLAS 內核 maxwell_sgemmBatched_64x64_raggedMn_tn 也出現相同情況。
顯而易見，cuBLAS batch 矩陣相乘的執行效率很低。基于這個原因，我們在 NMT 中使用 TVM 生成高效的 batch 矩陣相乘內核。
batch 矩陣相乘計算
在 TVM 中，普通 batch 矩陣相乘計算聲明如下：
# computation representation
A = tvm.placeholder((batch, M, K), name='A')
B = tvm.placeholder((batch, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute((batch, M, N),
lambda b, y, x: tvm.sum(A[b, y, k] * B[b, k, x], axis = k),
name = 'C')
調度優化
在聲明計算之后，我們需要仔細地設計調度來更好地發揮性能。
調節 block/線程數的參數
# thread indices
block_y = tvm.thread_axis("blockIdx.y")
block_x = tvm.thread_axis("blockIdx.x")
thread_y = tvm.thread_axis((0, num_thread_y), "threadIdx.y")
thread_x = tvm.thread_axis((0, num_thread_x), "threadIdx.x")
thread_yz = tvm.thread_axis((0, vthread_y), "vthread", name="vy")
thread_xz = tvm.thread_axis((0, vthread_x), "vthread", name="vx")
# block partitioning
BB, FF, MM, PP = s[C].op.axis
BBFF = s[C].fuse(BB, FF)
MMPP = s[C].fuse(MM, PP)
by, ty_block = s[C].split(BBFF, factor = num_thread_y * vthread_y)
bx, tx_block = s[C].split(MMPP, factor = num_thread_x * vthread_x)
s[C].bind(by, block_y)
s[C].bind(bx, block_x)
vty, ty = s[C].split(ty_block, nparts = vthread_y)
vtx, tx = s[C].split(tx_block, nparts = vthread_x)
s[C].reorder(by, bx, vty, vtx, ty, tx)
s[C].reorder(by, bx, ty, tx)
s[C].bind(ty, thread_y)
s[C].bind(tx, thread_x)
s[C].bind(vty, thread_yz)
s[C].bind(vtx, thread_xz)
我們融合了 batch 矩陣相乘的外部維度，例如 op 維的 BB 和 FF 在 batch 矩陣相乘計算中通常稱為「batch」維，我們用一個因子 (number_thread * vthread) 分割外部和內部維度。
在 batch 矩陣相乘中不需要 Strided 模式，因此將虛擬線程數（vthready 和 vthreadx）都設置為 1。
找到 number_thread 的最佳組合
下面的結果是基于 NVIDIA M40 GPU（CUDA 8.0）。
基于過去的經驗，找到 num_thread_y 和 num_thread_x 最佳組合的方法是通過暴力搜索（brute-force search）。經過暴力搜索后，可以找到當前形狀的最佳組合，在當前的計算中，num_thread_y = 8，num_thread_x = 32。
將 batch 矩陣相乘與其他運算融合
現有的「黑盒」cuBLAS 庫調用一般會作為常用的「op 融合」優化策略的邊界。然而，利用生成的高效 batch 矩陣相乘內核，融合邊界極易被打破，將不僅僅是各個元素之間的融合，因此可以獲得更好的性能改進。
從計算圖中可以看出，batch 矩陣相乘之后總是伴隨著廣播加法運算或轉置運算。
通過將「加法」或「轉置」運算與 batch 矩陣相乘融合，可以減少內核啟動開銷和冗余內存訪問時間。
batch 矩陣相乘和廣播加法融合計算的聲明如下：
# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
# the shape of B is (N, K) other than (K, N) is because B is transposed is this fusion pattern
B = tvm.placeholder((batch_size, features, N, K), name='B')
ENTER = tvm.placeholder((batch_size, 1, M, N), name = 'ENTER')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
(batch_size, features, M, N),
lambda yb, yf, m, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, x, k], axis = k),
name = 'C')
D = topi.broadcast_add(C, ENTER)
batch 矩陣相乘和轉置融合計算的聲明如下：
# computation representation
A = tvm.placeholder((batch_size, features, M, K), name='A')
B = tvm.placeholder((batch_size, features, K, N), name='B')
k = tvm.reduce_axis((0, K), 'k')
C = tvm.compute(
(batch_size, M, features, N),
lambda yb, m, yf, x: tvm.sum(A[yb, yf, m, k] * B[yb, yf, k, x], axis = k),
name = 'C')
融合內核性能
測試生成代碼性能時，形狀選擇為 [batch=64, heads=8, M=1, N=17, K=128]。選擇 17 作為序列長度是因為它是我們生產中的平均輸入長度。
tf-r1.4 BatchMatmul: 513.9 us
tf-r1.4 BatchMatmul + Transpose (separate): 541.9 us
TVM BatchMatmul: 37.62 us
TVM BatchMatmul + Transpose (fused): 38.39 us
內核融合優化帶來了 1.7 倍的加速。
集成 TensorFlow
在我們的工作負載中，batch 矩陣相乘的輸入形狀是有限的，易于提前枚舉。有了這些預定義的形狀，我們可以提前生成高度優化的 CUDA 內核（固定形狀的計算可以帶來最佳優化潛能）。同時，還將生成一個適合大多數形狀的通用 batch 矩陣相乘內核，為沒有提前生成內核的形狀提供回退機制。
我們將生成的針對特定形狀的高效內核和回退機制集成到 Tensorflow 中。我們開發了一些融合操作，例如 BatchMatMulTranspose 或 BatchMatMulAdd——使用 TVM runtime API 為確定輸入形狀啟動特定生成的內核或調用回退內核。
通過執行圖優化 pass，可以利用融合操作自動替換原始batch matmul + add/transpose。同時，通過結合更激進的圖優化 pass，我們嘗試利用 TVM 為長尾操作模式生成更高效的融合內核，以進一步提升端到端性能。
總結
在阿里巴巴，我們發現 TVM 是非常有效的開發高性能 GPU 內核的工具，可以滿足我們的內部需求。
在本博客中，我們以 Transformer 模型為例，說明了我們利用 TVM 的優化策略。
首先，我們通過第一原理分析確定了 Transformer 模型的關鍵問題。然后，我們使用 TVM 生成高度優化的 CUDA 內核來取代 cuBLAS 版本（此時達到 13 倍的加速）。
接下來，利用 TVM 的內核融合機制來融合 batch 矩陣相乘的前/后操作，以帶來進一步的性能改進（性能提升 1.7 倍）。端到端的性能改善達到 1.4 倍。基于這些生成的內核，我們開發了一個圖優化 pass，自動用 TVM 融合內核替換掉原有的計算模式，確保優化對終端用戶透明。
最后，所有這些優化都以松散耦合的方式集成到 TensorFlow 中，這展示了將 TVM 與不同深度學習框架集成的潛在方式。
目前我們還有一項正在進行的工作——將 TVM 整合為 TensorFlow 的 codegen 后端。我們希望將來與社群分享更多成果。
雷鋒網 AI 研習社編譯整理。
雷峰網版權文章，未經授權禁止轉載。詳情見轉載須知。
