编写 CUDA 内核

CUDA 内置目标弃用通知

Numba 内置的 CUDA 目标已被弃用，后续开发已移至 NVIDIA numba-cuda 包。请参阅内置 CUDA 目标的弃用和维护状态。

简介

CUDA 拥有一种不同于传统用于 CPU 编程的顺序执行模型。在 CUDA 中，你编写的代码将由多个线程同时执行（通常是数百或数千个）。你的解决方案将通过定义一个包含网格、块和线程的线程层次结构来建模。

Numba 对 CUDA 的支持提供了声明和管理此线程层次结构的功能。这些功能与 NVIDIA 的 CUDA C 语言所提供的功能大体相似。

Numba 还提供了三种 GPU 内存：全局设备内存（连接到 GPU 本身的大容量、相对较慢的片外内存）、片上共享内存和局部内存。对于除最简单的算法之外的所有情况，仔细考虑如何使用和访问内存以最大程度地减少带宽需求和争用非常重要。

内核声明

内核函数 是一种 GPU 函数，旨在从 CPU 代码 (*) 中调用。它有两个基本特征：

内核不能显式返回值；所有结果数据都必须写入传递给函数的数组中（如果计算标量，你可能需要传递一个单元素数组）；
内核在调用时会显式声明它们的线程层次结构：即线程块的数量和每个块中的线程数量（请注意，尽管内核只编译一次，但可以使用不同的块大小或网格大小多次调用它）。

乍一看，使用 Numba 编写 CUDA 内核非常类似于为 CPU 编写一个 JIT 函数

@cuda.jit
def increment_by_one(an_array):
    """
    Increment all array elements by one.
    """
    # code elided here; read further for different implementations

(*) 注意：较新的 CUDA 设备支持设备端内核启动；此功能称为 动态并行，但 Numba 目前不支持它）

内核调用

内核通常通过以下方式启动

threadsperblock = 32
blockspergrid = (an_array.size + (threadsperblock - 1)) // threadsperblock
increment_by_one[blockspergrid, threadsperblock](an_array)

我们注意到这里有两个步骤

实例化内核本身，通过指定块的数量（或“每网格的块”）以及每个块的线程数量。两者的乘积将给出启动的线程总数。内核实例化是通过获取已编译的内核函数（此处为 increment_by_one）并使用整数元组对其进行索引来完成的。
运行内核，通过向其传递输入数组（如果需要，还可以传递任何单独的输出数组）。内核异步运行：启动操作会将它们的执行排队到设备上，然后立即返回。你可以使用 cuda.synchronize() 来等待所有先前的内核启动完成执行。

注意

传递位于主机内存中的数组将隐式导致数据复制回主机，这将是同步的。在这种情况下，内核启动将直到数据复制回主机后才返回，因此看起来是同步执行的。

选择块大小

在声明内核所需的线程数量时，拥有两级层次结构可能看起来很奇怪。块大小（即每个块的线程数）通常至关重要

在软件方面，块大小决定了有多少线程共享给定区域的共享内存。
在硬件方面，块大小必须足够大以完全占用执行单元；建议可以在 CUDA C 编程指南中找到。

多维块和网格

为了帮助处理多维数组，CUDA 允许你指定多维块和网格。在上面的示例中，你可以将 blockspergrid 和 threadsperblock 设置为包含一、二或三个整数的元组。与等效大小的一维声明相比，这不会改变生成代码的效率或行为，但可以帮助你以更自然的方式编写算法。

线程定位

当运行内核时，内核函数的代码会由每个线程执行一次。因此，它必须知道自己位于哪个线程中，才能知道它负责哪些数组元素（复杂算法可能会定义更复杂的职责，但基本原则是相同的）。

一种方法是让线程确定其在网格和块中的位置，并手动计算相应的数组位置

@cuda.jit
def increment_by_one(an_array):
    # Thread id in a 1D block
    tx = cuda.threadIdx.x
    # Block id in a 1D grid
    ty = cuda.blockIdx.x
    # Block width, i.e. number of threads per block
    bw = cuda.blockDim.x
    # Compute flattened index inside the array
    pos = tx + ty * bw
    if pos < an_array.size:  # Check array boundaries
        an_array[pos] += 1

注意

除非你确定块大小和网格大小是数组大小的除数，否则你必须检查边界，如上所示。

threadIdx、blockIdx、blockDim 和 gridDim 是 CUDA 后端提供的特殊对象，其唯一目的是了解线程层次结构的几何形状以及当前线程在该几何形状中的位置。

这些对象可以是 1D、2D 或 3D，具体取决于内核是如何调用的。要访问每个维度上的值，请分别使用这些对象的 x、y 和 z 属性。

numba.cuda.threadIdx: 当前线程块中的线程索引。对于 1D 块，索引（由 x 属性给出）是一个整数，范围从 0（包含）到 numba.cuda.blockDim （不包含）。当使用多个维度时，每个维度都有类似的规则。

numba.cuda.blockDim: 线程块的形状，在实例化内核时声明。此值对于给定内核中的所有线程都是相同的，即使它们属于不同的块（即每个块都是“满的”）。

numba.cuda.blockIdx: 在启动内核的线程网格中的块索引。对于 1D 网格，索引（由 x 属性给出）是一个整数，范围从 0（包含）到 numba.cuda.gridDim （不包含）。当使用多个维度时，每个维度都有类似的规则。

numba.cuda.gridDim: 块网格的形状，即此内核调用所启动的块总数，在实例化内核时声明。

绝对位置

简单算法通常会像上述示例一样始终以相同方式使用线程索引。Numba 提供了额外的功能来自动化此类计算

numba.cuda.grid(ndim): 返回当前线程在整个块网格中的绝对位置。ndim 应与实例化内核时声明的维度数量相对应。如果 ndim 为 1，则返回一个整数。如果 ndim 为 2 或 3，则返回给定数量整数的元组。

numba.cuda.gridsize(ndim): 返回整个块网格在线程中的绝对大小（或形状）。ndim 的含义与上述 grid() 中的相同。

使用这些函数，增量示例可以变为

@cuda.jit
def increment_by_one(an_array):
    pos = cuda.grid(1)
    if pos < an_array.size:
        an_array[pos] += 1

对于 2D 数组和线程网格，相同的示例将是

@cuda.jit
def increment_a_2D_array(an_array):
    x, y = cuda.grid(2)
    if x < an_array.shape[0] and y < an_array.shape[1]:
       an_array[x, y] += 1

请注意，实例化内核时的网格计算仍需手动完成，例如

threadsperblock = (16, 16)
blockspergrid_x = math.ceil(an_array.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(an_array.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
increment_a_2D_array[blockspergrid, threadsperblock](an_array)

延伸阅读

请参阅 CUDA C 编程指南以获取有关 CUDA 编程的详细讨论。