NBEP 7: CUDA 外部内存管理插件

作者:

格雷厄姆·马克尔，英伟达

贡献者:

Thomson Comer, Peter Entschev, Leo Fang, John Kirkham, Keith Kraus

日期:

2020年3月

状态:

最终

背景与目标

The CUDA Array Interface 使得访问CUDA设备的不同的Python库之间可以共享数据。然而，每个库都独立管理自己的内存。例如：

• Numba 内部管理用于创建设备和映射主机数组的内存。

• RAPIDS 库 (cuDF, cuML 等) 使用 Rapids 内存管理器 来分配设备内存。

• CuPy 包括一个用于设备和固定内存的 内存池实现。

本NBEP的目标是描述一个插件接口，该接口允许用户用外部内存管理器替换Numba的内部内存管理。当使用插件接口时，Numba在创建数组时不再直接分配或释放任何内存，而是通过外部管理器请求分配和释放。

要求

在 Numba 中提供一个 外部内存管理器 (EMM) 接口。

• 当使用EMM时，Numba将使用EMM进行所有内存分配。它永远不会直接调用诸如``CuMemAlloc``、cuMemFree等函数。

• 在不使用 外部内存管理器（EMM） 的情况下，Numba 的当前行为保持不变（在撰写本文时，当前版本是 0.48 版本）。

如果使用 EMM，它将在程序执行期间完全替换 Numba 的内部内存管理。将提供一个用于设置内存管理器的接口。

设备 vs. 主机内存

EMM 将始终负责管理设备内存。然而，并非所有 CUDA 内存管理库都支持管理主机内存，因此将提供一个设施，使 Numba 在将设备内存控制权移交给 EMM 的同时继续管理主机内存。

释放策略

Numba 的内部内存管理使用了一种 延迟释放策略，旨在通过推迟释放操作直到有大量待处理释放时来提高效率。它还提供了一种机制，在关键部分使用 defer_cleanup() 上下文管理器来完全防止释放操作。

• 当EMM不使用时，defer_cleanup 的释放策略和操作保持不变。

• 当使用 EMM 时，释放策略由 EMM 实现，Numba 的内部释放机制不被使用。例如：

  • EMM 可以通过类似 Numba 的策略来实现，或者

  • 已释放的内存可能会立即返回到内存池。

• defer_cleanup 上下文管理器在 EMM 中的行为可能有所不同 - EMM 应附带关于 defer_cleanup 上下文管理器在使用时的行为说明。

  • 例如，一个池分配器即使在上下文管理器使用时也可以立即将内存返回给池，但可以选择在 defer_cleanup 未使用时不释放空池。

其他对象的管理

除了内存，Numba 还管理 事件、流 和模块（模块是一个编译对象，由 @cuda.jit 修饰的函数生成）的分配和释放。流、事件和模块的管理不应因 EMM 的存在或不存在而改变。

异步分配 / 释放

一个异步内存管理器可能会提供一个功能，使得分配或释放操作可以接受一个CUDA流并异步执行。对于释放操作，这不太可能引起问题，因为它在Python层之下运行，但对于分配操作，如果用户试图从这种异步内存分配中在默认流上启动内核，这可能会引起问题。

本提案中描述的接口不需要支持异步分配和释放，因此这些用例将不再进一步考虑。然而，本提案中的内容不应妨碍在未来版本的接口中直接添加异步操作。

非需求

为了最小化复杂性并将此提案限制在合理的范围内，以下内容将不予支持：

• 在不同的上下文中使用不同的内存管理器实现。所有上下文将使用相同的内存管理器实现——无论是 Numba 内部实现还是外部实现。

• 一旦执行开始，更改内存管理器。更改内存管理器并保留所有分配是不切实际的。清理整个状态然后切换到不同的内存分配器（而不是启动新进程）似乎是一个相当小众的用例。

• 对 __cuda_array_interface__ 的任何更改以进一步定义其语义，例如在 Numba Issue #4886 中讨论的获取/释放内存 - 这些都是独立的，可以作为单独的提案来处理。

