當前位置：
首頁 > 知識 > 發掘 ARM GPU 的全部深度學習性能，TVM 優化帶來高達 2 倍性能提升

發掘 ARM GPU 的全部深度學習性能，TVM 優化帶來高達 2 倍性能提升

知識 01-27

隨著深度學習取得了巨大成功，在移動設備上部署深度學習神經網路模型的需求也在迅速增長。與我們在桌面端平台所做的相類似，在移動設備上使用 GPU 可以同時實現加速推理計算和節約電能。但是現有的大多數深度學習框架並不能很好地支持移動端 GPU。問題的難點在於移動端 GPU 和桌面端 GPU 存在架構上的差異，這意味著需要投入更多專門的工作來實現移動端 GPU 的優化。正是這些額外的工作最終導致了大多數深度學習框架對移動端 GPU 的支持不足。

TVM 通過引入統一的 IR 棧來解決在不同硬體上的部署難題，通過這個 IR 棧可以輕鬆完成針對不同硬體的優化。在這篇文章中，我們展示了如何使用 TVM（http://t.cn/RC2VOMM）/NNVM（http://t.cn/RHHUkzw） 為 ARM Mali GPU 生成高效的內核，並進行端到端的編譯（End-to-end compilation）。在我們基於 Mali-T860 MP4 的測試中，與 Arm Compute Library（http://t.cn/R664UgG） 相比，我們的方法在 VGG-16 上快了 1.4 倍，在 MobileNet 上快 2.2 倍。圖形級別（Graph-level）和操作級別（Operator-level）的優化共同促進了這種加速。

在不同底層上測試 ImageNet 的推理速度

Mali Midgrad GPU

我們將使用帶有 Mali-T860 MP4 的 Firefly-RK3399 作為我們的測試環境，所以我們下面主要關注 Mali T8xx。

架構

圖 1 是 T860 和 T880 上的 Mali 架構圖。GPU 可擴展到 16 個連通著色器核心（Coherent shader cores）。在每個著色器內核中，有 2 或 3 條運算流水線（Arithmetic pipelines），1 條載入 / 存儲流水線（所謂的 TriPipe）。每個運算流水線中的 ALU 有四個 128 位向量單元和一個標量單元。我們使用 OpenCL 進行 GPU 計算。映射到 OpenCL 模型時，每個著色器核心負責執行一個或多個工作組。並且每個著色器核心最多支持 384 個並發執行的線程。OpenCL 中的每個工作項通常映射到 Mali GPU 上的單個線程。Mali GPU 使用 VLIW（超長指令字，Very Long Instruction Word）架構。每個指令字包含多個操作。Mali GPU 也可以使用 SIMD，因此大多數運算指令會在多個數據元素單元（Multiple data elements）上同時運行。[1]

圖 1. Mali T860 和 T880（來源

[2]）

與英偉達 GPU 相比的不同點

與英偉達 GPU 相比，下面是我們在為 Mali GPU 編寫 OpenCL 代碼時需要關注的一些區別點。

Mali GPU 使用統一的全局內存。在英偉達的 GPU 中，我們通常會將數據複製到共享內存中，因為英偉達的 GPU 在物理層面上將全局內存、共享內存和寄存器區分開了。在 Mali，這個複製操作並不會提高計算性能，因此可以移除這項操作。另外，Mali GPU 通常與 CPU 共享全局內存，所以 CPU 和 GPU 之間不需要數據的轉移複製。

Mali Midgrad GPU 是基於 SIMD（單指令多數據）而設計的，並且需要顯性地進行向量化。在英偉達的 CUDA 中，並行性是通過 SIMT（單指令多線程）實現的，不需要顯性地進行向量化。但是也要注意，較新的 Mali Bitfrost GPU 是基於四式矢量（Quad-style vectorization），並不需要顯性地進行向量化。

Mali GPU 中的所有線程都有獨立的程序計數器。這意味著warp的大小為 1，所以分支發散（Branch divergence）不是一個大問題。


優化：以卷積操作為例

卷積層是大多數深度神經網路的核心，並且佔用了大部分的計算時間。所以我們以卷積為例，說明如何在 TVM 中應用打包（Packing）、平鋪（Tiling）、展開（Unrolling）和向量化（Vectorization）等常用技術。

使用 GEMM 實現 Im2Col

