TVM 优化 ARM GPU 上的移动深度学习

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可以双峰计算和减半内存消耗，使速度翻倍。需要良好的输入形式，以延长矢量化和微调一些参数。

移动设备的进一步工作

应该承认，还有一些改进的余地，主要是在图形水平，如模型压缩和权重排布。

源代码

• 端到端基准
• 卷积和深度卷积时间表