GPU加速的編程思想,圖解,和經典案例,NVIDIA Python Numba CUDA大法好

0 前言

2018年7月到9月,我做一個項目,Python編程實現。Python程序寫出來了,但是很慢。Python的for loop真是龜速呀。這個程序的瓶頸部分,就是一個雙層for loop,內層for loop裏是矩陣乘法。於是乎想到了numba來給瓶頸部分做優化。簡單的@numba.jit可以加速幾十倍,但是很奇怪無法和joblib配合使用。

最終解決方案是使用@numba.cuda.jit,他可以輕鬆加速數千倍 — 這篇博客就帶你入門GPU編程,本文出了闡述我對於GPU編程的理解和小結,還引用了一些非常好的學習資料。我這裏說的GPU,專門指的是NVIDIA GPU的CUDA編程。


1 GPU 編程思想

傳統意義上來講,大部分程序是運行在CPU上的,GPU只是玩遊戲的時候派上用場。然而現在GPU的重要性大幅度提升,NVIDIA的股票也是高速上漲,因爲GPU除了遊戲,增加了一個killer app:機器學習。另外,比特幣挖礦也是要用GPU的。

  1. 編寫CPU的程序,有兩個特點:單線程思想,面向對象思想。
  2. 編寫GPU的程序,有兩個特點:千線程思想,面向數組思想

1.1 千線程思想

CPU當然也有多線程程序,比如Macbook Pro是雙核四線程,可以同時跑四個線程(有點不準確,不過意思)。但是CPU的多線程和GPU的多線程有兩點本質區別:1)CPU的“多”,規模是十,然而GPU的“多”,規模是;2)CPU多線程之間的並行,是多個function之間的並行,然而GPU多線程的並行,是一個function內部的並行

圖1:本圖和下文很多圖片,來自這個鏈接

進一步解釋第二點,假設一個functionfunction內部是一個雙層for loop i := 0 \cdots 999,程序需要調用四次functionfunction。那麼CPU的程序會同時搞出四個線程,每個線程調用一次function,每次順序執行for loop100021000^2次。而GPU的騷操作是,順序調用四次function,每次搞出100021000^2個線程一下子解決這個雙層for loop。當然了,你需要改寫程序,改爲function_gpufunction\_gpu(GPU裏面可以同時執行幾千的線程,其他線程處於等待狀態,不過你儘管搞出上百萬個線程都沒問題)。當然了,你可以CPU和GPU同時配合使用,CPU搞出四個線程,每個線程調用function_gpufunction\_gpu,這樣子可以增大GPU的利用率。

這裏申明很重要一點,不是所有的functionfunction能(適合)改寫爲function_gpufunction\_gpu。不過很多機器學習相關的算法,很多計算密集行算法,是可以改寫的。會把functionfunction改寫爲function_gpufunction\_gpu是一種當下少數人掌握的技能。在本文chapter 3將講解幾乎是的GPU編程的Hello world。

1.2 面向數組思想

面向對象的編程思想是當下主流思想:一個對象有若干個屬性,一個容器裝很多個對象,如果想獲取對象的屬性,需要獲取對象然後進行點操作。面向數組則完全不同。在做data science(DS) 和 machine learning(ML) 項目的時候,都是面向數組的思想。numpy.ndarray,pandas.DataFrame,和torch.Tensor都是設計的非常棒的”超級數組",其中torch.Tensor原生支持GPU。

在DS和ML項目裏面,數組之所以作爲唯一欽定的數據結構,我認爲是數組能夠完美勝任DS和ML 項目裏面組織和管理數據的工作。DS和ML項目完全不需要object-oriented的那一套封裝和繼承的思想,也不需要鏈表、棧、隊列、哈希表、二叉樹、B-樹、圖等其他數據結構。這背後的原因,我認爲是DS和ML項目和那種企業級開發存在天然的區別。比如企業級開發需要處理複雜的業務邏輯,DS和ML項目就沒有複雜的業務邏輯,只有對數組裏的數據的反覆“暴力計算”,這種對一個數組裏的數據進行反覆的暴力計算,正好是GPU說擅長的東西,SIMD(single instruction multiple data)既視感。

