原文: http://numba.pydata.org/numba-doc/latest/cuda/kernels.html
与用于编程 CPU 的传统顺序模型不同,CUDA 具有执行模型。在 CUDA 中,您编写的代码将由多个线程同时执行(通常为数百或数千)。您的解决方案将通过定义 grid , blocks 和 threads 的线程层次结构来建模。
Numba 的 CUDA 支持公开了用于声明和管理这种线程层次结构的工具。这些设施与 NVidia 的 CUDA C 语言大致相似。
Numba 还暴露了三种 GPU 内存:全局设备内存(连接到 GPU 本身的大型,相对较慢的片外内存),片上共享内存和 []本地记忆](memory.html#cuda-local-memory)。除了最简单的算法外,您必须仔细考虑如何使用和访问内存,以最大限度地减少带宽需求和争用。
_ 内核函数 _ 是一个 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 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 到包括 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
请注意,实例化内核时的网格计算仍必须手动完成,例如:
from __future__ import division # for Python 2
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 编程的详细讨论,请参见 CUDA C 编程指南。