使用 TVM 优化 ARM GPU 上的移动端深度学习
随着深度学习的巨大成功，将深度神经网络部署到移动设备的需求正在快速增长。与我们在桌面平台上的做法类似，利用移动设备中的 GPU 可以提高推理速度和能源效率。然而，大多数现有的深度学习框架对移动 GPU 的支持都不是很好。困难在于移动 GPU 架构和桌面 GPU 架构之间的差异。这意味着需要在移动 GPU 上进行专门的优化工作。这种额外的非平凡工作最终导致大多数深度学习框架对移动 GPU 的支持不佳。
TVM 通过引入统一的 IR 堆栈解决了为不同硬件部署的难题，通过该堆栈可以轻松完成针对不同硬件的优化。在这篇文章中，我们将展示如何使用 TVM/NNVM 为 ARM Mali GPU 生成高效的内核并进行端到端编译。在我们于 Mali-T860 MP4 上的测试中，与 Arm Compute Library 相比，我们的方法在 VGG-16 上快 1.4 倍，在 MobileNet 上快 2.2 倍。图级和算子级优化都有助于这种加速。
Mali Midgrad 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]
与 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（单指令多线程）实现的，这不需要显式向量化。但也要注意，较新的 Mali Bitfrost GPU 基于四路式向量化，不需要显式向量化。
- Mali GPU 中的所有线程都有单独的程序计数器。这意味着
warp size为 1，因此分支发散不是主要问题。
优化：以卷积为例
卷积层是大多数深度神经网络的核心，并且占据了大部分计算时间。因此，我们以卷积层为例，演示如何在 TVM 中应用常见的优化技术，如 packing（打包）、tiling（平铺）、unrolling（展开）和 vectorization（向量化）。
Im2Col 与 GEMM
卷积层的一种众所周知的算法是 im2col，它将小的 3D 输入立方体转换为矩阵列并执行 GEMM。这种方法的优点是易于利用高度优化的 BLAS 库。然而，内存冗余（对于 3x3 内核，内存增加 9 倍）非常糟糕。
空间 Packing（打包）
相反，我们采用一种方法来计算卷积，并逐步应用优化技术。VGG-16 中的卷积层用作调优案例，其配置如下所示。我们假设批量大小为 1 用于推理。
| 输入形状 | 输出形状 | 内核大小 | 步幅 | 填充 |
|---|---|---|---|---|
| 56x56x256 | 56x56x256 | 3x3 | (1, 1) | (1, 1) |
作为基线，我们还列出了此层在 Arm Compute Library 中的性能。
| 内核 | 成本（秒） | GFLOPS |
|---|---|---|
| ARMComputeLib 中的 GEMM 方法 | 0.1821 | 20.3111 |
声明计算：tiling（平铺）和 packing（打包）
Tiling（平铺）和 packing（打包）是旨在改善内存访问的两种方法。Tiling（平铺）将整个计算分成小块，以获得更好的数据重用。Packing（打包）根据 tiling（平铺）重新布局输入矩阵，以便我们可以顺序访问内存，从而降低缓存未命中率。
我们在输入图像的宽度维度和滤波器矩阵的 CO 维度上进行 tiling（平铺）。这由 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')
我们可以通过以下方式检查定义的 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)]))
}
}
}
}
}
}
}
}
}
内核 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)
通过此调度，我们的代码现在可以运行，但性能很差。
| 内核 | 成本（秒） | GFLOPS | 加速 |
|---|---|---|---|
| ARMComputeLib 中的 GEMM 方法 | 0.1821 | 20.3111 | 1 倍 |
| 内核 1：简单绑定 | 5.6154 | 0.6588 | 0.03 倍 |
内核 2：unrolling（展开）
循环 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)
| 内核 | 成本（秒） | GFLOPS | 加速 |
|---|---|---|---|
| ARMComputeLib 中的 GEMM 方法 | 0.1821 | 20.3111 | 1 倍 |
| 内核 1：简单绑定 | 5.6154 | 0.6588 | 0.03 倍 |
| 内核 2：+ unrolling（展开） | 0.3707 | 9.9796 | 0.49 倍 |
内核 3：vectorization（向量化）
如前所述，我们需要显式地进行 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)
| 内核 | 成本（秒） | GFLOPS | 加速 |
|---|---|---|---|
| ARMComputeLib 中的 GEMM 方法 | 0.1821 | 20.3111 | 1 倍 |
| 内核 1：简单绑定 | 5.6154 | 0.6588 | 0.03 倍 |
| 内核 2：+ unrolling（展开） | 0.3707 | 9.9796 | 0.49 倍 |
| 内核 3：+ vectorization（向量化） | 0.1304 | 28.3679 | 1.40 倍 |
如何设置可调参数
对于上面的可调参数，有些是可以计算的。对于向量化维度 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 代码太长，无法在此处粘贴，并且由于大量的 unrolling（展开），它很难阅读。如果您有兴趣，可以在 此处 查看。
端到端基准测试
在本节中，我们将比较一些流行的深度神经网络上不同后端之间的综合性能。我们的测试环境是
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 进行端到端编译。
性能
如图 2 所示，我们在 ImageNet 上测试了推理速度。在 Firefly-RK3399 上，Mali GPU 可以比 6 核 big.LITTLE CPU 快 2 倍到 4 倍。我们的端到端 pipeline 比 Arm Compute Library 快 1.4 倍到 2.2 倍。我们尝试了 Arm Compute Library 中卷积层的 GEMM 方法和直接方法，在这些测试用例中，GEMM 方法始终比直接方法更快，因此我们仅绘制了 GEMM 方法的结果。
图 2 中缺少一些结果，例如 Arm Compute Library 上的 resnet18。这是因为 Arm Compute Library 的图运行时目前不支持 skip connection，并且深度卷积的 neon 实现效果不佳。这也反映了 NNVM 软件栈的优势。
半精度性能
深度神经网络中的精度不是很重要，特别是对于移动设备上的推理。使用低精度算术可以使推理速度更快。我们还在 Mali GPU 上测试了半精度浮点数。
| 模型 | 后端 | 每张图像的时间成本（秒） | 相对于 FP32 的加速 |
|---|---|---|---|
| vgg16 | ACM-mali | 0.9694 | 1.69 |
| vgg16 | TVM-mali | 0.6896 | 1.87 倍 |
| MobileNet 1.0 | TVM-mali | 0.0479 | 1.60 倍 |
| ResNet18 | TVM-mali | 0.1183 | 1.73 倍 |
理论上，FP16 可以使峰值计算能力加倍，内存消耗减半，从而使速度加倍。但这需要良好的输入形状以实现更长的向量化，并微调一些参数。
移动设备上的进一步工作
我们应该承认，仍然有一些改进空间，主要是在图级别，例如模型压缩和权重预布局。NNVM 的进一步改进将尝试解决这些问题。
代码展示
作者简介 & 致谢
Lianmin Zheng 是上海交通大学 Apex 实验室的本科生。他对机器学习和构建计算机系统感兴趣。
作者非常感谢 Tianqi Chen 的有益建议和 Yizhi Liu 的早期工作。