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

Question

目标/问题

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

在以前的 SO 溢出帖子中 [ Cupy 内存不足错误时尝试在内存映射模式下加载更大尺寸的 .npy 文件，但 np.load 工作正常 ]

在提到这一点的地方，可以使用 CUDA 固定的"零拷贝"内存。此外，这种方法似乎是由此人开发的[ cuda - 零拷贝内存，内存映射文件 ]尽管那个人在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:

有人提到，

看起来 Cpy.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 gig。我正在使用Google colab，其中我的CPU Ram是25 gigs，GPU ram是12 gigs。所以看起来整个文件托管在主机内存中后，它会检查它是否可以放入设备内存，当它看到它只有 12 个所需的 16 个演出时，它抛出了一个错误(我最好的猜测)。

所以，现在我正在尝试找出一种方法来使用固定的"零拷贝"内存来处理内存映射文件，该文件将数据馈送到 GPU。

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

Answer 1

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

但是，似乎可以创建一个。 这应被视为实验性代码。 并且存在一些与它的使用相关的问题。

基本思想是，我们将用我们自己的内存分配器替换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)将创建不需要的主机分配和不需要的复制操作。 我认为解决这个问题需要的不仅仅是我建议的这个小改变。 但至少对于您的测试用例，这种额外的开销可能是可管理的。 您似乎想从磁盘加载一次数据，然后将其保留在那里。 对于这种类型的活动，这应该是可管理的，特别是因为您要将其分解为块。 正如我们将看到的，处理四个 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设备内存的其他文件大小对其进行了测试，它似乎可以工作。

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

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

更新：cupy 现在为固定分配器提供支持，请参阅此处。 这个答案只能用于历史参考。