外部内存管理 (EMM) 插件接口

CUDA 内置目标弃用通知

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

CUDA 数组接口使得访问 CUDA 设备的各种 Python 库之间能够共享数据。然而，每个库都独立于其他库管理自己的内存。例如

默认情况下，Numba 通过与 CUDA 驱动 API 交互来在 CUDA 设备上分配内存，调用诸如 cuMemAlloc 和 cuMemFree 等函数，这适用于许多用例。
RAPIDS 库（cuDF、cuML 等）使用 RAPIDS 内存管理器 (RMM) 来分配设备内存。
CuPy 包含一个针对设备内存和页面锁定内存的内存池实现。

当多个支持 CUDA 的库一起使用时，Numba 最好将内存管理工作交给另一个库。EMM 插件接口促成了这一点，它使 Numba 能够使用另一个支持 CUDA 的库进行所有分配和释放操作。

EMM 插件用于促进使用外部库进行内存管理。EMM 插件可以是外部库的一部分，也可以作为独立的库实现。

外部内存管理概述

当 EMM 插件启用时（参见设置 EMM 插件），Numba 将通过该插件进行内存分配和释放。它将永远不会直接调用诸如 cuMemAlloc、cuMemFree 等函数。

EMM 插件始终负责管理设备内存。然而，并非所有支持 CUDA 的库都支持管理主机内存，因此 Numba 提供了一个机制，允许 Numba 继续管理主机内存，同时将设备内存的控制权交给 EMM（参见仅主机 CUDA 内存管理器）。

对释放策略的影响

Numba 内部的释放行为旨在通过推迟释放直到累积大量待处理项来提高效率。它还提供了一种机制，使用 defer_cleanup() 上下文管理器，以在关键部分完全阻止释放。

当 EMM 插件启用时，释放策略由 EMM 实现，Numba 内部的释放机制不被使用。EMM 插件可以实现

类似于 Numba 释放行为的策略，或者
更适合该插件的策略——例如，已释放的内存可能会立即返回到内存池。

defer_cleanup 上下文管理器在与 EMM 插件一起使用时可能表现不同——EMM 插件应附带在使用时 defer_cleanup 上下文管理器行为的文档。例如，池分配器即使在使用上下文管理器时也可以始终立即将内存返回到池中，但可以选择在 defer_cleanup 不再使用之前不释放空池。

其他对象的管理

除了内存之外，Numba 还管理事件、流和模块（模块是从 @cuda.jit 修饰的函数生成的编译对象）的分配和释放。事件、流和模块的管理不受 EMM 插件使用的影响。

异步分配和释放

当前的 EMM 插件接口不支持异步分配和释放。这可能会在接口的未来版本中添加。

实现 EMM 插件

EMM 插件通过继承 BaseCUDAMemoryManager 来实现。以下是实现时的一些注意事项总结：

Numba 为每个上下文实例化一个 EMM 插件类的实例。如果需要，可以通过 self.context 访问拥有 EMM 插件对象的上下文。
EMM 插件对于任何使用 Numba 的代码都是透明的——它的所有方法都由 Numba 调用，无需被使用 Numba 的代码调用。
分配方法 memalloc、memhostalloc 和 mempin 应该使用底层库来分配和/或固定设备或主机内存，并构造一个表示要返回给 Numba 的内存的内存指针实例。这些方法始终在当前 CUDA 上下文是拥有 EMM 插件实例的上下文时被调用。
initialize 方法在 EMM 插件对象首次用于某个上下文之前由 Numba 调用。此方法应执行任何必要的准备工作，以使底层库在当前上下文中进行分配。此方法可能会被多次调用，并且在调用时不得使先前的状态失效。
reset 方法在上下文中所有分配都被清理时调用。它甚至可能在 initialize 之前被调用，EMM 插件实现需要对此进行防范。
为了支持 GPU 间通信，get_ipc_handle 方法应为给定的 IpcHandle 实例提供一个 MemoryPointer。此方法是 EMM 接口的一部分（而不是由 Numba 内部处理），因为分配的基地址仅由底层库知晓。关闭 IPC 句柄在 Numba 内部处理。
从 get_memory_info 方法提供内存信息是可选的，该方法提供上下文中设备上的总内存和可用内存计数。最好实现此方法，但这对于所有分配器可能不切实际。如果未提供内存信息，此方法应引发 RuntimeError。
defer_cleanup 方法应返回一个上下文管理器，该管理器在激活时确保避免昂贵的清理操作。这其中的细微差别因插件而异，因此插件文档应包含关于推迟清理如何影响释放以及整体性能的解释。
interface_version 属性用于确保插件版本与 Numba 版本提供的接口匹配。目前，此值应始终为 1。

