内存管理#

CuPy 默认使用 内存池 进行内存分配。内存池通过减少内存分配和 CPU/GPU 同步的开销，显著提高了性能。

CuPy 中有两种不同的内存池：

设备内存池（GPU 设备内存），用于 GPU 内存分配。
固定内存池（不可交换的 CPU 内存），用于在 CPU 到 GPU 数据传输期间使用。

注意

当你监控内存使用情况时（例如，使用 nvidia-smi 监控GPU内存或使用 ps 监控CPU内存），你可能会注意到即使数组实例超出作用域后，内存也没有被释放。这是一个预期行为，因为默认的内存池会“缓存”已分配的内存块。

有关内存管理API的详细信息，请参阅低级 CUDA 支持。

为了更方便地使用固定内存，我们在 cupyx 命名空间中还提供了一些高级API，包括 cupyx.empty_pinned()、cupyx.empty_like_pinned()、cupyx.zeros_pinned() 和 cupyx.zeros_like_pinned()。它们返回由固定内存支持的NumPy数组。如果正在使用CuPy的固定内存池，则从该池中分配固定内存。

备注

CuPy v8 及以上版本提供了一个 FFT 计划缓存，如果在使用 FFT 及相关函数时，可能会使用一部分设备内存。可以通过缩小或禁用缓存来释放占用的内存。

内存池操作#

内存池实例提供有关内存分配的统计信息。要访问默认的内存池实例，请使用 cupy.get_default_memory_pool() 和 cupy.get_default_pinned_memory_pool()。您还可以释放内存池中所有未使用的内存块。有关详细信息，请参阅下面的示例代码：

import cupy
import numpy

mempool = cupy.get_default_memory_pool()
pinned_mempool = cupy.get_default_pinned_memory_pool()

# Create an array on CPU.
# NumPy allocates 400 bytes in CPU (not managed by CuPy memory pool).
a_cpu = numpy.ndarray(100, dtype=numpy.float32)
print(a_cpu.nbytes)                      # 400

# You can access statistics of these memory pools.
print(mempool.used_bytes())              # 0
print(mempool.total_bytes())             # 0
print(pinned_mempool.n_free_blocks())    # 0

# Transfer the array from CPU to GPU.
# This allocates 400 bytes from the device memory pool, and another 400
# bytes from the pinned memory pool.  The allocated pinned memory will be
# released just after the transfer is complete.  Note that the actual
# allocation size may be rounded to larger value than the requested size
# for performance.
a = cupy.array(a_cpu)
print(a.nbytes)                          # 400
print(mempool.used_bytes())              # 512
print(mempool.total_bytes())             # 512
print(pinned_mempool.n_free_blocks())    # 1

# When the array goes out of scope, the allocated device memory is released
# and kept in the pool for future reuse.
a = None  # (or `del a`)
print(mempool.used_bytes())              # 0
print(mempool.total_bytes())             # 512
print(pinned_mempool.n_free_blocks())    # 1

# You can clear the memory pool by calling `free_all_blocks`.
mempool.free_all_blocks()
pinned_mempool.free_all_blocks()
print(mempool.used_bytes())              # 0
print(mempool.total_bytes())             # 0
print(pinned_mempool.n_free_blocks())    # 0

详情请参阅 cupy.cuda.MemoryPool 和 cupy.cuda.PinnedMemoryPool。

限制 GPU 内存使用#

你可以通过使用 CUPY_GPU_MEMORY_LIMIT 环境变量来硬限制可以分配的GPU内存量（详情请参阅环境变量）。

# Set the hard-limit to 1 GiB:
#   $ export CUPY_GPU_MEMORY_LIMIT="1073741824"

# You can also specify the limit in fraction of the total amount of memory
# on the GPU. If you have a GPU with 2 GiB memory, the following is
# equivalent to the above configuration.
#   $ export CUPY_GPU_MEMORY_LIMIT="50%"

import cupy
print(cupy.get_default_memory_pool().get_limit())  # 1073741824

你也可以使用 cupy.cuda.MemoryPool.set_limit() 来设置限制（或覆盖通过环境变量指定的值）。通过这种方式，你可以为每个GPU设备设置不同的限制。

import cupy

mempool = cupy.get_default_memory_pool()

with cupy.cuda.Device(0):
    mempool.set_limit(size=1024**3)  # 1 GiB

with cupy.cuda.Device(1):
    mempool.set_limit(size=2*1024**3)  # 2 GiB

备注

CUDA 在内存池之外分配了一些 GPU 内存（例如 CUDA 上下文、库句柄等）。根据使用情况，此类内存可能占用一到几百 MiB。这不会计入限制中。

更改内存池#

你可以通过将内存分配函数传递给 cupy.cuda.set_allocator() / cupy.cuda.set_pinned_memory_allocator() 来使用自己的内存分配器，而不是默认的内存池。内存分配器函数应接受1个参数（以字节为单位的请求大小）并返回 cupy.cuda.MemoryPointer / cupy.cuda.PinnedMemoryPointer。