• 托管内存 / UVM 不受支持。目前 Numba 不支持 UVM - 有关支持的讨论，请参见 Numba Issue #4362。

插件开发者接口

新类和函数将被添加到 numba.cuda.cudadrv.driver 中：

• BaseCUDAMemoryManager 和 HostOnlyCUDAMemoryManager：EMM 插件实现的基类。

• set_memory_manager: 用于向 Numba 注册外部内存管理器的方法。

这些将通过公共API暴露，位于 numba.cuda 模块中。此外，一些已经是 driver 模块一部分的类也将作为公共API的一部分暴露：

• MemoryPointer: 用于封装设备内存指针的信息。

• MappedMemory：用于保存关于映射到设备地址空间的主机内存信息（MemoryPointer 的子类）。

• PinnedMemory：用于保存关于主机内存的信息，这些内存是固定的（mviewbuf.MemAlloc 的子类，Numba 内部的类）。

除了调用 set_memory_manager 函数外，还可以使用环境变量来设置内存管理器。环境变量的值应为包含内存管理器的模块名称，该模块在全局范围内命名为 _numba_memory_manager:

export NUMBA_CUDA_MEMORY_MANAGER="<module>"

当设置此变量时，Numba 将自动使用指定模块中的内存管理器。对 set_memory_manager 的调用将发出警告，但会被忽略。

插件基类

EMM 插件通过继承 BaseCUDAMemoryManager 类来实现，该类定义如下：

class BaseCUDAMemoryManager(object, metaclass=ABCMeta):
    @abstractmethod
    def memalloc(self, size):
        """
        Allocate on-device memory in the current context. Arguments:

        - `size`: Size of allocation in bytes

        Returns: a `MemoryPointer` to the allocated memory.
        """

    @abstractmethod
    def memhostalloc(self, size, mapped, portable, wc):
        """
        Allocate pinned host memory. Arguments:

        - `size`: Size of the allocation in bytes
        - `mapped`: Whether the allocated memory should be mapped into the CUDA
                    address space.
        - `portable`: Whether the memory will be considered pinned by all
                      contexts, and not just the calling context.
        - `wc`: Whether to allocate the memory as write-combined.

        Returns a `MappedMemory` or `PinnedMemory` instance that owns the
        allocated memory, depending on whether the region was mapped into
        device memory.
        """

    @abstractmethod
    def mempin(self, owner, pointer, size, mapped):
        """
        Pin a region of host memory that is already allocated. Arguments:

        - `owner`: An object owning the memory - e.g. a `DeviceNDArray`.
        - `pointer`: The pointer to the beginning of the region to pin.
        - `size`: The size of the region to pin.
        - `mapped`: Whether the region should also be mapped into device memory.

        Returns a `MappedMemory` or `PinnedMemory` instance that refers to the
        allocated memory, depending on whether the region was mapped into device
        memory.
        """

    @abstractmethod
    def initialize(self):
        """
        Perform any initialization required for the EMM plugin to be ready to
        use.
        """

    @abstractmethod
    def get_memory_info(self):
        """
        Returns (free, total) memory in bytes in the context
        """

    @abstractmethod
    def get_ipc_handle(self, memory):
        """
        Return an `IpcHandle` from a GPU allocation. Arguments:

        - `memory`: A `MemoryPointer` for which the IPC handle should be created.
        """

    @abstractmethod
    def reset(self):
        """
        Clear up all memory allocated in this context.
        """

    @abstractmethod
    def defer_cleanup(self):
        """
        Returns a context manager that ensures the implementation of deferred
        cleanup whilst it is active.
        """

    @property
    @abstractmethod
    def interface_version(self):
        """
        Returns an integer specifying the version of the EMM Plugin interface
        supported by the plugin implementation. Should always return 1 for
        implementations described in this proposal.
        """

EMM 插件的所有方法都在 Numba 内部调用 - 它们永远不需要由 Numba 用户直接调用。

