PyCUDA：线程边界、逐元素运算与设备内存

学习教程更新于 2026-10-03编辑本文

PyCUDA 管理哪一层

PyCUDA 让 Python 调用 CUDA 驱动接口、管理设备内存并启动自己编写的 CUDA kernel。它适合需要定制 GPU 运算的场景；使用现成 ONNX 模型推理时，通常先用 ONNX Runtime 或 TensorRT，不必重新实现整套算子。

CPU 数组位于主机内存；GPUArray 的数据位于设备内存。to_gpu 上传，.get() 下载。Python 持有的对象不等于设备上的数据已经搬回 CPU。GPUArray 文档

向量加法：线程编号与边界

一个一维 grid 中，全局线程编号为：

i=bB+t,

其中 $b$ 是块编号、 $B$ 是每块线程数、 $t$ 是块内线程编号。处理 $N$ 个元素需要：

G=\left\lceil\frac{N}{B}\right\rceil =\left\lfloor\frac{N+B-1}{B}\right\rfloor

个块。最后一块可能有多余线程，所以 kernel 内还要判断 $i<N$ 。只做向下整除会漏掉尾部；只改成向上取整而不补边界判断，则可能越界。

下面示例是独立 PyCUDA 程序，要求可用的 NVIDIA CUDA 环境和编译工具。它支持空向量，并故意用 7 个元素检验“不整除一个块”的情况。

import numpy as np
import pycuda.autoinit  # 独立示例：自动创建并管理一个 CUDA 上下文
import pycuda.gpuarray as gpuarray
from pycuda.compiler import SourceModule

module = SourceModule(r"""
__global__ void vector_add(const float *a, const float *b,
                           float *c, unsigned int n) {
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        c[i] = a[i] + b[i];
    }
}
""")
kernel = module.get_function("vector_add")

def vector_add(a, b):
    a, b = np.asarray(a), np.asarray(b)
    if a.dtype != np.float32 or b.dtype != np.float32:
        raise TypeError("both inputs must be float32")
    if a.ndim != 1 or a.shape != b.shape:
        raise ValueError("inputs must be equal-length vectors")
    if a.size > np.iinfo(np.int32).max:
        raise ValueError("this example is limited to 32-bit indexing")
    if a.size == 0:
        return np.empty(0, dtype=np.float32)
    a, b = np.ascontiguousarray(a), np.ascontiguousarray(b)
    da, db = gpuarray.to_gpu(a), gpuarray.to_gpu(b)
    dc = gpuarray.empty_like(da)
    block = 256
    grid = (a.size + block - 1) // block
    kernel(da, db, dc, np.uint32(a.size),
           block=(block, 1, 1), grid=(grid, 1, 1))
    return dc.get()

a = np.arange(7, dtype=np.float32)
b = np.full(7, 2, dtype=np.float32)
np.testing.assert_allclose(vector_add(a, b), a + b)
assert vector_add(a[:0], b[:0]).size == 0

np.uint32 与 kernel 的 unsigned int 参数配对。数组 dtype、参数宽度、缓冲区大小必须一致；把 Python 任意精度整数随意传给底层函数并不等于已声明好 CUDA 参数类型。这里的上限只保护示例索引，实际最大 grid 尺寸和可分配内存仍受设备限制。

逐元素乘法不是矩阵乘法

原笔记的 C[i] = A[i] * B[i] 计算 Hadamard 积：

C_{ij}=A_{ij}B_{ij}.

通常说的矩阵乘法则是：

D_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}.

前者同位置相乘；后者是行与列做内积。ElementwiseKernel 在每个位置执行一个标量表达式，下面因此命名为 elementwise_product：ElementwiseKernel 文档

import numpy as np
import pycuda.autoinit
import pycuda.gpuarray as gpuarray
from pycuda.elementwise import ElementwiseKernel

a = np.array([[1, 2], [3, 4]], dtype=np.float32)
b = np.array([[5, 6], [7, 8]], dtype=np.float32)
da, db = gpuarray.to_gpu(a), gpuarray.to_gpu(b)
dc = gpuarray.empty_like(da)
multiply = ElementwiseKernel(
    "float *a, float *b, float *c", "c[i] = a[i] * b[i]",
    "elementwise_product",
)
multiply(da, db, dc)
np.testing.assert_array_equal(dc.get(), a * b)
print("elementwise:", a * b)  # [[5,12],[21,32]]
print("matrix:", a @ b)       # [[19,22],[43,50]]

真正的高性能矩阵乘法还涉及数据复用、分块、共享内存和硬件特性；常规任务优先考虑适合的成熟线性代数库。

上下文、异步操作与计时

pycuda.autoinit 会自动创建上下文；pycuda.autoprimaryctx 则保留设备的 primary context。和其他 CUDA 库混用时，应先约定谁持有上下文、流和内存，不能假设来自不同上下文的地址可以直接交换。初始化与上下文说明

kernel 提交通常是异步的，Python 调用返回不代表设备运算结束。读取结果、复用缓冲区或结束计时前，必须有正确的同步依赖。上面通过 .get() 取得结果；对 kernel 调用前后直接使用主机时钟，而不等待 GPU 完成，测到的可能主要是提交耗时。

评估性能时至少分别记录：首次编译、内存分配、上传、设备计算、下载。四个或七个元素的教学样例不适合证明 GPU 加速收益，传输与启动开销可能远大于计算本身。

本次复核用 CPU 枚举验证了网格覆盖公式，并验证了两种乘法的数值差别；CUDA 示例通过 Python 语法检查，但没有在 GPU 上编译或运行。它们保留为具备对应环境时的复核程序。