NVIDIA CUDA Python 常用 API 及详细教程

一、CUDA Python 简介

CUDA Python 是 NVIDIA 推出的用于从 Python 访问 CUDA 平台的工具集，包含多个子组件，覆盖从低级别API到高级别Pythonic接口，适用于GPU加速计算、并行算法实现等场景。核心组件包括：

• cuda.bindings：低级别Python绑定，直接映射CUDA C API
• cuda.core：高级别Pythonic接口，简化CUDA Runtime和核心功能使用
• cuda.pathfinder：CUDA组件定位工具
• cuda.cccl.*：基于CCCL库的并行算法和设备原语
• numba.cuda：Numba的CUDA目标，直接编译Python代码为CUDA核函数

二、安装指南

通过pip安装主要组件：

# 安装元包（包含多个子组件）
pip install cuda-python

# 单独安装子组件（按需选择）
pip install cuda-core[cu12]  # 高级别Pythonic接口
pip install cuda-bindings    # 低级别API绑定
pip install cuda-pathfinder  # 组件定位工具

三、核心组件常用API及示例

1. cuda.bindings：低级别CUDA API绑定

提供与CUDA C API几乎1:1的映射，适合需要精细控制的场景。核心模块包括driver（驱动API）、runtime（运行时API）、nvrtc（运行时编译）等。

（1）nvrtc模块：运行时编译PTX

用于动态编译CUDA C代码为PTX（并行线程执行）中间代码。

常用API：

• nvrtcCreateProgram：创建编译程序对象
• nvrtcCompileProgram：编译程序（生成PTX）
• nvrtcGetPTX：获取编译后的PTX代码
• nvrtcGetProgramLog：获取编译日志（用于调试错误）

示例：编译并获取PTX

from cuda.bindings import nvrtc
import numpy as np

# 定义CUDA核函数代码（SAXPY：y = a*x + y）
kernel_code = """
extern "C" __global__ void saxpy(float a, float *x, float *y, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] = a * x[i] + y[i];
}
"""

# 1. 创建程序对象
prog = nvrtc.nvrtcProgram()
err, = nvrtc.nvrtcCreateProgram(
    prog,               # 输出：程序对象
    kernel_code,        # 源代码
    "saxpy.cu",         # 文件名（用于日志）
    0, None, None       # 其他参数（宏定义等）
)
assert err == nvrtc.nvrtcResult.NVRTC_SUCCESS, "创建程序失败"

# 2. 编译程序（指定GPU架构，如compute_70）
compile_opts = ["--gpu-architecture=compute_70"]
err, = nvrtc.nvrtcCompileProgram(prog, len(compile_opts), compile_opts)
if err != nvrtc.nvrtcResult.NVRTC_SUCCESS:
    # 获取错误日志
    log_size = nvrtc.nvrtcGetProgramLogSize(prog)[1]
    log = nvrtc.nvrtcGetProgramLog(prog, log_size)[1]
    raise RuntimeError(f"编译失败：{log.decode()}")

# 3. 获取PTX代码
ptx_size = nvrtc.nvrtcGetPTXSize(prog)[1]
ptx = nvrtc.nvrtcGetPTX(prog, ptx_size)[1].decode()
print("PTX代码：\n", ptx[:500])  # 打印前500字符

（2）driver模块：设备管理与核函数执行

用于管理GPU设备、内存分配、加载PTX并执行核函数。

常用API：

• cuInit：初始化CUDA驱动
• cuDeviceGet：获取设备
• cuCtxCreate：创建CUDA上下文
• cuMemAlloc/cuMemFree：设备内存分配/释放
• cuMemcpyHtoD/cuMemcpyDtoH：主机与设备内存拷贝
• cuModuleLoadData：从PTX加载模块
• cuModuleGetFunction：从模块获取核函数
• cuLaunchKernel：执行核函数

示例：执行SAXPY核函数

from cuda.bindings import driver
import numpy as np

# 初始化驱动
driver.cuInit(0)

# 获取设备并创建上下文
device = driver.CUdevice()
driver.cuDeviceGet(device, 0)  # 获取第0号GPU
ctx = driver.CUcontext()
driver.cuCtxCreate(ctx, 0, device)

# 准备主机数据
n = 1024
a = 2.0
x = np.random.rand(n).astype(np.float32)
y = np.random.rand(n).astype(np.float32)
expected = a * x + y  # 预期结果

# 分配设备内存
x_dev = driver.CUdeviceptr()
y_dev = driver.CUdeviceptr()
driver.cuMemAlloc(x_dev, n * np.dtype(np.float32).itemsize)
driver.cuMemAlloc(y_dev, n * np.dtype(np.float32).itemsize)

# 主机到设备的数据拷贝
driver.cuMemcpyHtoD(x_dev, x.ctypes.data, n * np.dtype(np.float32).itemsize)
driver.cuMemcpyHtoD(y_dev, y.ctypes.data, n * np.dtype(np.float32).itemsize)

# 从PTX加载模块（使用前一步生成的ptx）
module = driver.CUmodule()
driver.cuModuleLoadData(module, ptx.encode())

