譯者 | 布加迪
審校 | 重樓
Python 速度提升80倍？探究如何用一行代碼將你的代碼變成GPU猛獸！
GPU非常適合處理需要對不同數(shù)據(jù)執(zhí)行相同操作的任務。這種方法名為單指令多數(shù)據(jù)（SIMD）。與只有幾個強大核心的CPU不同，GPU擁有數(shù)千個較小的核心，它們可以同時運行這些重復性操作。你會在機器學習中經(jīng)常看到這種模式，比如在對大型向量進行加法或乘法時，因為每個計算都是獨立的。這是使用GPU通過并行機制加快處理任務的理想場景。
英偉達創(chuàng)建了CUDA，以便開發(fā)者編寫在GPU上而不是CPU上運行的程序。它基于C語言，允許你編寫名為內(nèi)核的特殊函數(shù)，這些函數(shù)可以同時運行多個操作。問題在于，用C 或C++編寫CUDA對初學者來說并不友好，必須處理諸如手動分配內(nèi)存、協(xié)調(diào)線程以及理解GPU底層工作原理之類的問題。這可能會讓初學者感到不知所措，尤其是如果習慣用 Python編寫代碼的話。
這時候Numba可以助你一臂之力。它允許借助Python編寫CUDA內(nèi)核，并使用LLVM（低級虛擬機）編譯器基礎架構，將Python代碼直接編譯成與CUDA兼容的內(nèi)核。借助即時（JIT）編譯，你可以使用裝飾器注釋函數(shù)，其余所有工作交由Numba來處理。
我們在本文中將使用一個向量加法的常見示例，并使用Numba將簡單的CPU代碼轉換成CUDA內(nèi)核。向量加法是并行機制的理想例子，因為跨單個索引的加法與其他索引無關。這是完美的SIMD場景，因此所有索引可以同時相加，從而在一次運算中完成向量加法。
請注意，你需要一個CUDA GPU才能遵循本文操作。你可以使用Colab的免費T4 GPU或已安裝英偉達工具包和NVCC的本地GPU。
搭建環(huán)境并安裝Numba
Numba以Python軟件包的形式提供，你可以使用pip來安裝。此外，我們將使用numpy用于向量運算。使用以下命令搭建Python環(huán)境：
python3 -m venv venv
source venv/bin/activate
pip install numba-cuda numpyCPU上的向量加法
我們舉一個簡單的向量加法例子。對于兩個給定的向量，我們將每個索引對應的值相加以獲得最終值。我們將使用numpy生成隨機的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萬個隨機浮點數(shù)。
- 我們還創(chuàng)建一個空向量c來存儲結果。
- vector_add_cpu函數(shù)只是循環(huán)遍歷每個索引，并將a和b中的元素相加，將結果存儲在c中。
這是一個串行操作；每個加法操作都是一個接一個進行。雖然這種方法運行良好，但它并非最高效的方法，尤其是對大型數(shù)據(jù)集而言。由于每個加法彼此獨立，因此非常適合在GPU上并行執(zhí)行。
在下一節(jié)中，你將看到如何使用Numba轉換這個相同的操作以便在GPU上運行。通過將每個元素級加法分布到數(shù)千個GPU線程上，我們可以顯著加快任務完成速度。
借助Numba在GPU上進行向量加法
現(xiàn)在你將使用Numba定義一個可在CUDA上運行的Python函數(shù)，并在Python中執(zhí)行它。我們在執(zhí)行相同的向量加法運算，但現(xiàn)在它可以針對Numpy數(shù)組的每個索引并行運行，從而提高執(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不妨細述一下上面發(fā)生的操作。
理解GPU函數(shù)
@cuda.jit裝飾器告訴Numba將以下函數(shù)視為CUDA內(nèi)核；這是一個將跨GPU上的多個線程并行運行的特殊函數(shù)。在運行時，Numba會將此函數(shù)編譯成與CUDA兼容的代碼，并為你處理C-API轉譯。
@cuda.jit
defvector_add_gpu(a, b, c):
...該函數(shù)將同時在數(shù)千個線程上運行。但我們需要一種方法來確定每個線程應該處理數(shù)據(jù)的哪個部分。這就是接下來幾行代碼的作用：
- tx 是線程在其塊中的 ID。
- bx 是塊在網(wǎng)格中的 ID。
- bw 是塊中有多少個線程。
我們將這些數(shù)據(jù)組合起來計算出獨特的位置，該位置告訴每個線程應該添加數(shù)組中的哪個元素。請注意，線程和塊可能并不總是提供有效的索引，因為它們以2的冪次方進行操作。當向量長度不符合底層架構時，這可能會導致無效索引。因此，我們在執(zhí)行向量加法之前添加了一個保護條件來驗證索引。這可以防止訪問數(shù)組時出現(xiàn)任何越界運行時錯誤。
一旦我們知道了這個獨特位置，現(xiàn)在可以像在CPU實現(xiàn)中一樣添加值。以下代碼行與CPU實現(xiàn)一致：
c[position] = a[position] + b[position]啟動內(nèi)核
gpu_add函數(shù)負責進行設置：
- 它定義了要使用的線程和塊的數(shù)量。你可以嘗試塊和線程大小的不同值，并在GPU內(nèi)核中打印輸出相應的值。這可以幫助你理解底層GPU索引的工作原理。
- 它將輸入數(shù)組（a、b 和 c）從CPU內(nèi)存復制到GPU內(nèi)存，以便可以在GPU RAM中訪問這些向量。
- 它使用vector_add_gpu[blocks_per_grid,threads_per_block]運行GPU內(nèi)核。
- 最后，它將結果從GPU復制回到c數(shù)組，以便我們可以在CPU上訪問這些值。
比較實現(xiàn)和潛在加速
我們已有了CPU和GPU版本的向量加法，是時候比較一下它們了。驗證結果以及使用CUDA并行機制所能獲得的執(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")首先，我們運行兩種實現(xiàn)，檢查它們的結果是否一致。這對于確保我們的GPU代碼正常運行且輸出結果與CPU版本一致至關重要。
接下來，我們使用Python內(nèi)置的timeit模塊來測量每個版本的運行時間。我們運行每個函數(shù)幾次，取平均值以獲得可靠的時間。最后，我們計算GPU版本比CPU版本快多少倍。你應該會看到顯著的差異，因為GPU可以同時執(zhí)行多項操作，而CPU在循環(huán)中一次只處理一項操作。
以下是Colab上英偉達T4 GPU 的預期輸出。請注意，具體的加速效果可能因CUDA版本和底層硬件而異。
Results match: True
CPU implementation: 4.033822 seconds
GPU implementation: 0.047736 seconds
GPU speedup: 84.50x這個簡單的測試有助于展示GPU加速的強大功能，以及它為何對涉及大量數(shù)據(jù)和并行工作的任務如此有用。
結語
就是這樣。你現(xiàn)在已經(jīng)借助Numba編寫了第一個CUDA內(nèi)核，無需實際編寫任何C或CUDA代碼。Numba提供了一個簡單的接口，可以通過Python使用GPU，這使得Python工程師更容易上手CUDA編程。
現(xiàn)在，你可以使用相同的模板來編寫高級CUDA算法，這些算法在機器學習和深度學習中非常流行。如果你發(fā)現(xiàn)遵循SIMD范式存在問題，使用GPU來提升執(zhí)行速度始終是個好主意。
完整的代碼可以在Colab筆記本上找到，可以點擊此處訪問。你可以隨意測試并進行一些簡單的更改，以更好地理解CUDA索引和執(zhí)行的內(nèi)部工作原理。
原文標題：Writing Your First GPU Kernel in Python with Numba and CUDA，作者：Kanwal Mehreen
