rmm.mr (Memory Resources)

class rmm.mr.ArenaMemoryResource(DeviceMemoryResource upstream_mr, arena_size=None, bool dump_log_on_failure=False)

Bases: UpstreamResourceAdaptor

A suballocator that emphasizes fragmentation avoidance and scalable concurrency support.

Parameters:

upstream_mrDeviceMemoryResource

The DeviceMemoryResource from which to allocate memory for arenas.

arena_sizeint, optional

Size in bytes of the global arena. Defaults to half of the available memory on the current device.

dump_log_on_failurebool, optional

Whether to dump the arena on allocation failure.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.BinningMemoryResource(DeviceMemoryResource upstream_mr, int8_t min_size_exponent=-1, int8_t max_size_exponent=-1)

Bases: UpstreamResourceAdaptor

Allocates memory from a set of specified “bin” sizes based on a specified allocation size.

If min_size_exponent and max_size_exponent are specified, initializes with one or more FixedSizeMemoryResource bins in the range [2**min_size_exponent, 2**max_size_exponent].

Call add_bin() to add additional bin allocators.

Parameters:

upstream_mrDeviceMemoryResource

The memory resource to use for allocations larger than any of the bins.

min_size_exponentsize_t

The base-2 exponent of the minimum size FixedSizeMemoryResource bin to create.

max_size_exponentsize_t

The base-2 exponent of the maximum size FixedSizeMemoryResource bin to create.

Attributes:

bin_mrs

BinningMemoryResource.bin_mrs: list

upstream_mr

Methods

add_bin(self, size_t allocation_size, ...)	Adds a bin of the specified maximum allocation size to this memory resource.
allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)

add_bin(self, size_t allocation_size, DeviceMemoryResource bin_resource=None)

Adds a bin of the specified maximum allocation size to this memory resource. If specified, uses bin_resource for allocation for this bin. If not specified, creates and uses a FixedSizeMemoryResource for allocation for this bin.

Allocations smaller than allocation_size and larger than the next smaller bin size will use this fixed-size memory resource.

Parameters:

allocation_sizesize_t

The maximum allocation size in bytes for the created bin

bin_resourceDeviceMemoryResource

The resource to use for this bin (optional)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

bin_mrs

BinningMemoryResource.bin_mrs: list

Get the list of binned memory resources.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.CallbackMemoryResource(allocate_func, deallocate_func)

Bases: DeviceMemoryResource

A memory resource that uses the user-provided callables to do memory allocation and deallocation.

CallbackMemoryResource should really only be used for debugging memory issues, as there is a significant performance penalty associated with using a Python function for each memory allocation and deallocation.

Parameters:

allocate_func: callable

The allocation function must accept two arguments. An integer representing the number of bytes to allocate and a Stream on which to perform the allocation, and return an integer representing the pointer to the allocated memory.

deallocate_func: callable

The deallocation function must accept three arguments. an integer representing the pointer to the memory to free, a second integer representing the number of bytes to free, and a Stream on which to perform the deallocation.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

Examples

>>> import rmm
>>> base_mr = rmm.mr.CudaMemoryResource()
>>> def allocate_func(size, stream):
...     print(f"Allocating {size} bytes")
...     return base_mr.allocate(size, stream)
...
>>> def deallocate_func(ptr, size, stream):
...     print(f"Deallocating {size} bytes")
...     return base_mr.deallocate(ptr, size, stream)
...
>>> rmm.mr.set_current_device_resource(
    rmm.mr.CallbackMemoryResource(allocate_func, deallocate_func)
)
>>> dbuf = rmm.DeviceBuffer(size=256)
Allocating 256 bytes
>>> del dbuf
Deallocating 256 bytes

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.CudaAsyncMemoryResource

Bases: DeviceMemoryResource

Memory resource that uses cudaMallocAsync/cudaFreeAsync for allocation/deallocation.

Parameters:

initial_pool_sizeint | str, optional

Initial pool size in bytes. By default, half the available memory on the device is used. A string argument is parsed using parse_bytes.

