numpy - 如何将 CUDA 固定 "zero-copy"内存用于内存映射文件？

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:
有人提到

it appears that cupy.load will require that the entire file fit first in host memory, then in device memory.

还有人提到

CuPy can't handle mmap memory. So, CuPy uses GPU memory directly in default. https://docs-cupy.chainer.org/en/stable/reference/generated/cupy.cuda.MemoryPool.html#cupy.cuda.MemoryPool.malloc You can change default memory allocator if you want to use Unified Memory.

我尝试使用
cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.memory.malloc_managed).malloc)
但这似乎并没有什么不同。出现错误时，我的 CPU 内存约为 16 gigs，但我的 GPU ram 为 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.cuda.set_allocator 用我们自己的替换cupy 的默认设备内存分配器。正如已经向您建议的那样。我们需要自行更换 BaseMemory用作 cupy.cuda.memory.MemoryPointer 的存储库的类.这里的主要区别在于我们将使用固定内存分配器而不是设备分配器。这是PMemory的要点下面的课。
需要注意的其他一些事项:

在使用固定内存(分配)完成您需要的操作后，您可能应该恢复 cupy分配器为其默认值。不幸的是，不像cupy.cuda.set_allocator ，我没有找到对应的cupy.cuda.get_allocator ，这在我看来是 cupy 的不足之处，这似乎也值得向我提出一个cupy issue。然而，对于这个演示，我们将只恢复到 None选择，它使用默认设备内存分配器之一(但不是池分配器)。

通过提供这种简约的固定内存分配器，我们仍然建议丘比特这是普通设备内存。这意味着它不能从主机代码直接访问(实际上是这样，但cupy 不知道)。因此，各种操作(例如 cupy.load )将创建不需要的主机分配和不需要的复制操作。我认为要解决这个问题需要的不仅仅是我建议的这个小改动。但至少对于您的测试用例，这种额外的开销可能是可以管理的。看来您想从磁盘加载一次数据，然后将其留在那里。对于这种类型的事件，这应该是可以管理的，尤其是当您将其分解成块时。正如我们将看到的，处理 4 个 5GB 块对于 25GB 的主机内存来说太多了。我们需要为四个 5GB 块(实际上是固定的)分配主机内存，我们还需要额外的空间来容纳一个额外的 5GB“开销”缓冲区。所以 25GB 是不够的。但出于演示目的，如果我们将您的缓冲区大小减少到 4GB (5x4GB = 20GB)，我认为它可能适合您的 25GB 主机 RAM 大小。

与cupy 的默认设备内存分配器相关联的普通设备内存与特定设备相关联。 pinned memory 不需要有这样的关联，但是我们对 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 .这个答案应该只用于历史引用。