TVM在ARM GPU上優(yōu)化移動(dòng)深度學(xué)習(xí)
隨著深度學(xué)習(xí)的巨大成功，將深度神經(jīng)網(wǎng)絡(luò)部署到移動(dòng)設(shè)備的需求正在迅速增長。與在臺(tái)式機(jī)平臺(tái)上所做的類似，在移動(dòng)設(shè)備中使用GPU可以提高推理速度和能源效率。但是，大多數(shù)現(xiàn)有的深度學(xué)習(xí)框架都不能很好地支持移動(dòng)GPU。困難在于移動(dòng)GPU架構(gòu)和臺(tái)式機(jī)GPU架構(gòu)之間的差異。這意味著在移動(dòng)GPU上進(jìn)行優(yōu)化需要付出特殊的努力。繁瑣的額外工作最終導(dǎo)致大多數(shù)深度學(xué)習(xí)框架中對移動(dòng)GPU的支持不佳。
TVM通過引入統(tǒng)一的IR堆棧解決了部署不同硬件的困難，通過該IR堆?？梢暂p松完成針對不同硬件的優(yōu)化。本文展示了如何使用 TVM / NNVM為ARM Mali GPU生成有效的內(nèi)核并進(jìn)行端到端編譯。在對Mali-T860 MP4的測試中，與Arm Compute Library相比 ，的方法在VGG-16上快1.4倍，在MobileNet上快2.2倍。圖形級和算子級優(yōu)化都有助于加快速度。

ImageNet上不同后端的推理速度圖
MALI中級GPU
使用帶有Mali-T860 MP4的Firefly-RK3399作為的測試環(huán)境，主要關(guān)注下面的Mali T8xx。
建筑學(xué)
圖1是T860和T880上的Mali體系結(jié)構(gòu)的概述。GPU最多可擴(kuò)展到16個(gè)一致的著色器內(nèi)核。在每個(gè)著色器內(nèi)核內(nèi)部，有2或3條算術(shù)管道，1條加載/存儲(chǔ)管道和1條紋理管道（所謂的TriPipe）。每個(gè)算術(shù)流水線中的ALU具有四個(gè)128位向量單元和一個(gè)標(biāo)量單元。
使用OpenCL進(jìn)行GPU計(jì)算。映射到OpenCL模型時(shí)，每個(gè)著色器內(nèi)核將執(zhí)行一個(gè)或幾個(gè)工作組。每個(gè)著色器內(nèi)核最多支持384個(gè)并發(fā)執(zhí)行的線程。OpenCL中的每個(gè)工作項(xiàng)通常都映射到Mali GPU上的單個(gè)線程。Mali GPU使用VLIW（超長指令字）架構(gòu)。每個(gè)指令字包含多個(gè)操作。Mali GPU還使用SIMD，大多數(shù)算術(shù)指令可同時(shí)對多個(gè)數(shù)據(jù)元素進(jìn)行操作。

圖1. Mali T860和T880
與NVIDIA GPU的不同
與為NVIDIA GPU編寫代碼相比，在為Mali GPU編寫OpenCL代碼時(shí)，需要注意一些差異。
? Mali GPU使用統(tǒng)一的全局內(nèi)存。在NVIDIA的GPU中，通常將數(shù)據(jù)復(fù)制到共享內(nèi)存中，因?yàn)镹VIDIA的GPU具有物理上獨(dú)立的全局內(nèi)存，共享內(nèi)存和寄存器。在Mali，此副本不會(huì)提高性能，可以刪除。此外，Mali GPU通常與CPU共享全局內(nèi)存，無需在CPU和GPU之間進(jìn)行復(fù)制。
? Mali Midgrad GPU基于SIMD（單指令多數(shù)據(jù)），并且需要顯式矢量化。在NVIDIA CUDA中，并行性是通過SIMT（單指令多線程）實(shí)現(xiàn)的，而SIMT不需要顯式矢量化。注意，較新的Mali Bitfrost GPU基于四邊形矢量化，不需要顯式矢量化。
? Mali GPU中的所有線程都有單獨(dú)的程序計(jì)數(shù)器。這意味著warp size，因此分支分歧不是主要問題。
優(yōu)化：以卷積為例
卷積層是大多數(shù)深度神經(jīng)網(wǎng)絡(luò)的內(nèi)核，占用大部分計(jì)算時(shí)間。以卷積層為例來說明如何在TVM中應(yīng)用諸如打包，平鋪，展開和矢量化之類的常見優(yōu)化技術(shù)。
Im2Col與GEMM
卷積層的一種著名算法是im2col，將小3D輸入多維數(shù)據(jù)集轉(zhuǎn)換為矩陣的列并執(zhí)行GEMM。方法的優(yōu)點(diǎn)是易于利用高度優(yōu)化的BLAS庫。內(nèi)存冗余（3x3內(nèi)核為9x內(nèi)存）非常糟糕。
空間批處理
相反，采用一種方法來計(jì)算卷積，并逐步應(yīng)用優(yōu)化技術(shù)。VGG-16中的卷積層用作調(diào)整案例，其配置在下面列出。假設(shè)批處理大小為1以便進(jìn)行推斷。

