[FEA] Need a Python interface for accessing the C++ cuda_async_view_memory_resource

[lilohuang]

Hi @leofang and all,

Currently, the librmm library provides a CudaAsyncMemoryResource Python interface (https://docs.rapids.ai/api/rmm/stable/python_api/#rmm.mr.CudaAsyncMemoryResource), which lets users access the C++ cuda_async_memory_resource class. However, there is no equivalent Python interface for accessing the C++ cuda_async_view_memory_resource class.

import rmm
import cupy
from rmm.allocators.cupy import rmm_cupy_allocator

rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource(1024, 1024))
cupy.cuda.set_allocator(rmm_cupy_allocator)

The CudaAsyncMemoryResource always requires creating a new memory pool through cudaMemPoolCreate (https://github.com/rapidsai/rmm/blob/branch-24.08/include/rmm/mr/device/cuda_async_memory_resource.hpp#L107), which is problematic when integrating librmm and cuDF usage into an existing GPU-accelerated application.

An existing GPU-accelerated application needs to use the default CUDA memory pool for stream-ordered memory allocation and deallocation through cudaMallocAsync and cudaFreeAsync. If the librmm library always creates a separate pool, the memory pool allocated by librmm can only be used by itself or by cuPy and cuDF, causing some GPU memory resource wastage.

We hope that librmm can provide a CudaAsyncViewMemoryResource Python interface, similar to the one shown below, to access the C++ cuda_async_view_memory_resource class. This would allow us to pass the default memory pool handle (obtained from cudaDeviceGetDefaultMemPool) to librmm:

import rmm
import cupy
from rmm.allocators.cupy import rmm_cupy_allocator

rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncViewMemoryResource(pool_handle))
cupy.cuda.set_allocator(rmm_cupy_allocator)

Alternatively, introduce rmm.mr.CudaAsyncDefaultMemoryResource(), which automatically obtains the pool handle from cudaDeviceGetDefaultMemPool:

import rmm
import cupy
from rmm.allocators.cupy import rmm_cupy_allocator

rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncDefaultMemoryResource())
cupy.cuda.set_allocator(rmm_cupy_allocator)

Thanks,
Lilo

[wence-]

Thanks, this should be quite doable. We need to think about what type the pool_handle should have in python, probably a cuda-python cudart.cudaMemPool_t?

[lilohuang]

@wence- The cuda-python cudart.cudaMemPool_t appears to be a viable option for me. Thanks.