如何为内存映射文件使用 CUDA 固定的“零拷贝”内存？

Question

目标/问题

在 Python 中，我正在寻找一种从内存映射文件读取/写入数据到 GPU 的快速方法。

在之前的 SO 溢出帖子中 [Cupy OutOfMemoryError when trying to cupy.load larger dimension .npy files in memory map mode, but np.load works fine]

提到这可以使用 CUDA 固定的“零拷贝”内存。此外，这个方法似乎是由这个人开发的[ cuda - Zero-copy memory, memory-mapped file ] 尽管那个人正在使用 C++ 工作。

我之前的尝试是使用 Cupy，但我对任何 cuda 方法持开放态度。

到目前为止我所尝试的

我提到了我如何尝试使用 Cupy，它允许您以内存映射模式打开 numpy 文件。

import os
import numpy as np
import cupy

#Create .npy files. 
for i in range(4):
    numpyMemmap = np.memmap( 'reg.memmap'+str(i), dtype='float32', mode='w+', shape=( 2200000 , 512))
    np.save( 'reg.memmap'+str(i) , numpyMemmap )
    del numpyMemmap
    os.remove( 'reg.memmap'+str(i) )

# Check if they load correctly with np.load.
NPYmemmap = []
for i in range(4):
    NPYmemmap.append( np.load( 'reg.memmap'+str(i)+'.npy' , mmap_mode = 'r+' )  )
del NPYmemmap

# Eventually results in memory error. 
CPYmemmap = []
for i in range(4):
    print(i)
    CPYmemmap.append( cupy.load( 'reg.memmap'+str(i)+'.npy' , mmap_mode = 'r+' )  )

我尝试过的结果

我的尝试导致OutOfMemoryError:

有人提到

看来 cupy.load 将要求整个文件首先适合主机内存，然后是设备内存。

还有人提到

CuPy 无法处理 mmap 内存。因此，CuPy 默认直接使用 GPU 内存。 https://docs-cupy.chainer.org/en/stable/reference/generated/cupy.cuda.MemoryPool.html#cupy.cuda.MemoryPool.malloc 如果您想使用统一内存，可以更改默认内存分配器。

我尝试过使用

cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.memory.malloc_managed).malloc)

但这似乎并没有什么不同。发生错误时，我的 CPU 内存约为 16 gigs，但我的 GPU 内存为 0.32 gigs。我正在使用 Google colab，我的 CPU Ram 是 25 gigs，GPU ram 是 12 gigs。所以看起来在整个文件托管在主机内存中之后，它检查它是否适合设备内存，当它看到它只有所需的 16 个演出中的 12 个时，它抛出了一个错误（我最好的猜测）。

所以，现在我正在尝试找出一种方法来使用固定的“零拷贝”内存来处理内存映射文件，该文件会将数据提供给 GPU。

如果重要的话，我尝试传输的数据类型是浮点数组。通常，对于只读数据，二进制文件会加载到 GPU 内存中，但我正在处理数据，我在每一步都尝试读取和写入。

Answer 1

在我看来，目前cupy 没有提供可用于代替通常设备内存分配器的固定分配器，即可以用作cupy.ndarray 的支持。如果这对您很重要，您可以考虑提交cupy issue。

但是，似乎可以创建一个。这应该被认为是实验代码。还有一些与它的使用相关的问题。

基本的想法是我们将用我们自己的替换cupy的默认设备内存分配器，使用cupy.cuda.set_allocator，正如已经向您建议的那样。我们将需要为用作cupy.cuda.memory.MemoryPointer 的存储库的BaseMemory 类提供我们自己的替换。这里的关键区别在于我们将使用固定内存分配器而不是设备分配器。这是下面PMemory 类的要点。

需要注意的其他一些事项：

