譯者 | 布加迪

審校 | 重樓

Python 速度提升80倍？探究如何用一行代碼將你的代碼變成GPU猛獸！

GPU非常適合處理需要對不同數(shù)據(jù)執(zhí)行相同操作的任務(wù)。這種方法名為單指令多數(shù)據(jù)（SIMD）。與只有幾個強(qiáng)大核心的CPU不同，GPU擁有數(shù)千個較小的核心，它們可以同時運(yùn)行這些重復(fù)性操作。你會在機(jī)器學(xué)習(xí)中經(jīng)?？吹竭@種模式，比如在對大型向量進(jìn)行加法或乘法時，因?yàn)槊總€計(jì)算都是獨(dú)立的。這是使用GPU通過并行機(jī)制加快處理任務(wù)的理想場景。

英偉達(dá)創(chuàng)建了CUDA，以便開發(fā)者編寫在GPU上而不是CPU上運(yùn)行的程序。它基于C語言，允許你編寫名為內(nèi)核的特殊函數(shù)，這些函數(shù)可以同時運(yùn)行多個操作。問題在于，用C 或C++編寫CUDA對初學(xué)者來說并不友好，必須處理諸如手動分配內(nèi)存、協(xié)調(diào)線程以及理解GPU底層工作原理之類的問題。這可能會讓初學(xué)者感到不知所措，尤其是如果習(xí)慣用 Python編寫代碼的話。

這時候Numba可以助你一臂之力。它允許借助Python編寫CUDA內(nèi)核，并使用LLVM（低級虛擬機(jī)）編譯器基礎(chǔ)架構(gòu)，將Python代碼直接編譯成與CUDA兼容的內(nèi)核。借助即時（JIT）編譯，你可以使用裝飾器注釋函數(shù)，其余所有工作交由Numba來處理。

我們在本文中將使用一個向量加法的常見示例，并使用Numba將簡單的CPU代碼轉(zhuǎn)換成CUDA內(nèi)核。向量加法是并行機(jī)制的理想例子，因?yàn)榭鐔蝹€索引的加法與其他索引無關(guān)。這是完美的SIMD場景，因此所有索引可以同時相加，從而在一次運(yùn)算中完成向量加法。

請注意，你需要一個CUDA GPU才能遵循本文操作。你可以使用Colab的免費(fèi)T4 GPU或已安裝英偉達(dá)工具包和NVCC的本地GPU。

搭建環(huán)境并安裝Numba

Numba以Python軟件包的形式提供，你可以使用pip來安裝。此外，我們將使用numpy用于向量運(yùn)算。使用以下命令搭建Python環(huán)境：

python3 -m venv venv
source venv/bin/activate
pip install numba-cuda numpy

CPU上的向量加法

我們舉一個簡單的向量加法例子。對于兩個給定的向量，我們將每個索引對應(yīng)的值相加以獲得最終值。我們將使用numpy生成隨機(jī)的float32向量，并使用for循環(huán)生成最終輸出。

import numpy as np 
N = 10_000_000 # 10 million elements 
a = np.random.rand(N).astype(np.float32) 
b = np.random.rand(N).astype(np.float32) 
c = np.zeros_like(a) # Output array 
def vector_add_cpu(a, b, c): 
 """Add two vectors on CPU""" 
 for i in range(len(a)): 
 c[i] = a[i] + b[i]

代碼分解如下：

• 初始化兩個向量，每個向量包含1000萬個隨機(jī)浮點(diǎn)數(shù)。
• 我們還創(chuàng)建一個空向量c來存儲結(jié)果。
• vector_add_cpu函數(shù)只是循環(huán)遍歷每個索引，并將a和b中的元素相加，將結(jié)果存儲在c中。

這是一個串行操作；每個加法操作都是一個接一個進(jìn)行。雖然這種方法運(yùn)行良好，但它并非最高效的方法，尤其是對大型數(shù)據(jù)集而言。由于每個加法彼此獨(dú)立，因此非常適合在GPU上并行執(zhí)行。

在下一節(jié)中，你將看到如何使用Numba轉(zhuǎn)換這個相同的操作以便在GPU上運(yùn)行。通過將每個元素級加法分布到數(shù)千個GPU線程上，我們可以顯著加快任務(wù)完成速度。

借助Numba在GPU上進(jìn)行向量加法

現(xiàn)在你將使用Numba定義一個可在CUDA上運(yùn)行的Python函數(shù)，并在Python中執(zhí)行它。我們在執(zhí)行相同的向量加法運(yùn)算，但現(xiàn)在它可以針對Numpy數(shù)組的每個索引并行運(yùn)行，從而提高執(zhí)行速度。

以下是編寫內(nèi)核的代碼：

from numba import config
# Required for newer CUDA versions to enable linking tools. 
# Prevents CUDA toolkit and NVCC version mismatches.
config.CUDA_ENABLE_PYNVJITLINK = 1
from numba import cuda, float32
@cuda.jit
def vector_add_gpu(a, b, c):
 """Add two vectors using CUDA kernel"""
 # Thread ID in the current block
 tx = cuda.threadIdx.x
 # Block ID in the grid
 bx = cuda.blockIdx.x
 # Block width (number of threads per block)
 bw = cuda.blockDim.x
 # Calculate the unique thread position
 position = tx + bx * bw
 # Make sure we don't go out of bounds
 if position < len(a):
 c[position] = a[position] + b[position]
def gpu_add(a, b, c):
 # Define the grid and block dimensions
 threads_per_block = 256
 blocks_per_grid = (N + threads_per_block - 1) // threads_per_block
 # Copy data to the device
 d_a = cuda.to_device(a)
 d_b = cuda.to_device(b)
 d_c = cuda.to_device(c)
 # Launch the kernel
 vector_add_gpu[blocks_per_grid, threads_per_block](d_a, d_b, d_c)
 # Copy the result back to the host
 d_c.copy_to_host(c)