initialize 方法在 Numba 请求任何内存分配之前被调用。这为 EMM 提供了一个机会来初始化其正常操作所需的任何数据结构等。该方法在程序的生命周期内可能会被多次调用 - 后续调用不应使 EMM 的状态失效或重置。

memalloc、memhostalloc和 mempin 方法在 Numba 需要分配设备或主机内存，或固定主机内存时被调用。设备内存应在当前上下文中始终分配。

get_ipc_handle 在需要数组的IPC句柄时被调用。请注意，没有关闭IPC句柄的方法 - 这是因为由 get_ipc_handle 构造的 IpcHandle 对象在Numba中包含了 close() 方法，该方法通过调用 cuIpcCloseMemHandle 来关闭句柄。预计这对于一般用例已经足够，因此EMM插件接口不提供自定义关闭IPC句柄的功能。

get_memory_info 可以在 initialize 之后的任何时间调用。

reset 作为重置上下文的一部分被调用。Numba 通常不会自发调用 reset，但它可能会应用户的要求被调用。对 reset 的调用甚至可能发生在 initialize 被调用之前，因此插件应对此情况具有鲁棒性。

defer_cleanup 是在用户代码中使用 numba.cuda.defer_cleanup 上下文管理器时调用的。

interface_version 在设置内存管理器时被 Numba 调用，以确保插件实现的接口版本与正在使用的 Numba 版本兼容。

表示指针

设备内存

MemoryPointer 类用于表示指向内存的指针。虽然其实现细节多种多样，但与EMM插件开发相关的唯一方面是其初始化。__init__ 方法具有以下接口：

class MemoryPointer:
    def __init__(self, context, pointer, size, owner=None, finalizer=None):

• context：指针被分配的上下文。

• pointer：一个 ctypes 指针（例如 ctypes.c_uint64），保存内存地址。

• size：分配的大小，以字节为单位。

• owner：所有者有时由类的内部设置，或用于Numba的内部内存管理，但不需要由EMM插件的编写者提供——默认值``None``应该总是足够。

• finalizer：当 MemoryPointer 对象的最后一个引用被释放时调用的方法。通常，这将调用外部内存管理库，通知它该内存不再需要，并且可能可以被释放（尽管内存管理库不要求立即释放它）。

主机内存

映射到CUDA地址空间（当使用 mapped=True 调用 memhostalloc 或 mempin 方法时创建）的内存由 MappedMemory 类管理：

class MappedMemory(AutoFreePointer):
    def __init__(self, context, pointer, size, owner, finalizer=None):

• context：指针被分配的上下文。

• pointer：一个 ctypes 指针（例如 ctypes.c_void_p），保存分配内存的地址。

• size：分配的内存大小，单位为字节。

• owner：拥有内存的Python对象，例如 DeviceNDArray 实例。

• finalizer：当对 MappedMemory 对象的最后一个引用被释放时调用的方法。例如，该方法可以调用 cuMemFreeHost 来释放指针所指向的内存。

请注意，从 AutoFreePointer 继承是一个实现细节，EMM 插件的开发者无需关注 - MemoryPointer 在 MappedMemory 的 MRO 中处于更高位置。

仅位于主机地址空间且已被固定的内存由 PinnedMemory 类表示：

class PinnedMemory(mviewbuf.MemAlloc):
    def __init__(self, context, pointer, size, owner, finalizer=None):

• context：指针被分配的上下文。

• pointer：一个 ctypes 指针（例如 ctypes.c_void_p），保存着固定内存的地址。

• size：固定区域的大小，以字节为单位。

• owner：拥有内存的Python对象，例如 DeviceNDArray 实例。

• finalizer：当对 PinnedMemory 对象的最后一个引用被释放时调用的方法。此方法例如可以调用 cuMemHostUnregister 来立即解除对内存的固定。

仅提供设备内存管理

一些外部内存管理器将支持设备内存的管理，但不支持主机内存。为了便于使用这些管理器之一实现EMM插件，Numba将提供一个内存管理器类，其中包含 memhostalloc 和 mempin 方法的实现。以下是该类的简要定义：