基础类的完整文档如下：

class numba.cuda.BaseCUDAMemoryManager(*args, **kwargs)

外部内存管理 (EMM) 插件的抽象基类。

abstract memalloc(size)

在当前上下文中分配设备内存。

参数: size (int) – 分配的字节大小
返回: 拥有已分配内存的内存指针实例
返回类型: MemoryPointer

abstract memhostalloc(size, mapped, portable, wc)

分配页面锁定主机内存。

参数

size (int) – 分配的字节大小
mapped (bool) – 是否将分配的内存映射到 CUDA 地址空间。
portable (bool) – 内存是否会被所有上下文（而不仅仅是调用上下文）视为页面锁定。
wc (bool) – 是否将内存分配为写入合并模式。

返回

拥有已分配内存的内存指针实例。返回类型取决于该区域是否被映射到设备内存中。

返回类型

MappedMemory 或 PinnedMemory

abstract mempin(owner, pointer, size, mapped)

固定已分配的主机内存区域。

参数

owner – 拥有该内存的对象。
pointer (int) – 要固定区域起点的指针。
size (int) – 区域的字节大小。
mapped (bool) – 该区域是否也应映射到设备内存中。

返回

引用已分配内存的内存指针实例。

返回类型

MappedMemory 或 PinnedMemory

abstract initialize()

执行 EMM 插件实例准备就绪所需的任何初始化。

返回: None

abstract get_ipc_handle(memory)

从 GPU 分配中返回一个 IPC 句柄。

参数: memory (MemoryPointer) – 应为其创建 IPC 句柄的内存。
返回: 分配的 IPC 句柄
返回类型: IpcHandle

abstract get_memory_info()

返回上下文中以字节为单位的 (可用内存, 总内存)。如果返回此类信息不切实际（例如对于池分配器），则可能引发 NotImplementedError。

返回: 内存信息
返回类型: MemoryInfo

abstract reset()

清理此上下文中所有已分配的内存。

返回: None

abstract defer_cleanup()

返回一个上下文管理器，该管理器在其活动期间确保实现延迟清理。

返回: 上下文管理器

abstract property interface_version: 返回一个整数，指定插件实现支持的 EMM 插件接口版本。对于此版本规范的实现，应始终返回 1。

仅主机 CUDA 内存管理器

一些外部内存管理器将支持设备内存管理，但不支持主机内存。为了使用这些内存管理器实现 EMM 插件，Numba 提供了一个插件的部分实现，该插件实现了主机端的分配和固定。要使用它，请从 HostOnlyCUDAMemoryManager 而不是 BaseCUDAMemoryManager 派生。使用此类的指南如下：

仅主机内存管理器实现了 memhostalloc 和 mempin——EMM 插件仍应实现 memalloc。
如果重写了 reset，它还必须调用 super().reset() 以允许清理主机分配。
如果重写了 defer_cleanup，它必须持有 super().defer_cleanup() 返回的一个活动上下文管理器，以确保主机端清理也被延迟。

HostOnlyCUDAMemoryManager 类方法的文档如下：

class numba.cuda.HostOnlyCUDAMemoryManager(*args, **kwargs)

仅实现设备上分配的外部内存管理 (EMM) 插件的基类。子类无需实现 memhostalloc 和 mempin 方法。

此类还实现了 reset 和 defer_cleanup（参见 numba.cuda.BaseCUDAMemoryManager），用于其自身的内部状态管理。如果基于此类的 EMM 插件也实现了这些方法，则其实现也必须调用 super() 中的方法，以便 HostOnlyCUDAMemoryManager 能够对其管理的主机分配执行必要的工作。

此类未实现 interface_version，因为它将始终与其实现所在的 Numba 版本保持一致。子类化此类的 EMM 插件应改为实现 interface_version。

memhostalloc(size, mapped=False, portable=False, wc=False)

实现页面锁定主机内存的分配。

建议 EMM 插件实现不要重写此方法——而是使用 BaseCUDAMemoryManager。

mempin(owner, pointer, size, mapped=False)

实现主机内存的固定。

建议 EMM 插件实现不要重写此方法——而是使用 BaseCUDAMemoryManager。

