DOCSEARCHbyAlgolia
Submit Docs

RMM (RAPIDS Memory Manager)

repository·main·Indexed 20 days ago

https://github.com/rapidsai/rmm

RMM provides a common interface for customizing memory allocation on both host and device, supporting high-performance GPU workflows via pinned host memory and device memory pool sub-allocators. It includes various memory resources such as cuda_memory_resource, pool_memory_resource, and managed_memory_resource, along with C++ RAII classes like device_buffer and device_uvector, and Python APIs for explicit device memory management and CuPy integration.

Tokens
16.9K
Snippets
52
Records
76
Agent score
72%

What's inside rmm

  1. Overview of rmm.pylibrmm low-level bindings

    main

    The rmm.pylibrmm module contains the low-level Cython bindings for RMM, wrapping its C++ functionality. While some components are re-exported to the top-level rmm module for ease of use, others must be accessed directly via rmm.pylibrmm.

    Key components include:

    • DeviceBuffer: GPU memory buffers (re-exported as rmm.DeviceBuffer).
    • memory_resource: Implementations of memory resources (re-exported as rmm.mr).
    • Logging: Utilities available through the top-level rmm module.
    • CUDA Stream Wrappers: Low-level stream management classes available exclusively in rmm.pylibrmm.stream.

  2. Overview of the RAPIDS Memory Manager (RMM)

    main

    The RAPIDS Memory Manager (RMM) is designed to optimize performance in GPU-centric workflows by providing customizable memory allocation strategies for both host and device memory.

    Key capabilities include:

    • Customizable Allocation: A common interface to control how host and device memory are allocated.
    • Specialized Implementations: A collection of allocators, such as using "pinned" host memory for faster asynchronous host-to-device transfers, or device memory pool sub-allocators to minimize the overhead of frequent dynamic device memory allocations.
    • Integrated Data Structures: A suite of data structures that leverage the RMM interface for efficient memory management.

  3. Use the RMM Python API

    main

    RMM provides a Python interface for memory management, accessible by importing the rmm module. The API is organized into several submodules:

    • rmm: The main entry point for RMM functionality.
    • mr: Memory Resources (used for managing memory pools and device memory).
    • allocators: Python wrappers for RMM allocators.
    • statistics: Tools for gathering memory usage statistics.
    • pylibrmm and librmm: Low-level bindings and library interfaces.
    import rmm
# Access submodules via rmm or direct imports
from rmm import mr
from rmm import allocators

  4. Use rmm.statistics for statistical computations

    main
    The rmm.statistics module provides a suite of statistical functions designed to work with RMM-managed memory. It is intended for performing high-performance statistical operations on data structures compatible with the RAPIDS ecosystem.

  5. Taking ownership of C++ objects from Python

    main

    When interacting with a C++ library that uses RMM from Python, it is best practice to use C++ APIs that accept an explicit memory resource argument. This allows Python callers to pass a specific resource instead of relying on the process-global current-resource state.

    For example, a C++ function returning an rmm::device_buffer can accept an owning type-erased resource using cuda::mr::any_resource<cuda::mr::device_accessible>.

    On the Python side, the DeviceBuffer class provides a Cython function c_from_unique_ptr to construct a DeviceBuffer from a unique_ptr<rmm::device_buffer>, effectively taking ownership of it. When constructing the Python wrapper, pass the same Python memory resource object if you want the wrapper to maintain that resource association.

    std::unique_ptr<rmm::device_buffer> allocate(
  std::size_t size,
  cuda::mr::any_resource<cuda::mr::device_accessible> mr =
    rmm::mr::get_current_device_resource_ref())
{
    return std::make_unique<rmm::device_buffer>(size, rmm::cuda_stream_default, std::move(mr));
}

  6. How RMM memory resources work

    main

    RMM provides a common interface for device memory allocation that follows CCCL's memory resource concepts.

    Memory Resource Interface

    A device memory resource provides stream-ordered allocation and deallocation via these methods:

    void* allocate(cuda::stream_ref stream,
               std::size_t bytes,
               std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT);

