BEP 7：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”)
因此，mylib.numba_utils僅在Numba已經存在的情況下才導入包含EMM插件實現的。如果Numba不可用，則將永遠不會導入mylib.numba_utils（必須是import numba）。
建議所有帶有EMM插件的庫至少包括一些帶有Numba的環境，以便在使用中使用EMM插件進行測試，以及一些不帶Numba的環境，以避免引入意外的Numba依賴關系。
示例實現-RAPIDS內存管理器（RMM）插件
本節概述了Rapids Memory Manager（RMM）中EMM插件的實現。這旨在顯示實現的概述，以支持上述說明，并說明如何使用插件API-可以為產品就緒的實現做出不同的選擇。
插件實現包括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.pydef 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 specificationreturn 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())
運行此命令將導致輸出類似于以下內容：
注意，在RMM中有一些改進的余地，可以檢測發生分配/空閑的行號，但這超出了本提議示例的范圍。
通過環境設置內存管理器
除了調用rmm.use_rmm_for_numba()上面的示例之外，還可以將內存管理器設置為在環境變量中全局使用RMM，因此將調用Python解釋器以如下方式運行示例：
Numba內部變更
本部分主要面向Numba開發人員-對實現EMM插件的外部接口感興趣的開發人員可以選擇跳過本部分。
當前模型/實施
目前，Context該類中實現了內存管理 。它維護分配和釋放的清單：
? allocations是numba.core.utils.UniqueDict在上下文創建時創建的。
? deallocations是_PendingDeallocs該類的實例，并在Context.prepare_for_use()調用時創建。
這些用于跟蹤以下對象的分配和解除分配：
? 設備內存
? 固定內存
? 映射內存
? 流
? 記錄
? 模組
_PendingDeallocs類實現延遲釋放的策略-清除功能（如cuMemFree），用于上述項目由表示分配對象的最后，加到其未決解除分配的列表。當Python解釋器對擁有這些終結器的對象進行垃圾回收時，將運行這些終結器。當向取消分配列表添加新的清除功能，導致掛起的取消分配的數量或大小超過配置的比率時，該_PendingDeallocs對象將為其已知的所有項目運行解除分配器，然后清除其內部掛起的列表。
有關此實現的更多詳細信息，請參見解除分配。
變更
本節概述了為支持EMM插件API而進行的主要更改-為了適應這些更改，需要對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該類如何創建它們和成員。這些用于管理固定和映射的host內存的分配和重新分配。
? NumbaCUDAMemoryManager：子類HostOnlyCUDAMemoryManager，其中還包含memalloc基于Context類中當前存在的實現。這是默認的內存管理器，在添加EMM插件接口之前，它的使用會保留Numba的行為-也就是說，Numba數組的所有內存分配和釋放都在Numba中處理。
o 此類與其父類共享allocations和deallocations成員HostOnlyCUDAMemoryManager，并將它們用于管理其分配的設備內存。
? 該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。
o 請參閱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())
看到以下輸出：
此輸出類似于上面提供的示例用法的預期輸出（盡管請注意，指針地址和時間戳與示例相比有所不同），并提供了示例用法的一些驗證。
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插件具有一些原始記錄，因此看到了輸出：
Numba CUDA單元測試
除了提供正確的簡單示例執行之外，所有相關的Numba CUDA單元測試也都通過了原型分支，用于內部內存管理器和RMM EMM插件。
RMM
單元測試套件可以與具有以下功能的RMM EMM插件一起運行：
單元測試套件輸出的摘要是：
與內置的Numba內存管理一起運行時，輸出為：
也就是說，使用外部內存管理器進行的更改不會破壞內置的Numba內存管理。另外還有6個跳過的測試，來自：
? TestDeallocation：由于專門測試Numba的內部釋放策略而被跳過。
? TestDeferCleanup：由于專門測試Numba的延遲清除實施而被跳過。
? TestCudaArrayInterface.test_ownership：由于使用EMM插件時Numba不擁有內存而被跳過，但是此測試用例假定擁有所有權。
CuPy
可以使用CuPy插件運行測試套件，方法是：
目前，該插件實現比RMM實現更原始，并且會導致單元測試套件出現一些錯誤：
這8個錯誤是由于缺少get_ipc_handleCuPy EMM插件實現中的實現。預計將重新訪問并完成此實現，以便將來可以將CuPy穩定用作Numba的分配器。

總結

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

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