2.2. CUDA Python 简介

2.2.1.1. 在 Python 中使用 CUDA 库

CUDA C++ 拥有丰富的库生态系统，无需直接编写核函数或 GPU 代码即可实现 GPU 加速。CUDA C++ 于 2006 年推出时，库的数量很少，开发者基本上必须自行编写 GPU 核函数代码。此后，人们开发了大量库，使开发者无需编写太多 GPU 代码，甚至完全不必编写 GPU 代码，就能在 C++ 中利用 GPU 计算。

CUDA Python 生态系统则从另一个方向演进而来：在能够直接使用 Python 语法和语义编写自定义核函数之前，CuPy 等 Python 库就已经为 Python 开发者提供了计算和算法的 GPU 加速实现。这些库中的许多库为使用 CUDA C++ 实现的 GPU 代码提供了 Python 绑定。

在当今的 CUDA 时代，如果 GPU 加速库能够提供满足您需求所需的表达能力，那么几乎总是建议使用这些库。这些库中的许多实现都经过 GPU 计算专家的调优。当没有可用的库或库无法满足需求时，在 Python 中也可以像在 C++ 中一样，直接编写 GPU 核函数和设备函数。

2.2.1.2. 本章范围

虽然开发者应在可能的情况下优先使用库，但本章其余部分介绍的是当开发者需要用 Python 编写自定义 GPU 代码时该如何进行。本章采用与第 2.1 节介绍 C++ 时相同的方式，讲解如何在 GPU 上编写和运行自定义 Python 代码，首先介绍如何指定 GPU 核函数，然后介绍如何使用 cuPy 提供的 GPU 加速的ndarray在 GPU 上分配内存，并实现 CPU 与 GPU 之间的通信。

2.2.1.3. 环境设置

一般来说，CUDA Python 生态系统中的大多数组件都可在 PyPi 上获取，并可使用pip命令或任何流行的 Python 包管理器进行安装。所有包都要求系统上安装最新的 NVIDIA 驱动程序。编写或运行 CUDA Python 应用程序通常不需要 CUDA Toolkit。

有关在不同平台上安装和配置 CUDA Python 的信息，请参阅NVIDIA 开发者专区中的 CUDA Python。

2.2.1.4. 运行 CUDA Python 应用程序

CUDA Python 应用程序无论是使用 CUDA 加速库，还是包含用户编写的 GPU 代码，其运行方式都与常规 Python 应用程序相同。在本节中，示例始终通过命令行运行，具体做法是如下所示调用python3来执行名为cuda-python-app.py的程序。

$ python3 cuda-python-app.py

2.2.2.1. 指定核函数

在 CUDA Python 中指定核函数之前，必须导入numba.cuda包。通常采用如下所示的方式完成此操作。

from numba import cuda

这会导入numba.cuda包，并允许我们使用该包提供的cuda命名空间中的组件。

要在 CUDA Python 中将一个函数指定为核函数，请如下所示将装饰器@cuda.jit放在函数定义的上一行。

from numba import cuda

@cuda.jit
def function(input_array, output_array):
    ...

这将使核函数在第一次启动时，针对当前活动的 GPU 进行 JIT 编译。如果未指定其他 GPU，则使用默认的 CUDA 设备，本节展示的示例就是如此。

2.2.2.2. 启动核函数

执行核函数的线程数量在核函数启动时指定，是启动的一部分。这称为执行配置。每次核函数调用都可以具有独有的执行配置，例如不同的块大小或线程块数量。

kernel_name[number_of_thread_blocks, threads_per_block](arguments, ...)

from numba import cuda

@cuda.jit
def my_kernel(input, output):
    ...

## launch the kernel
my_kernel[num_thread_blocks, threads_per_block](in_array, out_array)

from numba import cuda

@cuda.jit
def function(input, output):
    ...

## launch the kernel
function[(gridX, gridY), (blockX, blockY)](in_array, out_array)

2.2.2.3. 线程和网格索引内建函数

第 1.2.2.1 节介绍了线程和网格，第 2.2.2.2 节展示了如何为核函数启动指定网格和线程块的大小。在核函数内部，每个线程都可以访问执行配置的参数，以及线程的索引和线程块在网格中的索引。