class HostOnlyCUDAMemoryManager(BaseCUDAMemoryManager):
    # Unimplemented methods:
    #
    # - memalloc
    # - get_memory_info

    def memhostalloc(self, size, mapped, portable, wc):
        # Implemented.

    def mempin(self, owner, pointer, size, mapped):
        # Implemented.

    def initialize(self):
        # Implemented.
        #
        # Must be called by any subclass when its initialize() method is
        # called.

    def reset(self):
        # Implemented.
        #
        # Must be called by any subclass when its reset() method is
        # called.

    def defer_cleanup(self):
        # Implemented.
        #
        # Must be called by any subclass when its defer_cleanup() method is
        # called.

一个类可以继承 HostOnlyCUDAMemoryManager ，然后它只需要添加设备内存方法的实现。任何子类必须遵守以下规则：

• 如果子类实现了 __init__，那么它还必须调用 HostOnlyCUDAMemoryManager.__init__，因为这是用于初始化其部分数据结构（self.allocations 和 self.deallocations）。

• 子类必须实现 memalloc 和 get_memory_info。

• initialize 和 reset 方法用于初始化 HostOnlyCUDAMemoryManager 使用的结构。

  • 如果子类在初始化（可能）或重置（不太可能）时没有需要执行的操作，那么它不需要实现这些方法。

  • 然而，如果它实现了这些方法，那么它也必须在其实现中调用 HostOnlyCUDAMemoryManager 的方法。

• 同样地，如果实现了 defer_cleanup ，它应该在 yield 之前（或在 __enter__ 方法中）进入由 HostOnlyCUDAManager.defer_cleanup() 提供的环境，并在退出之前（或在 __exit__ 方法中）释放它。

导入顺序

Numba 和实现 EMM 插件的库的顺序不应有影响。例如，如果 rmm 实现了并注册了一个 EMM 插件，那么：

from numba import cuda
import rmm

和

import rmm
from numba import cuda

是等价的 - 这是因为 Numba 在第一次调用 CUDA 函数之前不会初始化 CUDA 或分配任何内存 - 无论是实例化并注册一个 EMM 插件，还是导入 numba.cuda 都不会导致调用 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")

以便仅在Numba已存在的情况下导入包含EMM插件实现的``mylib.numba_utils``。如果Numba不可用，那么``mylib.numba_utils``（其必然导入``numba``）将永远不会被导入。

建议任何带有 EMM 插件的库至少包含一些使用 Numba 进行测试的环境，以及一些不使用 Numba 的环境，以避免意外引入 Numba 依赖。

示例实现 - RAPIDS 内存管理器 (RMM) 插件

本节概述了在 Rapids Memory Manager (RMM) 中实现 EMM 插件的方案。这旨在展示实现概述，以支持上述描述并说明如何使用插件接口 - 对于生产就绪的实现，可能会做出不同的选择。

插件实现包括对 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())

运行此操作应产生类似于以下的输出：

Event Type,Device ID,Address,Stream,Size (bytes),Free Memory,Total Memory,Current Allocs,Start,End,Elapsed,Location
Alloc,0,0x7fae06600000,0,80,0,0,1,1.10549,1.1074,0.00191666,<path>/numba/numba/cuda/cudadrv/driver.py:683
Free,0,0x7fae06600000,0,0,0,0,0,1.10798,1.10921,0.00122238,<path>/numba/numba/utils.py:678

需要注意的是，RMM 在检测分配/释放发生的行号方面有一些改进空间，但这超出了本提案中示例的范围。

通过环境设置内存管理器

在上面的例子中，可以通过环境变量将内存管理器设置为全局使用RMM，而不是调用 rmm.use_rmm_for_numba()，这样Python解释器在运行示例时会这样调用：

NUMBA_CUDA_MEMORY_MANAGER="rmm.RMMNumbaManager" python example.py

Numba 内部变化

本节主要面向Numba开发者 - 对实现EMM插件的外部接口感兴趣的开发者可以选择跳过本节。

当前模型 / 实现

