如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

本文作者：奕欣

2017-08-24 05:16

導(dǎo)語：高效的深度學(xué)習(xí)運算符是深度學(xué)習(xí)系統(tǒng)的核心。TVM 作為一種端到端的張量 IR / DSL 堆棧，能夠讓整個過程變得更加簡單。

數(shù)天前，陳天奇團隊宣布推出 TVM，在微博上表示，「我們今天發(fā)布了 TVM，和 NNVM 一起組成深度學(xué)習(xí)到各種硬件的完整優(yōu)化工具鏈，支持手機，cuda, opencl, metal, javascript 以及其它各種后端。歡迎對于深度學(xué)習(xí)，編譯原理，高性能計算，硬件加速有興趣的同學(xué)一起加入 dmlc 推動領(lǐng)導(dǎo)開源項目社區(qū) 。」

據(jù)雷鋒網(wǎng)AI科技評論了解，大多數(shù)現(xiàn)有系統(tǒng)針對窄范圍的服務(wù)器級 GPU 進(jìn)行優(yōu)化，且需要在包括手機、IOT 設(shè)備及專用加速器上部署大量工作。而 TVM 是一種將深度學(xué)習(xí)工作負(fù)載部署到硬件的端到端 IR（中間表示）堆棧。也就是說，這類解決方案能夠把深度學(xué)習(xí)模型分發(fā)到各種硬件設(shè)備上、實現(xiàn)端到端的調(diào)優(yōu)。

它存在三大特點：

能夠優(yōu)化 CPU、GPU 和其他專業(yè)化硬件在常規(guī)深度學(xué)習(xí)任務(wù)上的計算量；
能夠自動轉(zhuǎn)換計算圖，使得內(nèi)存利用率最小化，優(yōu)化數(shù)據(jù)布局，將計算模式進(jìn)行融合。
提供從現(xiàn)有前端到裸機硬件的端到端編譯，到瀏覽器可執(zhí)行 Javascripts。

雷鋒網(wǎng)AI科技評論了解到，TVM 的首篇博客是這樣介紹的：

「在 TVM 的幫助之下，開發(fā)者只需要少量的額外工作，便可輕易在手機端、嵌入式設(shè)備甚至瀏覽器上運行深度學(xué)習(xí)任務(wù)。TVM 還為多硬件平臺上的深度學(xué)習(xí)工作負(fù)載提供統(tǒng)一優(yōu)化框架，包括依賴全新計算原語的專用加速器?！?/p>

而在今天，陳天奇在微博上發(fā)布了新的動態(tài)，以圖森未來胡玉煒的教程介紹著重推介 TVM 的深度學(xué)習(xí) op 優(yōu)化。

「深度學(xué)習(xí) op 優(yōu)化是非常重要但是困難的問題。來自圖森未來的胡玉煒寫了一個教程介紹了如何利用 TVM 來優(yōu)化深度學(xué)習(xí)的 gpu op，通過幾十行 python 代碼獲得比已有 tf 實現(xiàn)兩三倍的提升?！?/p>

這一文章目前也同步更新到 TVM 博客上，雷鋒網(wǎng) AI 科技評論第一時間做了覆蓋和報道。

胡玉煒，北京航空航天大學(xué)電子工程碩士，目前 Gap 一年，現(xiàn)在圖森未來 HPC 小組實習(xí)。這篇文章題為《Optimize Deep Learning GPU Operators with TVM: A Depthwise Convolution Example》（以 Depthwise Convolution 為例，采用 TVM 優(yōu)化深度學(xué)習(xí) GPU 運算符）

高效的深度學(xué)習(xí)運算符是深度學(xué)習(xí)系統(tǒng)的核心。通常這些運算符很難優(yōu)化，需要 HPC 專家付出非常多的努力。TVM 作為一種端到端的張量 IR / DSL 堆棧，能夠讓整個過程變得更加簡單。

這篇文章提供了一個很好的參考，教會開發(fā)者如何在 TVM 的幫助下編寫高性能 GPU 運算符內(nèi)核。團隊采用的是 Depthwise Convolution（即 topi.nn.depthwise_conv2d_nchw）作為示例，并演示了如何可以改進(jìn)已經(jīng)手動優(yōu)化的 TensorFlow 中的 CUDA 內(nèi)核。

根據(jù)文章的描述，采用 TVM 的最終版本比不同工作負(fù)載下的 tf-1.2 中的優(yōu)化內(nèi)核快了 2-4 倍，運算符融合速度提升了 3 倍-7 倍。以下是在 GTX1080 下的測試結(jié)果， filter size= [1，256,3,3]，stride = [1，1]，padding ='SAME'：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

