CUDA 主机 API

CUDA 内置目标弃用通知

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

设备管理

设备检测与查询

以下函数可用于查询可用硬件

numba.cuda.is_available()

返回一个布尔值，指示 CUDA GPU 的可用性。

如果驱动尚未初始化，此操作将初始化驱动。

numba.cuda.detect()

检测支持的 CUDA 硬件并打印检测到的硬件摘要。

返回一个布尔值，指示是否检测到任何受支持的设备。

上下文管理

CUDA Python 函数在 CUDA 上下文中执行。系统中的每个 CUDA 设备都有一个关联的 CUDA 上下文，Numba 目前每个线程只允许一个上下文。有关 CUDA 上下文的更多详细信息，请参阅 CUDA 驱动 API 上下文管理文档 和 CUDA C 编程指南上下文文档。CUDA 上下文是 Context 类的实例

class numba.cuda.cudadrv.driver.Context(device, handle)

此对象封装了一个 CUDA 上下文资源。

用户代码不应直接构造上下文。

get_memory_info()

返回上下文中（空闲，总计）的内存字节数。

pop()

将此上下文从当前 CPU 线程中弹出。请注意，此上下文必须位于上下文堆栈的顶部，否则将发生错误。

push()

将此上下文压入当前 CPU 线程。

reset()

清理此上下文中所有拥有的资源。

以下函数可用于获取或选择上下文

numba.cuda.current_context(devnum=None)

获取当前设备或按设备编号使用设备，并返回 CUDA 上下文。

numba.cuda.require_context(fn)

一个装饰器，确保在执行 *fn* 时 CUDA 上下文可用。

注意：函数 *fn* 不能切换 CUDA 上下文。

以下函数影响当前上下文

numba.cuda.synchronize()

同步当前上下文。

numba.cuda.close()

明确清除当前线程中的所有上下文，如果当前线程是主线程，则销毁所有上下文。

设备管理

Numba 维护一个支持 CUDA 的设备列表

numba.cuda.gpus

一个可索引的受支持 CUDA 设备列表。此列表通过整数设备 ID 进行索引。

另外，可以获取当前设备

numba.cuda.gpus.current

当前选定的设备。

通过 numba.cuda.gpus 获取设备总是提供 numba.cuda.cudadrv.devices._DeviceContextManager 的实例，它充当所选设备的上下文管理器

class numba.cuda.cudadrv.devices._DeviceContextManager(device)

提供了一个上下文管理器，用于在所选设备的上下文中执行。此类型实例的正常用法来自 numba.cuda.gpus。例如，要在设备 2 上执行

with numba.cuda.gpus[2]:
    d_a = numba.cuda.to_device(a)

将数组 *a* 复制到设备 2 上，由 *d_a* 引用。

还可以使用以下三个函数选择上下文和设备或获取当前设备

numba.cuda.select_device(device_id)

将与设备 *device_id* 关联的上下文设为当前上下文。

返回一个 Device 实例。

出错时引发异常。

numba.cuda.get_current_device()

获取与当前线程关联的当前设备

numba.cuda.list_devices()

返回所有检测到的设备列表

numba.cuda.cudadrv.driver.Device 类可用于查询所选设备的功能

class numba.cuda.cudadrv.driver.Device

与特定上下文关联的设备。

compute_capability

一个元组，*(major, minor)*，指示支持的计算能力。

id

设备的整数 ID。

name

设备的名称（例如“GeForce GTX 970”）。

uuid

设备的 UUID（例如“GPU-e6489c45-5b68-3b03-bab7-0e7c8e809643”）。

reset()

删除设备的上下文。这将销毁上下文中创建的所有内存分配、事件和流。

supports_float16

如果设备支持 float16 操作则返回 True，否则返回 False。

编译

Numba 提供了一个入口点，用于编译 Python 函数而无需调用任何驱动程序 API。这对于以下情况很有用：

• 生成将内联到其他 PTX 代码（例如来自 Numba / Python 生态系统外部）中的 PTX。

• 生成 PTX 或 LTO-IR，以便与非 Python 翻译单元中的对象链接。

• 在没有设备存在时生成代码。

• 在 fork 之前生成代码而不初始化 CUDA。

注意

用户有责任管理因编译到 PTX / LTO-IR 而产生的任何 ABI 问题。传递 abi="c" 关键字参数可以解决大多数可能出现的问题 - 请参阅 使用 C ABI。

