V2AS问路
TVM 优化 ARM GPU 上的移动深度学习
阅读原文时间：2021年07月16日阅读：1

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

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

移动设备的进一步工作

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

源代码

手机扫一扫

移动阅读更方便

阿里云服务器
腾讯云服务器
七牛云服务器

你可能感兴趣的文章