GPU 编程实战——使用 PyCUDA 与 CuPy 功能

PyCUDA 中的动态 Kernel 编译

到目前为止，我们都是在 Python 脚本里以静态字符串的形式定义 CUDA C kernel，手动编写、一次性编译然后按需调用。这对于固定不变的 kernel 足够，但随着项目规模增大，我们往往需要更多灵活性：根据运行时参数生成代码、调优编译选项，或在同一框架里支持多种算法变体。

动态编译概述

动态编译允许我们在 Python 会话中即时生成、修改并编译 CUDA kernel，而无需手动管理临时文件或离开 Python 环境。它在以下场景中特别有用：

根据运行时输入调整编译时参数（块大小、数据类型、展开因子、数学常量等）。
构建支持多种算法变体的框架或工具，而无需为每种情况手写 kernel。
快速原型开发、测试或基准测量，要求快速迭代和自动化。
硬件细节或用户偏好不同，需要动态切换不同优化实现。

有了 PyCUDA，我们只需将 kernel 源码作为 Python 字符串，插入参数，然后在运行时用 SourceModule 编译即可。

示例：可参数化的向量缩放 Kernel

假设我们要将一个大数组的每个元素乘以一个缩放因子，并希望在编译时灵活调整展开（unroll）因子，而在运行时设置缩放常数。

数据准备

javascript 复制代码

import numpy as np
import pycuda.autoinit
import pycuda.gpuarray as gpuarray
from pycuda.compiler import SourceModule

N            = 2_000_000
host_array   = np.random.rand(N).astype(np.float32)
device_array = gpuarray.to_gpu(host_array)

生成带展开参数的 Kernel

ini 复制代码

def generate_scaling_kernel(unroll):
    return f"""
    __global__ void scale_unroll(float *data, float scale, int n)
    {{
        int idx = blockDim.x * blockIdx.x + threadIdx.x;
        #pragma unroll
        for (int i = 0; i < {unroll}; ++i) {{
            int offset = idx + i * gridDim.x * blockDim.x;
            if (offset < n)
                data[offset] *= scale;
        }}
    }}
    """

unroll_factor = 4
kernel_code    = generate_scaling_kernel(unroll_factor)
mod            = SourceModule(kernel_code)
scale_unroll   = mod.get_function("scale_unroll")

启动动态编译的 Kernel

scss 复制代码

threads_per_block = 256
blocks_per_grid   = (N + threads_per_block * unroll_factor - 1) \
                    // (threads_per_block * unroll_factor)

scale_unroll(
    device_array, np.float32(2.5), np.int32(N),
    block=(threads_per_block, 1, 1),
    grid =(blocks_per_grid,   1   )
)

验证并继续实验

ini 复制代码

result_host = device_array.get()
expected    = host_array * 2.5
print("结果是否匹配？", np.allclose(result_host, expected))

若要测试不同的展开因子或数学操作，只需重新调用 generate_scaling_kernel 生成新代码、再编译即可------无需离开 Python 也不必复制粘贴大量代码。这样我们的代码库更精简，实验更快捷，GPU 程序也能灵活适应各种场景和需求。

使用 CuPy 实现自定义操作

自定义通用函数与原生内核

到目前为止，我们一直在使用 CuPy 的 ufunc（通用函数）进行各种快速的逐元素操作，就像在 NumPy 中那样。CuPy 的 ufunc 已针对 GPU 做了优化，向量化操作直接映射到底层硬件，使加法、乘法、三角函数、指数函数等都能用简洁熟悉的语法完成。这些 ufunc 是高性能逐元素计算的关键。

然而，实际的科学、工程或数据驱动场景常常需要超出内置算术或数学运算的自定义逻辑，比如分段函数、阈值非线性变换，或数学与逻辑步骤的组合。在 NumPy 中我们可能会用 np.vectorize 或 np.frompyfunc 包装自定义逻辑，但这些方式无法在 GPU 上加速，也不被 CuPy 原生支持。标准 ufunc 虽强大，但对需要 GPU 级速度的高级工作流而言往往不足。

CuPy 提供了 ElementwiseKernel 和 RawKernel 接口，让我们能定义自己的通用函数。通过这些工具，可以直接用 CUDA C 编写专用逐元素 kernel，或使用 CuPy 灵活的 Python API，然后像调用内置 ufunc 一样，在数组表达式中使用它们。这使我们能跳出标准 ufunc 的限制，以 GPU 速度加速任意复杂的自定义逻辑。

定义自定义逐元素通用函数

假设我们要实现机器学习中常用的"带泄漏的 ReLU"（leaky ReLU）变换：

ini 复制代码

import cupy as cp

leaky_relu = cp.ElementwiseKernel(
    'float32 x, float32 slope',  # 输入参数
    'float32 y',                  # 输出参数
    'y = x > 0 ? x : slope * x;',# C 语法的操作
    'leaky_relu'                  # 函数名
)

