雷鋒網(wǎng) AI 研習(xí)社按，日前，阿里機(jī)器翻譯團(tuán)隊(duì)和 PAI 團(tuán)隊(duì)發(fā)表博文，闡述將 TVM 引入 TensorFlow，可以帶來至少 13 倍的 batch 矩陣相乘（matmul）加速。雷鋒網(wǎng) AI 研習(xí)社將原文編譯整理如下：

背景

神經(jīng)機(jī)器翻譯（NMT）是一種端到端的自動翻譯方法，可能克服傳統(tǒng)的基于短語的翻譯系統(tǒng)的缺點(diǎn)。最近，阿里巴巴集團(tuán)正致力于在全球電子商務(wù)中部署 NMT 服務(wù)。

目前，我們將 Transformer（https://arxiv.org/pdf/1706.03762.pdf） 作為 NMT 系統(tǒng)的核心組成。相較于傳統(tǒng)基于 RNN/LSTM 的方法，它更適合于高效的離線訓(xùn)練，有著相同或更高的精度。

Transformer 在時(shí)間步長中打破了相關(guān)性，對離線訓(xùn)練更友好，但在在線推理上，它并沒有那么高效。我們在生產(chǎn)環(huán)境中發(fā)現(xiàn)初版 Transformer 的推理速度大約比 LSTM 版本慢 1.5 倍到 2 倍。為了提高推理性能，我們已經(jīng)進(jìn)行了一些優(yōu)化，包括圖級別的 op 融合、循環(huán)不變節(jié)點(diǎn)外提（loop invariant node motion）。我們觀察到一個(gè)特殊問題：batch 矩陣相乘是 Transformer 中的一個(gè)關(guān)鍵問題，目前它在 cuBLAS 中的實(shí)現(xiàn)并未得到很好的優(yōu)化。

圖1：Transformer 模型架構(gòu)

下圖表明，通過 TVM 生成的內(nèi)核可以帶來至少 13 倍的 batch 矩陣相乘加速，伴隨算子融合，速度將更快。

batch 矩陣相乘

為什么選擇利用 batch 矩陣相乘

在 Transformer 中，batch 矩陣相乘被廣泛應(yīng)用于 multi-head attention 的計(jì)算。利用 batch 矩陣相乘，可以并行運(yùn)行 attention 層中的 multiple heads，這有助于提高硬件的計(jì)算效率。

圖2：左圖為 Scaled Dot-Product Attention，右圖為并行運(yùn)行若干 attention 層的 Multi-Head Attention

我們在推理階段對 Transformer 模型進(jìn)行了全面分析，結(jié)果表明，batch 矩陣相乘計(jì)算的開銷達(dá)到 GPU 內(nèi)核執(zhí)行時(shí)間的 30%。當(dāng)使用 nvprof 對 cuBLAS batch 矩陣相乘內(nèi)核做一些第一原理（first-principle）分析，很明顯，這種方法的表現(xiàn)并不好，同時(shí)我們還發(fā)現(xiàn)幾個(gè)有趣的現(xiàn)象。

什么是 batch 矩陣相乘

通常，batch 矩陣相乘計(jì)算會在一批矩陣上執(zhí)行矩陣-矩陣乘法。batch 被認(rèn)為是「統(tǒng)一的」，即所有實(shí)例都具有相同的維度（M，N，K）、leading 維度 (lda，ldb，ldc) 和它們各自的 A、B、C 矩陣的轉(zhuǎn)置。

batch 矩陣相乘計(jì)算具體可以描述如下：

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 矩陣相乘形狀

在語言翻譯任務(wù)中，batch 矩陣相乘的形狀比在其他工作負(fù)載下的常規(guī)矩陣相乘計(jì)算要小得多。Transformer 的形狀與輸入語句的長度和解碼器步長有關(guān)。一般來說小于 30。

至于 batch 維度，當(dāng)給定推理 batch 大小時(shí)，它是固定數(shù)字。例如，如果 batch size 是 16，beam size 是 4，batch 維度是 16 * 4 * #head (在 multi-headattention 中 head 的數(shù)目，通常為 8)。矩陣 M、K、N 的范圍在 [1, max decode length] 或 [1, max encode length] 內(nèi)。

batch 矩陣相乘的性能問題

首先，我們在理論上對 batch 矩陣相乘內(nèi)核進(jìn)行了 FLOP 分析。結(jié)果非常有趣：所有 batch 矩陣相乘的計(jì)算強(qiáng)度都是受限的（TFLOP 數(shù)少于 1）。

然后，我們通過 nvprof 描述了多形狀 batch 矩陣相乘的 cuBLAS 性能。下面的表格中是使用 NVIDIA M40 GPU（CUDA 8.0）得到的一些指標(biāo)。

即使形狀不同（在 M、N、K 間變化），所有 maxwell_sgemmBatched_128x128_raggedMn_tn 調(diào)用執(zhí)行的都是相同的 FLOP 數(shù)，這比理論值大得多。從中可以推斷，所有這些不同的形狀最終都會被填充成確定的形狀。在所有的形狀中，即使在最好的情況下，理論 FLOP 只占實(shí)際執(zhí)行 FLOP 的 2.74%，因此大多數(shù)計(jì)算都是多余的。類似地，調(diào)用另一個(gè) cuBLAS 內(nèi)核 maxwell_sgemmBatched_64x64_raggedMn_tn 也出現(xiàn)相同情況。

顯而易見，cuBLAS batch 矩陣相乘的執(zhí)行效率很低?；谶@個(gè)原因，我們在 NMT 中使用 TVM 生成高效的 batch 矩陣相乘內(nèi)核。

