rmm.pylibrmm#

This module contains the low-level Cython bindings for RMM. Some components from this module are re-exported through the top-level rmm module for convenience, while others are available only through this module.

Overview#

rmm.pylibrmm provides the Cython layer that wraps RMM’s C++ functionality:

DeviceBuffer - GPU memory buffer (available as rmm.DeviceBuffer)
memory_resource - Memory resource implementations (available as rmm.mr)
Logging utilities (available through rmm)
CUDA stream wrappers (documented below)

CUDA Stream#

The stream classes are available only through rmm.pylibrmm and provide low-level CUDA stream management.

rmm.pylibrmm.stream#

class rmm.pylibrmm.stream.CudaStreamFlags(*values)#

Bases: IntEnum

Enumeration of CUDA stream creation flags.

Attributes:

SYNC_DEFAULTint: Created stream synchronizes with the default stream.
NON_BLOCKINGint: Created stream does not synchronize with the default stream.

Methods

`as_integer_ratio`(/)	Return a pair of integers, whose ratio is equal to the original int.
`bit_count`(/)	Number of ones in the binary representation of the absolute value of self.
`bit_length`(/)	Number of bits necessary to represent self in binary.
`conjugate`(/)	Returns self, the complex conjugate of any int.
`from_bytes`(/, bytes[, byteorder, signed])	Return the integer represented by the given array of bytes.
`is_integer`(/)	Returns True.
`to_bytes`(/[, length, byteorder, signed])	Return an array of bytes representing an integer.

NON_BLOCKING = 1#

SYNC_DEFAULT = 0#

class rmm.pylibrmm.stream.Stream(obj=None, flags=None)#

Bases: object

Methods

`is_default`(self)	Check if we are the default CUDA stream
`synchronize`(self)	Synchronize the CUDA stream

is_default(self)#: Check if we are the default CUDA stream

synchronize(self)#: Synchronize the CUDA stream

CUDA Stream Pool#

class rmm.pylibrmm.cuda_stream_pool.CudaStreamPool#

Bases: object

A pool of CUDA streams for efficient stream management.

Provides thread-safe access to a collection of CUDA stream objects. Successive calls may return views of identical streams.

Parameters:

pool_sizeint, optional: Number of streams in the pool. Defaults to 16.
flagsCudaStreamFlags, optional: Flags used to create each stream. Defaults to CudaStreamFlags.SYNC_DEFAULT.

Methods

`get_pool_size`(self)	Get the pool size.
`get_stream`(self, int stream_id)	Get a Stream from the pool (optionally by ID).

get_pool_size(self) → int#

Get the pool size.

Returns:

int: The number of streams in the pool

get_stream( self, int stream_id: Optional[int] = None, ) → Stream#

Get a Stream from the pool (optionally by ID).

Parameters:

stream_idOptional[int], optional: The ID of the stream to get. If None, the next stream from the pool is returned.

Returns:

Stream: A non-owning Stream object from the pool

Device Buffer Functions#

rmm.pylibrmm.device_buffer.copy_device_to_ptr( uintptr_t d_src, uintptr_t d_dst, size_t count, Stream stream=DEFAULT_STREAM, ) → void#

Copy from a device pointer to a device pointer

Parameters:

d_srcint: Pointer to device data to copy.
d_dstint: Pointer to device memory to write into.
countint: Number of bytes to copy.
streamStream, optional: CUDA stream to use for copying. Defaults to the default stream.

Examples

>>> import rmm
>>> db = rmm.DeviceBuffer(size=5)
>>> db2 = rmm.DeviceBuffer.to_device(b"abc")
>>> rmm.pylibrmm.device_buffer.copy_device_to_ptr(db2.ptr, db.ptr, db2.size)
>>> hb = db.copy_to_host()
>>> hb
array([97, 98, 99,  0,  0], dtype=uint8)

rmm.pylibrmm.device_buffer.copy_host_to_ptr( const unsigned char[::1] hb, uintptr_t db, Stream stream=DEFAULT_STREAM, ) → void#

Copy from a host pointer to a device pointer

Parameters:

hbbuffer: Host data to copy.
dbint: Pointer to device memory to write into.
streamStream, optional: CUDA stream to use for copying. Defaults to the default stream.

Notes

If stream is the default stream, it is synchronized after the copy. If a non-default stream is provided, the copy is asynchronous. The caller must keep hb alive and unmodified until the stream is synchronized.

Examples

>>> import rmm
>>> db = rmm.DeviceBuffer(size=10)
>>> hb = b"abc"
>>> rmm.pylibrmm.device_buffer.copy_host_to_ptr(hb, db.ptr)
>>> hb = db.copy_to_host()
>>> print(hb)
array([97, 98, 99,  0,  0,  0,  0,  0,  0,  0], dtype=uint8)

rmm.pylibrmm.device_buffer.copy_ptr_to_host( uintptr_t db, unsigned char[::1] hb, Stream stream=DEFAULT_STREAM, ) → void#

Copy from a device pointer to a buffer on host

Parameters:

dbint: Pointer to device data to copy.
hbwritable buffer: Host buffer to write into.
streamStream, optional: CUDA stream to use for copying. Defaults to the default stream.

Notes

If stream is the default stream, it is synchronized after the copy. However if a non-default stream is provided, this function is fully asynchronous.

Examples

>>> import rmm
>>> db = rmm.DeviceBuffer.to_device(b"abc")
>>> hb = bytearray(db.nbytes)
>>> rmm.pylibrmm.device_buffer.copy_ptr_to_host(db.ptr, hb)
>>> print(hb)
bytearray(b'abc')

rmm.pylibrmm.device_buffer.to_device( const unsigned char[::1] b, Stream stream=DEFAULT_STREAM, ) → DeviceBuffer#

Return a new DeviceBuffer with a copy of the data.

Parameters:

bbuffer: Host data to copy to device.
streamStream, optional: CUDA stream to use for copying. Defaults to the default stream.

Returns:

DeviceBuffer: Device buffer containing a copy of the host data.

Notes

If stream is the default stream, it is synchronized after the copy. If a non-default stream is provided, the copy is asynchronous. The caller must keep b alive and unmodified until the stream is synchronized.

Examples

>>> import rmm
>>> db = rmm.pylibrmm.device_buffer.to_device(b"abc")
>>> print(bytes(db))
b'abc'