CUDA 内核 API

内核声明

@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即时编译Python函数。

参数:

• func_or_sig – 一个用于JIT编译的函数，或编译的函数的*签名*。如果提供了一个函数，则返回一个 调度器。否则，func_or_sig 可能是一个签名或签名列表，并返回一个函数。返回的函数接受另一个函数，它将编译该函数并返回一个 调度器。有关传递签名的更多信息，请参见 JIT 函数。 .. 注意:: 内核不能有任何返回值。

• device (bool) – 指示这是否是一个设备函数。

• link (list) – 包含PTX或CUDA C/C++源代码的文件列表，用于与函数链接

• debug – 如果为 True，则在执行内核时检查抛出的异常。由于这会降低性能，因此仅应用于调试目的。如果设置为 True，则 opt 应设置为 False。默认为 False。（默认值可以通过设置环境变量 NUMBA_CUDA_DEBUGINFO=1 来覆盖。）

• fastmath – 当为 True 时，启用 CUDA Fast Math 文档 中概述的 fastmath 优化。

• max_registers – 请求内核每个线程最多使用此数量的寄存器。如果ABI要求的寄存器数量大于请求的数量，则可能不会遵守此限制。这对于提高占用率很有用。

• opt (bool) – 是否在启用优化的情况下从LLVM IR编译为PTX。当为``True``时，-opt=3``传递给NVVM。当为``False``时，-opt=0``传递给NVVM。默认为``True``。

• lineinfo (bool) – 如果为真，生成源代码与汇编代码之间的行映射。这使得可以在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。

需要在 PATH 中提供 nvdisasm。

inspect_types(file=None)

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

specialize(*args)

为给定的 args 创建此调度程序的新实例。

property specialized

如果调度器已被特化，则为真。

内在属性和函数

