BEP 7:CUDA外部内存管理插件(上)
阅读原文时间:2023年07月10日阅读:2

BEP 7:CUDA外部内存管理插件(上)

背景和目标

CUDA阵列接口使得能够共享不同的Python之间的数据库的访问CUDA设备。但是,每个库都与其它库区别对待。例如:

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

手机扫一扫

移动阅读更方便

阿里云服务器
腾讯云服务器
七牛云服务器

你可能感兴趣的文章