release_threshold: int, optional

Release threshold in bytes. If the pool size grows beyond this value, unused memory held by the pool will be released at the next synchronization point.

enable_ipc: bool, optional

If True, enables export of POSIX file descriptor handles for the memory allocated by this resource so that it can be used with CUDA IPC.

enable_fabric: bool, optional

If True, enables export of fabric handles for the memory allocated by this resource.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.CudaAsyncViewMemoryResource

Bases: DeviceMemoryResource

Memory resource that uses cudaMallocAsync/cudaFreeAsync for allocation/deallocation with an existing CUDA memory pool.

This resource uses an existing CUDA memory pool handle (such as the default pool) instead of creating a new one. This is useful for integrating with existing GPU applications that already use a CUDA memory pool, or customizing the flags used by the memory pool.

The memory pool passed in must not be destroyed during the lifetime of this memory resource.

Parameters:

pool_handlecudaMemPool_t or CUmemoryPool

Handle to a CUDA memory pool which will be used to serve allocation requests.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
pool_handle(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

pool_handle(self)

class rmm.mr.CudaMemoryResource

Bases: DeviceMemoryResource

Memory resource that uses cudaMalloc/cudaFree for allocation/deallocation.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.DeviceMemoryResource

Bases: object

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.FailureCallbackResourceAdaptor(DeviceMemoryResource upstream_mr, callback)

Bases: UpstreamResourceAdaptor

Memory resource that call callback when memory allocation fails.

Parameters:

upstreamDeviceMemoryResource

The upstream memory resource.

callbackcallable

Function called when memory allocation fails.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.FixedSizeMemoryResource(DeviceMemoryResource upstream_mr, size_t block_size=0x100000, size_t blocks_to_preallocate=128)

Bases: UpstreamResourceAdaptor

Memory resource which allocates memory blocks of a single fixed size.

Parameters:

upstream_mrDeviceMemoryResource

The DeviceMemoryResource from which to allocate blocks for the pool.

block_sizeint, optional

The size of blocks to allocate (default is 1MiB).

blocks_to_preallocateint, optional

The number of blocks to allocate to initialize the pool.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)

Notes

Supports only allocations of size smaller than the configured block_size.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.LimitingResourceAdaptor(DeviceMemoryResource upstream_mr, size_t allocation_limit)

Bases: UpstreamResourceAdaptor

Memory resource that limits the total allocation amount possible performed by an upstream memory resource.

Parameters:

upstream_mrDeviceMemoryResource

The upstream memory resource.

allocation_limitsize_t

