面向 Python 的 NumPy/SciPy 兼容数组库,用于 GPU 加速计算
CuPy 是一个开源的数组库,用于使用 Python 进行 GPU 加速计算。CuPy 利用 CUDA Toolkit 库(包括 cuBLAS、cuRAND、cuSOLVER、cuSPARSE、cuFFT、cuDNN 和 NCCL)充分发挥 GPU 架构的优势。
该图显示了 CuPy 相对于 NumPy 的加速效果。大多数操作在使用 CuPy 的 GPU 上开箱即用即可良好运行。CuPy 将某些操作的速度提高了 100 多倍。请阅读 RAPIDS AI Medium 博客上的原始基准测试文章 Single-GPU CuPy Speedups。
CuPy 的接口与 NumPy 和 SciPy 高度兼容;在大多数情况下,它可以作为直接替代品使用。您只需在 Python 代码中将 numpy
和 scipy
替换为 cupy
和 cupyx.scipy
即可。CuPy 基础教程有助于学习 CuPy 的第一步。
CuPy 支持各种方法、索引、数据类型、广播等。此比较表显示了 NumPy / SciPy API 及其对应的 CuPy 实现列表。
>>> import cupy as cp
>>> x = cp.arange(6).reshape(2, 3).astype('f')
>>> x
array([[ 0., 1., 2.],
[ 3., 4., 5.]], dtype=float32)
>>> x.sum(axis=1)
array([ 3., 12.], dtype=float32)
安装 CuPy 最简单的方法是使用 pip。CuPy 为 Linux 和 Windows 提供了 wheels(预编译的二进制包)。请阅读安装指南了解更多详情。
CuPy 也可以从 Conda-Forge 安装或从源代码安装。
# For CUDA 11.2 ~ 11.x
pip install cupy-cuda11x
# For CUDA 12.x
pip install cupy-cuda12x
# For AMD ROCm 4.3
pip install cupy-rocm-4-3
# For AMD ROCm 5.0
pip install cupy-rocm-5-0
如果您想让代码运行得更快,您可以轻松创建自定义 CUDA 内核,只需一小段 C++ 代码片段即可。CuPy 会自动封装并编译它以生成 CUDA 二进制文件。编译后的二进制文件会被缓存并在后续运行中重用。请阅读用户定义内核教程。
此外,您还可以通过原始模块 (Raw modules)使用原始 CUDA 内核。
>>> x = cp.arange(6, dtype='f').reshape(2, 3)
>>> y = cp.arange(3, dtype='f')
>>> kernel = cp.ElementwiseKernel(
... 'float32 x, float32 y', 'float32 z',
... '''
... if (x - 2 > y) {
... z = x * y;
... } else {
... z = x + y;
... }
... ''', 'my_kernel')
>>> kernel(x, y)
array([[ 0., 2., 4.],
[ 0., 4., 10.]], dtype=float32)