# 获取核函数
kernel = driver.CUfunction()
driver.cuModuleGetFunction(kernel, module, "saxpy")

# 配置核函数参数（a, x_dev, y_dev, n）
args = [
    np.float32(a),
    x_dev,
    y_dev,
    np.int32(n)
]
# 计算网格和块大小（每个块256线程）
block_dim = (256, 1, 1)
grid_dim = ((n + block_dim[0] - 1) // block_dim[0], 1, 1)

# 执行核函数
driver.cuLaunchKernel(
    kernel,
    grid_dim[0], grid_dim[1], grid_dim[2],  # 网格大小
    block_dim[0], block_dim[1], block_dim[2],  # 块大小
    0,  # 共享内存大小
    None,  # 流（默认）
    args,  # 参数列表
    None   # 额外参数
)
driver.cuCtxSynchronize()  # 等待核函数执行完成

# 设备到主机的数据拷贝（获取结果）
y_result = np.empty_like(y)
driver.cuMemcpyDtoH(y_result.ctypes.data, y_dev, n * np.dtype(np.float32).itemsize)

# 验证结果
assert np.allclose(y_result, expected, atol=1e-5), "结果不匹配"
print("SAXPY执行成功！")

# 释放资源
driver.cuMemFree(x_dev)
driver.cuMemFree(y_dev)
driver.cuModuleUnload(module)
driver.cuCtxDestroy(ctx)

2. cuda.core：高级别Pythonic接口

提供更简洁的API，封装了设备管理、内存操作、核函数编译等功能，适合快速开发。

常用类及方法：

• Device：GPU设备管理（如Device(0)获取第0号设备）
• Stream：流操作（异步执行）
• Event：事件计时（如记录核函数执行时间）
• Program：编译CUDA代码（封装NVRTC）
• LaunchConfig：核函数启动配置（网格/块大小）

示例：用cuda.core实现SAXPY

from cuda.core import Device, Program, LaunchConfig
import numpy as np

# 获取设备并设置为当前设备
device = Device(0)
device.use()

# 准备数据
n = 1024
a = 2.0
x = np.random.rand(n).astype(np.float32)
y = np.random.rand(n).astype(np.float32)
expected = a * x + y

# 设备内存分配（自动与主机数据同步）
x_dev = device.mem_alloc(x.nbytes)
y_dev = device.mem_alloc(y.nbytes)
x_dev.copy_from_host(x)
y_dev.copy_from_host(y)

# 编译核函数（自动处理NVRTC流程）
kernel_code = """
extern "C" __global__ void saxpy(float a, float *x, float *y, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] = a * x[i] + y[i];
}
"""
program = Program(kernel_code, "saxpy.cu")
program.compile(["--gpu-architecture=compute_70"])  # 编译

# 配置核函数启动参数
launch_config = LaunchConfig(
    grid=(n + 255) // 256,  # 网格大小
    block=(256,),           # 块大小
)

# 执行核函数（自动处理参数传递）
program.saxpy[launch_config](a, x_dev, y_dev, n)
device.synchronize()  # 等待完成

# 获取结果
y_result = np.empty_like(y)
y_dev.copy_to_host(y_result)

# 验证
assert np.allclose(y_result, expected, atol=1e-5)
print("cuda.core SAXPY执行成功！")

3. cuda.pathfinder：CUDA组件定位

用于查找系统中的CUDA动态库、头文件等组件，适合跨环境部署。

常用API：

• find_cuda_library：查找CUDA动态库（如cudart）
• find_cuda_include_dirs：查找CUDA头文件目录

示例：

from cuda.pathfinder import find_cuda_library, find_cuda_include_dirs

# 查找CUDA运行时库
cudart_path = find_cuda_library("cudart")
print(f"CUDA Runtime库路径：{cudart_path}")

# 查找CUDA头文件目录
include_dirs = find_cuda_include_dirs()
print(f"CUDA头文件目录：{include_dirs}")

四、进阶场景：使用cuda.cccl.parallel进行并行算法

cuda.cccl.parallel封装了CCCL库的高效并行算法（如排序、归约），可直接在主机调用。

示例：并行归约（求和）

from cuda.cccl.parallel import reduce
import numpy as np

# 生成随机数据
data = np.random.rand(1024 * 1024).astype(np.float32)
expected_sum = data.sum()

# 使用CCCL的reduce进行并行求和
result = reduce(data, 0.0, lambda a, b: a + b)  # 初始值0.0，求和操作

assert np.isclose(result, expected_sum), "归约结果不匹配"
print(f"并行求和结果：{result}（预期：{expected_sum}）")

五、总结

CUDA Python 提供了从低级别到高级别的完整接口：

• 需精细控制时使用cuda.bindings（直接映射CUDA C API）
• 追求开发效率时使用cuda.core（Pythonic封装）
• 并行算法可借助cuda.cccl.parallel
• 设备定位用cuda.pathfinder

更多细节可参考官方文档：

• cuda.core 文档
• cuda.bindings 文档