def time_gpu():
 c_gpu = np.zeros_like(a)
 gpu_add(a, b, c_gpu)
 return c_gpu

不妨細(xì)述一下上面發(fā)生的操作。

理解GPU函數(shù)

@cuda.jit裝飾器告訴Numba將以下函數(shù)視為CUDA內(nèi)核；這是一個將跨GPU上的多個線程并行運(yùn)行的特殊函數(shù)。在運(yùn)行時，Numba會將此函數(shù)編譯成與CUDA兼容的代碼，并為你處理C-API轉(zhuǎn)譯。

@cuda.jit
defvector_add_gpu(a, b, c):
...

該函數(shù)將同時在數(shù)千個線程上運(yùn)行。但我們需要一種方法來確定每個線程應(yīng)該處理數(shù)據(jù)的哪個部分。這就是接下來幾行代碼的作用：

• tx 是線程在其塊中的 ID。
• bx 是塊在網(wǎng)格中的 ID。
• bw 是塊中有多少個線程。

我們將這些數(shù)據(jù)組合起來計(jì)算出獨(dú)特的位置，該位置告訴每個線程應(yīng)該添加數(shù)組中的哪個元素。請注意，線程和塊可能并不總是提供有效的索引，因?yàn)樗鼈円?/span>2的冪次方進(jìn)行操作。當(dāng)向量長度不符合底層架構(gòu)時，這可能會導(dǎo)致無效索引。因此，我們在執(zhí)行向量加法之前添加了一個保護(hù)條件來驗(yàn)證索引。這可以防止訪問數(shù)組時出現(xiàn)任何越界運(yùn)行時錯誤。

一旦我們知道了這個獨(dú)特位置，現(xiàn)在可以像在CPU實(shí)現(xiàn)中一樣添加值。以下代碼行與CPU實(shí)現(xiàn)一致：

c[position] = a[position] + b[position]

啟動內(nèi)核

gpu_add函數(shù)負(fù)責(zé)進(jìn)行設(shè)置：

• 它定義了要使用的線程和塊的數(shù)量。你可以嘗試塊和線程大小的不同值，并在GPU內(nèi)核中打印輸出相應(yīng)的值。這可以幫助你理解底層GPU索引的工作原理。
• 它將輸入數(shù)組（a、b 和 c）從CPU內(nèi)存復(fù)制到GPU內(nèi)存，以便可以在GPU RAM中訪問這些向量。
• 它使用vector_add_gpu[blocks_per_grid,threads_per_block]運(yùn)行GPU內(nèi)核。
• 最后，它將結(jié)果從GPU復(fù)制回到c數(shù)組，以便我們可以在CPU上訪問這些值。

比較實(shí)現(xiàn)和潛在加速

我們已有了CPU和GPU版本的向量加法，是時候比較一下它們了。驗(yàn)證結(jié)果以及使用CUDA并行機(jī)制所能獲得的執(zhí)行提升非常重要。

import timeit
c_cpu = time_cpu()
c_gpu = time_gpu()
print("Results match:", np.allclose(c_cpu, c_gpu))
cpu_time = timeit.timeit("time_cpu()", globals=globals(), number=3) / 3
print(f"CPU implementation: {cpu_time:.6f} seconds")
gpu_time = timeit.timeit("time_gpu()", globals=globals(), number=3) / 3
print(f"GPU implementation: {gpu_time:.6f} seconds")
speedup = cpu_time / gpu_time
print(f"GPU speedup: {speedup:.2f}x")

首先，我們運(yùn)行兩種實(shí)現(xiàn)，檢查它們的結(jié)果是否一致。這對于確保我們的GPU代碼正常運(yùn)行且輸出結(jié)果與CPU版本一致至關(guān)重要。

接下來，我們使用Python內(nèi)置的timeit模塊來測量每個版本的運(yùn)行時間。我們運(yùn)行每個函數(shù)幾次，取平均值以獲得可靠的時間。最后，我們計(jì)算GPU版本比CPU版本快多少倍。你應(yīng)該會看到顯著的差異，因?yàn)?/span>GPU可以同時執(zhí)行多項(xiàng)操作，而CPU在循環(huán)中一次只處理一項(xiàng)操作。

以下是Colab上英偉達(dá)T4 GPU 的預(yù)期輸出。請注意，具體的加速效果可能因CUDA版本和底層硬件而異。

Results match: True
CPU implementation: 4.033822 seconds
GPU implementation: 0.047736 seconds
GPU speedup: 84.50x

這個簡單的測試有助于展示GPU加速的強(qiáng)大功能，以及它為何對涉及大量數(shù)據(jù)和并行工作的任務(wù)如此有用。

結(jié)語

就是這樣。你現(xiàn)在已經(jīng)借助Numba編寫了第一個CUDA內(nèi)核，無需實(shí)際編寫任何C或CUDA代碼。Numba提供了一個簡單的接口，可以通過Python使用GPU，這使得Python工程師更容易上手CUDA編程。

現(xiàn)在，你可以使用相同的模板來編寫高級CUDA算法，這些算法在機(jī)器學(xué)習(xí)和深度學(xué)習(xí)中非常流行。如果你發(fā)現(xiàn)遵循SIMD范式存在問題，使用GPU來提升執(zhí)行速度始終是個好主意。

完整的代碼可以在Colab筆記本上找到，可以點(diǎn)擊此處訪問。你可以隨意測試并進(jìn)行一些簡單的更改，以更好地理解CUDA索引和執(zhí)行的內(nèi)部工作原理。

原文標(biāo)題：Writing Your First GPU Kernel in Python with Numba and CUDA，作者：Kanwal Mehreen