reset()

清理当前上下文中所有主机内存（已映射和/或已固定）。

重写此方法的 EMM 插件必须调用 super().reset() 以确保主机分配也得到清理。

defer_cleanup()

返回一个上下文管理器，该管理器在其活动期间禁用当前上下文中已映射或已固定主机内存的清理。

重写此方法的 EMM 插件在让出控制权之前必须从此方法获取上下文管理器，以确保主机分配的清理也被延迟。

IPC 句柄混合类

get_ipc_handle() 函数的一个实现是在 GetIpcHandleMixin 类中提供的。它使用驱动 API 来确定分配的基地址，以便打开 IPC 句柄。如果此实现适用于 EMM 插件，可以通过混合 GetIpcHandleMixin 类来添加它。

class numba.cuda.GetIpcHandleMixin

提供 get_ipc_handle() 默认实现的类。

get_ipc_handle(memory): 通过使用 cuMemGetAddressRange 确定分配的基指针来打开一个 IPC 内存句柄。一个类型为 cu_ipc_mem_handle 的 IPC 句柄被构造并使用 cuIpcGetMemHandle 初始化。返回一个 numba.cuda.IpcHandle，其中填充了底层的 ipc_mem_handle。

返回对象的类和结构

本节概述了 EMM 插件需要构造的类和结构。

内存指针

EMM 插件应构造表示其分配的内存指针实例，以返回给 Numba。每个方法中要使用的适当内存指针类是：

MemoryPointer: 从 memalloc 返回
MappedMemory: 当主机内存映射到设备内存空间时，从 memhostalloc 或 mempin 返回。
PinnedMemory: 当主机内存未映射到设备内存空间时，从 memhostalloc 或 mempin 返回。

内存指针可以带有一个终结器（finalizer），该终结器是一个在不再需要缓冲区时调用的函数。通常，终结器会调用内存管理库（Numba 内部的，或如果是 EMM 插件分配的则是外部的）来告知它不再需要该内存，并且该内存可能被释放和/或解除固定。内存管理器可以选择在终结器运行后的任何时间延迟实际清理内存——它不要求立即释放缓冲区。

内存指针类的文档如下。

class numba.cuda.MemoryPointer(context, pointer, size, owner=None, finalizer=None)

一个拥有缓冲区并带有可选终结器的内存指针。内存指针提供引用计数，实例初始化时引用计数为 1。

基本的 MemoryPointer 类不使用引用计数来管理缓冲区生命周期。相反，缓冲区生命周期与内存指针实例的生命周期绑定。

当实例被删除时，将调用终结器。
当引用计数降至 0 时，不执行任何操作。

MemoryPointer 的子类可能会修改这些语义，例如将缓冲区生命周期与引用计数绑定，以便在没有更多引用时释放缓冲区。

参数

context (Context) – 指针分配所在的上下文。
pointer (ctypes.c_void_p) – 缓冲区的地址。
size (int) – 分配的字节大小。
owner (NoneType) – owner 有时由该类的内部机制设置，或用于 Numba 的内部内存管理。它不应由 MemoryPointer 类的外部用户（例如来自 EMM 插件内部）提供；默认值 None 应始终足够。
finalizer (function) – 在缓冲区将要释放时调用的函数。

AutoFreePointer 类不需要直接使用，但由于它被 numba.cuda.MappedMemory 子类化，因此在此处进行了说明。

class numba.cuda.cudadrv.driver.AutoFreePointer(*args, **kwargs)

修改 MemoryPointer 的所有权语义，使实例生命周期直接与引用计数绑定。

当引用计数达到零时，将调用终结器。

构造函数参数与 MemoryPointer 的相同。

class numba.cuda.MappedMemory(context, pointer, size, owner=None, finalizer=None)

指向主机上已映射到设备内存中的缓冲区的内存指针。

参数

context (Context) – 指针映射所在的上下文。
pointer (ctypes.c_void_p) – 缓冲区的地址。
size (int) – 缓冲区的字节大小。
owner (NoneType) – owner 有时由该类的内部机制设置，或用于 Numba 的内部内存管理。它不应由 MappedMemory 类的外部用户（例如来自 EMM 插件内部）提供；默认值 None 应始终足够。
finalizer (function) – 在缓冲区将要释放时调用的函数。

class numba.cuda.PinnedMemory(context, pointer, size, owner=None, finalizer=None)

指向主机上已固定缓冲区的指针。

