CUDA 核函数 API

CUDA 内置目标弃用通知

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

核函数声明

@cuda.jit 装饰器用于创建可配置和启动的 CUDA 调度器对象。

numba.cuda.jit(func_or_sig=None, device=False, inline=False, link=[], debug=None, opt=True, lineinfo=False, cache=False, **kws)

为 CUDA GPU JIT 编译 Python 函数。

参数

func_or_sig –
要进行 JIT 编译的函数，或要编译函数的签名。如果提供函数，则返回 Dispatcher。否则，func_or_sig 可以是单个签名或签名列表，并返回一个函数。返回的函数接受另一个函数，该函数将被编译并返回一个 Dispatcher。有关传递签名的更多信息，请参阅 JIT 函数。

注意

核函数不能有任何返回值。
device (bool) – 指示这是否为设备函数。
link (list) – 包含 PTX 或 CUDA C/C++ 源文件的列表，用于与函数链接。
debug – 如果为 True，则在执行核函数时检查抛出的异常。由于这会降低性能，因此应仅用于调试目的。如果设置为 True，则 opt 应设置为 False。默认为 False。（默认值可以通过设置环境变量 NUMBA_CUDA_DEBUGINFO=1 来覆盖。）
fastmath – 当为 True 时，启用 CUDA Fast Math 文档中概述的快速数学优化。
max_registers – 请求将核函数限制为每个线程最多使用此数量的寄存器。如果 ABI 要求比请求更多的寄存器，则此限制可能不被遵守。对于提高占用率很有用。
opt (bool) – 是否启用优化从 LLVM IR 编译到 PTX。当为 True 时，-opt=3 传递给 NVVM。当为 False 时，-opt=0 传递给 NVVM。默认为 True。
lineinfo (bool) – 如果为 True，则生成源代码和汇编代码之间的行映射。这使得在 NVIDIA 分析工具中检查源代码并与程序计数器采样相关联成为可能。
cache (bool) – 如果为 True，则为此函数启用基于文件的缓存。

调度器对象

使用下标配置调度器以进行启动的常用语法如下，其参数如下：

# func is some function decorated with @cuda.jit
func[griddim, blockdim, stream, sharedmem]

griddim 和 blockdim 参数指定网格和线程块的大小，可以是整数或长度最多为 3 的元组。stream 参数是可选的，表示核函数将在此流上启动，而 sharedmem 参数指定动态共享内存的大小（以字节为单位）。

对调度器进行下标操作会返回一个配置对象，该对象可以通过核函数参数调用

configured = func[griddim, blockdim, stream, sharedmem]
configured(x, y, z)

然而，更惯用的做法是在一个语句中配置和调用核函数

func[griddim, blockdim, stream, sharedmem](x, y, z)

这与 CUDA C/C++ 中的启动配置类似

func<<<griddim, blockdim, sharedmem, stream>>>(x, y, z)

注意

Numba 中 stream 和 sharedmem 的顺序与 CUDA C/C++ 中相反。

调度器对象还提供了几个用于检查和创建专门实例的实用方法

class numba.cuda.dispatcher.CUDADispatcher(py_func, targetoptions, pipeline_class=<class 'numba.cuda.compiler.CUDACompiler'>)

CUDA 调度器对象。配置并调用后，调度器将根据给定的参数（如果尚无合适的专用版本）和计算能力进行专门化，并在与当前上下文关联的设备上启动。

调度器对象不由用户构造，而是使用 numba.cuda.jit() 装饰器创建。

property extensions

一个对象列表，这些对象必须具有 prepare_args 函数。当调用专用核函数时，每个参数都将通过 prepare_args（从列表中最后一个对象到第一个对象）传递。传递给 prepare_args 的参数是：

ty 参数的 numba 类型
val 参数值本身
stream 用于当前核函数调用的 CUDA 流
retr 一个零参数函数列表，你可以向其追加调用后的清理工作。

prepare_args 函数必须返回一个元组 (ty, val)，该元组将依次传递给下一个最右侧的 extension。在所有扩展都被调用后，结果 (ty, val) 将传递给 Numba 的默认参数调度逻辑。

forall(ntasks, tpb=0, stream=0, sharedmem=0)

返回针对给定任务数量的 1D 配置调度器。

这假定：

核函数将全局线程 ID cuda.grid(1) 与任务一对一映射。
核函数检查全局线程 ID 不超过 ntasks，如果不符合则不执行任何操作。

参数

ntasks – 任务数量。
tpb – 块的大小。如果未提供此参数，则选择适当的值。
stream – 配置好的调度器将在此流上启动。
sharedmem – 核函数所需的动态共享内存字节数。

返回

一个配置好的调度器，准备好在一组参数上启动。

get_const_mem_size(signature=None)

返回此核函数在当前上下文的设备上使用的常量内存大小（以字节为单位）。

