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

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.py

 

    def 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 specification

        return 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中处理。
    • 此类与其父类共享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插件时，将需要跳过它们。

原型/实验实施

已经产生了一些原型/实验性实施方式来说明本文档中介绍的设计。当前的实现可以在以下位置找到：

目前的实施状态

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的分配器。

posted @ 2020-12-26 17:06  吴建明wujianming  阅读(151)  评论(0)  编辑  收藏  举报