发掘 ARM GPU 的全部深度学习性能，TVM 优化带来高达2倍性能提升

雷锋网 AI 科技评论按：本文是由来自上海交通大学 Apex 实验室的本科生 Lianmin Zheng 发表于 TVM 的一篇博客，文中阐述了如何使用 TVM 优化移动端上的 ARM GPU 的深度学习。雷锋网 AI 科技评论对原文进行了编译。

随着深度学习取得了巨大成功，在移动设备上部署深度学习神经网络模型的需求也在迅速增长。与我们在桌面端平台所做的相类似，在移动设备上使用 GPU 可以同时实现加速推理计算和节约电能。但是现有的大多数深度学习框架并不能很好地支持移动端 GPU。问题的难点在于移动端 GPU 和桌面端 GPU 存在架构上的差异，这意味着需要投入更多专门的工作来实现移动端 GPU 的优化。正是这些额外的工作最终导致了大多数深度学习框架对移动端 GPU 的支持不足。

TVM 通过引入统一的 IR 栈来解决在不同硬件上的部署难题，通过这个 IR 栈可以轻松完成针对不同硬件的优化。在这篇文章中，我们展示了如何使用TVM/NNVM为 ARM Mali GPU 生成高效的内核，并进行端到端的编译（End-to-end compilation）。在我们基于 Mali-T860 MP4 的测试中，与Arm Compute Library相比，我们的方法在 VGG-16 上快了 1.4 倍，在 MobileNet 上快 2.2 倍。图形级别（Graph-level）和操作级别（Operator-level）的优化共同促进了这种加速。

在不同底层上测试 ImageNet 的推理速度

Mali Midgrad GPU

我们将使用带有 Mali-T860 MP4 的 Firefly-RK3399 作为我们的测试环境，所以我们下面主要关注 Mali T8xx。

架构

图 1 是 T860 和 T880 上的 Mali 架构图。GPU 可扩展到 16 个连通着色器核心（Coherent shader cores）。在每个着色器内核中，有 2 或 3 条运算流水线（Arithmetic pipelines），1 条加载/存储流水线（所谓的 TriPipe）。每个运算流水线中的 ALU 有四个 128 位向量单元和一个标量单元。我们使用 OpenCL 进行 GPU 计算。映射到 OpenCL 模型时，每个着色器核心负责执行一个或多个工作组。并且每个着色器核心最多支持 384 个并发执行的线程。OpenCL 中的每个工作项通常映射到 Mali GPU 上的单个线程。Mali GPU 使用 VLIW（超长指令字，Very Long Instruction Word）架构。每个指令字包含多个操作。Mali GPU 也可以使用 SIMD，因此大多数运算指令会在多个数据元素单元（Multiple data elements）上同时运行。

图1. Mali T860 和 T880（来源）

与英伟达 GPU 相比的不同点

与英伟达 GPU 相比，下面是我们在为 Mali GPU 编写 OpenCL 代码时需要关注的一些区别点。

Mali GPU 使用统一的全局内存。在英伟达的 GPU 中，我们通常会将数据复制到共享内存中，因为英伟达的 GPU 在物理层面上将全局内存、共享内存和寄存器区分开了。在 Mali，这个复制操作并不会提高计算性能，因此可以移除这项操作。另外，Mali GPU 通常与 CPU 共享全局内存，所以 CPU 和 GPU 之间不需要数据的转移复制。

Mali Midgrad GPU 是基于 SIMD（单指令多数据）而设计的，并且需要显性地进行向量化。在英伟达的 CUDA 中，并行性是通过 SIMT（单指令多线程）实现的，不需要显性地进行向量化。但是也要注意，较新的 Mali Bitfrost GPU 是基于四式矢量（Quad-style vectorization），并不需要显性地进行向量化。

Mali GPU 中的所有线程都有独立的程序计数器。这意味着 的大小为 1，所以分支发散（Branch divergence）不是一个大问题。

优化：以卷积操作为例

卷积层是大多数深度神经网络的核心，并且占用了大部分的计算时间。所以我们以卷积为例，说明如何在 TVM 中应用打包（Packing）、平铺（Tiling）、展开（Unrolling）和向量化（Vectorization）等常用技术。

使用 GEMM 实现 Im2Col

众所周知的卷积层算法是 im2col，它的原理是将小的 3D 输入立方体转换成矩阵的列并执行 GEMM 算法。这么做的优点在于，转化为矩阵运算之后可以使用高度优化的 BLAS 库。但是内存冗余问题（3x3 卷积存在 9 倍的内存冗余）也是相当可怕。

空间填充（Spatial Packing）

相反，我们采用另一种方法来计算卷积，并逐步应用一些优化技术。使用 VGG-16 中的卷积层作为微调样例，其配置如下所示。这里我们假设批量的大小为 1。

作为基准，我们还列出了 Arm Compute Library 中该层的性能。

声明计算过程：平铺和打包

平铺（Tiling）和打包（Packing）操作是用于更好地实现内存访问的两种方法。平铺操作将整个计算分成多个小块，以获得更好的数据重用（Data reuse）性能。包装操作则根据平铺重新排列输入矩阵，以便我们可以顺序地访问存储器，从而降低缓存未命中率。

我们在输入图像的宽度维度和滤波器矩阵的 CO 维度上进行平铺操作。这由代码 进行声明。

# set tiling factor

VH = 1VW = 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）

下面的代码简单地将坐标轴（axes）绑定到 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：展开操作

循环展开（Loop unrolling）可以减少循环控制的指令，减少分支惩罚并隐藏内存读取的延迟。在 TVM 中，可以通过调用 来实现。

# 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 parame

ternum_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 con

v_, 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 来实现端到端编译。

性能

图2. 在不同底层上测试 ImageNet 的推理速度

如图2所示，我们在 ImageNet 上测试推理速度。在 Firefly-RK3399 上，Mali GPU 可以比 6 核 big.LITTLE 的 CPU 快 2 至 4 倍。我们的端到端流水线比 Arm Compute Library 快 1.4 至 2.2 倍。在 Arm Compute Library 中，我们尝试使用 GEMM 和直接卷积的方法，在这些测试用例中 GEMM 方法总是比直接方法快，所以我们只绘制了 GEMM 方法的结果。

图中缺失了一些结果，比如 Arm Compute Library 上的 resnet18，这是因为 Arm Compute Library 的图形运行时还暂时不支持跳转连接（Skip connection）操作，并且深度卷积（Depthwise convolution）的实现效果较差。这也反映了 NNVM 软件栈的优势。

半精度性能

深度神经网络的精度不是很重要，特别是对移动设备的推理过程而言。使用低精度算术可以使得推理速度更快。我们还测试了 Mali GPU 上的半精度浮点数。

表1. ImageNet 上 FP16 的推理速度

从理论上讲，FP16 既可以使得峰值计算加倍又可以使得内存开销减半，从而使速度提高一倍。但是对于较长的向量化和调优某些参数，它则需要更好的输入形状（Input shape）。

在移动设备上的更多工作

我们承认还有一些改进空间，它们主要是在图形层面。比如模型压缩和权重预布局。NNVM 的下一步改进将试图解决这些问题。

代码传送门

End-to-End benchmark

Convolution and Depthwise Convolution Schedule

引用

[1]ARM Mali GPU OpenCL Developer Guide

[2]ARM Developer

ViaOptimizing Mobile Deep Learning on ARM GPU with TVM，由雷锋网 AI 科技评论编译。