如何利用 TVM 優化深度學習GPU op？教你用幾十行Python代碼實現2-3倍提升

數天前，陳天奇團隊宣布推出 TVM，在微博上表示，「我們今天發布了 TVM，和 NNVM 一起組成深度學習到各種硬體的完整優化工具鏈，支持手機，cuda, opencl, metal, javascript 以及其它各種後端。歡迎對於深度學習，編譯原理，高性能計算，硬體加速有興趣的同學一起加入 dmlc 推動領導開源項目社區 。」

據雷鋒網AI科技評論了解，大多數現有系統針對窄範圍的伺服器級 GPU 進行優化，且需要在包括手機、IOT 設備及專用加速器上部署大量工作。而 TVM 是一種將深度學習工作負載部署到硬體的端到端 IR（中間表示）堆棧。也就是說，這類解決方案能夠把深度學習模型分發到各種硬體設備上、實現端到端的調優。

它存在三大特點：

• 能夠優化 CPU、GPU 和其他專業化硬體在常規深度學習任務上的計算量；

• 能夠自動轉換計算圖，使得內存利用率最小化，優化數據布局，將計算模式進行融合。

• 提供從現有前端到裸機硬體的端到端編譯，到瀏覽器可執行 Javascripts。

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

「在 TVM 的幫助之下，開發者只需要少量的額外工作，便可輕易在手機端、嵌入式設備甚至瀏覽器上運行深度學習任務。TVM 還為多硬體平台上的深度學習工作負載提供統一優化框架，包括依賴全新計算原語的專用加速器。」

而在今天，陳天奇在微博上發布了新的動態，以圖森未來胡玉煒的教程介紹著重推介 TVM 的深度學習 op 優化。

「深度學習 op 優化是非常重要但是困難的問題。來自圖森未來的胡玉煒寫了一個教程介紹了如何利用 TVM 來優化深度學習的 gpu op，通過幾十行 python 代碼獲得比已有 tf 實現兩三倍的提升。」

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

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

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

這篇文章提供了一個很好的參考，教會開發者如何在 TVM 的幫助下編寫高性能 GPU 運算符內核。團隊採用的是 Depthwise Convolution（即 topi.nn.depthwise_conv2d_nchw）作為示例，並演示了如何可以改進已經手動優化的 TensorFlow 中的 CUDA 內核。

根據文章的描述，採用 TVM 的最終版本比不同工作負載下的 tf-1.2 中的優化內核快了 2-4 倍，運算符融合速度提升了 3 倍-7 倍。以下是在 GTX1080 下的測試結果， filter size= [1，256,3,3]，stride = [1，1]，padding ="SAME"：

Depthwise Convolution 是一種構建模型的基本思想，能夠有效降低深度神經網路的計算複雜度，包括谷歌的 Xception 和 MobileNet 都屬於 Depthwise Convolution。

在 TVM 環境下，運行 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 = pad_left, j - pad_left 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 優化指南

胡玉煒在文章中提到了優化 CUDA 代碼時通常需要注意的三大問題，即數據重用（data reuse）、共享內存（shared memory）和訪問衝突（bank conflicts）。

在現代計算架構中，從從內存中載入數據的成本遠高於單個浮點計算。因此，我們希望在輸入數據被載入到寄存器或 cache 後能夠再次使用。

在 depthwise convolution 中有兩種形式的數據重用：過濾器重用和輸入重用，前者發生在輸入通道上滑過並計算多次時，後者在平鋪時發生，以 3x3 的 depthwise convolution 為例：

不平鋪情況下，每個線程計算 1 個輸出元素並載入 3x3 輸入數據。16 線程共有 9x16 個負載。

平鋪情況下，每個線程分別計算 2x2 個輸出元素並載入 4x4 輸入數據。4 線程共有 16x4 個負載。

共享內存可以看作 GPU 中的緩存，且是片上的，速度較快。通常的做法是，將數據從全局內存載入到共享內存中，然後塊中的所有線程都從共享內存中讀取數據。

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

詳細內容可參考https://devblogs.nvidia.com/parallelforall/using-shared-memory-cuda-cc/

具體的優化過程

計算填充輸入內嵌以節省內存分配

Padding 被明確聲明為一個單獨的階段。通過計算內聯以避免冗餘內存分配：

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

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

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

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)

下圖為測試結果，在 GTX 1080 上 1000 次運行的平均時間成本，並與 depthwise_conv2d in tensorflow進行比較。

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

在一個 cuda 塊中實現 32 x 32 的線程，如下：

num_thread_y 和 num_thread_x 這兩個參數如何調才能得到最優解？在 Filter = [256, 1, 3, 3] and stride = [1, 1]下：

通過測試，團隊得到以下結果：

• 大規模平鋪對於數據重用是有好處的，但不利於本地存儲器讀取。

• num_thread_y 和 num_thread_x 對訪問衝突的影響不同。

• num_thread_y 和 num_thread_x 的最佳組合，需要實現高效共享內存訪問（避免存儲區衝突），數據重用和本地內存讀取的平衡。

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

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

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

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

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

運算符融合是一個典型優化深度學習網路的方法，在 TVM 中，考慮到原有的模式 depthwise_conv2d + scale_shift + relu，可以稍作如下修改：

生成 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) }

}

}

}

}

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 的結果寫入全局內存之前，會計算 scale_shift 和 relu。融合運算符與單個 depthwise_conv2d 一樣快。以下是 Input = [1，256，96，96]，Filter = [256，1,3,3]，stride = [1，1]，padding =『SAME"的結果：

• tf-1.2 depthwise_conv2d：251.6 us

• tf-1.2 depthwise_conv2d + scale_shift + relu（separate）：419.9 us

• TVM depthwise_conv2d：90.9 us

• TVM depthwise_conv2d + scale_shift + relu（fusion）：91.5 us

更多優化代碼可參考以下鏈接：

Declare: https://github.com/dmlc/tvm/blob/master/topi/python/topi/nn/convolution.py

Schedule: https://github.com/dmlc/tvm/blob/master/topi/python/topi/cuda/depthwise_conv2d.py

Test: https://github.com/dmlc/tvm/blob/master/topi/recipe/conv/depthwise_conv2d_test.py

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！