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

The CUDA Array Interface 使得访问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 在将设备内存控制权移交给 EM 的同时继续管理主机内存（参见主机专用 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 方法由 Numba 在首次使用 EMM 插件对象进行上下文操作之前调用。此方法应执行任何必要操作，以准备当前上下文中的底层库进行分配。此方法可能会被多次调用，并且在调用时不得使之前的状态失效。
当上下文中的所有分配需要清理时，会调用 reset 方法。它可能在 initialize 之前被调用，因此 EMM 插件实现需要对此进行防护。
为了支持GPU间通信，get_ipc_handle 方法应为给定的 MemoryPointer 实例提供一个 IpcHandle。此方法属于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插件实例所需的任何初始化，以便其准备好使用。

返回:: 无

abstract get_ipc_handle(memory)[源代码]

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

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

abstract get_memory_info()[源代码]

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

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

abstract reset()[源代码]

清除在此上下文中分配的所有内存。

返回:: 无

abstract defer_cleanup()[源代码]

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

返回:: 上下文管理器

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

主机专用 CUDA 内存管理器

一些外部内存管理器将支持设备内存的管理，但不支持主机内存。为了使用这些内存管理器实现EMM插件，提供了一个实现主机端分配和固定的插件的部分实现。要使用它，请从 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 初始化。返回一个填充了底层 ipc_mem_handle 的 numba.cuda.IpcHandle。

返回对象的类和结构

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

内存指针

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

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

内存指针可以有一个终结器，这是一个在缓冲区不再需要时调用的函数。通常，终结器会调用内存管理库（无论是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) – 所有者有时由该类的内部设置，或用于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) – 所有者有时由该类的内部设置，或用于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插件实现，默认的 None 对于在 memhostalloc 中分配的内存是足够的 - 对于 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 获得。
可以通过调用 self.context.device.get_device_identity() 来创建 source_info 的适当值。
如果底层内存不指向由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 插件的类。
返回:: 无

重置内存管理器

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

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

备注

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