掘金 人工智能 08月15日
GPU 编程实战——使用 PyCUDA 与 CuPy 功能
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

PyCUDA允许动态编译CUDA kernel,根据运行时参数调整编译选项,支持多种算法变体,适用于快速原型开发。CuPy通过ElementwiseKernel和RawKernel实现自定义逐元素函数,突破ufunc限制。CuPy的广播和高级索引功能简化多维数组操作。PyCUDA与CuPy可通过设备指针共享内存,实现混合GPU工作流,提升性能与开发便捷性。

💡PyCUDA动态编译:可在Python中即时生成、修改并编译CUDA kernel,根据运行时参数调整编译选项(如块大小、数据类型),支持多种算法变体,无需手动管理临时文件,无需离开Python环境,适用于快速原型开发和灵活的GPU编程。

🔧CuPy自定义ufunc:通过ElementwiseKernel使用Python API定义简单逐元素函数,通过RawKernel使用CUDA C编写复杂逻辑,两者均支持广播、任意形状,完全在GPU上执行,突破标准ufunc限制,实现GPU加速的任意自定义逻辑。

📈CuPy广播与索引:支持从数组末尾维度开始对比形状的广播规则,自动扩展较小数组形状以匹配较大数组,无需手动调整或复制数据;支持整数数组索引、布尔掩码索引和切片对象等高级索引方式,配合广播实现复杂数据操作,无需显式循环,代码简洁高效。

🔗PyCUDA与CuPy数据共享:通过交换设备指针并创建视图,PyCUDA的GPUArray和CuPy的ndarray可共享底层GPU内存,无需数据复制,实现高效数据交换,允许混合使用自定义kernel、底层内存管理与高级数值运算,构建高性能混合GPU工作流。

🚀混合GPU工作流:结合PyCUDA的动态内核编译和精细控制,与CuPy的类似NumPyAPI和高级数组操作,构建稳定灵活的GPU应用,满足科研、工程或生产环境中对性能与灵活性的双重需求,提升开发效率和计算性能。

PyCUDA 中的动态 Kernel 编译

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

动态编译概述

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

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

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

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

数据准备

import numpy as npimport pycuda.autoinitimport pycuda.gpuarray as gpuarrayfrom pycuda.compiler import SourceModuleN            = 2_000_000host_array   = np.random.rand(N).astype(np.float32)device_array = gpuarray.to_gpu(host_array)

生成带展开参数的 Kernel

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 = 4kernel_code    = generate_scaling_kernel(unroll_factor)mod            = SourceModule(kernel_code)scale_unroll   = mod.get_function("scale_unroll")

启动动态编译的 Kernel

threads_per_block = 256blocks_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   ))

验证并继续实验

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

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

使用 CuPy 实现自定义操作

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

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

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

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

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

假设我们要实现机器学习中常用的“带泄漏的 ReLU”(leaky ReLU)变换:

import cupy as cpleaky_relu = cp.ElementwiseKernel(    'float32 x, float32 slope',  # 输入参数    'float32 y',                  # 输出参数    'y = x > 0 ? x : slope * x;',# C 语法的操作    'leaky_relu'                  # 函数名)# 在 GPU 数组上使用自定义 ufunca = cp.linspace(-5, 5, 10_000, dtype=cp.float32)slope = 0.1b = leaky_relu(a, slope)

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

使用 RawKernel 实现更复杂逻辑

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

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 = 256blocks_per_grid   = (a.size + threads_per_block - 1) // threads_per_block# 启动 RawKernelpiecewise_scale(    (blocks_per_grid,), (threads_per_block,),    (a, b, a.size))

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

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

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

数组广播与索引

CuPy 的广播规则

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

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

示例:CuPy 数组的广播

不同形状的逐元素算术

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

import cupy as cprows, cols = 512, 128matrix = 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 张量的最后一个维度:

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 支持多种索引方式:

整数索引

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]

布尔掩码

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

多维切片

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

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

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 数组类管理设备内存:

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

从 PyCUDA 转为 CuPy ndarray

假设我们已有一个 PyCUDA GPUArray

import pycuda.gpuarray as gpuarrayimport pycuda.autoinitimport numpy as npimport cupy as cparr_host    = np.arange(10_000, dtype=np.float32)arr_gpu_py  = gpuarray.to_gpu(arr_host)  # PyCUDA GPUArray

可直接将其包装为 CuPy ndarray

ptr   = arr_gpu_py.gpudata            # 原始设备指针shape = arr_gpu_py.shapedtype = 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 数组:

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

可同样创建 PyCUDA 视图:

import pycuda.driver as drvimport pycuda.gpuarray as gpuarrayimport numpy as npptr   = arr_cupy.data.ptr             # 设备指针shape = arr_cupy.shapedtype = 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 编程的工具集。

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

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

PyCUDA 动态编译 CuPy 自定义ufunc 广播 高级索引 数据共享
相关文章