void deallocate(cuda::stream_ref stream,
                void* ptr,
                std::size_t bytes,
                std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept;

    Resource Types

    • rmm::device_async_resource_ref: A lightweight, non-owning reference to a device resource (alias for cuda::mr::resource_ref<cuda::mr::device_accessible>).
    • cuda::mr::any_resource<cuda::mr::device_accessible>: An owning, type-erased resource.

    Resources with non-trivial state are value types with shared ownership; copying them is inexpensive and keeps the underlying state alive.

    // Example of the expected interface for a device memory resource
void* ptr = resource.allocate(stream, bytes, alignment);
resource.deallocate(stream, ptr, bytes, alignment);

  7. Available RMM Device Memory Resources

    main

    RMM provides several device memory resources for different allocation strategies:

    • cuda_memory_resource: Uses standard cudaMalloc and cudaFree.
    • managed_memory_resource: Uses cudaMallocManaged and cudaFree. Note: NVIDIA vGPU does not support this by default unless Unified Memory is enabled.
    • pool_memory_resource: A coalescing, best-fit pool sub-allocator for fast dynamic allocation.
    • fixed_size_memory_resource: Allocates a single fixed size with constant cost for allocation and deallocation.
    • binning_memory_resource: Uses multiple upstream resources for different bin sizes (e.g., multiple fixed_size_memory_resources for small bins and a pool_memory_resource for larger ones).

  8. Understand stream-ordered memory allocation semantics

    main

    RMM memory resources use stream-ordered allocation. This allows the allocator to reuse memory deallocated on the same stream without requiring heavy synchronization.

    Rules and Constraints

    • Validity: A pointer returned by resource.allocate(stream_a, bytes) is only guaranteed to be valid for use on stream_a. Using it on stream_b without prior synchronization (e.g., cudaStreamSynchronize(stream_a) or a CUDA event) is Undefined Behavior.
    • Deallocation: The stream passed to deallocate should be the stream on which the memory was last used. This allows the resource to manage available memory with minimal synchronization.
    • Stream Destruction: It is Undefined Behavior to destroy a CUDA stream that is currently being passed to deallocate. If a stream is about to be destroyed, synchronize it first and pass a different stream (like the default stream) to deallocate.

    These semantics also apply to RMM device data structures like rmm::device_buffer and rmm::device_uvector.

  9. Understand the RMM namespaces

    main

    The RAPIDS Memory Manager (RMM) organizes its public API within two primary namespaces:

    • rmm: The main namespace containing the core memory management utilities, allocators, and pool managers.
    • rmm::mr: A sub-namespace typically used for memory resource abstractions or specialized memory management components.

  10. Python-specific memory and resource management considerations

    main

    When working with the RMM Python layer, developers must ensure proper management of GPU memory and resource lifecycles. Key considerations include:

    • Memory Resource Lifecycle: Use context managers for resource management where appropriate and ensure proper cleanup in __del__ methods. Ownership semantics between Python and C++ layers must be clearly documented.
    • Cython Bindings: Implement proper memory management using __dealloc__, handle exceptions correctly across the Python/C++ boundary, and ensure correct GIL handling for CUDA operations.
    • Array Interfaces: Support __cuda_array_interface__ to ensure interoperability with libraries like CuPy and Numba-CUDA. Ensure shape, strides, and data pointers are correctly handled and that the data pointer remains valid for the object's lifetime.
    • Stream Handling: Ensure correct CUDA stream semantics by propagating stream parameters and managing synchronization correctly in Python bindings.

  11. Stream management and asynchronous memory operations

    main

    Because RMM handles GPU memory, managing CUDA streams is critical to prevent race conditions and use-after-free errors.

    Core Requirements

    • Async Allocation Signature: Asynchronous allocation functions must accept a cuda_stream_view.
    • Stream Synchronization: You must ensure that asynchronous operations are complete before memory is returned to a pool or deallocated. Use stream.synchronize_no_throw(*this) before calling upstream_->deallocate(...) if the memory is still in use by an async task.
    • Lifecycle Management: Avoid using destroyed streams and ensure proper stream/event cleanup to prevent leaks.
    // Example: Synchronize before deallocating to prevent use-after-free
stream.synchronize_no_throw(*this);
upstream_->deallocate(ptr, size, stream);

  12. Memory management and RAII patterns in RMM

    main

    RMM provides RAII-based device memory management to prevent leaks and ensure resource safety, especially during exception paths.

    Best Practices

    • Use RAII Wrappers: Prefer rmm::device_uvector and rmm::device_buffer over raw pointers to manage device memory lifecycle.
    • Resource Inheritance: All custom memory resources must derive from rmm::device_memory_resource.
    • Upstream Delegation: When implementing memory resources (like pools or arenas), ensure you properly forward allocations and deallocations to the upstream resource, including correct size, alignment, and error path handling.
    • Avoid Raw Pointers: Avoid raw pointer allocations without an RAII wrapper to prevent leaks when exceptions are thrown.