Depthwise Convolution 是一種構(gòu)建模型的基本思想，能夠有效降低深度神經(jīng)網(wǎng)絡(luò)的計算復(fù)雜度，包括谷歌的 Xception 和 MobileNet 都屬于 Depthwise Convolution。

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

在 TVM 環(huán)境下，運行 Depthwise Convolution 的代碼如下：

# padding stagePaddedInput = tvm.compute(
(batch, in_channel, height_after_pad, width_after_pad),
lambda b, c, i, j: tvm.select(
tvm.all(i >= pad_top, i - pad_top < in_height, j >= pad_left, j - pad_left < in_width),
Input[b, c, i - pad_top, j - pad_left], tvm.const(0.0)),
name="PaddedInput")# depthconv stagedi = tvm.reduce_axis((0, filter_height), name='di')dj = tvm.reduce_axis((0, filter_width), name='dj')Output = tvm.compute(
(batch, out_channel, out_height, out_width),
lambda b, c, i, j: tvm.sum(
PaddedInput[b, c/channel_multiplier, i*stride_h + di, j*stride_w + dj] * Filter[c/channel_multiplier, c%channel_multiplier, di, dj],
axis=[di, dj]),
name='DepthwiseConv2d')

通用 GPU 優(yōu)化指南

胡玉煒在文章中提到了優(yōu)化 CUDA 代碼時通常需要注意的三大問題，即數(shù)據(jù)重用（data reuse）、共享內(nèi)存（shared memory）和訪問沖突（bank conflicts）。

在現(xiàn)代計算架構(gòu)中，從從內(nèi)存中加載數(shù)據(jù)的成本遠(yuǎn)高于單個浮點計算。因此，我們希望在輸入數(shù)據(jù)被加載到寄存器或 cache 后能夠再次使用。

在 depthwise convolution 中有兩種形式的數(shù)據(jù)重用：過濾器重用和輸入重用，前者發(fā)生在輸入通道上滑過并計算多次時，后者在平鋪時發(fā)生，以 3x3 的 depthwise convolution 為例：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

不平鋪情況下，每個線程計算 1 個輸出元素并加載 3x3 輸入數(shù)據(jù)。16 線程共有 9x16 個負(fù)載。

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

平鋪情況下，每個線程分別計算 2x2 個輸出元素并加載 4x4 輸入數(shù)據(jù)。4 線程共有 16x4 個負(fù)載。

共享內(nèi)存和訪問沖突

共享內(nèi)存可以看作 GPU 中的緩存，且是片上的，速度較快。通常的做法是，將數(shù)據(jù)從全局內(nèi)存加載到共享內(nèi)存中，然后塊中的所有線程都從共享內(nèi)存中讀取數(shù)據(jù)。

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

而為了避免訪問沖突，連續(xù)的線程最好訪問連續(xù)的內(nèi)存地址，如下所示（每種顏色代表一個共享內(nèi)存庫）：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

詳細(xì)內(nèi)容可參考https://devblogs.nvidia.com/parallelforall/using-shared-memory-cuda-cc/

具體的優(yōu)化過程

計算填充輸入內(nèi)嵌以節(jié)省內(nèi)存分配

Padding 被明確聲明為一個單獨的階段。通過計算內(nèi)聯(lián)以避免冗余內(nèi)存分配：

s = tvm.create_schedule(Output.op)s[PaddedInput].compute_inline()

將一個大通道分成較小的塊

一個簡單的做法是一個 CUDA 塊處理一個輸入通道和相應(yīng)的過濾器，加載到共享存儲器后計算：

IS = s.cache_read(PaddedInput, "shared", [DepthwiseConv2d])
FS = s.cache_read(Filter, "shared", [DepthwiseConv2d])
block_y = tvm.thread_axis("blockIdx.y")
block_x = tvm.thread_axis("blockIdx.x")
# bind the dimension of batch (N in NCHW) with block_y
s[Output].bind(Output.op.axis[0], block_y)
# bind the dimension of channel (C in NCHW) with block_x
s[Output].bind(Output.op.axis[1], block_x)

下圖為測試結(jié)果，在 GTX 1080 上 1000 次運行的平均時間成本，并與 depthwise_conv2d in tensorflow進(jìn)行比較。

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

如果是 21 x 21 或者 32 x 32 的通道大小，則性能表現(xiàn)良好，但如果是 64 x 64，那么性能會大大下降。如果做一些修改，那么效果會提升很多：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

