librmm-cu12

import rmm from rmm.mr import PoolMemoryResource, set_current_device_resource from rmm.device_buffer import DeviceBuffer import cupy as cp # 1. Configure RMM with a memory resource (e.g., a memory pool) pool_size_bytes = 2 * 1024**3 # 2 GB max_pool_size_bytes = 4 * 1024**3 # 4 GB # Create a PoolMemoryResource # Note: As of v26.02.00, host-only memory resources were removed. # This example uses a device-backed pool. mr = PoolMemoryResource(initial_pool_size=pool_size_bytes, maximum_pool_size=max_pool_size_bytes) set_current_device_resource(mr) print(f"Current RMM memory resource set to: {rmm.mr.get_current_device_resource()}") # 2. Allocate device memory directly with RMM db = DeviceBuffer(size=1024, dtype='uint8') print(f"Allocated DeviceBuffer of size {db.size} bytes: {db}") # 3. Use CuPy with RMM integration (CuPy will automatically use RMM) a = cp.arange(10**6, dtype=cp.float32) b = a * 2 print(f"CuPy array created using RMM: {a.shape}, dtype={a.dtype}") # Clean up (optional, as RMM resources are typically global and managed) del db del a, b # Note: The PoolMemoryResource itself will be deallocated when 'mr' goes out of scope # or when the program exits, releasing its managed memory.