CuPy 提供了两种这样的分配器，用于在 GPU 上使用托管内存和流顺序内存，详情分别参见 cupy.cuda.malloc_managed() 和 cupy.cuda.malloc_async()。要启用由托管内存支持的内存池，您可以使用其分配器设置为 malloc_managed() 来构造一个新的 MemoryPool 实例，如下所示

import cupy

# Use managed memory
cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.malloc_managed).malloc)

请注意，如果你直接将 malloc_managed() 传递给 set_allocator() 而不构建 MemoryPool 实例，当内存被释放时，它将立即释放回系统，这可能或可能不是期望的行为。

Stream Ordered Memory Allocator 是自 CUDA 11.2 起新增的一个新特性。CuPy 为其提供了一个 实验性 接口。与 CuPy 的内存池类似，Stream Ordered Memory Allocator 也以流顺序的方式从/向内存池中异步分配/释放内存。关键区别在于，它是由 NVIDIA 在 CUDA 驱动中实现的固有特性，因此同一进程中的其他 CUDA 应用程序可以轻松地从同一池中分配内存。

要启用一个管理流顺序内存的内存池，你可以构造一个新的 MemoryAsyncPool 实例：

import cupy

# Use asynchronous stream ordered memory
cupy.cuda.set_allocator(cupy.cuda.MemoryAsyncPool().malloc)

# Create a custom stream
s = cupy.cuda.Stream()

# This would allocate memory asynchronously on stream s
with s:
    a = cupy.empty((100,), dtype=cupy.float64)

请注意，在这种情况下我们不使用 MemoryPool 类。MemoryAsyncPool 的输入参数与 MemoryPool 不同，以指示使用哪个池。有关更多详细信息，请参阅 MemoryAsyncPool 的文档。

请注意，如果你直接将 malloc_async() 传递给 set_allocator() 而不构造一个 MemoryAsyncPool 实例，将使用设备的当前内存池。

在使用流顺序内存时，重要的是您自己使用例如 Stream 和 Event API 来维护正确的流语义（详见流和事件）；CuPy 不会试图为您智能处理。在释放时，内存会在其分配的流上（第一次尝试）或任何当前的 CuPy 流上（第二次尝试）异步释放。允许在所有分配在该流上的内存被释放之前，销毁分配内存的流。

此外，应用程序/库内部使用 cudaMalloc``（CUDA 的默认同步分配器）可能会与流顺序内存分配器产生意外的相互作用。具体来说，释放到内存池的内存可能不会立即对 ``cudaMalloc 可见，从而导致潜在的内存不足错误。在这种情况下，您可以调用 free_all_blocks() 或手动执行（事件/流/设备）同步，然后重试。

目前，MemoryAsyncPool 接口是 实验性的。特别是，虽然它的 API 与 MemoryPool 的 API 大致相同，但由于 CUDA 的限制，池的几个方法需要足够新的驱动程序（当然，还需要支持的硬件、CUDA 版本和平台）。

你甚至可以通过以下代码禁用默认的内存池。请确保在进行任何其他 CuPy 操作之前执行此操作。

import cupy

# Disable memory pool for device memory (GPU)
cupy.cuda.set_allocator(None)

# Disable memory pool for pinned memory (CPU).
cupy.cuda.set_pinned_memory_allocator(None)

统一内存编程 (UMP) 支持 (实验性！)#

在启用异构内存管理（HMM）或地址转换服务（ATS）的系统上，例如NVIDIA Grace Hopper超级芯片，可以使NumPy和CuPy使用/共享系统分配的内存。要激活此功能，目前您需要：

安装 numpy_allocator
设置环境变量 CUPY_ENABLE_UMP=1
为CuPy创建一个内存池，以分配系统内存（malloc_system），例如：

import cupy as cp
cp.cuda.set_allocator(cp.cuda.MemoryPool(cp.cuda.memory.malloc_system).malloc)

切换到 NumPy 的对齐分配器以绘制系统内存

import cupy._core.numpy_allocator as ac
import numpy_allocator
import ctypes
lib = ctypes.CDLL(ac.__file__)
class my_allocator(metaclass=numpy_allocator.type):
    _calloc_ = ctypes.addressof(lib._calloc)
    _malloc_ = ctypes.addressof(lib._malloc)
    _realloc_ = ctypes.addressof(lib._realloc)
    _free_ = ctypes.addressof(lib._free)
my_allocator.__enter__()  # change the allocator globally

通过此设置更改，所有数据移动API，如 get()、asnumpy() 和 asarray() 变为无操作（不进行复制），并且以下代码得到加速：

# a, b, c are np.ndarray, d is cp.ndarray
a = np.random.random(100)
b = np.random.random(100)
c = np.add(a, b)
d = cp.matmul(cp.asarray(a), cp.asarray(c))

本质上，CPU/GPU 内存空间的区别已经消失，np/cp 用于表示执行空间（代码应在CPU还是GPU上运行）。

除了为 NumPy/CuPy 设置配置外，不需要更改用户代码。