batch 矩陣相乘計(jì)算

在 TVM 中，普通 batch 矩陣相乘計(jì)算聲明如下：

# 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')

調(diào)度優(yōu)化

在聲明計(jì)算之后，我們需要仔細(xì)地設(shè)計(jì)調(diào)度來更好地發(fā)揮性能。

調(diào)節(jié) block/線程數(shù)的參數(shù)

# 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)

我們?nèi)诤狭?batch 矩陣相乘的外部維度，例如 op 維的 BB 和 FF 在 batch 矩陣相乘計(jì)算中通常稱為「batch」維，我們用一個(gè)因子 (number_thread * vthread) 分割外部和內(nèi)部維度。

在 batch 矩陣相乘中不需要 Strided 模式，因此將虛擬線程數(shù)（vthready 和 vthreadx）都設(shè)置為 1。

找到 number_thread 的最佳組合

下面的結(jié)果是基于 NVIDIA M40 GPU（CUDA 8.0）。

基于過去的經(jīng)驗(yàn)，找到 num_thread_y 和 num_thread_x 最佳組合的方法是通過暴力搜索（brute-force search）。經(jīng)過暴力搜索后，可以找到當(dāng)前形狀的最佳組合，在當(dāng)前的計(jì)算中，num_thread_y = 8，num_thread_x = 32。

將 batch 矩陣相乘與其他運(yùn)算融合

現(xiàn)有的「黑盒」cuBLAS 庫調(diào)用一般會作為常用的「op 融合」優(yōu)化策略的邊界。然而，利用生成的高效 batch 矩陣相乘內(nèi)核，融合邊界極易被打破，將不僅僅是各個(gè)元素之間的融合，因此可以獲得更好的性能改進(jìn)。

從計(jì)算圖中可以看出，batch 矩陣相乘之后總是伴隨著廣播加法運(yùn)算或轉(zhuǎn)置運(yùn)算。

通過將「加法」或「轉(zhuǎn)置」運(yùn)算與 batch 矩陣相乘融合，可以減少內(nèi)核啟動開銷和冗余內(nèi)存訪問時(shí)間。

batch 矩陣相乘和廣播加法融合計(jì)算的聲明如下：

# 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 矩陣相乘和轉(zhuǎn)置融合計(jì)算的聲明如下：

# 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')

融合內(nèi)核性能

測試生成代碼性能時(shí)，形狀選擇為 [batch=64, heads=8, M=1, N=17, K=128]。選擇 17 作為序列長度是因?yàn)樗俏覀兩a(chǎn)中的平均輸入長度。

• 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

內(nèi)核融合優(yōu)化帶來了 1.7 倍的加速。

集成 TensorFlow

在我們的工作負(fù)載中，batch 矩陣相乘的輸入形狀是有限的，易于提前枚舉。有了這些預(yù)定義的形狀，我們可以提前生成高度優(yōu)化的 CUDA 內(nèi)核（固定形狀的計(jì)算可以帶來最佳優(yōu)化潛能）。同時(shí)，還將生成一個(gè)適合大多數(shù)形狀的通用 batch 矩陣相乘內(nèi)核，為沒有提前生成內(nèi)核的形狀提供回退機(jī)制。

我們將生成的針對特定形狀的高效內(nèi)核和回退機(jī)制集成到 Tensorflow 中。我們開發(fā)了一些融合操作，例如 BatchMatMulTranspose 或 BatchMatMulAdd——使用 TVM runtime API 為確定輸入形狀啟動特定生成的內(nèi)核或調(diào)用回退內(nèi)核。

通過執(zhí)行圖優(yōu)化 pass，可以利用融合操作自動替換原始batch matmul + add/transpose。同時(shí)，通過結(jié)合更激進(jìn)的圖優(yōu)化 pass，我們嘗試?yán)?TVM 為長尾操作模式生成更高效的融合內(nèi)核，以進(jìn)一步提升端到端性能。

總結(jié)

在阿里巴巴，我們發(fā)現(xiàn) TVM 是非常有效的開發(fā)高性能 GPU 內(nèi)核的工具，可以滿足我們的內(nèi)部需求。

在本博客中，我們以 Transformer 模型為例，說明了我們利用 TVM 的優(yōu)化策略。

首先，我們通過第一原理分析確定了 Transformer 模型的關(guān)鍵問題。然后，我們使用 TVM 生成高度優(yōu)化的 CUDA 內(nèi)核來取代 cuBLAS 版本（此時(shí)達(dá)到 13 倍的加速）。

接下來，利用 TVM 的內(nèi)核融合機(jī)制來融合 batch 矩陣相乘的前/后操作，以帶來進(jìn)一步的性能改進(jìn)（性能提升 1.7 倍）。端到端的性能改善達(dá)到 1.4 倍。基于這些生成的內(nèi)核，我們開發(fā)了一個(gè)圖優(yōu)化 pass，自動用 TVM 融合內(nèi)核替換掉原有的計(jì)算模式，確保優(yōu)化對終端用戶透明。

最后，所有這些優(yōu)化都以松散耦合的方式集成到 TensorFlow 中，這展示了將 TVM 與不同深度學(xué)習(xí)框架集成的潛在方式。

目前我們還有一項(xiàng)正在進(jìn)行的工作——將 TVM 整合為 TensorFlow 的 codegen 后端。我們希望將來與社群分享更多成果。

via：http://www.tvmlang.org

雷鋒網(wǎng) AI 研習(xí)社編譯整理。

雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。