線程數(shù)調(diào)參

在一個 cuda 塊中實現(xiàn) 32 x 32 的線程，如下：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

num_thread_y 和 num_thread_x 這兩個參數(shù)如何調(diào)才能得到最優(yōu)解？在 Filter = [256, 1, 3, 3] and stride = [1, 1]下：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

通過測試，團隊得到以下結(jié)果：

大規(guī)模平鋪對于數(shù)據(jù)重用是有好處的，但不利于本地存儲器讀取。
num_thread_y 和 num_thread_x 對訪問沖突的影響不同。
num_thread_y 和 num_thread_x 的最佳組合，需要實現(xiàn)高效共享內(nèi)存訪問（避免存儲區(qū)沖突），數(shù)據(jù)重用和本地內(nèi)存讀取的平衡。

通過強力搜索，在 TVM 中我們可以將 num_thread_y 和 num_thread_x 作為參數(shù)傳遞給 schedule 函數(shù)，并嘗試所有可能的組合來找到最優(yōu)組合。

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

Vthread（virtual thread）與 Strided Patterns

在 TVM 中，Vthread 能夠有效支持 Strided Patterns。

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

在 Filter = [256, 1, 3, 3], stride = [1, 1], blocking_h = 32, blocking_w = 32 的情況下，結(jié)果如下：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

case 2 比 case 1 更快，因為在 case 2 num_thread_x = 8 和 num_vthread_x = 4 的情況下，確保了連續(xù)的線程訪問連續(xù)的存儲器地址，從而避免訪問沖突，如下所示（每個顏色表示一個線程的工作負(fù)載）：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

再來回顧一下和 Tensorflow 的對比：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

運算符融合

運算符融合是一個典型優(yōu)化深度學(xué)習(xí)網(wǎng)絡(luò)的方法，在 TVM 中，考慮到原有的模式 depthwise_conv2d + scale_shift + relu，可以稍作如下修改：

如何利用 TVM 優(yōu)化深度學(xué)習(xí)GPU op？教你用幾十行Python代碼實現(xiàn)2-3倍提升

生成 IR 如下：

/* Input = [1, 1, 32, 32], Filter = [1, 1, 3, 3], stride = [1, 1], padding = 'SAME' */produce Relu {
// attr [iter_var(blockIdx.y, , blockIdx.y)] thread_extent = 1 // attr [DepthwiseConv2d] storage_scope = "local" allocate DepthwiseConv2d[float32 * 1 * 1 * 4 * 4]
// attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 1 // attr [iter_var(threadIdx.y, Range(min=0, extent=8), threadIdx.y)] thread_extent = 8 // attr [iter_var(threadIdx.x, Range(min=0, extent=8), threadIdx.x)] thread_extent = 8 produce DepthwiseConv2d {
for (i, 0, 4) {
for (j, 0, 4) {
DepthwiseConv2d[((i*4) + j)] = 0.000000f
for (di, 0, 3) {
for (dj, 0, 3) {
DepthwiseConv2d[((i*4) + j)] = (DepthwiseConv2d[((i*4) + j)] + (tvm_if_then_else(((((((1 - di) - i) <= (((blockIdx.x*8) + threadIdx.y)*4)) && ((((blockIdx.x*8) + threadIdx.y)*4) < ((33 - di) - i))) && (((1 - dj) - j) <= (threadIdx.x*4))) && ((threadIdx.x*4) < ((33 - dj) - j))), Input[(((((((((((blockIdx.y + blockIdx.x)*8) + threadIdx.y)*32) + threadIdx.x)*4) + (i*32)) + j) + (di*32)) + dj) + -33)], 0.000000f)*Filter[((di*3) + dj)]))
}
}
}
}
}
for (i2.inner.inner.inner, 0, 4) {
for (i3.inner.inner.inner, 0, 4) {
Relu[((((((((blockIdx.y + blockIdx.x)*8) + threadIdx.y)*32) + threadIdx.x)*4) + (i2.inner.inner.inner*32)) + i3.inner.inner.inner)] = max(((DepthwiseConv2d[((i2.inner.inner.inner*4) + i3.inner.inner.inner)]*Scale[0]) + Shift[0]), 0.000000f)
}
}}

可以看到，每個線程在將 depthwise_conv2d 的結(jié)果寫入全局內(nèi)存之前，會計算 scale_shift 和 relu。融合運算符與單個 depthwise_conv2d 一樣快。以下是 Input = [1，256，96，96]，F(xiàn)ilter = [256，1,3,3]，stride = [1，1]，padding =‘SAME'的結(jié)果：