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的內部重新分配機制。例如：
o EMM可以實施與Numba相似的策略，或者
o 釋放的內存可能會立即返回到內存池。
? 該defer_cleanup上下文管理器會有與EMM不同的表現-EMM是在由執行文檔附加 defer_cleanup上下文管理器使用。
o 例如，即使在使用上下文管理器，池分配器也始終可以立即將內存返回給池，但是可以選擇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 bytesReturns: 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 CUDAaddress space.- `portable`: Whether the memory will be considered pinned by allcontexts, and not just the calling context.- `wc`: Whether to allocate the memory as write-combined.Returns a `MappedMemory` or `PinnedMemory` instance that owns theallocated memory, depending on whether the region was mapped intodevice 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 theallocated memory, depending on whether the region was mapped into devicememory."""@abstractmethod
def initialize(self):"""Perform any initialization required for the EMM plugin to be ready touse."""@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 deferredcleanup whilst it is active."""@property
@abstractmethod
def interface_version(self):"""Returns an integer specifying the version of the EMM Plugin interfacesupported by the plugin implementation. Should always return 1 forimplementations 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。
o 如果子類與初始化（可能）或復位（不太可能）無關，則無需實現這些方法。
o 但是，如果它確實實現了這些方法，那么它還必須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函數的調用。

總結

以上是生活随笔為你收集整理的BEP 7：CUDA外部内存管理插件（上）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。