numba.cuda.compile(pyfunc, sig, debug=False, lineinfo=False, device=True, fastmath=False, cc=None, opt=True, abi='c', abi_info=None, output='ptx')

将 Python 函数编译为给定参数类型的 PTX 或 LTO-IR。

参数

• pyfunc – 要编译的 Python 函数。

• sig – 表示函数输入和输出类型的签名。如果这是一个不包含返回类型的参数类型元组，则此函数将返回推断的返回类型。如果传递了包含返回类型的签名，则编译后的代码将包含从推断返回类型到指定返回类型的转换，并且此函数将返回指定的返回类型。

• debug (bool) – 是否在编译后的代码中包含调试信息。

• lineinfo (bool) – 是否包含从编译代码到源代码的行映射。通常，这与优化代码一起使用（因为调试模式会自动包含此功能），因此我们希望在 LLVM IR 中包含调试信息，但最终输出中只包含行映射。

• device (bool) – 是否编译设备函数。

• fastmath (bool) – 是否启用快速数学标志（ftz=1, prec_sqrt=0, prec_div=, and fma=1）

• cc (tuple) – 要编译的计算能力，以元组 (MAJOR, MINOR) 的形式。默认为 (5, 0)。

• opt (bool) – 启用优化。默认为 True。

• abi (str) – 编译函数的 ABI - "numba" 或 "c"。请注意，Numba ABI 不被认为是稳定的。目前 C ABI 仅支持设备函数。

• abi_info (dict) – 一个包含 ABI 特定选项的字典。 "c" ABI 支持一个选项，"abi_name"，用于提供包装函数的名称。 "numba" ABI 没有选项。

• output (str) – 要生成的输出类型，可以是 "ptx" 或 "ltoir"。

返回

(code, resty)：编译后的代码和推断的返回类型

返回类型

tuple

环境变量 NUMBA_CUDA_DEFAULT_PTX_CC 可以设置来控制 compile 目标默认计算能力 - 请参阅 GPU 支持。如果需要当前设备的计算能力代码，可以使用 compile_for_current_device 函数

numba.cuda.compile_for_current_device(pyfunc, sig, debug=False, lineinfo=False, device=True, fastmath=False, opt=True, abi='c', abi_info=None, output='ptx')

将 Python 函数编译为 PTX 或 LTO-IR，针对当前设备的计算能力给出签名。这会调用 compile() 并为当前设备提供适当的 cc 值。

Numba 还提供了两个函数，可在专门只编译为 PTX 的旧代码中使用

numba.cuda.compile_ptx(pyfunc, sig, debug=False, lineinfo=False, device=False, fastmath=False, cc=None, opt=True, abi='numba', abi_info=None)

将 Python 函数编译为给定签名的 PTX。参见 compile()。此函数的默认设置是使用 Numba ABI 编译内核，而不是 compile() 默认的用 C ABI 编译设备函数。

numba.cuda.compile_ptx_for_current_device(pyfunc, sig, debug=False, lineinfo=False, device=False, fastmath=False, opt=True, abi='numba', abi_info=None)

将 Python 函数编译为给定签名的 PTX，针对当前设备的计算能力。参见 compile_ptx()。

测量

性能分析

NVidia Visual Profiler 可直接用于执行 CUDA Python 代码——无需在用户代码中插入对这些函数的调用。然而，这些函数可用于允许对代码的特定部分进行选择性性能分析。有关性能分析的更多信息，请参阅 NVidia Profiler 用户指南。

numba.cuda.profile_start()

在当前上下文中启用性能分析数据收集。

numba.cuda.profile_stop()

在当前上下文中禁用性能分析数据收集。

numba.cuda.profiling()

一个上下文管理器，在进入时启用性能分析，在退出时禁用性能分析。

事件

事件可用于监视执行进度并记录到达特定点的时间戳。事件创建立即返回，并且可以查询创建的事件以确定是否已到达。有关更多信息，请参阅 CUDA C 编程指南事件部分。

以下函数用于创建事件和测量事件之间的时间

numba.cuda.event(timing=True)

创建一个 CUDA 事件。只有在创建时设置 timing=True，事件才会记录计时数据。

numba.cuda.event_elapsed_time(evtstart, evtend)

计算两个事件之间经过的时间（以毫秒为单位）。

事件是 numba.cuda.cudadrv.driver.Event 类的实例

class numba.cuda.cudadrv.driver.Event(context, handle, finalizer=None)

query()

如果最近一次记录之前的所有工作都已完成，则返回 True；否则返回 False。