眾所周知的卷積層演算法是 im2col，它的原理是將小的 3D 輸入立方體轉換成矩陣的列並執行 GEMM 演算法。這麼做的優點在於，轉化為矩陣運算之後可以使用高度優化的 BLAS 庫。但是內存冗餘問題（3x3 卷積存在 9 倍的內存冗餘）也是相當可怕。

空間填充（Spatial Packing）

相反，我們採用另一種方法來計算卷積，並逐步應用一些優化技術。使用 VGG-16 中的卷積層作為微調樣例，其配置如下所示。這裡我們假設批量的大小為 1。

作為基準，我們還列出了 Arm Compute Library 中該層的性能。

聲明計算過程：平鋪和打包

平鋪（Tiling）和打包（Packing）操作是用於更好地實現內存訪問的兩種方法。平鋪操作將整個計算分成多個小塊，以獲得更好的數據重用（Data reuse）性能。包裝操作則根據平鋪重新排列輸入矩陣，以便我們可以順序地訪問存儲器，從而降低緩存未命中率。

我們在輸入圖像的寬度維度和濾波器矩陣的 CO 維度上進行平鋪操作。這由代碼tvm.compute進行聲明。

# set tiling factor

VH = 1VW = VC = 4

# get input shape

_, CI, IH, IW = data.shape

CO, CI, KH, KW = kernel.shape

TH = IH + 2 * H_PAD

TW = IW + 2 * W_PAD

# calc output shape

OH = (IH + 2*H_PAD - KH) // H_STR + 1

OW = (IW + 2*W_PAD - KW) // W_STR + 1

# data shape after packing

