TVM 優化 ARM GPU 上的移動深度學習
隨着深度學習的巨大成功，將深度神經網絡部署到移動設備的需求正在迅速增長。與桌面平台上所做的類似，在移動設備中使用 GPU 既有利於推理速度，也有利於能源效率。但是，大多數現有的深度學習框架並不很好地支持移動 GPU。難點在於移動 GPU 架構和桌面 GPU 架構之間的區別。這意味着在移動 GPU 上進行優化需要特別努力。非平凡的額外工作最終導致移動 GPU 在大多數深度學習框架中支持不力。
TVM 通過引入統一的 IR 堆棧，解決為不同硬件部署的困難，從而輕松完成對不同硬件的優化。本文展示了如何使用TVM/NNVM為ARMMaliGPU生成高效的內核，並進行端到端編譯。在Mali-T860 MP4的測試中，與ARM計算庫相比，方法在VGG-16上快1.4倍，在Mobilet上快2.2倍。圖形級別和算子級別優化都有助於加快速度。
Figure1. ImageNet上不同后端的推理速度
Mali Midgrad GPU
將使用螢火飛-RK3399與Mali-T860 MP4作為測試環境，所以主要專注於MaliT8xx。
架構
圖1是T860和T880Mali建築的概述。GPU 可擴展至 16 個連續的着色器內核。每個着色器內核有 2 或 3 條算術管道、1 條負載/存儲管道和 1 條紋理管線（稱為 TriPipe）。每個算術管道中的 ALU 有四個 128 位矢量單元和一個Mali單元。
使用開放CL進行GPU計算。映射到 OpenCL 模型時，每個着色器內核執行一個或多個工作組。每個着色器內核支持多達 384 個同步執行線程。OpenCL 中的每個工作項目通常映射到Mali GPU 上的單個線程。Mali GPU 使用 VLIW（很長的指令字）架構。每個指令字包含多個算子。Mali GPU 還使用 SIMD，以便大多數算術指令同時在多個數據元素上運行。
Figure 2. Mali T860 and T880
與 NVIDIA GPU 差異
以下是我們在為Mali GPU 編寫 OpenCL 代碼時應該關注的一些差異，而為 NVIDIA 的 GPU 編寫這些差異。
- Mali GPU 使用統一的通用內存。在 NVIDIA 的 GPU 中，通常將數據復制到共享內存中，因為 NVIDIA 的 GPU 具有物理上獨立的全局內存、共享內存和注冊。Mali副本不能提高性能，可以刪除。此外，Mali GPU 通常與 CPU 共享全局內存，無需在 CPU 和 GPU 之間復制。
- Mali MidGrad GPU基於SIMD（單一指令多重數據），需要確定的矢量化。在 NVIDIA CUDA 中，並行通過 SIMT（單指令多線程）實現，不需要確定矢量化。注意，較新的Mali Bitfrost GPU基於quad-style vectorization矢量化，不需要明確的矢量化。
- Mali GPU 的所有線程都有單獨的程序計數器。意思是是 1 ，所以
warp size分支發散不是大問題。
Optimization : Convolution as Example卷積為例
卷積層是最深神經網絡的核心，占用了大部分計算時間。以卷積層為例，演示在 TVM 中應用了packing, tiling, unrolling and vectorization等常見優化技術。
Im2Col with GEMM
im2col是卷積層的一個眾所周知的算法，將小3D輸入立方體轉換為矩陣的列，並在GEMM上執行。這種方法的優點是易於利用高度優化的BLAS庫。然而，內存冗余（3x3內核的9倍內存）是可怕的。
Spatial Packing
采用一種計算卷積的方法，逐步應用優化技術。VGG-16 中的卷積層用作調諧tuning case，其配置如下。假設批次大小為1作為推理。
| Input Shape |
Output Shape |
Kernel Size |
Stride |
Padding |
| 56x56x256 |
56x56x256 |
3x3 |
(1, 1) |
(1, 1) |
As a baseline, we also list the performance of this layer in Arm Compute Library.
| Kernel |
Cost (second) |
GFLOPS |
| GEMM method in ARMComputeLib |
0.1821 |
20.3111 |
Declare the computation: tiling and packing
Tiling and packing are two methods intended for better memory access. Tiling separates the whole computation into small blocks for better datareuse. Packing re-layouts the input matrices according to the tiling so that we can access the memory sequentially, which reduces cache miss rate.
平鋪和包裝是兩種用於較好訪問內存的方法。平鋪將整個計算分離成小塊，以便更好的重用數據。Packing根據平鋪重新布局輸入矩陣，以便能夠按順序訪問內存，從而降低緩存誤差率。
根據filter矩陣的輸入圖像和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 + 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')
We can inspect the defined IR by
print(tvm.lower(s, [data, kernel, output], simple_mode=True))
I pick the convolution part here.
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)]))
}
}
}
}
}
}
}
}
}
Kernel 1: bind thread
TVM中，首先聲明計算，然后調度。此機制將算法和實現詳細信息脫鈎。（這個想法來自Halid）。
以下調度表只需將軸與 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)
With this schedule, our code can run now, but the performance is terrible.
| Kernel |
Cost (second) |
GFLOPS |
speedup |
| GEMM method in ARMComputeLib |
0.1821 |
20.3111 |
1x |
| Kernel 1: simple bind |
5.6154 |
0.6588 |
0.03x |
Kernel 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)
| Kernel |
Cost (second) |
GFLOPS |
speedup |
| GEMM method in ARMComputeLib |
0.1821 |
20.3111 |
1x |
| Kernel 1: simple bind |
5.6154 |
0.6588 |
0.03x |
| Kernel 2: + unrolling |
0.3707 |
9.9796 |
0.49x |
Kernel3: vectorization
如前所述，需要進行解釋性向量化，以便在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)
| Kernel |
Cost (second) |
GFLOPS |
speedup |
| GEMM method in ARMComputeLib |
0.1821 |
20.3111 |
1x |
| Kernel 1: simple bind |
5.6154 |
0.6588 |
0.03x |
| Kernel 2: + unrolling |
0.3707 |
9.9796 |
0.49x |
| Kernel 3: + vectorization |
0.1304 |
28.3679 |
1.40x |
如何設置可調參數
至於上面的可調參數，可以計算一些。對於矢量維度，應該填寫128位寄存器，設置為128/32+4，用於VC 中float32和128/16=8用於float16。
更常見的情況是，由於runtime復雜，無法確定最佳值。在TVM中使用網格搜索。可以做到非常有效，在TVM的高水平IR，而不是直接OpenCL代碼中編寫python代碼。
生成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
We use NNVM and TVM to do end-to-end compilation.
Performance
圖3. ImageNet上不同后端的推斷速度
如圖 3 所示，測試 ImageNet上的推理速度。在Firefly-RK3399上，MaliGPU的速度可以是6核大的2倍~4倍，小端方式。端到端管道比ARM計算庫快 1.4 倍~2.2 倍。嘗試在ARM計算庫中同時采用GEMM 和直接卷積層的方法，在這些測試案例中，GEMM方法總是比直接方法快，所以只繪制GEMM 方法的結果。
圖3 中缺少某些結果，如ARM計算庫上的 resnet18。這是因為 Arm 計算庫的圖形runtime目前不支持跳轉連接， 並且具有深度卷積的neon implementation實施不良。這也反映了NNVM軟件堆棧的優勢。
半精度性能
深神經網絡的精度不是很重要，尤其是對於移動設備上的推理。使用低精度算術可以使推理更快。還在Mali GPU 上測試了半精度float。
| 型 |
后端 |
每張圖片的時間成本（秒） |
加速到FP32 |
| vgg16 |
阿姆Mali |
0.9694 |
1.69 |
| vgg16 |
電視 - Mali |
0.6896 |
1.87倍 |
| 移動網 1.0 |
電視 - Mali |
0.0479 |
1.60倍 |
| 雷斯網18 |
電視 - Mali |
0.1183 |
1.73倍 |
表1 . 圖像網上 FP16 的推理速度
從理論上講，FP16可以雙峰計算和減半內存消耗，使速度翻倍。需要良好的輸入形式，以延長矢量化和微調一些參數。
移動設備的進一步工作
應該承認，還有一些改進的余地，主要是在圖形水平，如模型壓縮和權重排布。
源代碼