CUDA 主机 API

设备管理

设备检测与查询

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

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 上，设备 2 由 d_a 引用。

也可以通过以下三种函数选择上下文和设备，或获取当前设备：

numba.cuda.select_device(device_id)

将关联到设备 device_id 的上下文设置为当前上下文。

返回一个设备实例。

在错误时引发异常。

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 代码中的 PTX（例如来自 Numba / Python 生态系统之外的代码）。

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

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

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

备注

用户需自行管理由于编译为PTX/LTO-IR而产生的任何ABI问题。传递``abi=”c”``关键字参数可以解决大多数可能出现的问题 - 参见 cuda-使用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=, 和 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"。

返回:

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

返回类型:

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 Stream 中的外部流的指针

运行时

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

class numba.cuda.cudadrv.runtime.Runtime

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

get_version()

返回 CUDA 运行时版本作为一个元组 (主要版本, 次要版本)。

is_supported_version()

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

property supported_versions

所有支持的 CUDA 工具包版本的元组。版本以 (主版本, 次版本) 的形式给出。

当前运行时是否与当前版本的 Numba 正式支持并经过测试，也可以进行查询：

numba.cuda.is_supported_version()

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

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

• 继续默默地，

• 发出警告，

• 生成错误或阻止使用CUDA。