参数

context (Context) – 指针映射所在的上下文。
owner – 拥有该内存的对象。对于 EMM 插件实现，这可以
pointer (ctypes.c_void_p) – 缓冲区的地址。
size (int) – 缓冲区的字节大小。
owner – 拥有已固定缓冲区的对象。对于 EMM 插件实现，memhostalloc 中分配的内存使用默认值 None 即可——对于 mempin，它应是传递给 mempin 方法的所有者。
finalizer (function) – 在缓冲区将要释放时调用的函数。

内存信息

如果 get_memory_info() 的实现要提供结果，则应返回 MemoryInfo 命名元组的实例。

class numba.cuda.MemoryInfo(free, total)

设备的可用内存和总内存。

free: 可用设备内存（字节）。

total: 总设备内存（字节）。

IPC

get_ipc_handle() 的实现必须返回 IpcHandle 的实例。

class numba.cuda.IpcHandle(base, handle, size, source_info=None, offset=0)

CUDA IPC 句柄。CUDA IPC 句柄对象的序列化在此处实现。

参数

base (MemoryPointer) – 对原始分配的引用，以保持其生命周期。
handle – CUDA IPC 句柄，作为 ctypes 字节数组。
size (int) – 原始分配的大小。
source_info (dict) – 打开 IPC 句柄的设备的标识。
offset (int) – 此 IPC 句柄所引用内存相对于底层分配的偏移量。

在实现 EMM 插件的上下文中构造 IPC 句柄的指南

传递给 EMM 插件的 get_ipc_handle 方法的 memory 参数可以作为 base 参数传递。
handle 的合适类型可以构造为 ctypes.c_byte * 64。 handle 的数据必须使用适合底层库的获取 CUDA IPC 句柄的方法来填充。
size 应与原始分配的大小匹配，这可以在 get_ipc_handle 中通过 memory.size 获取。
source_info 的适当值可以通过调用 self.context.device.get_device_identity() 来创建。
如果底层内存未指向 CUDA 驱动程序或运行时 API 返回的分配基地址（例如，如果正在使用池分配器），则必须提供相对于基地址的 offset。

设置 EMM 插件

默认情况下，Numba 使用其内部内存管理——如果要使用 EMM 插件，则必须对其进行配置。有两种机制可以配置 EMM 插件的使用：环境变量和函数。

环境变量

可以在环境变量 NUMBA_CUDA_MEMORY_MANAGER 中提供模块名称。如果设置了此环境变量，Numba 将尝试导入该模块，并使用其 _numba_memory_manager 全局变量作为内存管理器类。这主要用于使用 EMM 插件运行 Numba 测试套件，例如：

$ NUMBA_CUDA_MEMORY_MANAGER=rmm python -m numba.runtests numba.cuda.tests

函数

set_memory_manager() 函数可用于在运行时设置内存管理器。此函数应在任何上下文初始化之前调用，因为 EMM 插件实例是与上下文一起实例化的。

numba.cuda.set_memory_manager(mm_plugin)

配置 Numba 以使用外部内存管理 (EMM) 插件。如果 EMM 插件版本与此版本的 Numba 支持的版本不匹配，将引发 RuntimeError。

参数: mm_plugin (BaseCUDAMemoryManager) – 实现 EMM 插件的类。
返回: None

重置内存管理器

建议在开始使用任何 CUDA 功能之前设置一次内存管理器，并在执行的其余部分保持不变。可以多次设置内存管理器，但请注意以下几点：

在创建时，上下文会绑定到一个内存管理器实例，并在其生命周期内保持绑定。
更改内存管理器不会对现有上下文产生影响——只有在内存管理器更新后创建的上下文才会使用新内存管理器的实例。
设置内存管理器后，可以使用 numba.cuda.close() 来销毁上下文，以便使用新的内存管理器重新创建它们。
- 这将使上下文中拥有的任何数组、流、事件和模块失效。
- 尝试使用无效数组、流或事件可能会因驱动程序 API 函数返回 CUDA_ERROR_INVALID_CONTEXT 或 CUDA_ERROR_CONTEXT_IS_DESTROYED 错误代码而引发异常。
- 尝试使用无效模块将导致类似错误，或在某些情况下导致段错误/访问冲突。

注意

模块失效意味着在上下文销毁之前用 @cuda.jit 编译的所有函数都需要重新定义，因为它们底层的代码也将从 GPU 卸载。