圖2:截屏自這個YouTube視頻。

所以我在使用GPU加速的方法來對程序進行改造的時候,我給類寫了一個 to_array(self) 的方法,爲了把類的非numpy.array數據成員轉換成numpy.array數據成員


2 圖解GPU

圖3:CPU裏面有很大的緩存(黃色),GPU裏面緩存很少。CPU擅長複雜的程序控制(藍色),但是ALU算術運算單元(綠色)比較少。GPU最大的特點就是ALU很多,擅長算數計算。

圖4:GPU加速的方法是,把程序裏面計算密集型的CPU代碼,改寫爲GPU代碼。讓CPU做CPU擅長的事情,讓GPU做GPU擅長的事情,優勢互補。

圖5:GPU是如何和內存和CPU相配合的,分爲4個步驟。其中步驟1和4是很消耗時間的,實際編程的時候,需要考慮如何減少1和4。

圖6:CUDA是這樣子組織上千個線程的,若干線程匯聚爲一個block,若干block匯聚爲一個grid。這就是CUDA的two-level thread hierarchy。深刻理解這個two-level對於編寫CUDA程序十分重要。

圖7:GPU的內存模型。GPU裏面的內存分爲三種:per thread local memory, per block shared memory,和global memory。在實際編程的時候,需要考慮多用shared memory,少用global memory,因爲shared比global的訪問和存取速度快很多。


3 GPU的Hello world: 矩陣相乘

能夠讓人理解GPU編程的 Hello world 程序,便是矩陣相乘這個紐約大學的教程非常棒,詳細講解了如何編寫GPU程序進行矩陣相乘。我當時學習Numba和CUDA,這個教程發揮了很大的作用。

3.1介紹幾個名詞

首先,要弄清楚下面6個名詞的概念。編寫GPU程序,其實是一種CPU和GPU之間的“互動”,所謂的異構開發

  1. Host。代表CPU。
  2. Device。代表GPU。
  3. Host memory。RAM內存。
  4. Device memory。GPU上的存儲。
  5. Kernal function。GPU函數,執行在device上面,調用者是host。
  6. Device function。GPU函數,執行在device上面,調用者是kernal function或者device function。

3.2 GPU程序的執行流程

圖5可視化了GPU程序的流程:

  1. 把數組的數據(numpy.ndarray)從host memory拷貝到device memory上。
  2. 配置kernal function的參數,調用kernal function。參數有兩種:一種用中括號[ ],一種用小括號( )。中括號的參數是threadperblock和blockpergrid,參考圖6。小括號就是那種普通函數的參數。
  3. 幾千個線程同時調用同一個kernal function,在GPU裏面進行計算。kernal function的編寫,是一門技術。
  4. 把GPU裏面的運算結果,從device memory拷貝回host memory。THE END。

3.3 矩陣相乘源代碼

這份Python代碼,有6個函數,進行 C=A×B\mathbf{C} = \mathbf{A} \times \mathbf{B}

  1. cpu_mat_mul(A, B, C), 基於CPU的矩陣相乘,O(n3)O(n^3)
  2. cpu_mat_mul_jit(A, B, C),和cpu_mat_mul相比,唯一的區別就是加了裝飾器 @jit,這是一種編譯優化工具,可以加快幾十倍。
  3. mat_mul_naive_kernal(A, B, C),矩陣相乘的GPU函數,使用global memory。
  4. mat_mul_shared_kernal(A, B, C),矩陣相乘的GPU,使用了shared memory。shared memory比global memory訪問速度快不少。
  5. host_naive(A, B, C),mat_mul_naive_kernal的“配套設施”,kernal不能獨立存在的,前前後後需要一些代碼輔助。
  6. host_optimized(A, B, C),mat_mul_shared_kernal的“配套設施”。如果host_naive比cpu_mat_mul快三千倍的話,host_optimized可能快三千五百倍。