record(stream=0)

将事件的记录点设置为给定流中的当前点。

当在调用 record() 时在流中排队的所有工作都已完成时，该事件将被视为已发生。

synchronize()

同步主机线程以等待事件完成。

wait(stream=0)

所有未来提交到流中的工作都将等待事件完成。

流管理

流允许在给定上下文中的单个设备上并发执行。同一流中排队的工作项按顺序执行，但不同流中的工作项可以并发执行。大多数涉及 CUDA 设备的操作都可以使用流异步执行，包括数据传输和内核执行。有关流的更多详细信息，请参阅 CUDA C 编程指南流部分。

Numba 默认使用传统默认流作为默认流。通过将环境变量 NUMBA_CUDA_PER_THREAD_DEFAULT_STREAM 设置为 1，可以将每线程默认流设为默认流（请参阅 CUDA 环境变量部分）。无论此设置如何，都可以使用以下函数构造代表传统默认流和每线程默认流的对象。

流是 numba.cuda.cudadrv.driver.Stream 的实例

class numba.cuda.cudadrv.driver.Stream(context, handle, finalizer, external=False)

add_callback(callback, arg=None)

向计算流添加回调。一旦所有前面的流操作完成，用户提供的函数将从驱动程序线程中调用。

回调函数是从 CUDA 驱动程序线程调用的，而不是从调用 add_callback 的线程调用的。在回调函数内部不能调用任何 CUDA API 函数。

回调函数的持续时间应保持简短，因为回调会阻塞流中后续的工作，并可能阻止其他回调的执行。

注意：此方法底层的驱动函数已被标记为最终弃用，并可能在未来的 CUDA 版本中被替换。

参数

• callback – 带有参数 (stream, status, arg) 的回调函数。

• arg – 可选的用户数据，将传递给回调函数。

async_done() → Future

返回一个可等待对象，一旦所有前面的流操作完成即解析。可等待对象的结果是当前流。

auto_synchronize()

一个上下文管理器，在退出上下文时等待此流中的所有命令执行完毕并提交所有挂起的内存传输。

synchronize()

等待此流中的所有命令执行完毕。这将提交所有挂起的内存传输。

要创建新流

numba.cuda.stream()

创建一个 CUDA 流，它表示设备的命令队列。

要获取默认流

numba.cuda.default_stream()

获取默认 CUDA 流。通常，CUDA 语义规定默认流是传统默认流或每线程默认流，具体取决于正在使用的 CUDA API。在 Numba 中，传统默认流的 API 始终在使用中，但未来可能会提供使用每线程默认流 API 的选项。

要获取默认流并明确选择它是传统默认流还是每线程默认流

numba.cuda.legacy_default_stream()

获取传统默认 CUDA 流。

numba.cuda.per_thread_default_stream()

获取每线程默认 CUDA 流。

为了使用在其他地方分配的流构造 Numba Stream 对象，提供了 external_stream 函数。请注意，外部流的生命周期必须由用户管理——Numba 不会释放外部流，并且在 Numba Stream 对象使用期间，该流必须保持有效。

numba.cuda.external_stream(ptr)

为 Numba 外部分配的流创建 Numba 流对象。

参数

ptr (int) – 指向要在 Numba 流中封装的外部流的指针

运行时

Numba 通常使用 Driver API，但它提供了 Runtime API 的一个简单包装，以便可以查询正在使用的运行时版本。可以通过 cuda.runtime 访问，它是 numba.cuda.cudadrv.runtime.Runtime 类的一个实例

class numba.cuda.cudadrv.runtime.Runtime

延迟绑定运行时 API 函数的运行时对象。

get_version()

以元组 (major, minor) 形式返回 CUDA 运行时版本。

is_supported_version()

如果 CUDA 运行时是受支持的版本，则返回 True。

property supported_versions

一个包含所有支持的 CUDA 工具包版本的元组。版本以 (major_version, minor_version) 的形式给出。

还可以查询当前运行时是否被 Numba 当前版本官方支持和测试

numba.cuda.is_supported_version()

如果 CUDA 运行时是受支持的版本，则返回 True。

不支持的版本（例如，比 Numba 已知的版本更新的版本）可能仍然有效；此函数提供了一个工具来检查当前 Numba 版本是否已测试并已知与当前运行时版本兼容。如果当前版本不受支持，调用者可以决定如何处理。选项包括

• 静默继续，

• 发出警告，

• 生成错误或以其他方式阻止 CUDA 的使用。