编写 CUDA 内核

介绍

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

Numba 的 CUDA 支持提供了声明和管理这种线程层次结构的工具。这些工具在很大程度上类似于 NVidia 的 CUDA C 语言所提供的工具。

Numba 还暴露了三种类型的 GPU 内存：全局设备内存和本地内存。对于所有但最简单的算法，重要的是你要仔细考虑如何使用和访问内存，以最小化带宽需求和争用。

内核声明

一个 内核函数 是一个从 CPU 代码中调用的 GPU 函数（*）。它赋予了两个基本特性：

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

乍一看，使用 Numba 编写 CUDA 内核看起来非常像为 CPU 编写即时编译函数:

@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 设为包含一、二或三个整数的元组。与等效大小的 1D 声明相比，这不会改变生成代码的效率或行为，但可以帮助你以更自然的方式编写算法。

线程定位

当运行一个内核时，内核函数的代码会被每个线程执行一次。因此，它必须知道它处于哪个线程中，以便知道它负责哪个数组元素（复杂的算法可能定义更复杂的责任，但基本原理是相同的）。

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

@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（包含）到:attr:`numba.cuda.blockDim`（不包含）。当使用多个维度时，每个维度都存在类似的规则。

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

numba.cuda.blockIdx: 在启动内核的线程网格中，块索引。对于一维网格，索引（由 x 属性给出）是一个整数，范围从 0（包含）到 :attr:`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

对于二维数组和线程网格的相同示例将是:

@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 编程的详细讨论。