CUDA Python 入門 (Intro to CUDA Python)

重點總覽

本筆記是 C++ 版（2.1 節）的 Python 對應。CUDA programming model 的概念兩種語言相同，差別主要在語法與工具鏈。以下每節都會對照 C++ 寫法凸顯 Python 的等價作法。

CUDA Python Ecosystem

CUDA Python 是一整套讓 Python 進行 GPU 運算的工具與函式庫生態系，並非每個元件都是本章所需。它可分為三類：

主要元件（GPU 控制 / 執行 library 提供的 GPU code）

• cuda.core — CUDA 控制（記憶體、device 管理）的 Pythonic 介面；為 Python 提供 CUDA Runtime 為 CUDA C++ 提供的功能。
• cuda.compute — 提供 CCCL (CUDA Core Compute Library) 的 GPU-accelerated 函式。
• CuPy — 提供 Numpy routine 的 GPU 加速版本，以及 GPU 加速版的 ndarray 資料容器（本章記憶體操作主力）。

Kernel 撰寫元件

• cuda.lang — 用 Python 子集、以 SIMT 模型撰寫 kernel 與 device function 的 Python DSL (Domain-specific language)。
• cuda.coop — 提供 CCCL 的 device-callable primitive（以 cuda.lang 實作）。
• cuda.tile — 以 Tile programming model 撰寫 kernel/device function 的 Python DSL。

其他元件

• cuda.pathfinder — 定位 Python 環境中已安裝 CUDA 元件的工具。
• cuda.bindings — CUDA 函式庫的低階 Python bindings（CUDA Driver API、Runtime API、NVRTC、NVVM 等）；功能與 cuda.core 相同，但它是C 語言 API 的 Python wrapper，不是 native 的 Pythonic 介面。

CUDA Python Ecosystem
├─ 控制 / 跑現成 GPU code ─ cuda.core | cuda.compute | CuPy
├─ 寫 kernel / device fn ── cuda.lang(SIMT) | cuda.coop | cuda.tile(Tile)
└─ 其他工具 ────────────── cuda.pathfinder | cuda.bindings(低階 C wrapper)

在 Python 使用 CUDA libraries

• CUDA C++ 自 2006 年問世時函式庫很少，開發者大多得自己寫 kernel；之後累積了大量 library，讓 C++ 開發者幾乎不用寫 GPU code 就能用 GPU。
• CUDA Python 則是反方向演進：CuPy 等 Python library 先提供了 GPU 加速的演算法實作（許多是 CUDA C++ code 的 Python binding），之後才有用 Python 語法/語義直接寫 custom kernel 的能力。

Getting Setup 與執行

• 多數 CUDA Python 元件發佈在 PyPi，可用 pip 或任何主流套件管理器安裝。
• 所有套件都需要系統安裝最新的 NVIDIA Driver。
• CUDA Toolkit 通常不需要就能寫/跑 CUDA Python 應用程式（與 C++ 需要 nvcc 工具鏈不同）。
• 執行方式與一般 Python 程式完全相同：

python3 cuda-python-app.py

SIMT Kernels in Python

在 GPU 上執行、可從 host 呼叫的函式稱為 kernel。CUDA 提供兩種模型：SIMT 與 CUDA tile。SIMT kernel 由大量平行 thread 同時執行——這個概念在 CUDA Python 與 CUDA C++ 完全一致。本章以 SIMT kernel 介紹 CUDA Python。

指定 kernel（Specifying Kernels）

先 import numba.cuda，再用 @cuda.jit 裝飾器把函式標記成 kernel：

from numba import cuda

@cuda.jit
def function(input_array, output_array):
    ...

• @cuda.jit 會讓 kernel 在第一次 launch 時對當前 GPU 進行 JIT 編譯。
• 未指定其他 GPU 時，使用 default CUDA device。

Launching Kernels

執行 kernel 的 thread 數量在 launch 時指定，稱為 execution configuration；每次呼叫可有不同的 block size 或 block 數量。execution configuration 放在 kernel 名稱後、引數前的方括號 [ ] 內，引數順序與 C++ 三角括號相同：

from numba import cuda

@cuda.jit
def my_kernel(input, output):
    ...

## launch the kernel
my_kernel[num_thread_blocks, threads_per_block](in_array, out_array)

C++/Python launch 語法對照：

多維 grid 與 thread block

grid 與 thread block 可為 1、2 或 3 維。1 維時用整數；2/3 維時用 tuple：

