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

原文地址：https://tvm.ai/2018/01/16/opt-mali-gpu.html

       随着深度学习的巨大成功，对移动设备部署深度神经网络的需求正在迅速增长。与我们在桌面平台上的工作类似，在移动设备中使用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

我们将使用Firefly-RK3399和Mali-T860 MP4作为我们的测试环境，因此我们主要关注下面的Mali T8xx。

架构

图1是T860和T880上Mali Architecture的概述。GPU可扩展至16个相干着色器核心。在每个着色器核心内部，有2或3个算术管道，1个加载/存储管道和1个纹理管道（所谓的TriPipe）。每个算术流水线中的ALU具有四个128位向量单元和一个标量单元。

我们使用OpenCL进行GPU计算。映射到OpenCL模型时，每个着色器核心都会执行一个或多个工作组。每个着色器核心最多支持384个并发执行的线程。OpenCL中的每个工作项通常映射到Mali GPU上的单个线程。Mali GPU使用VLIW（超长指令字）架构。每个指令字包含多个操作。Mali GPU也使用SIMD，因此大多数算术指令同时对多个数据元素进行操作。

与NVIDIA的GPU不同

与为NVIDIA的GPU编写代码相比，在为Mali GPU编写OpenCL代码时，我们应该关注一些差异。

• Mali GPU使用统一的全局内存。在NVIDIA的GPU中，我们通常将数据复制到共享内存，因为NVIDIA的GPU具有物理上独立的全局内存，共享内存和寄存器。在马里，此副本不会提高性能，可以删除。此外，Mali GPU通常与CPU共享全局内存，因此不需要在CPU和GPU之间进行复制。
• Mali Midgrad GPU基于SIMD（单指令多数据）并需要显式矢量化。在NVIDIA CUDA中，并行性是通过SIMT（单指令多线程）实现的，不需要显式矢量化。但另请注意，较新的Mali Bitfrost GPU基于四元矢量化，不需要显式矢量化。
• Mali GPU中的所有线程都有单独的程序计数器。它意味着warp size是1，因此分支差异不是主要问题。

优化：卷积为例

卷积层是大多数深度神经网络的核心，占用了大部分的计算时间。因此，我们以卷积层为例，演示如何在TVM中应用包装，平铺，展开和矢量化等常用优化技术。

Im2Col与GEMM

众所周知的卷积层算法是im2col，它将小的3D输入立方体转换为矩阵的列并执行GEMM。该方法的优点是易于利用高度优化的BLAS库。但是，内存冗余（3x3内核的9倍内存）非常糟糕。

空间包装

相反，我们采用一种方法来计算卷积，并逐步应用优化技术。VGG-16中的卷积层用作调谐情况，其配置如下所示。我们假设批量大小为1用于推断。

input shape	output shape	kernel size	stride	pad
56x56x256	56x56x256	3×3	（1,1）	（1,1）

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

kernel	cost(second)	GFLOPS
ARM Compute Lib中的GEMM方法	0.1821	20.3111

声明计算：平铺和打包(tiling and packing)

平铺和打包是两种旨在更好地访问内存的方法。平铺将整个计算分成小块以获得更好的数据使用。打包根据平铺重新排列输入矩阵，以便我们可以顺序访问内存，从而降低缓存未命中率。

我们对输入图像的宽度尺寸和滤波器矩阵的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')

我们可以检查定义的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	1X
内核1：简单绑定	5.6154	0.6588	0.03X

内核2：展开

循环展开可以减少循环控制的指令，减少分支惩罚并隐藏读取内存中的延迟。在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	1X
内核1：简单绑定	5.6154	0.6588	0.03X
内核2：+展开	0.3707	9.9796	0.49x

Kernel3：矢量化

如前所述，我们需要明确地进行矢量化，以便在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	1X
内核1：简单绑定	5.6154	0.6588	0.03X
内核2：+展开	0.3707	9.9796	0.49x
内核3：+矢量化	0.1304	28.3679	1.40x

如何设置可调参数

至于上面的可调参数，可以计算一些。对于矢量化维度VC，我们应该填充128位寄存器，因此对于float32可以设置为128/32 = 4，对于float16，可以设置为128/16 = 8。

但由于运行时间复杂，我们更常无法确定最佳值。我们在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

我们使用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，在图2中缺失。这是因为Arm Compute Library的图形运行时当前不支持跳过连接，并且深度卷积的霓虹灯实现很差。这也反映了NNVM软件堆栈的优势。

半精度性能

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

模型	后端	每张图像的时间成本（秒）	加速到FP32
vgg16	ACM-马里	0.9694	1.69
vgg16	TVM-马里	0.6896	1.87x
MobileNet 1.0	TVM-马里	0.0479	1.60x
ResNet18	TVM-马里	0.1183	1.73x

表1. ImageNet上FP16的推理速度

 

从理论上讲，FP16可以使峰值计算加倍，并使内存消耗减半，从而使速度加倍。但它需要良好的输入形状，以实现更长的矢量化和微调某些参数。