• 在使用固定内存（分配）完成所需操作后，您可能应该将cupy 分配器恢复为其默认值。不幸的是，与cupy.cuda.set_allocator 不同，我没有找到对应的cupy.cuda.get_allocator，这让我觉得cupy 存在缺陷，这似乎也值得向我提出一个杯子问题。但是，对于本演示，我们将仅恢复到 None 选项，它使用默认设备内存分配器之一（但不是池分配器）。
• 通过提供这种简约的固定内存分配器，我们仍然建议 cupy 这是普通的设备内存。这意味着它不能直接从主机代码访问（实际上是，但 cupy 不知道）。因此，各种操作（例如cupy.load）将创建不需要的主机分配和不需要的复制操作。我认为解决这个问题需要的不仅仅是我建议的这个小改动。但至少对于您的测试用例，这种额外的开销可能是可控的。看来您想从磁盘加载一次数据，然后将其保留在那里。对于那种类型的活动，这应该是可管理的，特别是因为您将其分解成块。正如我们将看到的，处理 4 个 5GB 的块对于 25GB 的主机内存来说太多了。我们将需要为四个 5GB 块（实际上是固定的）分配主机内存，并且我们还需要为一个额外的 5GB“开销”缓冲区分配额外的空间。所以 25GB 是不够的。但出于演示目的，如果我们将您的缓冲区大小减少到 4GB (5x4GB = 20GB)，我认为它可能适合您的 25GB 主机 RAM 大小。
• 与 cupy 的默认设备内存分配器关联的普通设备内存与特定设备有关联。固定内存不需要有这样的关联，但是我们用相似类替换BaseMemory 意味着我们建议cupy 这个“设备”内存，就像所有其他普通设备内存一样，具有特定的设备关联。在像您这样的单一设备设置中，这种区别是没有意义的。但是，这不适用于稳定的多设备使用固定内存。为此，建议再次对cupy 进行更强大的更改，可能是通过提交问题。

这是一个例子：

import os
import numpy as np
import cupy



class PMemory(cupy.cuda.memory.BaseMemory):
    def __init__(self, size):
        self.size = size
        self.device_id = cupy.cuda.device.get_device_id()
        self.ptr = 0
        if size > 0:
            self.ptr = cupy.cuda.runtime.hostAlloc(size, 0)
    def __del__(self):
        if self.ptr:
            cupy.cuda.runtime.freeHost(self.ptr)

def my_pinned_allocator(bsize):
    return cupy.cuda.memory.MemoryPointer(PMemory(bsize),0)

cupy.cuda.set_allocator(my_pinned_allocator)

#Create 4 .npy files, ~4GB each
for i in range(4):
    print(i)
    numpyMemmap = np.memmap( 'reg.memmap'+str(i), dtype='float32', mode='w+', shape=( 10000000 , 100))
    np.save( 'reg.memmap'+str(i) , numpyMemmap )
    del numpyMemmap
    os.remove( 'reg.memmap'+str(i) )

# Check if they load correctly with np.load.
NPYmemmap = []
for i in range(4):
    print(i)
    NPYmemmap.append( np.load( 'reg.memmap'+str(i)+'.npy' , mmap_mode = 'r+' )  )
del NPYmemmap

# allocate pinned memory storage
CPYmemmap = []
for i in range(4):
    print(i)
    CPYmemmap.append( cupy.load( 'reg.memmap'+str(i)+'.npy' , mmap_mode = 'r+' )  )
cupy.cuda.set_allocator(None)

我还没有在具有这些文件大小的 25GB 主机内存的设置中对此进行测试。但是我已经用超过我的 GPU 设备内存的其他文件大小对其进行了测试，它似乎可以工作。

再次，实验代码，未经彻底测试，您的里程可能会有所不同，最好通过提交cupy github问题来实现此功能。而且，正如我之前提到的，从设备代码访问这种“设备内存”通常比普通的cupy 设备内存要慢得多。

最后，这并不是真正的“内存映射文件”，因为所有文件内容都将加载到主机内存中，此外，这种方法会“耗尽”主机内存。如果要访问 20GB 的文件，则需要 20GB 以上的主机内存。只要您“加载”了这些文件，就会使用 20GB 的主机内存。

更新：cupy 现在提供对固定分配器的支持，请参阅here。此答案仅供历史参考。