本节中其余的属性和函数只能在 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（包含）到 :attr:`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)

在CUDA内核的本地内存空间中创建一个具有给定 shape 和 dtype 的数组。

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

备注

同一个线程块中的所有线程看到的是同一个数组。

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

在CUDA内核的本地内存空间中创建一个具有给定 shape 和 dtype 的数组。

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

备注

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

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 的维度数量匹配。

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.sub(array, idx, value)

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

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.and_(array, idx, value)

执行 array[idx] &= value。仅支持 int32、uint32、int64 和 uint64。idx 参数可以是整数或用于多维数组索引的整数索引元组。idx 中的元素数量必须与 array 的维数匹配。

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.or_(array, idx, value)

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

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.xor(array, idx, value)

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

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.exch(array, idx, value)

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

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.inc(array, idx, value)

执行 array[idx] = (0 if array[idx] >= value else array[idx] + 1)。仅支持 uint32 和 uint64。idx 参数可以是整数或用于索引多维数组的整数元组。idx 中的元素数量必须与 array 的维度数量匹配。

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.dec(array, idx, value)

执行 array[idx] = (value if (array[idx] == 0) or (array[idx] > value) else array[idx] - 1)。仅支持 uint32 和 uint64。idx 参数可以是整数或用于索引多维数组的整数索引元组。idx 中的元素数量必须与 array 的维度数量匹配。

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.max(array, idx, value)

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

返回存储新值之前的 array[idx] 的值。行为类似于原子加载。

numba.cuda.atomic.cas(array, idx, old, value)

执行 if array[idx] == old: array[idx] = value。仅支持 int32、int64、uint32、uint64 索引。idx 参数可以是整数或用于多维数组索引的整数索引元组。idx 中的元素数量必须与 array 的维度数量匹配。

返回存储新值之前的 array[idx] 的值。行为类似于原子比较和交换。

numba.cuda.syncthreads()

同步同一线程块中的所有线程。此函数实现了与传统多线程编程中的屏障相同的模式：此函数等待块中的所有线程调用它，此时它将控制权返回给所有调用者。

numba.cuda.syncthreads_count(predicate)

对 numba.cuda.syncthreads 的扩展，其中返回值是 predicate 为真的线程数。

numba.cuda.syncthreads_and(predicate)

对 numba.cuda.syncthreads 的扩展，如果所有线程的 predicate 为真，则返回 1，否则返回 0。

numba.cuda.syncthreads_or(predicate)

对 numba.cuda.syncthreads 的扩展，如果 predicate 对任何线程为真则返回 1，否则返回 0。

警告

所有 syncthreads 函数必须由线程块中的每个线程调用。未能这样做可能会导致未定义的行为。

协作组

numba.cuda.cg.this_grid()

获取当前网格组。

返回:

当前网格组

返回类型:

numba.cuda.cg.GridGroup

class numba.cuda.cg.GridGroup

一个网格组。用户不应直接构造 GridGroup - 而是使用 cg.this_grid() 获取当前网格组。

sync()

同步当前网格组。

内存栅栏

内存栅栏用于保证内存操作的效果对同一线程块、同一GPU设备以及同一系统（跨GPU的全局内存）中的其他线程可见。优化过程保证内存加载和存储不会跨越内存栅栏。

警告

内存栅栏被认为是高级API，大多数用例应使用线程屏障（例如 syncthreads()）。

numba.cuda.threadfence()

设备级别的内存栅栏（在GPU内部）。

numba.cuda.threadfence_block()

线程块级别的内存屏障。

numba.cuda.threadfence_system()

系统级别的内存栅栏（跨GPU）。

Warp 内部函数

参数 membermask 是一个 32 位整数掩码，每个位对应于 warp 中的一个线程，1 表示该线程在函数调用中的线程子集中。如果 GPU 计算能力低于 7.x，则 membermask 必须全为 1。

numba.cuda.syncwarp(membermask)

同步一个warp中线程的掩码子集。

numba.cuda.all_sync(membermask, predicate)

如果 predicate 对掩码线程块中的所有线程都为真，则返回一个非零值，否则返回0。

numba.cuda.any_sync(membermask, predicate)

如果在掩码的warp中任何线程的 predicate 为真，则返回一个非零值，否则返回0。

numba.cuda.eq_sync(membermask, predicate)

如果在掩码线程束中所有线程的布尔 predicate 相同，则返回非零值，否则返回0。

numba.cuda.ballot_sync(membermask, predicate)

返回在给定掩码内，其 predicate 为真的线程束中所有线程的掩码。

numba.cuda.shfl_sync(membermask, value, src_lane)

在掩码的warp中打乱 value 并返回来自 src_lane 的 value 。如果超出warp范围，则返回给定的 value 。

numba.cuda.shfl_up_sync(membermask, value, delta)

在掩码的warp中打乱 value ，并返回来自 laneid - delta 的 value 。如果超出warp范围，则返回给定的 value 。

numba.cuda.shfl_down_sync(membermask, value, delta)

在掩码的warp中打乱 value ，并返回来自 laneid + delta 的 value 。如果超出warp范围，则返回给定的 value 。

numba.cuda.shfl_xor_sync(membermask, value, lane_mask)

在掩码的warp中打乱 value ，并返回来自 laneid ^ lane_mask 的 value 。

numba.cuda.match_any_sync(membermask, value, lane_mask)

返回一个掩码，该掩码表示在掩码的线程束中具有与给定 value 相同的 value 的线程。

numba.cuda.match_all_sync(membermask, value, lane_mask)

返回一个元组 (mask, pred)，其中 mask 是一个线程掩码，表示在掩码 warp 内具有与给定 value 相同 value 的线程，如果它们都具有相同的值，否则为 0。而 pred 是一个布尔值，表示掩码 warp 中的所有线程是否具有相同的 warp。

numba.cuda.activemask()

返回调用 warp 中所有当前活动线程的 32 位整数掩码。当调用 activemask() 时，如果 warp 中的第 N 个通道是活动的，则第 N 位被设置。不活动的线程在返回的掩码中用 0 位表示。已经退出内核的线程总是被标记为不活动的。

numba.cuda.lanemask_lt()

返回一个32位整数掩码，包含所有ID小于当前通道的通道（包括非活动通道）。

整数内在函数

CUDA Math API 的一部分整数内在函数是可用的。如需进一步的文档，包括语义，请参阅 CUDA Toolkit 文档。

numba.cuda.popc(x)

返回 x 中设置的位数。

numba.cuda.brev(x)

返回 x 的位模式反转。例如，0b10110110 变为 0b01101101。

numba.cuda.clz(x)

返回 x 中前导零的数量。

numba.cuda.ffs(x)

返回 x 中第一个（最低有效）位设置为 1 的位置，其中最低有效位位置为 1。ffs(0) 返回 0。

浮点数内部函数

CUDA Math API 的一部分浮点内建函数是可用的。如需进一步的文档，包括语义，请参阅 CUDA Toolkit 文档中的 单精度 和 双精度 部分。

numba.cuda.fma()

执行融合乘加操作。以C API中的 fma 和 fmaf 命名，但映射到PTX指令 fma.rn.f32 和 fma.rn.f64 （四舍五入到最近偶数）。

numba.cuda.cbrt(x)

执行立方根操作，x ** (1/3)。以C API中的函数 cbrt 和 cbrtf 命名。仅支持 float32 和 float64 参数。

16位浮点内建函数

cuda.fp16 模块中的函数用于对16位浮点操作数进行操作。这些函数返回一个16位浮点结果。

要确定在当前配置下 Numba 是否支持编译使用 float16 类型的代码，请使用：

numba.cuda.is_float16_supported()

如果支持16位浮点数，则返回 True ，否则返回 False 。

要检查设备是否支持 float16，请使用其 supports_float16 属性。

numba.cuda.fp16.hfma(a, b, c)

在舍入到最近模式下，对16位浮点参数执行融合乘加操作 (a * b) + c。映射到 fma.rn.f16 PTX 指令。

返回融合乘加运算的16位浮点结果。

numba.cuda.fp16.hadd(a, b)

在四舍五入到最近的模式下，对16位浮点参数执行加法操作 a + b。对应于 add.f16 PTX 指令。

返回加法的16位浮点结果。

numba.cuda.fp16.hsub(a, b)

在四舍五入到最近的模式下，对16位浮点参数执行减法操作 a - b。对应于 sub.f16 PTX 指令。

返回减法操作的16位浮点结果。

numba.cuda.fp16.hmul(a, b)

在四舍五入到最近的模式下，对16位浮点参数执行乘法操作 a * b。映射到 mul.f16 PTX 指令。

返回乘法的16位浮点结果。

numba.cuda.fp16.hdiv(a, b)

在四舍五入模式下对16位浮点参数执行除法操作 a / b。

返回除法的16位浮点结果。

numba.cuda.fp16.hneg(a)

对16位浮点参数执行取反操作 -a 。对应于 neg.f16 PTX 指令。

返回否定结果的16位浮点数。

numba.cuda.fp16.habs(a)

对16位浮点参数执行绝对值操作 |a| 。

返回绝对值操作的16位浮点结果。

numba.cuda.fp16.hsin(a)

计算16位浮点数参数的三角正弦函数。

返回正弦运算的16位浮点结果。

numba.cuda.fp16.hcos(a)

计算16位浮点数参数的三角余弦函数。

返回余弦操作的16位浮点结果。

numba.cuda.fp16.hlog(a)

计算16位浮点参数的自然对数。

返回自然对数运算的16位浮点结果。

numba.cuda.fp16.hlog10(a)

计算16位浮点参数的以10为底的对数。

返回以10为底的对数运算的16位浮点结果。

numba.cuda.fp16.hlog2(a)

计算16位浮点参数的以2为底的对数。

返回以2为底的对数运算的16位浮点结果。

numba.cuda.fp16.hexp(a)

计算16位浮点数参数的自然指数运算。

返回指数运算的16位浮点结果。

numba.cuda.fp16.hexp10(a)

计算16位浮点参数的以10为底数的指数。

返回指数运算的16位浮点结果。

numba.cuda.fp16.hexp2(a)

计算16位浮点参数的以2为底的指数。

返回指数运算的16位浮点结果。

numba.cuda.fp16.hfloor(a)

计算16位浮点参数上的地板操作，即小于或等于``a``的最大整数。

返回 floor 操作的 16 位浮点结果。

numba.cuda.fp16.hceil(a)

计算16位浮点参数上的天花板操作，即大于或等于``a``的最小整数。

返回 ceil 操作的 16 位浮点结果。

numba.cuda.fp16.hsqrt(a)

计算16位浮点参数的平方根运算。

返回平方根运算的16位浮点结果。

numba.cuda.fp16.hrsqrt(a)

计算16位浮点参数的平方根的倒数。

返回倒数平方根操作的16位浮点结果。

numba.cuda.fp16.hrcp(a)

计算16位浮点参数的倒数。

返回倒数的16位浮点结果。

numba.cuda.fp16.hrint(a)

将输入的16位浮点数参数四舍五入到最近的整数值。

返回四舍五入后的16位浮点结果。

numba.cuda.fp16.htrunc(a)

将输入的16位浮点数参数截断为不超过输入参数大小的最接近的整数。

返回截断后的16位浮点结果。

numba.cuda.fp16.heq(a, b)

对16位浮点参数执行比较操作 a == b。

返回一个布尔值。

numba.cuda.fp16.hne(a, b)

对16位浮点参数执行比较操作 a != b。

返回一个布尔值。

numba.cuda.fp16.hgt(a, b)

对16位浮点参数执行比较操作 a > b。

返回一个布尔值。

numba.cuda.fp16.hge(a, b)

对16位浮点参数执行比较操作 a >= b。

返回一个布尔值。

numba.cuda.fp16.hlt(a, b)

对16位浮点参数执行比较操作 a < b。

返回一个布尔值。

numba.cuda.fp16.hle(a, b)

对16位浮点参数执行比较操作 a <= b。

返回一个布尔值。

numba.cuda.fp16.hmax(a, b)

执行操作 a if a > b else b.

返回一个16位浮点数值。

numba.cuda.fp16.hmin(a, b)

执行操作 a if a < b else b.

返回一个16位浮点数值。

控制流指令

CUDA 控制流指令的一个子集可以直接作为内部函数使用。避免分支是提高 CUDA 性能的关键方法，使用这些内部函数意味着你不必依赖 nvcc 优化器来识别和移除分支。有关进一步的文档，包括语义，请参阅 相关的 CUDA 工具包文档。

numba.cuda.selp()

根据第一个参数的值，在两个表达式之间进行选择。类似于LLVM的``select``指令。

计时器内在函数

numba.cuda.nanosleep(ns)

将线程暂停一个接近于指定延迟 ns 的睡眠时间，单位为纳秒。