發掘 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]

圖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 庫。但是內存冗餘問題(3x3 卷積存在 9 倍的內存冗餘)也是相當可怕。

空間填充(Spatial Packing)

相反,我們採用另一種方法來計算卷積,並逐步應用一些優化技術。使用 VGG-16 中的卷積層作為微調樣例,其配置如下所示。這裡我們假設批量的大小為 1。

作為基準,我們還列出了 Arm Compute Library 中該層的性能。

聲明計算過程:平鋪和打包

平鋪(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)

有了這些代碼后,我們的代碼就可以運行了,但是性能卻是非常糟糕的。

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

內核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 的下一步改進將試圖解決這些問題。

代碼傳送門

引用

[1] ARM Mali GPU OpenCL Developer Guide

[2] ARM Developer

Via Optimizing Mobile Deep Learning on ARM GPU with TVM,由雷鋒網 AI 科技評論編譯。


想在手機閱讀更多博客及網站收入資訊?下載【香港矽谷】Android應用
分享到Facebook
技術平台: Nasthon Systems