BEP 7：CUDA外部内存管理插件（下）
Numba依赖
向库中添加EMM插件的实现自然会使Numba成为库的依赖项，而以前可能没有。为了使依赖关系可选，如果需要的话，可以有条件地实例化并注册EMM插件，如下所示：
try:
import numba
from mylib.numba_utils import MyNumbaMemoryManager
numba.cuda.cudadrv.driver.set_memory_manager(MyNumbaMemoryManager)
except:
print(“Numba not importable – not registering EMM Plugin”)
因此，mylib.numba_utils仅在Numba已经存在的情况下才导入包含EMM插件实现的。如果Numba不可用，则将永远不会导入mylib.numba_utils（必须是import numba）。
建议所有带有EMM插件的库至少包括一些带有Numba的环境，以便在使用中使用EMM插件进行测试，以及一些不带Numba的环境，以避免引入意外的Numba依赖关系。
示例实现-RAPIDS内存管理器（RMM）插件
本节概述了Rapids Memory Manager（RMM）中EMM插件的实现。这旨在显示实现的概述，以支持上述说明，并说明如何使用插件API-可以为产品就绪的实现做出不同的选择。
插件实现包括python / rmm / rmm.py的附加内容：
# New imports:
from contextlib import context_manager
# RMM already has Numba as a dependency, so these imports need not be guarded
# by a check for the presence of numba.
from numba.cuda import (HostOnlyCUDAMemoryManager, MemoryPointer, IpcHandle,
set_memory_manager)
# New class implementing the EMM Plugin:
class RMMNumbaManager(HostOnlyCUDAMemoryManager):
def memalloc(self, size):
# Allocates device memory using RMM functions. The finalizer for the
# allocated memory calls back to RMM to free the memory.
addr = librmm.rmm_alloc(bytesize, 0)
ctx = cuda.current_context()
ptr = ctypes.c_uint64(int(addr))
finalizer = _make_finalizer(addr, stream)
return MemoryPointer(ctx, ptr, size, finalizer=finalizer)
def get_ipc_handle(self, memory):
“””
Get an IPC handle for the memory with offset modified by the RMM memory
pool.
“””
# This implementation provides a functional implementation and illustrates
# what get_ipc_handle needs to do, but it is not a very “clean”
# implementation, and it relies on borrowing bits of Numba internals to
# initialise ipchandle.
#
# A more polished implementation might make use of additional functions in
# the RMM C++ layer for initialising IPC handles, and not use any Numba
# internals.
ipchandle = (ctypes.c_byte * 64)() # IPC handle is 64 bytes
cuda.cudadrv.memory.driver_funcs.cuIpcGetMemHandle(
ctypes.byref(ipchandle),
memory.owner.handle,
)
source_info = cuda.current_context().device.get_device_identity()
ptr = memory.device_ctypes_pointer.value
offset = librmm.rmm_getallocationoffset(ptr, 0)
return IpcHandle(memory, ipchandle, memory.size, source_info,
offset=offset)
def get_memory_info(self):
# Returns a tuple of (free, total) using RMM functionality.
return get_info() # Function defined in rmm.py
def initialize(self):
# Nothing required to initialize RMM here, but this method is added
# to illustrate that the super() method should also be called.
super().initialize()
@contextmanager
def defer_cleanup(self):
# Does nothing to defer cleanup – a full implementation may choose to
# implement a different policy.
with super().defer_cleanup():
yield
@property
def interface_version(self):
# As required by the specification
return 1
# The existing _make_finalizer function is used by RMMNumbaManager:
def _make_finalizer(handle, stream):
“””
Factory to make the finalizer function.
We need to bind *handle* and *stream* into the actual finalizer, which
takes no args.
“””
def finalizer():
“””
Invoked when the MemoryPointer is freed
“””
librmm.rmm_free(handle, stream)
return finalizer
# Utility function register `RMMNumbaManager` as an EMM:
def use_rmm_for_numba():
set_memory_manager(RMMNumbaManager)
# To support `NUMBA_CUDA_MEMORY_MANAGER=rmm`:
_numba_memory_manager = RMMNumbaManager
用法示例
将Numba配置为使用RMM进行内存管理，并创建设备阵列的简单示例如下：
# example.py
import rmm
import numpy as np
from numba import cuda
rmm.use_rmm_for_numba()
a = np.zeros(10)
d_a = cuda.to_device(a)
del(d_a)
print(rmm.csv_log())
运行此命令将导致输出类似于以下内容：
注意，在RMM中有一些改进的余地，可以检测发生分配/空闲的行号，但这超出了本提议示例的范围。
通过环境设置内存管理器
除了调用rmm.use_rmm_for_numba()上面的示例之外，还可以将内存管理器设置为在环境变量中全局使用RMM，因此将调用Python解释器以如下方式运行示例：
Numba内部变更
本部分主要面向Numba开发人员-对实现EMM插件的外部接口感兴趣的开发人员可以选择跳过本部分。
当前模型/实施
目前，Context该类中实现了内存管理 。它维护分配和释放的清单：
- allocations是numba.core.utils.UniqueDict在上下文创建时创建的。
- deallocations是_PendingDeallocs该类的实例，并在Context.prepare_for_use()调用时创建。
这些用于跟踪以下对象的分配和解除分配：
- 设备内存
- 固定内存
- 映射内存
- 流
- 记录
- 模组
_PendingDeallocs类实现延迟释放的策略-清除功能（如cuMemFree），用于上述项目由表示分配对象的最后，加到其未决解除分配的列表。当Python解释器对拥有这些终结器的对象进行垃圾回收时，将运行这些终结器。当向取消分配列表添加新的清除功能，导致挂起的取消分配的数量或大小超过配置的比率时，该_PendingDeallocs对象将为其已知的所有项目运行解除分配器，然后清除其内部挂起的列表。
有关此实现的更多详细信息，请参见解除分配。
变更
本节概述了为支持EMM插件API而进行的主要更改-为了适应这些更改，需要对Numba的其它部分进行各种小的更改；没有提供这些的详尽列表。
上下文变化
该numba.cuda.cudadrv.driver.Context将不再直接分配和释放内存。相反，上下文将保留对内存管理器实例的引用，并且其内存分配方法将调用内存管理器，例如：
def memalloc(self, size):
return self.memory_manager.memalloc(size)
def memhostalloc(self, size, mapped=False, portable=False, wc=False):
return self.memory_manager.memhostalloc(size, mapped, portable, wc)
def mempin(self, owner, pointer, size, mapped=False):
if mapped and not self.device.CAN_MAP_HOST_MEMORY:
raise CudaDriverError(“%s cannot map host memory” % self.device)
return self.memory_manager.mempin(owner, pointer, size, mapped)
def prepare_for_use(self):
self.memory_manager.initialize()
def get_memory_info(self):
self.memory_manager.get_memory_info()
def get_ipc_handle(self, memory):
return self.memory_manager.get_ipc_handle(memory)
def reset(self):
# … Already-extant reset logic, plus:
self._memory_manager.reset()
在memory_manager创建上下文时构件被初始化。
该memunpin方法（上面未显示，但当前存在于Context 类中）从未实现过-它当前引发一个NotImplementedError。此方法可以说是不需要的-固定内存将由其终结器立即解除锁定，并且在终结器运行之前解除锁定将使PinnedMemory仍保留其引用的对象的状态无效 。建议在对Context类进行其它更改时将其删除。
该Context班将仍然实例化self.allocations，并 self.deallocations像以前一样-这些仍然会通过上下文来管理事件，流和模块，这不是由EMM插件处理的分配和释放。
driver模块的新组件
- BaseCUDAMemoryManager：一个抽象类，如上面的插件接口中所定义。
- HostOnlyCUDAMemoryManager：子类BaseCUDAMemoryManager，逻辑来自Context.memhostalloc和Context.mempin移入其中。该类还将创建自己的allocations和deallocations成员，类似于Context该类如何创建它们和成员。这些用于管理固定和映射的host内存的分配和重新分配。
- NumbaCUDAMemoryManager：子类HostOnlyCUDAMemoryManager，其中还包含memalloc基于Context类中当前存在的实现。这是默认的内存管理器，在添加EMM插件接口之前，它的使用会保留Numba的行为-也就是说，Numba数组的所有内存分配和释放都在Numba中处理。
- 此类与其父类共享allocations和deallocations成员HostOnlyCUDAMemoryManager，并将它们用于管理其分配的设备内存。
- 该set_memory_manager函数设置指向内存管理器类的全局指针。该全局变量最初成立NumbaCUDAMemoryManager（默认）。
分级IPC
分级IPC不应拥有其分配的内存的所有权。使用默认的内部内存管理器时，已为临时阵列分配的内存已拥有。使用EMM插件时，获取内存所有权是不合法的。
可以通过应用以下小补丁来进行此更改，该小补丁已经过测试，对CUDA测试套件没有影响：
diff –git a/numba/cuda/cudadrv/driver.py b/numba/cuda/cudadrv/driver.py
index 7832955..f2c1352 100644
— a/numba/cuda/cudadrv/driver.py
+++ b/numba/cuda/cudadrv/driver.py
@@ -922,7 +922,11 @@ class _StagedIpcImpl(object):
with cuda.gpus[srcdev.id]:
impl.close()
– return newmem.own()
+ return newmem
测试
除了为新功能添加适当的测试外，还将对所需的现有测试进行一些重构，但是这些更改并不重要。重新分配策略的检查（如TestDeallocation， TestDeferCleanup），就需要进行修改，以确保正在研究组正确的释放操作的。使用EMM插件时，将需要跳过它们。
原型/实验实施
已经产生了一些原型/实验性实施方式来说明本文档中介绍的设计。当前的实现可以在以下位置找到：
- Numba分支：https : //github.com/gmarkall/numba/tree/grm-numba-nbep-7。
- RMM分支：https : //github.com/gmarkall/rmm/tree/grm-numba-nbep-7。
- CuPy实施：https : //github.com/gmarkall/nbep-7/blob/master/nbep7/cupy_mempool.py-使用未经修改的CuPy。
- 请参阅CuPy内存管理文档。
目前的实施状态
RMM插件
举一个最小的例子，简单的分配和免费使用RMM可以按预期工作。对于示例代码（类似于上面的RMM示例）：
import rmm
import numpy as np
from numba import cuda
rmm.use_rmm_for_numba()
a = np.zeros(10)
d_a = cuda.to_device(a)
del(d_a)
print(rmm.csv_log())
看到以下输出：
此输出类似于上面提供的示例用法的预期输出（尽管请注意，指针地址和时间戳与示例相比有所不同），并提供了示例用法的一些验证。
CuPy插件
from nbep7.cupy_mempool import use_cupy_mm_for_numba
import numpy as np
from numba import cuda
use_cupy_mm_for_numba()
a = np.zeros(10)
d_a = cuda.to_device(a)
del(d_a)
原型CuPy插件具有一些原始记录，因此看到了输出：
Numba CUDA单元测试
除了提供正确的简单示例执行之外，所有相关的Numba CUDA单元测试也都通过了原型分支，用于内部内存管理器和RMM EMM插件。
RMM
单元测试套件可以与具有以下功能的RMM EMM插件一起运行：
单元测试套件输出的摘要是：
与内置的Numba内存管理一起运行时，输出为：
也就是说，使用外部内存管理器进行的更改不会破坏内置的Numba内存管理。另外还有6个跳过的测试，来自：
- TestDeallocation：由于专门测试Numba的内部释放策略而被跳过。
- TestDeferCleanup：由于专门测试Numba的延迟清除实施而被跳过。
- TestCudaArrayInterface.test_ownership：由于使用EMM插件时Numba不拥有内存而被跳过，但是此测试用例假定拥有所有权。
CuPy
可以使用CuPy插件运行测试套件，方法是：
目前，该插件实现比RMM实现更原始，并且会导致单元测试套件出现一些错误：
这8个错误是由于缺少get_ipc_handleCuPy EMM插件实现中的实现。预计将重新访问并完成此实现，以便将来可以将CuPy稳定用作Numba的分配器。