我在學習Numba CUDA的時候,使用的Anaconda Python3.6,Numba 0.38, cudatoolkit 9.0。

'''
Matrix multiplication sample, some numba and CUDA testing code
'''
import math
import time
import numpy as np
from numba import cuda, jit, float64

TPB = 16 # thread per block

def cpu_mat_mul(A, B, C):
    '''matrix mulplication on cpu, O(n^3) implementation
    '''
    for i in range(C.shape[0]):
        for j in range(C.shape[1]):
            summation = 0
            for k in range(A.shape[1]):
                summation += A[i, k] * B[k, j]
            C[i, j] = summation

@jit
def cpu_mat_mul_jit(A, B, C):
    '''matrix mulplication on cpu O(n^3) implementation with @jit decocation
    '''
    for i in range(C.shape[0]):
        for j in range(C.shape[1]):
            summation = 0
            for k in range(A.shape[1]):
                summation += A[i, k] * B[k, j]
            C[i, j] = summation

@cuda.jit
def mat_mul_naive_kernal(A, B, C):
    '''matrix multiplication on gpu, naive method using global device memory
    '''
    i, j = cuda.grid(2)
    if i < C.shape[0] and j < C.shape[1]:
        summation = 0
        for k in range(A.shape[1]):
            summation += A[i, k] * B[k, j]
        C[i, j] = summation

@cuda.jit
def mat_mul_shared_kernal(A, B, C):
    '''matrix multiplication on gpu, optimized version using shared memory.
    '''
    s_A = cuda.shared.array((TPB, TPB), dtype=float64)  # s_ --> shared
    s_B = cuda.shared.array((TPB, TPB), dtype=float64)
    x, y = cuda.grid(2)
    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    bw = cuda.blockDim.x
    bh = cuda.blockDim.y
    #print((x, y), (tx, ty), (bx, by), (bw, bh))

    if x >= C.shape[0] or y >= C.shape[1]:
        return

    tmp = 0
    for i in range(int(A.shape[1]/TPB)):
        #print((x, y), (tx, ty), i)
        s_A[tx, ty] = A[x, ty + bw*i]
        s_B[tx, ty] = B[tx + bh*i, y]
        cuda.syncthreads()

        for j in range(TPB):
            tmp += s_A[tx, j] * s_B[j, ty]

        cuda.syncthreads()
    C[x, y] = tmp


def host_naive(A, B, C):
    '''host code for calling naive kernal
    '''
    d_A = cuda.to_device(A)  # d_ --> device
    d_B = cuda.to_device(B)
    d_C = cuda.device_array(C.shape, np.float64)

    threadsperblock = (TPB, TPB)
    blockspergrid_x = math.ceil(A.shape[0]/threadsperblock[0])
    blockspergrid_y = math.ceil(B.shape[1]/threadsperblock[1])
    blockspergrid = (blockspergrid_x, blockspergrid_y)

    mat_mul_naive_kernal[blockspergrid, threadsperblock](d_A, d_B, d_C)

    return d_C.copy_to_host()


def host_optimized(A, B, C):
    '''host code for calling naive kernal
    '''
    d_A = cuda.to_device(A)  # d_ --> device
    d_B = cuda.to_device(B)
    d_C = cuda.device_array(C.shape, np.float64)

    threadsperblock = (TPB, TPB)
    blockspergrid_x = math.ceil(A.shape[0]/threadsperblock[0])
    blockspergrid_y = math.ceil(B.shape[1]/threadsperblock[1])
    blockspergrid = (blockspergrid_x, blockspergrid_y)

    mat_mul_shared_kernal[blockspergrid, threadsperblock](d_A, d_B, d_C)

    return d_C.copy_to_host()