dvshape = (N, TH // (VH*H_STRIDE), TW // (VW*W_STRIDE), CI, VH*H_STRIDE+HCAT, VW*W_STRIDE+WCAT)

# kernel shape after packing

kvshape = (CO // VC, CI, KH, KW, VC)

ovshape = (N, CO // VC, OH // VH, OW // VW, VH, VW, VC)

oshape = (N, CO, OH, OW)

# define packing

data_vec = tvm.compute(dvshape, lambda n, h, w, ci, vh, vw:

data_pad[n][ci][h*VH*H_STRIDE+vh][w*VW*W_STRIDE+vw], name="data_vec")

kernel_vec = tvm.compute(kvshape, lambda co, ci, kh, kw, vc:

kernel[co*VC+vc][ci][kh][kw], name="kernel_vec")

# define convolution

ci = tvm.reduce_axis((0, CI), name="ci")

kh = tvm.reduce_axis((0, KH), name="kh")

kw = tvm.reduce_axis((0, KW), name="kw")

conv = tvm.compute(ovshape, lambda n, co, h, w, vh, vw, vc:

tvm.sum(data_vec[n, h, w, ci, vh*H_STRIDE+kh, vw*W_STRIDE+kw].astype(out_dtype) *

kernel_vec[co, ci, kh, kw, vc].astype(out_dtype),

axis=[ci, kh, kw]), name="conv")

# unpack to correct layout

output = tvm.compute(oshape, lambda n, co, h, w:

conv[n][co//VC][h/VH][w//VW][h%VH][w%VW][co%VC],

name="output_unpack", tag="direct_conv_output")

我們可以通過以下代碼查看定義的 IR。

print(tvm.lower(s, [data, kernel, output], simple_mode=True))

我在這裡選了卷積部分。

produce conv {

for (co, 0, 64) {

for (h, 0, 56) {

for (w, 0, 14) {

for (vw.init, 0, 4) {

for (vc.init, 0, 4) {

conv[((((((((co*56) + h)*14) + w)*4) + vw.init)*4) + vc.init)] = 0.000000f

}

}

for (ci, 0, 256) {

for (kh, 0, 3) {

for (kw, 0, 3) {

for (vw, 0, 4) {

for (vc, 0, 4) {

conv[((((((((co*56) + h)*14) + w)*4) + vw)*4) + vc)] = (conv[((((((((co*56) + h)*14) + w)*4) + vw)*4) + vc)] + (data_vec[(((((((((h*14) + w)*256) + ci)*3) + kh)*6) + kw) + vw)]*kernel_vec[((((((((co*256) + ci)*3) + kh)*3) + kw)*4) + vc)]))

內核 1：線程綁定

在 TVM 中，我們首先聲明計算，然後進行規劃。該機制可以將演算法和實現細節進行分離。（這個想法來自於 Halide：http://halide-lang.org/）

下面的代碼簡單地將坐標軸（axes）綁定到 GPU 線程，以便我們的代碼可以在 Mali GPU 上運行。

# helper function for binding thread

def tile_and_bind3d(s, tensor, z, y, x, z_factor=2, y_factor=None, x_factor=None):

"""tile and bind 3d"""

y_factor = y_factor or z_factor

x_factor = x_factor or y_factor

zo, zi = s[tensor].split(z, z_factor)

yo, yi = s[tensor].split(y, y_factor)

xo, xi = s[tensor].split(x, x_factor)

s[tensor].bind(zo, tvm.thread_axis("blockIdx.z"))

s[tensor].bind(zi, tvm.thread_axis("threadIdx.z"))

s[tensor].bind(yo, tvm.thread_axis("blockIdx.y"))

s[tensor].bind(yi, tvm.thread_axis("threadIdx.y"))

s[tensor].bind(xo, tvm.thread_axis("blockIdx.x"))

s[tensor].bind(xi, tvm.thread_axis("threadIdx.x"))

# set tunable parameter

num_thread = 8

# schedule data packing

_, h, w, ci, vh, vw = s[data_vec].op.axis

tile_and_bind3d(s, data_vec, h, w, ci, 1)

# schedule kernel packing

co, ci, kh, kw, vc = s[kernel_vec].op.axis

tile_and_bind(s, kernel_vec, co, ci, 1)

# schedule conv

_, c, h, w, vh, vw, vc = s[conv].op.axis

kc, kh, kw = s[conv].op.reduce_axis

s[conv].reorder(_, c, h, w, vh, kc, kh, kw, vw, vc)

tile_and_bind3d(s, conv, c, h, w, num_thread, 1, 1)

_, co, oh, ow = s[output].op.axis

tile_and_bind3d(s, output, co, oh, ow, num_thread, 1, 1)

有了這些代碼後，我們的代碼就可以運行了，但是性能卻是非常糟糕的。

內核 2：展開操作

循環展開（Loop unrolling）可以減少循環控制的指令，減少分支懲罰並隱藏內存讀取的延遲。在 TVM 中，可以通過調用s.unroll(axis)來實現。

# set tunable parameter

num_thread = 8

# schedule data packing

_, h, w, ci, vh, vw = s[data_vec].op.axis

tile_and_bind3d(s, data_vec, h, w, ci, 1)

"""!! ADD UNROLL HERE !!"""

s[data_vec].unroll(vw)

# schedule kernel packing

co, ci, kh, kw, vc = s[kernel_vec].op.axis

tile_and_bind(s, kernel_vec, co, ci, 1)

"""!! ADD UNROLL HERE !!"""

s[kernel_vec].unroll(kh)

s[kernel_vec].unroll(kw)

s[kernel_vec].unroll(vc)

# schedule conv

_, c, h, w, vh, vw, vc = s[conv].op.axis

kc, kh, kw = s[conv].op.reduce_axis

s[conv].reorder(_, c, h, w, vh, kc, kh, kw, vw, vc)

tile_and_bind3d(s, conv, c, h, w, num_thread, 1, 1)

"""!! ADD UNROLL HERE !!"""

s[conv].unroll(kh)

s[conv].unroll(kw)

s[conv].unroll(vw)

s[conv].unroll(vc)

_, co, oh, ow = s[output].op.axis

tile_and_bind3d(s, output, co, oh, ow, num_thread, 1, 1)

內核 3：向量化

如前所述，為了在 Mali GPU 上實現最佳性能，我們需要顯性地進行向量化。

# set tunable parame

ternum_thread = 8

# schedule data packing

_, h, w, ci, vh, vw = s[data_vec].op.axis

tile_and_bind3d(s, data_vec, h, w, ci, 1)

# unroll

s[data_vec].unroll(vw)

# schedule kernel packing

co, ci, kh, kw, vc = s[kernel_vec].op.axis

tile_and_bind(s, kernel_vec, co, ci, 1)

# unroll

s[kernel_vec].unroll(kh)

s[kernel_vec].unroll(kw)

"""!! VECTORIZE HERE !!"""

s[kernel_vec].vectorize(vc)

# schedule con

v_, c, h, w, vh, vw, vc = s[conv].op.axis

kc, kh, kw = s[conv].op.reduce_axis

s[conv].reorder(_, c, h, w, vh, kc, kh, kw, vw, vc)

tile_and_bind3d(s, conv, c, h, w, num_thread, 1, 1)

# unroll

s[conv].unroll(kh)

s[conv].unroll(kw)

s[conv].unroll(vw)

"""!! VECTORIZE HERE !!"""

s[conv].vectorize(vc)

_, co, oh, ow = s[output].op.axis

tile_and_bind3d(s, output, co, oh, ow, num_thread, 1, 1)

設置可調參數

至於上面的可調參數，有些可以被計算出來。對於向量化維度 VC，我們應該填充 128 位寄存器，所以 float32 可以設置為 128/32 = 4，float16 設置為 128/16 = 8。

但是由於運行過於複雜，我們很難去確定最佳超參數值。因此我們在 TVM 中使用網格搜索。由於我們在 TVM 的高級 IR 中編寫了 python 代碼，而不是直接使用 OpenCL 代碼，所以它可以做得非常有效。

生成 OpenCL 代碼

我們可以通過以下代碼，看到所生成的 OpenCL 代碼。

print(func.imported_modules[0].get_source())

由於 OpenCL 代碼太長，無法在這裡粘貼，而由於做了大量的展開，也很難以閱讀。如果你們感興趣可以到這裡（http://t.cn/RQsLkiD）查看。


端到端的基準測試

在本節中，我們將採用一些比較流行的深度學習網路，用來測試不同底層間的性能差異。我們的測試環境是：

Firefly-RK3399 4G

CPU: dual-core Cortex-A72 + quad-core Cortex-A53

GPU: Mali-T860MP4

Arm Compute Library : v17.12

MXNet: v1.0.1

Openblas: v0.2.18

我們使用 NNVM 和 TVM 來實現端到端編譯。


性能

圖 2. 在不同底層上測試 ImageNet 的推理速度

如圖 2 所示，我們在 ImageNet 上測試推理速度。在 Firefly-RK3399 上，Mali GPU 可以比 6 核 big.LITTLE 的 CPU 快 2 至 4 倍。我們的端到端流水線比 Arm Compute Library 快 1.4 至 2.2 倍。在 Arm Compute Library 中，我們嘗試使用 GEMM 和直接卷積的方法，在這些測試用例中 GEMM 方法總是比直接方法快，所以我們只繪製了 GEMM 方法的結果。

圖中缺失了一些結果，比如 Arm Compute Library 上的 resnet18，這是因為 Arm Compute Library 的圖形運行時還暫時不支持跳轉連接（Skip connection）操作，並且深度卷積（Depthwise convolution）的實現效果較差。這也反映了 NNVM 軟體棧的優勢。


半精度性能

深度神經網路的精度不是很重要，特別是對移動設備的推理過程而言。使用低精度算術可以使得推理速度更快。我們還測試了 Mali GPU 上的半精度浮點數。

表 1. ImageNet 上 FP16 的推理速度

從理論上講，FP16 既可以使得峰值計算加倍又可以使得內存開銷減半，從而使速度提高一倍。但是對於較長的向量化和調優某些參數，它則需要更好的輸入形狀（Input shape）。


在移動設備上的更多工作

我們承認還有一些改進空間，它們主要是在圖形層面。比如模型壓縮和權重預布局。NNVM 的下一步改進將試圖解決這些問題。


代碼傳送門

End-to-End benchmark（http://t.cn/RQSJQjB）

Convolution and Depthwise Convolution Schedule（http://t.cn/RQsy07B）


引用

[1] ARM Mali GPU OpenCL Developer Guide（http://t.cn/RQsyjYE）

[2] ARM Developer（https://developer.arm.com/）

NLP 工程師入門實踐班：基於深度學習的自然語言處理

三大模塊，五大應用，手把手快速入門 NLP

海外博士講師，豐富項目經驗

演算法 + 實踐，搭配典型行業應用

隨到隨學，專業社群，講師在線答疑

新人福利

關注 AI 研習社（okweiwu），回復1領取

【超過 1000G 神經網路 / AI / 大數據，教程，論文】

深度學習中如何選擇一款合適的 GPU 卡的一些經驗和建議分享


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

本站內容充實豐富，博大精深，小編精選每日熱門資訊，隨時更新，點擊「搶先收到最新資訊」瀏覽吧！


請您繼續閱讀更多來自 AI研習社 的精彩文章:

※不必再費心尋找，2017最全的開發乾貨就在這1067頁PDF里

TAG:AI研習社 |