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

How to use CUDA pinned "zero-copy" memory for a memory mapped file?

Objective/Problem

在 Python 中,我正在寻找一种从内存映射文件到 GPU 的 read/write 数据的快速方法。

在之前的 SO 溢出中 post [ ]

在提到的地方可以使用 CUDA 固定 "zero-copy" 内存。而且,这个方法好像是这个人开发的【 ] 虽然那个人在用 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 Ram 约为 16 gig,但我的 GPU ram 为 0.32 gig。我正在使用 Google colab,其中我的 CPU Ram 是 25 gig,GPU ram 是 12 gig。所以看起来整个文件都托管在主机内存中后,它检查它是否适合设备内存,当它发现它只有所需的 16 个演出中的 12 个时,它抛出了一个错误(我最好的猜测).

所以,现在我正在尝试找出一种方法来使用固定 'zero-copy' 内存来处理内存映射文件,该文件会将数据提供给 GPU。

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

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

但是,似乎可以创建一个。这应该被视为实验代码。它的使用也存在一些问题。

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

其他一些注意事项:

  • 在使用固定内存(分配)完成所需的操作后,您可能应该将 cupy 分配器恢复为默认值。不幸的是,与 cupy.cuda.set_allocator 不同,我没有找到相应的 cupy.cuda.get_allocator,这让我觉得 cupy 有缺陷,这似乎也值得向我提出一个 cupy 问题。然而,对于这个演示,我们将恢复到 None 选择,它使用默认设备内存分配器之一(但是不是池分配器)。
  • 通过提供这个简约的固定内存分配器,我们仍然向 cupy 建议这是普通的设备内存。这意味着它不能直接从主机代码访问(实际上是,但 cupy 不知道)。因此,各种操作(例如 cupy.load)将创建不需要的主机分配和不需要的复制操作。我认为解决这个问题需要的不仅仅是我建议的这个小改动。但至少对于您的测试用例,这种额外的开销可能是可控的。看来您想从磁盘加载数据一次,然后将其留在那里。对于那种类型的 activity,这应该是可管理的,尤其是当您将其分解成块时。正如我们将看到的,处理四个 5GB 的块对于 25GB 的主机内存来说太多了。我们需要为四个 5GB 块(实际上是固定的)分配主机内存,并且我们还需要额外的 space 用于一个额外的 5GB“开销”缓冲区。所以 25GB 是不够的。但出于演示目的,如果我们将您的缓冲区大小减少到 4GB (5x4GB = 20GB),我认为它可能适合您的 25GB 主机 RAM 大小。
  • 与 cupy 的默认设备内存分配器关联的普通设备内存与特定设备相关联。固定内存不需要有这样的关联,但是我们用相似的 class 替换 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。本回答仅供历史参考。