發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

本文作者：隔壁王大喵

編輯：楊曉凡

2018-01-28 20:25

導(dǎo)語：為移動GPU尋找更適合的計(jì)算內(nèi)核

雷鋒網(wǎng) AI 科技評論按：本文是由來自上海交通大學(xué) Apex 實(shí)驗(yàn)室的本科生 Lianmin Zheng 發(fā)表于 TVM 的一篇博客，文中闡述了如何使用 TVM 優(yōu)化移動端上的 ARM GPU 的深度學(xué)習(xí)。雷鋒網(wǎng) AI 科技評論對原文進(jìn)行了編譯。

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

隨著深度學(xué)習(xí)取得了巨大成功，在移動設(shè)備上部署深度學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)模型的需求也在迅速增長。與我們在桌面端平臺所做的相類似，在移動設(shè)備上使用 GPU 可以同時實(shí)現(xiàn)加速推理計(jì)算和節(jié)約電能。但是現(xiàn)有的大多數(shù)深度學(xué)習(xí)框架并不能很好地支持移動端 GPU。問題的難點(diǎn)在于移動端 GPU 和桌面端 GPU 存在架構(gòu)上的差異，這意味著需要投入更多專門的工作來實(shí)現(xiàn)移動端 GPU 的優(yōu)化。正是這些額外的工作最終導(dǎo)致了大多數(shù)深度學(xué)習(xí)框架對移動端 GPU 的支持不足。

TVM 通過引入統(tǒng)一的 IR 棧來解決在不同硬件上的部署難題，通過這個 IR ?？梢暂p松完成針對不同硬件的優(yōu)化。在這篇文章中，我們展示了如何使用 TVM/NNVM 為 ARM Mali GPU 生成高效的內(nèi)核，并進(jìn)行端到端的編譯（End-to-end compilation）。在我們基于 Mali-T860 MP4 的測試中，與 Arm Compute Library 相比，我們的方法在 VGG-16 上快了 1.4 倍，在 MobileNet 上快 2.2 倍。圖形級別（Graph-level）和操作級別（Operator-level）的優(yōu)化共同促進(jìn)了這種加速。

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

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

Mali Midgrad GPU

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

架構(gòu)

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

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

圖1. Mali T860 和 T880（來源^[2]）

與英偉達(dá) GPU 相比的不同點(diǎn)

與英偉達(dá) GPU 相比，下面是我們在為 Mali GPU 編寫 OpenCL 代碼時需要關(guān)注的一些區(qū)別點(diǎn)。

Mali GPU 使用統(tǒng)一的全局內(nèi)存。在英偉達(dá)的 GPU 中，我們通常會將數(shù)據(jù)復(fù)制到共享內(nèi)存中，因?yàn)橛ミ_(dá)的 GPU 在物理層面上將全局內(nèi)存、共享內(nèi)存和寄存器區(qū)分開了。在 Mali，這個復(fù)制操作并不會提高計(jì)算性能，因此可以移除這項(xiàng)操作。另外，Mali GPU 通常與 CPU 共享全局內(nèi)存，所以 CPU 和 GPU 之間不需要數(shù)據(jù)的轉(zhuǎn)移復(fù)制。
Mali Midgrad GPU 是基于 SIMD（單指令多數(shù)據(jù)）而設(shè)計(jì)的，并且需要顯性地進(jìn)行向量化。在英偉達(dá)的 CUDA 中，并行性是通過 SIMT（單指令多線程）實(shí)現(xiàn)的，不需要顯性地進(jìn)行向量化。但是也要注意，較新的 Mali Bitfrost GPU 是基于四式矢量（Quad-style vectorization），并不需要顯性地進(jìn)行向量化。
Mali GPU 中的所有線程都有獨(dú)立的程序計(jì)數(shù)器。這意味著 warp 的大小為 1，所以分支發(fā)散（Branch divergence）不是一個大問題。

優(yōu)化：以卷積操作為例

卷積層是大多數(shù)深度神經(jīng)網(wǎng)絡(luò)的核心，并且占用了大部分的計(jì)算時間。所以我們以卷積為例，說明如何在 TVM 中應(yīng)用打包（Packing）、平鋪（Tiling）、展開（Unrolling）和向量化（Vectorization）等常用技術(shù)。

使用 GEMM 實(shí)現(xiàn) Im2Col

眾所周知的卷積層算法是 im2col，它的原理是將小的 3D 輸入立方體轉(zhuǎn)換成矩陣的列并執(zhí)行 GEMM 算法。這么做的優(yōu)點(diǎn)在于，轉(zhuǎn)化為矩陣運(yùn)算之后可以使用高度優(yōu)化的 BLAS 庫。但是內(nèi)存冗余問題（3x3 卷積存在 9 倍的內(nèi)存冗余）也是相當(dāng)可怕。

空間填充（Spatial Packing）

相反，我們采用另一種方法來計(jì)算卷積，并逐步應(yīng)用一些優(yōu)化技術(shù)。使用 VGG-16 中的卷積層作為微調(diào)樣例，其配置如下所示。這里我們假設(shè)批量的大小為 1。

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

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

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

聲明計(jì)算過程：平鋪和打包

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

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

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

內(nèi)核1：線程綁定

在 TVM 中，我們首先聲明計(jì)算，然后進(jìn)行規(guī)劃。該機(jī)制可以將算法和實(shí)現(xiàn)細(xì)節(jié)進(jìn)行分離。（這個想法來自于 Halide）

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

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

有了這些代碼后，我們的代碼就可以運(yùn)行了，但是性能卻是非常糟糕的。

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

內(nèi)核2：展開操作

循環(huán)展開（Loop unrolling）可以減少循環(huán)控制的指令，減少分支懲罰并隱藏內(nèi)存讀取的延遲。在 TVM 中，可以通過調(diào)用 s.unroll(axis) 來實(shí)現(xiàn)。

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

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

內(nèi)核3：向量化

如前所述，為了在 Mali GPU 上實(shí)現(xiàn)最佳性能，我們需要顯性地進(jìn)行向量化。

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

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

設(shè)置可調(diào)參數(shù)

至于上面的可調(diào)參數(shù)，有些可以被計(jì)算出來。對于向量化維度 VC，我們應(yīng)該填充 128 位寄存器，所以 float32 可以設(shè)置為 128/32 = 4，float16 設(shè)置為 128/16 = 8。

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

生成 OpenCL 代碼

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

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

由于 OpenCL 代碼太長，無法在這里粘貼，而由于做了大量的展開，也很難以閱讀。如果你們感興趣可以到這里查看。

端到端的基準(zhǔn)測試

在本節(jié)中，我們將采用一些比較流行的深度學(xué)習(xí)網(wǎng)絡(luò)，用來測試不同底層間的性能差異。我們的測試環(huán)境是：

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 來實(shí)現(xiàn)端到端編譯。

性能

發(fā)掘 ARM GPU 的全部深度學(xué)習(xí)性能，TVM 優(yōu)化帶來高達(dá)2倍性能提升

圖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 方法的結(jié)果。

圖中缺失了一些結(jié)果，比如 Arm Compute Library 上的 resnet18，這是因?yàn)?Arm Compute Library 的圖形運(yùn)行時還暫時不支持跳轉(zhuǎn)連接（Skip connection）操作，并且深度卷積（Depthwise convolution）的實(shí)現(xiàn)效果較差。這也反映了 NNVM 軟件棧的優(yōu)勢。