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

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

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

随着深度学习取得了巨大成功,在移动设备上部署深度学习神经网络模型的需求也在迅速增长。与我们在桌面端平台所做的相类似,在移动设备上使用 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)的优化共同促进了这种加速。

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

在不同底层上测试 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]

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

图1. Mali T860 和 T880(来源[2]

与英伟达 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 中的所有线程都有独立的程序计数器。这意味着 warp 的大小为 1,所以分支发散(Branch divergence)不是一个大问题。

优化:以卷积操作为例

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

使用 GEMM 实现 Im2Col

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

空间填充(Spatial Packing)

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

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

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

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

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

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

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

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

有了这些代码后,我们的代码就可以运行了,但是性能却是非常糟糕的。

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

内核2:展开操作

循环展开(Loop 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)

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

内核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)

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

设置可调参数

至于上面的可调参数,有些可以被计算出来。对于向量化维度 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 来实现端到端编译。

性能

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

图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 上的半精度浮点数。

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

表1. ImageNet 上 FP16 的推理速度

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

在移动设备上的更多工作

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

代码传送门

引用

[1] ARM Mali GPU OpenCL Developer Guide

[2] ARM Developer

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

雷锋网版权文章,未经授权禁止转载。详情见。


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

原创文章,作者:ItWorker,如若转载,请注明出处:https://blog.ytso.com/129298.html

(0)
上一篇 2021年8月30日
下一篇 2021年8月30日

相关推荐

发表回复

登录后才能评论