@cuda.jit
def function(input, output):
    ...

## 2D grid 與 2D block：(gridX, gridY) 與 (blockX, blockY)
function[(gridX, gridY), (blockX, blockY)](in_array, out_array)

Thread/Grid Index Intrinsics

kernel 內可用下列變數判斷自己的身分（皆為 3 分量向量，具 .x / .y / .z）：

• 未在 launch 指定的維度：維度 (Dim) 預設 1，index (Idx) 預設 0。
• cuda.threadIdx 與 cuda.blockIdx zero-indexed：threadIdx.x 取值 0 ~ blockDim.x - 1，.y/.z 同理。

一個 element-wise 向量加法 kernel（C = A + B）：

@cuda.jit
def vecadd(A, B, C):
    idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    C[idx] = A[idx] + B[idx]

• idx 是 thread 在 grid 中的唯一 index，範圍 0 ~ N-1，其中 N = cuda.gridDim.x * cuda.blockDim.x。此 kernel 假設 1 維 block + 1 維 grid。
• 這個「global index」計算太常見，Numba 提供 shorthand cuda.grid(n)（n 為維度數）：

## 以下兩行等價
idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
idx = cuda.grid(1)

grid (1D)                      thread global index
┌──────────┬──────────┬──────────┐
│ block 0  │ block 1  │ block 2  │   idx = threadIdx.x
│ t0..t255 │ t0..t255 │ t0..t255 │       + blockIdx.x * blockDim.x
└──────────┴──────────┴──────────┘       = cuda.grid(1)
   blockDim.x = 256,  gridDim.x = 3,  N = 3 * 256 = 768

Memory in GPU computing

GPU 有自己連接的 DRAM。kernel 要用的資料陣列一般必須先位於 GPU DRAM 才能被 kernel 存取。在 Python 中，控制資料位置（在 CPU↔GPU 間搬移）是程式設計師的責任——這與 C++ 的 explicit memory management 情況相同。

在 GPU 上 instantiate arrays

CuPy 提供在 GPU 上建立指定型別/維度 ndarray 的函式，函式簽章與 Numpy 相似：

import cupy as cp
import numpy as np

## 在 GPU 上建立全零矩陣；未指定 datatype 時預設 float32
A_device = cp.zeros((1024, 1024))
## 在 GPU 上建立 2^20 個隨機 double
B_device = cp.rand.random((2**20), dtype=np.double)
## 建立與既有 array 同 shape、同 datatype 的全零 array
C_device = cp.zeros_like(A)

host↔GPU 複製

import cupy as cp
import numpy as np

A_host = np.zeros((1024, 1024))   # host memory (Numpy)
A_device = cp.array(A_host)       # H→D：複製到 GPU

B_device = cp.rand.random((1024, 1024))  # GPU memory
B_host = cp.asnumpy(B_device)            # D→H：複製回 host

C++/Python 記憶體操作對照：

ndarray 物件型別

• ndarray 只存在於 host 記憶體或 GPU 記憶體其中之一，不會同時兩邊都有。
• 把位於 host 的 array 傳給 kernel → error。
• 把位於 GPU 的 array 傳給一般 Python 函式（非 kernel）→ error。
• CuPy 不會隱式在 CPU↔GPU 間複製（複製昂貴、過度複製傷效能），要求程式設計師明確決定何時複製。
• 在 kernel 內使用 ndarray 的好處：array 攜帶各維度的 extent，因此可自動做邊界檢查；當所需 thread 總數略小於 block/grid 總大小時，kernel code 不需手動檢查越界。

        Numpy ndarray            CuPy ndarray
        (host memory)            (GPU/device memory)
            │                          │
   cp.array(host) ─────────────────────▶  (H→D)
            ◀───────────────────── cp.asnumpy(dev)  (D→H)

   傳給 kernel：只接受 GPU 上的 ndarray（host array → error）
   傳給一般 Python fn：只接受 host array（GPU array → error）

Synchronizing the CPU and the GPU

與 C++ 相同，CUDA Python 的 kernel launch 對 host thread 是非同步的：host code 在 launch 後即繼續往下執行，不保證 kernel 已完成、甚至已開始。要確保 GPU kernel 完成，host thread 必須做某種同步。

最簡單的形式是同步整個 GPU（device-wide synchronization），由 CUDA driver 提供，CuPy 與 numba.cuda 都以 synchronize() 方法暴露：

import cupy as cp
from numba import cuda

