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

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函数的调用。