# 在 GPU 数组上使用自定义 ufunc
a = cp.linspace(-5, 5, 10_000, dtype=cp.float32)
slope = 0.1
b = leaky_relu(a, slope)

此时，leaky_relu 已成为一流的 ufunc：支持广播、任意形状，并完全在 GPU 上执行。

使用 RawKernel 实现更复杂逻辑

若需更专门化的操作，比如根据不同区间调整缩放比：

ini 复制代码

raw_kernel_code = r'''
extern "C" __global__
void piecewise_scale(const float* x, float* y, int n)
{
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    if (idx < n) {
        float val = x[idx];
        if (val < 0)
            y[idx] = val * 0.5f;
        else if (val < 1)
            y[idx] = val * 2.0f;
        else
            y[idx] = val * 0.1f;
    }
}
'''
mod                = cp.RawModule(code=raw_kernel_code)
piecewise_scale    = mod.get_function('piecewise_scale')

a = cp.linspace(-2, 3, 100_000, dtype=cp.float32)
b = cp.empty_like(a)

threads_per_block = 256
blocks_per_grid   = (a.size + threads_per_block - 1) // threads_per_block

# 启动 RawKernel
piecewise_scale(
    (blocks_per_grid,), (threads_per_block,),
    (a, b, a.size)
)

在数组表达式中集成自定义 Kernel

定义完成后，这些 kernel 可像内置 ufunc 一样自然融入工作流：

ElementwiseKernel 函数支持广播和混合类型，便于与其他数组操作组合。
RawKernel 则使用标准 CUDA 启动语法，与 CuPy 数组无缝配合，管理内存与数据传输高效统一。

借助 ElementwiseKernel 与 RawKernel，我们能够在 GPU 上设计、优化和扩展任意复杂的自定义逻辑，满足各种科学与工程问题对性能和灵活性的双重需求。

数组广播与索引

CuPy 的广播规则

广播（Broadcasting）是 NumPy 和 CuPy 中最强大、最便捷的特性之一。它允许我们对形状不同的数组执行算术或逻辑运算，只要它们满足兼容的广播规则即可。这样，我们无需手动调整形状、平铺或循环遍历数组------CuPy 会在内部以高效的方式"拉伸"较小数组，使其在运算时与较大数组形状匹配，且不实际复制数据。广播规则让我们能够用简洁而富有表现力的高层代码来完成从简单到极其复杂的 GPU 加速数组运算，无需显式循环或手动管理内存。

其原理是从数组的末尾维度开始对比形状：如果某个维度相等，或其中一个为 1，则该维度可广播；CuPy 在计算时会在逻辑上将长度为 1 的维度扩展到所需大小，但不会在内存中重复存储数据。

示例：CuPy 数组的广播

不同形状的逐元素算术

假设有一个 2D 数组和一个 1D 数组，想将后者加到前者的每一行：

ini 复制代码

import cupy as cp

rows, cols = 512, 128
matrix = cp.random.rand(rows, cols).astype(cp.float32)
vector = cp.linspace(1, 2, cols, dtype=cp.float32)

# 向量形状 (128,) 在内部广播为 (512, 128)
result = matrix + vector

无需循环或手动扩展，CuPy 会自动完成广播。

多维数组上的广播

将一个 1D 数组应用到 3D 张量的最后一个维度：

ini 复制代码

tensor          = cp.random.rand(32, 64, 128).astype(cp.float32)
scaling_factors = cp.linspace(0.1, 1.0, 128, dtype=cp.float32)

# scaling_factors 在计算时广播为 (32, 64, 128)
scaled_tensor = tensor * scaling_factors

每个切片的最后 128 个元素都被对应的缩放因子逐一相乘。

高级索引

配合广播，高级索引能让我们方便地选择、修改或分析多维数组。CuPy 支持多种索引方式：

整数数组索引：使用索引数组挑选任意元素。
布尔掩码索引：选出满足条件的元素。
切片对象：提取整行、整列或任意子数组。

整数索引

ini 复制代码

data        = cp.random.rand(1024, 1024)
row_indices = cp.array([10, 200, 400])
col_indices = cp.array([5, 100, 800])

# 选出 (10,5)、(200,100)、(400,800) 三个位置的元素
selected_elements = data[row_indices, col_indices]

布尔掩码

ini 复制代码

mask     = matrix > 0.5
filtered = matrix[mask]  # 返回所有大于 0.5 的元素

多维切片

ini 复制代码

# 切出子矩阵
sub_matrix = matrix[100:200, 50:90]

在许多实际任务中，广播和索引常常配合使用，例如沿某个轴归一化或中心化数据：

ini 复制代码

mean     = matrix.mean(axis=0)
centered = matrix - mean    # 广播将每列的均值从该列中减去