可以从核函数内部访问以下变量，以确定线程的身份：

• cuda.threadIdx.[xyz]给出线程在其线程块内的索引。一个线程块中的每个线程都有不同的索引。

• cuda.blockDim.[xyz]给出线程块的各维度大小，这些大小是在核函数启动的执行配置中指定的。

• cuda.blockIdx.[xyz]给出线程块在网格中的索引。每个线程块都有不同的索引。

• cuda.gridDim.[xyz]给出网格的各维度大小，这些大小是在核函数启动时的执行配置中指定的。

这些变量中的每一个都是一个三分量向量，具有.x、.y和.z成员。如果在核函数启动的执行配置中未指定某个维度，则对于维度大小，其默认值为 1；对于索引，其默认值为 0。

cuda.threadIdx和cuda.blockIdx的索引从零开始。也就是说，cuda.threadIdx.x的取值从 0 到cuda.blockDim.x - 1，包含上限值。.y和.z在各自的维度中采用相同的方式。

下面展示了一个简单的向量加法核函数的代码，它将两个向量逐元素相加。该函数接受三个数组A、B和C，以实现逐元素向量加法C = A + B

# C = A + B vector addition
@cuda.jit
def vecadd(A, B, C):
    idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    C[idx] = A[idx] + B[idx]

该核函数首先计算线程在网格中的唯一索引。此核函数假定它是在一维网格中以一维线程块启动的。idx变量是一个从 0 到N-1的唯一索引，其中 N 是网格中的线程总数，即N = cuda.gridDim.x * cuda.blockDim.x。

上面代码块所示的线程索引计算模式非常常见，因此 Numba 为此操作提供了一种简写语法：cuda.grid(n)，其中n是维数。在上面的示例中，下列这一行

idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x

可以替换为简单得多的

idx = cuda.grid(1)

此核函数一个值得注意的方面是，它不检查对A、B或C的越界访问。在本章中，我们假定它们是由 cuPy 创建的ndarrays，稍后将在第 2.2.3.3 节介绍。使用 cuPyndarrays时，边界检查由数组类型隐式实现。

2.2.3.1. 在 GPU 上实例化数组

CuPy 提供了在 GPU 上创建指定类型和维度的ndarray对象的函数，也提供了在 CPU 与 GPU 之间复制数据的函数。cuPy 中许多函数的函数签名与 Numpy 中用于创建ndarrays的函数相似。下面展示了几个示例，说明如何使用 CuPy 在 GPU 内存中创建和填充数组。

import cupy as cp
import numpy as np

## create a matrix of zeros on the GPU
## when a datatype is not specified, float32 is used by default
A_device = cp.zeros((1024, 1024))

## create an array of 2^20 random doubles on the GPU
B_device = cp.rand.random((2**20), dtype=np.double)

## create an array of zeroes with the same shape and datatype as an existing array
C_device = cp.zeros_like(A)

2.2.3.2. 在主机内存与 GPU 内存之间复制数组

CuPy 还可用于将数据从位于 CPU 内存中的 Numpyndarrays复制到位于 GPU 内存中的 CuPy 数组。

import cupy as cp
import numpy as np

## Create an array in host memory
A_host = np.zeros((1024, 1024))
## Copy the array to the GPU
A_device = cp.array(A_host)

## Create an array in GPU memory
B_device = cp.rand.random((1024, 1024))
## copy the array to host memory
B_host = cp.asnumpy(B_device)

2.2.3.3. ndarray 对象类型

上一节展示的ndarray对象要么存在于主机内存中，要么存在于 GPU 内存中，但不会同时存在于两者中。将驻留在主机上的数组作为实参传递给核函数会导致错误。将驻留在 GPU 内存中的数组传递给普通 Python 函数，即非核函数，也会导致错误。CuPy 不会隐式执行 CPU 与 GPU 之间的复制，因为这些操作可能开销很大，而且过多的数据复制会损害性能。因此，CuPy 要求程序员有意识地决定何时在 CPU 与 GPU 之间复制数据。

在 GPU 核函数中使用ndarray类型的一个优点是，数组自身携带其各维度的范围。如第 2.2.2.3 节所示，边界检查会自动完成，当所需的线程总数略小于执行块或网格的总大小时，核函数代码无需检查越界访问。