TVM在ARM GPU上优化移动深度学习
TVM在ARM GPU上優化移動深度學習
隨著深度學習的巨大成功，將深度神經網絡部署到移動設備的需求正在迅速增長。與在臺式機平臺上所做的類似，在移動設備中使用GPU可以提高推理速度和能源效率。但是，大多數現有的深度學習框架都不能很好地支持移動GPU。困難在于移動GPU架構和臺式機GPU架構之間的差異。這意味著在移動GPU上進行優化需要付出特殊的努力。繁瑣的額外工作最終導致大多數深度學習框架中對移動GPU的支持不佳。
TVM通過引入統一的IR堆棧解決了部署不同硬件的困難，通過該IR堆棧可以輕松完成針對不同硬件的優化。本文展示了如何使用 TVM / NNVM為ARM Mali GPU生成有效的內核并進行端到端編譯。在對Mali-T860 MP4的測試中，與Arm Compute Library相比 ，的方法在VGG-16上快1.4倍，在MobileNet上快2.2倍。圖形級和算子級優化都有助于加快速度。
ImageNet上不同后端的推理速度圖
MALI中級GPU
使用帶有Mali-T860 MP4的Firefly-RK3399作為的測試環境，主要關注下面的Mali T8xx。
建筑學
圖1是T860和T880上的Mali體系結構的概述。GPU最多可擴展到16個一致的著色器內核。在每個著色器內核內部，有2或3條算術管道，1條加載/存儲管道和1條紋理管道（所謂的TriPipe）。每個算術流水線中的ALU具有四個128位向量單元和一個標量單元。
使用OpenCL進行GPU計算。映射到OpenCL模型時，每個著色器內核將執行一個或幾個工作組。每個著色器內核最多支持384個并發執行的線程。OpenCL中的每個工作項通常都映射到Mali GPU上的單個線程。Mali GPU使用VLIW（超長指令字）架構。每個指令字包含多個操作。Mali GPU還使用SIMD，大多數算術指令可同時對多個數據元素進行操作。
圖1. Mali T860和T880
與NVIDIA GPU的不同
與為NVIDIA GPU編寫代碼相比，在為Mali GPU編寫OpenCL代碼時，需要注意一些差異。
? Mali GPU使用統一的全局內存。在NVIDIA的GPU中，通常將數據復制到共享內存中，因為NVIDIA的GPU具有物理上獨立的全局內存，共享內存和寄存器。在Mali，此副本不會提高性能，可以刪除。此外，Mali GPU通常與CPU共享全局內存，無需在CPU和GPU之間進行復制。
? Mali Midgrad GPU基于SIMD（單指令多數據），并且需要顯式矢量化。在NVIDIA CUDA中，并行性是通過SIMT（單指令多線程）實現的，而SIMT不需要顯式矢量化。注意，較新的Mali Bitfrost GPU基于四邊形矢量化，不需要顯式矢量化。
? Mali GPU中的所有線程都有單獨的程序計數器。這意味著warp size，因此分支分歧不是主要問題。
優化：以卷積為例
卷積層是大多數深度神經網絡的內核，占用大部分計算時間。以卷積層為例來說明如何在TVM中應用諸如打包，平鋪，展開和矢量化之類的常見優化技術。
Im2Col與GEMM
卷積層的一種著名算法是im2col，將小3D輸入多維數據集轉換為矩陣的列并執行GEMM。方法的優點是易于利用高度優化的BLAS庫。內存冗余（3x3內核為9x內存）非常糟糕。
空間批處理
相反，采用一種方法來計算卷積，并逐步應用優化技術。VGG-16中的卷積層用作調整案例，其配置在下面列出。假設批處理大小為1以便進行推斷。
作為基準，還在Arm Compute庫中列出了該層的性能。
聲明計算：平鋪和打包
平鋪和打包是旨在更好地訪問內存的兩種方法。平鋪將整個計算分成小塊，以實現更好的數據重用。打包根據平鋪對輸入矩陣進行重新布局，以便可以順序訪問內存，從而降低了緩存未命中率。
對輸入圖像的寬度尺寸和濾鏡矩陣的CO尺寸進行平鋪。通過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)]))
}
}
}
}
}
}
}
}
}
內核1：綁定線程
在TVM中，首先聲明計算，然后調度。這種機制使算法和實現細節脫鉤。（這個想法來自Halide）。
以下調度僅將軸綁定到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：展開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 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)
如何設置可調參數
至于上面的可調參數，可以計算一些。對于矢量化維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代碼太長，無法粘貼到此處，由于展開太重而難以閱讀。
端到端基準測試
比較一些流行的深度神經網絡上不同后端之間的綜合性能。測試環境是
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進行端到端編譯。
性能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方法的結果。
圖2中缺少一些結果，例如Arm Compute Library上的resnet18，因為Arm Compute Library的圖形運行時當前不支持跳過連接，并且深度卷積的霓虹燈實現較差。這也反映了NNVM軟件堆棧的優勢。
半精度性能
深度神經網絡的精度不是很重要，特別是對于移動設備的推斷而言。使用低精度算術可以使推理更快。還在Mali GPU上測試了半精度浮點數。
從理論上講，FP16既可以使峰值計算加倍，又可以使內存消耗減半，從而使速度加倍。需要良好的輸入形狀，以實現更長的矢量化和微調一些參數。
在移動設備上的進一步工作
還有一些改進的空間，主要是在圖形級別，例如模型壓縮和權重布局。NNVM的進一步改進將嘗試解決這些問題。
總結
以上是生活随笔為你收集整理的TVM在ARM GPU上优化移动深度学习的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： NNVM AI框架编译器
- 下一篇： TensorFlow+TVM优化NMT神