# 将矩阵中所有负值置零
matrix[matrix < 0] = 0

所有这些操作无需显式 Python 循环，代码既可读又执行极其高效。广播与高级索引使我们能够在几行代码内表达复杂的数据操作，并借助 GPU 并行加速满足真实的科学、工程和分析需求。

PyCUDA 与 CuPy 之间的数据交换

为什么需要数据共享？

有时仅靠 PyCUDA 或仅靠 CuPy 都无法满足所有需求。PyCUDA 提供了精细的控制、动态内核编译和对 CUDA 原始特性的直接访问；而 CuPy 则带来类似 NumPy 的表达式 API，用于快速的逐元素数组操作和丰富的高级实用功能。某些工作流可能用 PyCUDA 编写自定义 kernel 或分配内存，然后切换到 CuPy 进行便捷的切片、广播或高级数学运算；也有场景是外部库或遗留代码基于 PyCUDA 构建，但项目其余部分使用 CuPy。

高效地在 PyCUDA 和 CuPy 之间共享数据，对于性能优化和开发便捷性至关重要。我们希望避免额外的设备 ↔ 主机拷贝，将所有数据留在 GPU 上，并在两个库之间无缝传递数据所有权或视图。

CuPy 与 PyCUDA 的互操作性

两者都通过自己的 GPU 数组类管理设备内存：

PyCUDA ：pycuda.gpuarray.GPUArray
CuPy ：cupy.ndarray

它们底层都封装了 CUDA 设备指针与内存池，但 Python 接口不同。幸运的是，两者都支持基于对方的原始设备指针创建数组视图，无需复制数据。

从 PyCUDA 转为 CuPy ndarray

假设我们已有一个 PyCUDA GPUArray：

python 复制代码

import pycuda.gpuarray as gpuarray
import pycuda.autoinit
import numpy as np
import cupy as cp

arr_host    = np.arange(10_000, dtype=np.float32)
arr_gpu_py  = gpuarray.to_gpu(arr_host)  # PyCUDA GPUArray

可直接将其包装为 CuPy ndarray：

ini 复制代码

ptr   = arr_gpu_py.gpudata            # 原始设备指针
shape = arr_gpu_py.shape
dtype = arr_gpu_py.dtype

# 用 UnownedMemory 封装 PyCUDA 内存，不发生复制
unowned = cp.cuda.UnownedMemory(int(ptr), arr_gpu_py.nbytes, arr_gpu_py)
memptr   = cp.cuda.MemoryPointer(unowned, 0)

# 构造 CuPy ndarray 视图
arr_cupy = cp.ndarray(shape, dtype=dtype, memptr=memptr)

这样，CuPy 数组共享了 PyCUDA 的底层缓冲，只要任一对象存在，内存就保持有效。

从 CuPy 转为 PyCUDA GPUArray

若已有一个 CuPy 数组：

ini 复制代码

arr_cupy = cp.arange(10_000, dtype=cp.float32)

可同样创建 PyCUDA 视图：

ini 复制代码

import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
import numpy as np

ptr   = arr_cupy.data.ptr             # 设备指针
shape = arr_cupy.shape
dtype = np.dtype(str(arr_cupy.dtype))

# 构造 DeviceAllocation 对象
gpudata  = drv.DeviceAllocation(ptr)

# 构造 PyCUDA GPUArray 视图
arr_gpu_py = gpuarray.GPUArray(shape, dtype, gpudata)

此时两个对象都指向同一段 GPU 缓冲，无需数据复制。

通过交换设备指针并创建视图，我们可以在 PyCUDA 与 CuPy 之间高效桥接，大幅减少重复的设备--主机传输，并充分发挥双方库的优势，打造流畅、高性能的混合 GPU 工作流。

总结

简而言之，我们系统地了解了 PyCUDA 与 CuPy，并扩展了用于 Python 中高、低级 GPU 编程的工具集。

在 PyCUDA 中，我们学习了动态内核编译，可以在运行时从 Python 环境生成并调优 CUDA 代码，无需离开脚本或手动管理临时文件。
在 CuPy 中，我们通过 ElementwiseKernel 和 RawKernel 自定义逐元素内核，突破了标准 ufunc 的限制，将任意复杂逻辑以 GPU 速度加速运行。
随后，我们使用 CuPy 的广播和高级索引功能，在多维数组上完成了复杂操作而无需显式循环，大幅简化代码并实现并行加速。
最后，我们演示了如何通过指针交换在 PyCUDA 与 CuPy 之间高效共享设备内存，混合使用自定义 kernel、底层内存管理与高级数值运算，避免不必要的数据拷贝或主机--设备传输。

凭借以上技能，我们已具备构建稳定灵活的 GPU 应用所需的基础，能够将 PyCUDA 的精细控制与 CuPy 的高层便捷结合起来，满足科研、工程或生产环境中对性能与灵活性的双重需求。