概述

CuPy 是一个与 NumPy/SciPy 兼容的数组库，用于使用 Python 进行 GPU 加速计算。CuPy 可以作为直接替代品，用于在 NVIDIA CUDA 或 AMD ROCm 平台上运行现有的 NumPy/SciPy 代码。

CuPy 为 GPU 设备提供了 ndarray、稀疏矩阵以及相关的例程，它们都具有与 NumPy 和 SciPy 相同的 API

• N 维数组 (ndarray): cupy.ndarray

  • 数据类型 (dtypes)：布尔型 (bool_)、整型 (int8, int16, int32, int64, uint8, uint16, uint32, uint64)、浮点型 (float16, float32, float64) 和复数型 (complex64, complex128)

  • 支持与 numpy.ndarray 相同的语义，包括基本/高级索引和广播

• 稀疏矩阵: cupyx.scipy.sparse

  • 二维稀疏矩阵: csr_matrix, coo_matrix, csc_matrix, 和 dia_matrix

• NumPy 例程

  • 模块级函数 (cupy.*)

  • 线性代数函数 (cupy.linalg.*)

  • 快速傅里叶变换 (cupy.fft.*)

  • 随机数生成器 (cupy.random.*)

• SciPy 例程

  • 离散傅里叶变换 (cupyx.scipy.fft.* 和 cupyx.scipy.fftpack.*)

  • 高级线性代数 (cupyx.scipy.linalg.*)

  • 多维图像处理 (cupyx.scipy.ndimage.*)

  • 稀疏矩阵 (cupyx.scipy.sparse.*)

  • 稀疏线性代数 (cupyx.scipy.sparse.linalg.*)

  • 特殊函数 (cupyx.scipy.special.*)

  • 信号处理 (cupyx.scipy.signal.*)

  • 统计函数 (cupyx.scipy.stats.*)

这些例程由 CUDA 库（cuBLAS、cuFFT、cuSPARSE、cuSOLVER、cuRAND）、Thrust、CUB 和 cuTENSOR 支持，以提供最佳性能。

还可以使用以下方法轻松实现与 ndarray 配合使用的自定义 CUDA 内核：

• 内核模板：将逐元素操作和归约操作快速定义为单个 CUDA 内核

• 原始内核：导入现有的 CUDA C/C++ 代码

• 即时转译器 (JIT)：从 Python 源代码生成 CUDA 内核

• 内核融合：将多个 CuPy 操作融合到单个 CUDA 内核中

CuPy 可以在多 GPU 或集群环境中运行。分布式通信包（cupyx.distributed）由 NCCL 支持，为 ndarray 提供了集合通信和点对点通信原语。

对于需要更精细的性能控制的用户，可以访问低级 CUDA 特性：

• 流和事件：所有 API 都支持 CUDA 流和每线程默认流

• 内存池：带有内置内存池的可定制内存分配器

• 性能分析器：支持使用 CUDA Profiler 和 NVTX 进行代码性能分析

• 主机 API 绑定：直接从 Python 调用 CUDA 库的 API，例如 NCCL、cuDNN、cuTENSOR 和 cuSPARSELt

CuPy 实现了用于数据交换和互操作性的标准 API，例如 DLPack、CUDA 数组接口、__array_ufunc__ (NEP 13)、__array_function__ (NEP 18) 和 数组 API 标准。得益于这些协议，CuPy 可以轻松地与 NumPy、PyTorch、TensorFlow、MPI4Py 以及任何其他支持该标准的库集成。

在 AMD ROCm 环境下，CuPy 会自动将所有 CUDA API 调用转换为 ROCm HIP（hipBLAS、hipFFT、hipSPARSE、hipRAND、hipCUB、hipThrust、RCCL 等），从而使使用 CuPy 编写的代码无需任何修改即可在 NVIDIA 和 AMD GPU 上运行。

项目目标

CuPy 项目的目标是为 Python 用户提供 GPU 加速能力，而无需深入了解底层 GPU 技术。CuPy 团队专注于提供

• 完整的 NumPy 和 SciPy API 覆盖，成为一个全面的直接替代品，以及高级的 CUDA 功能以最大限度地提高性能。

• 成熟优质的库，作为所有需要加速的项目的基本包，从实验室环境到大规模集群。