目前，内存管理在 Context 类中实现。它维护了分配和释放的列表：

• allocations 是一个 numba.core.utils.UniqueDict，在上下文创建时生成。

• deallocations 是 _PendingDeallocs 类的一个实例，当调用 Context.prepare_for_use() 时创建。

这些用于跟踪以下内容的分配和释放：

• 设备内存

• 固定内存

• 映射内存

• 流

• 事件

• 模块

_PendingDeallocs 类实现了延迟释放策略——清理函数（如 cuMemFree）对于上述项目的清理函数会被对象的终结器添加到其待释放列表中。这些终结器在对象被 Python 解释器垃圾回收时运行。当向释放列表添加新的清理函数导致待释放的数量或大小超过配置的比例时，_PendingDeallocs 对象会运行其已知的所有项目的释放器，然后清空其内部的待释放列表。

有关此实现的更多详细信息，请参阅 释放行为。

提议的更改

本节概述了为支持EMM插件接口将进行的主要更改——为了适应这些更改，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 类创建它们的方式。这些成员用于管理固定和映射的主内存的分配和释放。

• NumbaCUDAMemoryManager：HostOnlyCUDAMemoryManager 的子类，其中还包含基于 Context 类中现有实现的 memalloc 实现。这是默认的内存管理器，其使用保留了在添加 EMM 插件接口之前 Numba 的行为——即，Numba 数组的所有内存分配和释放都在 Numba 内部处理。

  • 此类与其父类 HostOnlyCUDAMemoryManager 共享 allocations 和 deallocations 成员，并使用这些成员来管理其分配的设备内存。

• 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())

我们看到以下输出：

Event Type,Device ID,Address,Stream,Size (bytes),Free Memory,Total Memory,Current Allocs,Start,End,Elapsed,Location
Alloc,0,0x7f96c7400000,0,80,0,0,1,1.13396,1.13576,0.00180059,<path>/numba/numba/cuda/cudadrv/driver.py:686
Free,0,0x7f96c7400000,0,0,0,0,0,1.13628,1.13723,0.000956004,<path>/numba/numba/utils.py:678

此输出类似于上述示例用法中的预期输出（尽管请注意，指针地址和时间戳与示例中的不同），并提供了一些示例用例的验证。

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 插件的原型具有较为原始的日志记录功能，因此我们能看到输出：

Allocated 80 bytes at 7f004d400000
Freeing 80 bytes at 7f004d400000

Numba CUDA 单元测试

除了提供一个简单示例的正确执行外，所有相关的 Numba CUDA 单元测试在原型分支中也通过了测试，无论是内部内存管理器还是 RMM EMM 插件。

RMM

单元测试套件可以通过 RMM EMM 插件运行，使用以下命令：

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

单元测试套件输出的摘要如下：

Ran 564 tests in 142.211s

OK (skipped=11)

当使用内置的 Numba 内存管理运行时，输出为：

Ran 564 tests in 133.396s

OK (skipped=5)

即使用外部内存管理器的更改不会破坏内置的 Numba 内存管理。此外，还有 6 个跳过的测试，来自：

• TestDeallocation：已跳过，因为它专门测试 Numba 的内部释放策略。

• TestDeferCleanup：已跳过，因为它专门测试 Numba 的延迟清理实现。

• TestCudaArrayInterface.test_ownership：由于使用EMM插件时Numba不拥有内存，但此测试用例假设拥有内存，因此跳过。

CuPy

可以使用 CuPy 插件运行测试套件：

NUMBA_CUDA_MEMORY_MANAGER=nbep7.cupy_mempool python -m numba.runtests numba.cuda.tests

这个插件的实现目前比 RMM 的实现更为原始，并且在单元测试套件中导致了一些错误：

Ran 564 tests in 111.699s

FAILED (errors=8, skipped=11)

这8个错误是由于CuPy EMM插件实现中缺少``get_ipc_handle``的实现。预计这一实现将被重新审视并完成，以便CuPy未来能够稳定地作为Numba的分配器使用。