Maximum memory allowed for this allocator.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_allocated_bytes(self)	Query the number of bytes that have been allocated.
get_allocation_limit(self)	Query the maximum number of bytes that this allocator is allowed to allocate.
get_upstream(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_allocated_bytes(self) → size_t

Query the number of bytes that have been allocated. Note that this can not be used to know how large of an allocation is possible due to both possible fragmentation and also internal page sizes and alignment that is not tracked by this allocator.

get_allocation_limit(self) → size_t

Query the maximum number of bytes that this allocator is allowed to allocate. This is the limit on the allocator and not a representation of the underlying device. The device may not be able to support this limit.

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.LoggingResourceAdaptor(DeviceMemoryResource upstream_mr, log_file_name=None)

Bases: UpstreamResourceAdaptor

Memory resource that logs information about allocations/deallocations performed by an upstream memory resource.

Parameters:

upstreamDeviceMemoryResource

The upstream memory resource.

log_file_namestr

Path to the file to which logs are written.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
flush(self)
get_file_name(self)
get_upstream(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

flush(self)

get_file_name(self)

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.ManagedMemoryResource

Bases: DeviceMemoryResource

Memory resource that uses cudaMallocManaged/cudaFree for allocation/deallocation.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.PinnedHostMemoryResource

Bases: DeviceMemoryResource

Memory resource that uses cudaHostAlloc/cudaFreeHost for allocation/deallocation of pinned host memory.

Pinned (page-locked) host memory is accessible from both the host and device, enabling faster data transfers between host and device.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.PoolMemoryResource(DeviceMemoryResource upstream_mr, initial_pool_size=None, maximum_pool_size=None)

Bases: UpstreamResourceAdaptor

Coalescing best-fit suballocator which uses a pool of memory allocated from an upstream memory resource.

Parameters:

upstream_mrDeviceMemoryResource

The DeviceMemoryResource from which to allocate blocks for the pool.

initial_pool_sizeint | str, optional

Initial pool size in bytes. By default, half the available memory on the device is used.

maximum_pool_sizeint | str, optional

Maximum size in bytes, that the pool can grow to.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)
pool_size(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

pool_size(self)

upstream_mr

class rmm.mr.PrefetchResourceAdaptor(DeviceMemoryResource upstream_mr)

Bases: UpstreamResourceAdaptor

Memory resource that prefetches all allocations.

Parameters:

upstreamDeviceMemoryResource

The upstream memory resource.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

upstream_mr

class rmm.mr.SamHeadroomMemoryResource(size_t headroom)

Bases: DeviceMemoryResource

Memory resource that uses malloc/free for allocation/deallocation.

Parameters:

headroomsize_t

Size of the reserved GPU memory as headroom

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.StatisticsResourceAdaptor(DeviceMemoryResource upstream_mr)

Bases: UpstreamResourceAdaptor

Memory resource that tracks the current, peak and total allocations/deallocations performed by an upstream memory resource. Includes the ability to query these statistics at any time.

A stack of counters is maintained. Use push_counters() and pop_counters() to track statistics at different nesting levels.

Parameters:

upstreamDeviceMemoryResource

The upstream memory resource.

Attributes:

allocation_counts

StatisticsResourceAdaptor.allocation_counts: Statistics

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)
pop_counters(self)	Pop a counter pair (bytes and allocations) from the stack
push_counters(self)	Push a new counter pair (bytes and allocations) on the stack

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

allocation_counts

StatisticsResourceAdaptor.allocation_counts: Statistics

Gets the current, peak, and total allocated bytes and number of allocations.

The dictionary keys are current_bytes, current_count, peak_bytes, peak_count, total_bytes, and total_count.

Returns:

dict: Dictionary containing allocation counts and bytes.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

pop_counters(self) → Statistics

Pop a counter pair (bytes and allocations) from the stack

Returns:

The popped statistics

push_counters(self) → Statistics

Push a new counter pair (bytes and allocations) on the stack

Returns:

The statistics _before_ the push

upstream_mr

class rmm.mr.SystemMemoryResource

Bases: DeviceMemoryResource

Memory resource that uses malloc/free for allocation/deallocation.

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

class rmm.mr.TrackingResourceAdaptor(DeviceMemoryResource upstream_mr, bool capture_stacks=False)

Bases: UpstreamResourceAdaptor

Memory resource that logs tracks allocations/deallocations performed by an upstream memory resource. Includes the ability to query all outstanding allocations with the stack trace, if desired.

Parameters:

upstreamDeviceMemoryResource

The upstream memory resource.

capture_stacksbool

Whether or not to capture the stack trace with each allocation.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_allocated_bytes(self)	Query the number of bytes that have been allocated.
get_outstanding_allocations_str(self)	Returns a string containing information about the current outstanding allocations.
get_upstream(self)
log_outstanding_allocations(self)	Logs the output of get_outstanding_allocations_str() to the current RMM log file if enabled.

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_allocated_bytes(self) → size_t

Query the number of bytes that have been allocated. Note that this can not be used to know how large of an allocation is possible due to both possible fragmentation and also internal page sizes and alignment that is not tracked by this allocator.

get_outstanding_allocations_str(self) → str

Returns a string containing information about the current outstanding allocations. For each allocation, the address, size and optional stack trace are shown.

get_upstream(self) → DeviceMemoryResource

log_outstanding_allocations(self)

Logs the output of get_outstanding_allocations_str() to the current RMM log file if enabled.

upstream_mr

class rmm.mr.UpstreamResourceAdaptor

Bases: DeviceMemoryResource

Parent class for all memory resources that track an upstream.

Upstream resource tracking requires maintaining a reference to the upstream mr so that it is kept alive and may be accessed by any downstream resource adaptors.

Attributes:

upstream_mr

Methods

allocate(self, size_t nbytes, ...)	Allocate nbytes bytes of memory.
deallocate(self, uintptr_t ptr, ...)	Deallocate memory pointed to by ptr of size nbytes.
get_upstream(self)

allocate(self, size_t nbytes, Stream stream=DEFAULT_STREAM)

Allocate nbytes bytes of memory.

Parameters:

nbytessize_t

The size of the allocation in bytes

streamStream

Optional stream for the allocation

Raises:

MemoryError

If allocation fails.

Notes

On integrated memory systems, attempting to allocate more memory than available can cause the process to be killed by the operating system instead of raising a catchable MemoryError.

deallocate(self, uintptr_t ptr, size_t nbytes, Stream stream=DEFAULT_STREAM)

Deallocate memory pointed to by ptr of size nbytes.

Parameters:

ptruintptr_t

Pointer to be deallocated

nbytessize_t

Size of the allocation in bytes

streamStream

Optional stream for the deallocation

get_upstream(self) → DeviceMemoryResource

upstream_mr

rmm.mr.available_device_memory()

Returns a tuple of free and total device memory memory.

rmm.mr.disable_logging()

Disable logging if it was enabled previously using rmm.reinitialize() or rmm.enable_logging().

rmm.mr.enable_logging(log_file_name=None)

Enable logging of runtime events for all devices.

Parameters:

log_file_namestr, optional

Name of the log file. If not specified, the environment variable RMM_LOG_FILE is used. A ValueError is thrown if neither is available. A separate log file is produced for each device, and the suffix ".dev{id}" is automatically added to the log file name.

Notes

Note that if you use the environment variable CUDA_VISIBLE_DEVICES with logging enabled, the suffix may not be what you expect. For example, if you set CUDA_VISIBLE_DEVICES=1, the log file produced will still have suffix 0. Similarly, if you set CUDA_VISIBLE_DEVICES=1,0 and use devices 0 and 1, the log file with suffix 0 will correspond to the GPU with device ID 1. Use rmm.get_log_filenames() to get the log file names corresponding to each device.

rmm.mr.get_current_device_resource() → DeviceMemoryResource

Get the memory resource used for RMM device allocations on the current device.

If the returned memory resource is used when a different device is the active CUDA device, behavior is undefined.

rmm.mr.get_current_device_resource_type()

Get the memory resource type used for RMM device allocations on the current device.

rmm.mr.get_log_filenames()

Returns the log filename (or None if not writing logs) for each device in use.

Examples

>>> import rmm
>>> rmm.reinitialize(devices=[0, 1], logging=True, log_file_name="rmm.log")
>>> rmm.get_log_filenames()
{0: '/home/user/workspace/rapids/rmm/python/rmm.dev0.log',
 1: '/home/user/workspace/rapids/rmm/python/rmm.dev1.log'}

rmm.mr.get_per_device_resource(int device)

Get the default memory resource for the specified device.

If the returned memory resource is used when a different device is the active CUDA device, behavior is undefined.

Parameters:

deviceint

The ID of the device for which to get the memory resource.

rmm.mr.get_per_device_resource_type(int device)

Get the memory resource type used for RMM device allocations on the specified device.

Parameters:

deviceint

The device ID

rmm.mr.is_initialized()

Check whether RMM is initialized

rmm.mr.set_current_device_resource(DeviceMemoryResource mr)

Set the default memory resource for the current device.

Parameters:

mrDeviceMemoryResource

The memory resource to set. Must have been created while the current device is the active CUDA device.

rmm.mr.set_per_device_resource(int device, DeviceMemoryResource mr)

Set the default memory resource for the specified device.

Parameters:

deviceint

The ID of the device for which to get the memory resource.

mrDeviceMemoryResource

The memory resource to set. Must have been created while device was the active CUDA device.