## 等待所有先前發出的 GPU 工作完成（CuPy 介面）
cp.synchronize()
## 等待所有先前發出的 GPU 工作完成（numba.cuda 介面）
cuda.synchronize()

• device-wide 同步會阻塞 host thread，直到所有先前發出的 GPU 工作完成。
• 對照 C++：等同 cudaDeviceSynchronize()。

host thread:  launch kernel ──▶ (繼續執行 host code，不等 GPU)
                                       │  synchronize()
GPU:          [====== kernel 執行中 ======]
host thread:                            └──▶ 阻塞直到 GPU 全部完成，再往下走

Putting it All Together

完整的平行向量加法（Python 版），刻意讓 vector_size 不是 2 的次方、也不是 block_size 的倍數，以示範自動邊界檢查：

import numpy as np
from numba import cuda
import cupy as cp

@cuda.jit
def vecadd(A, B, C):
    work_index = cuda.grid(1)
    C[work_index] = A[work_index] + B[work_index]

vector_size = 2**24 + 11          # 非 2 次方、非 block_size 倍數
device = cp.cuda.Device()

## 直接在 GPU 上建立輸入（隨機 float32）與輸出（全零）
a = cp.random.uniform(-1, 1, vector_size)
b = cp.random.uniform(-1, 1, vector_size)
c = cp.zeros_like(a)

block_size = 256
grid_size = int(np.ceil(vector_size / block_size))
vecadd[grid_size, block_size](a, b, c)

## 顯式同步示範好習慣；其實下面的 asnumpy 複製也會隱式等待 kernel 完成
device.synchronize()

## 把 3 個 array 複製回 CPU 做驗證
a_np = cp.asnumpy(a)
b_np = cp.asnumpy(b)
c_np = cp.asnumpy(c)
expected = a_np + b_np
np.testing.assert_array_almost_equal(c_np, expected)
print("Test succeeded")

• grid_size = ceil(vector_size / block_size) 保證 thread 數 ≥ 元素數；多出的 thread 因 ndarray 自動邊界檢查而安全。
• device.synchronize() 此處純為示範好習慣；cp.asnumpy() 的 D→H 複製本身就會隱式等待 kernel 完成。
• A、B 在 GPU 上產生與初始化，最後才複製回 CPU，只是為了讓 CPU 也算一次以驗證 GPU 答案。

Error Checking in CUDA Python

任何影響 GPU 的操作（記憶體配置、複製、kernel launch）都可能引發錯誤；如同 C++（2.1.7 節），確認過程中未發生錯誤是 best practice。

在 Python，CUDA 錯誤會拋出 exception，若未捕捉會終止程式；可用一般 Python 的 try/except 捕捉。下例把 block_size 設成 2048（超過任何現行 GPU 的 1024 上限），使 kernel launch 失敗並拋 exception：

try:
    vector_size = 2**24 + 11
    device = cp.cuda.Device()
    a = cp.random.uniform(-1, 1, vector_size)
    b = cp.random.uniform(-1, 1, vector_size)
    c = cp.zeros_like(a)
    block_size = 2048                       # 對任何現行 GPU 都過大
    grid_size = int(np.ceil(vector_size / block_size))
    vecadd[grid_size, block_size](a, b, c)  # Error: 無效的 block size
    device.synchronize()
    print("Test did not encounter any errors")
except Exception as e:
    print(f"Exception occurred: {e}")

執行輸出：

$ python3 vecadd_error.py
Exception occurred: CUDA_ERROR_INVALID_VALUE: This indicates that one or more
of the parameters passed to the API call is not within an acceptable range of values.

• 程式捕捉到 exception 後正常結束。
• 若沒有 try/except，程式會異常結束並 dump traceback，顯示相同錯誤。

考試/測驗重點

Related Notes

• 02-Programming-GPUs/01-CUDA-Cpp-Kernels-and-Launch
• 02-Programming-GPUs/02-CUDA-Cpp-Memory-Management
• 02-Programming-GPUs/03-CUDA-Cpp-Sync-and-Workflow
• 02-Programming-GPUs/04-CUDA-Cpp-Errors-and-Specifiers
• 02-Programming-GPUs/06-SIMT-Basics-and-Thread-Hierarchy
• 02-Programming-GPUs/14-Async-Streams-and-Events
• 01-Introduction-to-CUDA/02-Execution-Model-and-SIMT
• 01-Introduction-to-CUDA/05-CUDA-Platform
• 02-Programming-GPUs/Practice-Programming-GPUs
• 00-Dashboard/Exam-Traps