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可以双峰计算和减半内存消耗，使速度翻倍。需要良好的输入形式，以延长矢量化和微调一些参数。
移动设备的进一步工作
应该承认，还有一些改进的余地，主要是在图形水平，如模型压缩和权重排布。
源代码