BEP 7：CUDA外部内存管理插件（上）
背景和目标
在CUDA阵列接口使得能够共享不同的Python之间的数据库的访问CUDA设备。但是，每个库都与其它库区别对待。例如：
- Numba在内部管理内存以创建设备和映射的host阵列。
- RAPIDS库（cuDF，cuML等）使用Rapids Memory Manager分配设备内存。
- CuPy包括 用于设备和固定内存的内存池实现。
该NBEP的目的是描述一个插件接口，该接口使Numba的内部存储器管理可由用户替换为外部存储器管理器。使用插件接口时，Numba在创建数组时不再直接分配或释放任何内存，而是通过外部管理器请求分配和释放。
需求
在Numba中提供一个外部内存管理器（EMM）API。
- 使用EMM时，Numba将使用EMM进行所有内存分配。它永远不会直接调用的功能，例如CuMemAlloc，cuMemFree等等。
- 不使用外部存储器管理器（EMM）时，Numba的当前行为不变（在撰写本文时，当前版本为0.48版本）。
如果要使用EMM，它将在程序执行期间完全取代Numba的内部内存管理。将提供用于设置内存管理器的API。
设备与host内存
EMM将始终负责设备内存的管理。但是，并非所有CUDA内存管理库都支持管理host内存，因此将为Numba提供一种在将设备内存的控制权转让给EMM的同时，继续管理host内存的功能。
解除分配策略
Numba的内部内存管理使用重新分配策略，该策略旨在通过将重新分配，推迟到有大量待处理状态来提高效率。它还使用defer_cleanup()上下文管理器提供了一种机制，可以在关键部分完全防止重新分配。
- 当不使用EMM时，重新分配策略及其操作 defer_cleanup保持不变。
- 使用EMM时，重新分配策略由EMM实施，并且不使用Numba的内部重新分配机制。例如：
- EMM可以实施与Numba相似的策略，或者
- 释放的内存可能会立即返回到内存池。
- 该defer_cleanup上下文管理器会有与EMM不同的表现-EMM是在由执行文档附加 defer_cleanup上下文管理器使用。
- 例如，即使在使用上下文管理器，池分配器也始终可以立即将内存返回给池，但是可以选择defer_cleanup，在不使用之前不释放空池。
其它对象的管理
除了内存，Numba还管理事件，流和模块的分配和取消分配 （模块是从@cuda.jit-ted函数生成的已编译对象）。流，事件和模块的管理应通过是否存在EMM来保持不变。
异步分配/解除分配
异步内存管理器可能会提供分配的功能，或者可以免费获取CUDA流并异步执行。对于释放，这不太可能引起问题，因为它在Python之下的一个层上运行，但是对于分配，如果用户尝试从此异步内存分配中的默认流上启动内核，则可能会出现问题。
本文描述的接口将不需要支持异步分配和解除分配，因此将不再考虑这些用例。但是，此建议中的任何内容，都不应排除在接口的未来版本中，直接添加异步操作的可能性。
不需要
为了最小化复杂性并将此提案限制在合理的范围内，将不支持以下内容：
- 针对不同的上下文使用不同的内存管理器实现。所有上下文将使用相同的内存管理器实现-Numba内部实现或外部实现。
- 执行开始后更改内存管理器。更改内存管理器并保留所有分配是不切实际的。清理整个状态然后更改为其它内存分配器（而不是启动新进程）似乎是一个相当小众的用例。
- 对任何更改以__cuda_array_interface__进一步定义其语义，例如，如Numba所述，用于获取/释放内存-这些更改是独立的，可以作为单独建议的一部分进行处理。
- 不支持托管内存/ UVM。目前，Numba不支持。
插件开发API
新的类和函数将添加到numba.cuda.cudadrv.driver：
- BaseCUDAMemoryManager和HostOnlyCUDAMemoryManager：EMM插件实现的基类。
- set_memory_manager：一种向Numba注册外部存储器管理器的方法。
这些将通过numba.cuda模块中的公共API开放。此外，某些已经属于驱动程序模块的类将作为公共API的一部分开放：
- MemoryPointer：用于封装有关指向设备内存的指针的信息。
- MappedMemory：用于保存有关host内存的信息，该信息已映射到设备地址空间（的子类MemoryPointer）。
- PinnedMemory：用于保存有关固定的host内存的信息（mviewbuf.MemAlloc子类，Numba内部的类）。
作为调用该set_memory_manager函数的替代方法，可以使用环境变量来设置内存管理器。环境变量的值应该是在其全局范围内包含内存管理器的模块的名称，名为_numba_memory_manager：
设置此变量后，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无效或重置EMM的状态。
memalloc，memhostalloc和mempin当需要Numba设备或host内存，或host存储器的分配方法被调用。设备内存应始终在当前上下文中分配。
get_ipc_handle当需要数组的IPC句柄时调用。注意，没有关闭IPC句柄的方法-这是因为IpcHandle构造的 对象在Numba中get_ipc_handle包含close()作为其定义一部分的方法，该方法通过调用来关闭句柄 cuIpcCloseMemHandle。预期这对于一般用例就足够了，因此EMM插件接口没有提供用于自定义IPC句柄关闭的功能。
get_memory_info可以在之后的任何时间调用initialize。
reset被称为重置上下文的一部分。Numba通常不会自发调用reset，但是可以根据用户的要求调用它。reset甚至可能在initialize调用之前发生to的 调用，因此插件应对此事件保持健壮。
defer_cleanupnumba.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释放对象的最后一个引用时调用的方法 。通常，这将调用外部内存管理库，以说明不再需要该内存，并且有可能将其释放（尽管不需要EMM立即释放它）。
host内存
使用以下类管理映射到CUDA地址空间（由调用memhostalloc或mempin方法时创建的地址）的内存 ：mapped=TrueMappedMemory
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在的MRO中较高MappedMemory。
用PinnedMemory类表示仅位于host地址空间中且已固定的内存：
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指针以立即取消固定内存。
仅提供设备内存管理
某些外部内存管理器将支持对设备上内存的管理，但不支持host内存。为了使使用这些管理器之一轻松实现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已实现，则应HostOnlyCUDAManager.defer_cleanup()在yield（或在__enter__方法中）输入之前提供的上下文，并在（或在方法中）退出之前释放它__exit__ 。
输入顺序
Numba和实现EMM插件的库的顺序无关紧要。例如，如果rmm要实施和注册EMM插件，则：
from numba import cuda
import rmm
和
import rmm
from numba import cuda
是等效的-这是因为Numba不会在第一次调用CUDA函数之前初始化CUDA或分配任何内存-既不实例化和注册EMM插件，也不通过导入numba.cuda导致对CUDA函数的调用。