参数: signature – 要获取常量内存用量的已编译核函数的签名。对于专用核函数，此项可以省略。
返回: 针对给定签名和当前设备的已编译核函数变体所分配的常量内存大小（以字节为单位）。

get_local_mem_per_thread(signature=None)

返回此核函数每个线程的局部内存大小（以字节为单位）。

参数: signature – 要获取局部内存用量的已编译核函数的签名。对于专用核函数，此项可以省略。
返回: 针对给定签名和当前设备的已编译核函数变体所分配的局部内存量。

get_max_threads_per_block(signature=None)

返回此核函数每个块允许的最大线程数。超过此阈值将导致核函数无法启动。

参数: signature – 要获取每个块最大线程数的已编译核函数的签名。对于专用核函数，此项可以省略。
返回: 针对给定签名和当前设备的已编译核函数变体允许的最大每个块线程数。

get_regs_per_thread(signature=None)

返回此核函数中每个线程在当前上下文的设备上使用的寄存器数量。

参数: signature – 要获取寄存器用量的已编译核函数的签名。对于专用核函数，此项可以省略。
返回: 针对给定签名和当前设备的已编译核函数变体使用的寄存器数量。

get_shared_mem_per_block(signature=None)

返回此核函数静态分配的共享内存大小（以字节为单位）。

参数: signature – 要获取共享内存用量的已编译核函数的签名。对于专用核函数，此项可以省略。
返回: 针对给定签名和当前设备的已编译核函数变体所分配的共享内存量。

inspect_asm(signature=None)

返回此核函数在当前上下文的设备上的 PTX 汇编代码。

参数: signature – 参数类型的元组。
返回: 给定签名的 PTX 代码，或所有先前遇到的签名的 PTX 代码字典。

inspect_llvm(signature=None)

返回此核函数的 LLVM IR。

参数: signature – 参数类型的元组。
返回: 给定签名的 LLVM IR，或所有先前遇到的签名的 LLVM IR 字典。

inspect_sass(signature=None)

返回此核函数在当前上下文的设备上的 SASS 汇编代码。

参数: signature – 参数类型的元组。
返回: 给定签名的 SASS 代码，或所有先前遇到的签名的 SASS 代码字典。

返回当前上下文设备上的 SASS。

需要 nvdisasm 在 PATH 中可用。

inspect_types(file=None): 生成此函数的 Python 源代码转储，并附带相应的 Numba IR 和类型信息。转储写入到 file，如果 file 为 None 则写入 sys.stdout。

specialize(*args): 创建此调度器的新实例，并针对给定 args 进行专门化。

property specialized: 如果调度器已专门化，则为 True。

内在属性和函数

本节其余属性和函数只能在 CUDA 核函数内部调用。

线程索引

numba.cuda.threadIdx: 当前线程块中的线程索引，通过属性 x, y 和 z 访问。每个索引都是一个整数，范围从 0（包含）到 numba.cuda.blockDim 中对应属性值（不包含）的范围。

numba.cuda.blockIdx: 线程块网格中的块索引，通过属性 x, y 和 z 访问。每个索引都是一个整数，范围从 0（包含）到 numba.cuda.gridDim 中对应属性值（不包含）的范围。

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

numba.cuda.gridDim: 块网格的形状，通过属性 x, y 和 z 访问。

numba.cuda.laneid: 当前 warp 中的线程索引，一个整数，范围从 0（包含）到 numba.cuda.warpsize（不包含）。

numba.cuda.warpsize: GPU 上一个 warp 的线程大小。目前始终为 32。

numba.cuda.grid(ndim)

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

第一个整数的计算方式如下：

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

其他两个索引的计算方式类似，但使用 y 和 z 属性。

numba.cuda.gridsize(ndim)

返回整个块网格的绝对大小（或形状）（以线程为单位）。ndim 应与实例化核函数时声明的维度数量相对应。

第一个整数的计算方式如下：

cuda.blockDim.x * cuda.gridDim.x

其他两个索引的计算方式类似，但使用 y 和 z 属性。

内存管理

numba.cuda.shared.array(shape, dtype)

使用给定的 shape 和 dtype 在 CUDA 核函数的局部内存空间中创建一个数组。

返回一个内容未初始化的数组。

注意

同一线程块中的所有线程都看到相同的数组。

numba.cuda.local.array(shape, dtype)

使用给定的 shape 和 dtype 在 CUDA 核函数的局部内存空间中创建一个数组。

返回一个内容未初始化的数组。

注意

每个线程看到一个唯一的数组。

numba.cuda.const.array_like(ary)

在编译时将 ary 复制到 CUDA 核函数的常量内存空间中。

返回一个类似于 ary 参数的数组。

注意

所有线程和块都看到相同的数组。

同步和原子操作

numba.cuda.atomic.add(array, idx, value)

执行 array[idx] += value。仅支持 int32、int64、float32 和 float64。idx 参数可以是整数，也可以是整数索引元组，用于多维数组索引。idx 中的元素数量必须与 array 的维度数量匹配。