BEP 7：CUDA外部内存管理插件（下）

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的分配器。