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

作者

Graham Markall, NVIDIA

贡献者

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

日期

2020 年 3 月

状态

最终版

背景和目标

CUDA 数组接口 允许在访问 CUDA 设备的各种 Python 库之间共享数据。然而，每个库都以与其他库不同的方式管理自己的内存。例如

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

• RAPIDS 库 (cuDF, cuML, etc.) 使用 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 问题 #4886 中讨论的内存获取/释放——这些是独立的，可以在单独的提案中解决。

• 不支持托管内存/UVM。目前 Numba 不支持 UVM——有关支持的讨论，请参阅 Numba 问题 #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 用户直接调用。

在请求任何内存分配之前，Numba 会调用 initialize 方法。这使得 EMM 有机会初始化其正常操作所需的任何数据结构等。该方法可以在程序生命周期内多次调用——随后的调用不应使 EMM 的状态失效或重置。

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

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

get_memory_info 可以在 initialize 之后随时调用。

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

当用户代码中使用 numba.cuda.defer_cleanup 上下文管理器时，会调用 defer_cleanup。

当设置内存管理器时，Numba 会调用 interface_version，以确保插件实现的接口版本与正在使用的 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 对象的最后一个引用被释放时调用的方法。通常，这将调用外部内存管理库，告知其内存不再需要，并且可能被释放（尽管 EMM 不要求立即释放）。

主机内存

映射到 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 的环境，以便在使用 EMM 插件的情况下进行测试，以及一些没有 Numba 的环境，以避免引入意外的 Numba 依赖。

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

本节概述了 Rapids 内存管理器 (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.use_rmm_for_numba() 之外，还可以通过环境变量全局设置内存管理器以使用 RMM，这样就可以调用 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）通过表示分配的对象的终结器添加到其挂起解分配列表中。这些终结器在拥有它们的 L对象被 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 的分配器。