cupy.cuda.MemoryPool

class cupy.cuda.MemoryPool(allocator=_malloc)

Memory pool for all GPU devices on the host.

A memory pool preserves any allocations even if they are freed by the user. Freed memory buffers are held by the memory pool as free blocks, and they are reused for further memory allocations of the same sizes. The allocated blocks are managed for each device, so one instance of this class can be used for multiple devices.

Note

When the allocation is skipped by reusing the pre-allocated block, it does not call cudaMalloc and therefore CPU-GPU synchronization does not occur. It makes interleaves of memory allocations and kernel invocations very fast.

Note

The memory pool holds allocated blocks without freeing as much as possible. It makes the program hold most of the device memory, which may make other CUDA programs running in parallel out-of-memory situation.

Parameters:allocator (function) – The base CuPy memory allocator. It is used for allocating new blocks when the blocks of the required size are all in use.

Methods

free_all_blocks(self, stream=None)

Release free blocks.

Parameters:stream (cupy.cuda.Stream) – Release free blocks in the arena of the given stream. The default releases blocks in all arenas.

free_all_free(self)

Release free blocks.

free_bytes(self)

Get the total number of bytes acquired but not used in the pool.

Returns:The total number of bytes acquired but not used in the pool. Return type:int

malloc(self, Py_ssize_t size) → MemoryPointer

Allocates the memory, from the pool if possible.

This method can be used as a CuPy memory allocator. The simplest way to use a memory pool as the default allocator is the following code:

set_allocator(MemoryPool().malloc)

Also, the way to use a memory pool of Managed memory (Unified memory) as the default allocator is the following code:

set_allocator(MemoryPool(malloc_managed).malloc)

Parameters:size (int) – Size of the memory buffer to allocate in bytes. Returns:Pointer to the allocated buffer. Return type:MemoryPointer

n_free_blocks(self)

Count the total number of free blocks.

Returns:The total number of free blocks. Return type:int

total_bytes(self)

Get the total number of bytes acquired in the pool.

Returns:The total number of bytes acquired in the pool. Return type:int

used_bytes(self)

Get the total number of bytes used.

Returns:The total number of bytes used. Return type:int