def main():
    '''main
    '''
    A = np.full((TPB*4, TPB*6), 0.5, dtype=np.float64)
    B = np.full((TPB*6, TPB*2), 2, dtype=np.float64)
    C = np.full((TPB*4, TPB*2), 0, dtype=np.float64)

    start = time.time()
    cpu_mat_mul(A, B, C)
    print('cpu mat mul:', time.time()-start)

    start = time.time()
    cpu_mat_mul_jit(A, B, C)
    print('cpu mat mul with numba.jit:', time.time()-start)

    start = time.time()
    ans = host_naive(A, B, C)
    print('gpu mat mul global:', time.time()-start)
    print(ans)
    
    start = time.time()
    ans = host_optimized(A, B, C)
    print('gpu mat mul shared:', time.time()-start)
    print(ans)

if __name__ == '__main__':
    main()

3.4 理解GPU的矩陣相乘算法的關鍵

C=A×B\mathbf{C} = \mathbf{A} \times \mathbf{B}
假設A.shape = (64, 96), B.shape = (96, 32),那麼C.shape = (64, 32),即C矩陣有2048個元素。

  1. CPU的算法,是1個線程,順序依次計算2048個元素的值。
  2. GPU的算法,是同時開出2048個線程,每個線程計算一個元素的值。線程與線程之間是並行的。

目前最新的英偉達GeForce RTX 2070,可以同時開出2304個線程,2048個線程簡直是輕鬆無壓力。

3.5 作圖:Y軸時間,X軸數組的長度

作圖代碼:GitHub鏈接(有時候GitHub無法加載Jupyter Notebook,此時點這個nbviewer
CPU, GPU 性能作圖

  1. 29=5122^9=512 ,此時 CPU 超過一分鐘了
  2. 211=20482^{11}=2048 ,此時 CPU+numba.jit 超過一分鐘了
  3. 214=163842^{14}=16384 ,此時 GPU 超過一分鐘了

一個(16384, 16384)的數組,在Python裏面已經超過2G的內存了,這樣大的數組在CPU和GPU之間傳輸是需要花不少時間的。

3.6 優化時間

shared memory 充其量是一個小優化,真正的大優化是減少CPU和GPU之間的通信時間(本質上是內存和GPU顯存之間的通信)。

我的親身例子:

  1. GPU裏,一個大數組的每個元素計算好了,把大數組從GPU傳輸到內存,使用CPU求和 numpy.sum()
  2. GPU裏, 一個大數組的每個元素計算好了,大數組原地不動,使用GPU求和(cuda.reduce),給CPU傳輸求和結果(一個float)

方法2 和方法1相比,GPU向CPU的傳輸降低了很多,CPU做求和和GPU做求和速度差不多。最終時間從1降低到0.5,時間降低50%,或者加速1倍。我發現,CPU和GPU之間的通信是非常耗費時間的。要降低CPU和GPU之間的通信,有時候是需要改變一下算法的。


4 Debug

  • 常見Error:與memory相關的報錯。
    • 常用解決方法:在行命令裏面輸入: watch -n 0.5 nvidia-smi, 這樣可以查看GPU的顯存佔用多少,有一些時候是GPU不夠用了。看看有沒有“殭屍”進程之類的,有的話,人工kill掉。
    • 本人之前嘗試在程序中同時使用 joblib 和 CUDA,也就是說CPU並行+GPU。發現比較容易出 memory 方面的 bug。解決方案是不用joblib,單線程操作GPU。

  \

  • GPU kernel 內部的邏輯有bug
    • 加一個環境變量 export NUMBA_ENABLE_CUDASIM=1. 在開發環境中(VS Code, PyCharm)加一個conditional breakpoint(條件斷點),斷點的位置在 i, j = cuda.grid(2) 這一行之後,condition 舉個例子: if i == 0 and j == 0。然後運行程序,觸發斷點,單步調試

5 後記

11/24/2019 是一個值得銘記的日子。今天有一人邀請我給他授課,兩個小時,價格真香(保密)。內容是用Python 做GPU編程。這個人是通過本文找到我的(在此感謝CSDN平臺)

最終授課完美結束,我授課娓娓道來,並且成功解答了對方的所有疑問,對方很滿意。

給我的啓示:好好寫博客,不僅能夠幫助他人,還能給自己經濟上的收入


The EndThe\ End

喜歡記得點贊,收藏,甚至打賞喲~


Last update: 11/24/2019

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章