作為基準(zhǔn)，還在Arm Compute庫中列出了該層的性能。
聲明計(jì)算：平鋪和打包
平鋪和打包是旨在更好地訪問內(nèi)存的兩種方法。平鋪將整個(gè)計(jì)算分成小塊，以實(shí)現(xiàn)更好的數(shù)據(jù)重用。打包根據(jù)平鋪對輸入矩陣進(jìn)行重新布局，以便可以順序訪問內(nèi)存，從而降低了緩存未命中率。
對輸入圖像的寬度尺寸和濾鏡矩陣的CO尺寸進(jìn)行平鋪。通過tvm.compute來描述。

set tiling factor

VH = 1
VW = 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 + 2H_PAD - KH) // H_STR + 1
OW = (IW + 2W_PAD - KW) // W_STR + 1

data shape after packing

dvshape = (N, TH // (VHH_STRIDE), TW // (VWW_STRIDE), CI, VHH_STRIDE+HCAT, VWW_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][hVHH_STRIDE+vh][wVWW_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, vhH_STRIDE+kh, vwW_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[((((((((co56) + h)*14) + w)*4) + vw)4) + vc)] = (conv[((((((((co56) + h)*14) + w)*4) + vw)4) + vc)] + (data_vec[(((((((((h14) + w)*256) + ci)*3) + kh)*6) + kw) + vw)]kernel_vec[((((((((co256) + ci)*3) + kh)*3) + kw)*4) + vc)]))
}
}
}
}
}
}
}
}
}
內(nèi)核1：綁定線程
在TVM中，首先聲明計(jì)算，然后調(diào)度。這種機(jī)制使算法和實(shí)現(xiàn)細(xì)節(jié)脫鉤。（這個(gè)想法來自Halide）。
以下調(diào)度僅將軸綁定到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)
有了這個(gè)時(shí)間表，的代碼現(xiàn)在可以運(yùn)行了，但是性能卻很糟糕。

內(nèi)核2：展開unrolling
循環(huán)展開可以減少循環(huán)控制的指令，減少分支懲罰并隱藏讀取內(nèi)存中的延遲。TVM通過調(diào)用以下命令輕松完成此操作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)

內(nèi)核3：矢量化
為了在Mali GPU上實(shí)現(xiàn)最佳性能，需要明確地進(jì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)

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

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)

如何設(shè)置可調(diào)參數(shù)
至于上面的可調(diào)參數(shù)，可以計(jì)算一些。對于矢量化維VC，應(yīng)該填充128位寄存器，因此對于float32可以將其設(shè)置為128/32 = 4，對于float16可以將其設(shè)置為128/16 = 8。
由于運(yùn)行時(shí)間復(fù)雜，常常無法確定最佳值。在TVM中使用網(wǎng)格搜索。在TVM的高級IR中編寫python代碼，不是直接編寫OpenCL代碼，可以非常有效地完成。
生成的OpenCL代碼
可以通過以下方式查看生成的OpenCL代碼：
print(func.imported_modules[0].get_source())
OpenCL代碼太長，無法粘貼到此處，由于展開太重而難以閱讀。
端到端基準(zhǔn)測試
比較一些流行的深度神經(jīng)網(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進(jìn)行端到端編譯。
性能Performance

圖2. ImageNet上不同后端的推理速度
如圖2所示，在ImageNet上測試推理速度。在Firefly-RK3399上，Mali GPU的速度可以比6核big.LITTLE CPU快2倍至4倍。端到端管道比Arm Compute庫快1.4倍至2.2倍。在Arm Compute Library中嘗試了GEMM和卷積層的直接方法，在這些測試案例中，GEMM方法總是比直接方法快，因此僅繪制GEMM方法的結(jié)果。
圖2中缺少一些結(jié)果，例如Arm Compute Library上的resnet18，因?yàn)锳rm Compute Library的圖形運(yùn)行時(shí)當(dāng)前不支持跳過連接，并且深度卷積的霓虹燈實(shí)現(xiàn)較差。這也反映了NNVM軟件堆棧的優(yōu)勢。
半精度性能
深度神經(jīng)網(wǎng)絡(luò)的精度不是很重要，特別是對于移動(dòng)設(shè)備的推斷而言。使用低精度算術(shù)可以使推理更快。還在Mali GPU上測試了半精度浮點(diǎn)數(shù)。

從理論上講，FP16既可以使峰值計(jì)算加倍，又可以使內(nèi)存消耗減半，從而使速度加倍。需要良好的輸入形狀，以實(shí)現(xiàn)更長的矢量化和微調(diào)一些參數(shù)。
在移動(dòng)設(shè)備上的進(jìn)一步工作
還有一些改進(jìn)的空間，主要是在圖形級別，例如模型壓縮和權(quán)重布局。NNVM的進(jìn)一步改進(jìn)將嘗試解決這些問題。

總結(jié)

以上是生活随笔為你收集整理的TVM在ARM GPU上优化移动深度学习的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。