Repository: rapidsai/rmm
Pull request #1864: Add async view memory resource bindings to Python.
Author: login: bdice
Labels: name: feature request
description: New feature or request
name: Python
description: Related to RMM Python API
name: non-breaking
description: Non-breaking change
name: cpp
description: Pertains to C++ code
Created: 2025-03-17T17:01:09Z | Merged: 2025-03-18T22:56:58Z

## Description
Closes #1611.

## Checklist
- [x] I am familiar with the [Contributing Guidelines](https://github.com/rapidsai/rmm/blob/HEAD/CONTRIBUTING.md).
- [x] New or existing tests cover these changes.
- [x] The documentation is up to date with these changes.


Linked issue #1611: [FEA] Need a Python interface for accessing the C++ cuda_async_view_memory_resource
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.

```python
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`:

```python
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`:

```python
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
Linked issue comments: user: login: wence-
created_at: 2024-07-17T09:25:54Z
body: 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`?
author_association: CONTRIBUTOR
user: login: lilohuang
created_at: 2024-07-17T10:48:50Z
body: @wence- The `cuda-python` `cudart.cudaMemPool_t` appears to be a viable option for me. Thanks.
author_association: CONTRIBUTOR
user: login: bdice
created_at: 2025-01-16T14:55:38Z
body: Thanks @lilohuang! I added this to our project board and @Matt711 is assigned.
author_association: COLLABORATOR
user: login: bdice
created_at: 2025-03-17T17:15:07Z
body: I started a PR for this in #1864. I am seeing some test failures that I need to debug, and may have to retarget to 25.06 depending on how long it takes me to get back to fixing those.
author_association: COLLABORATOR

Diff:
diff --git a/include/rmm/mr/device/cuda_async_view_memory_resource.hpp b/include/rmm/mr/device/cuda_async_view_memory_resource.hpp
index 92aea2072..0ef159c9f 100644
--- a/include/rmm/mr/device/cuda_async_view_memory_resource.hpp
+++ b/include/rmm/mr/device/cuda_async_view_memory_resource.hpp
@@ -47,13 +47,13 @@ class cuda_async_view_memory_resource final : public device_memory_resource {
    *
    * @throws rmm::logic_error if the CUDA version does not support `cudaMallocAsync`
    *
-   * @param valid_pool_handle Handle to a CUDA memory pool which will be used to
+   * @param pool_handle Handle to a CUDA memory pool which will be used to
    * serve allocation requests.
    */
-  cuda_async_view_memory_resource(cudaMemPool_t valid_pool_handle)
-    : cuda_pool_handle_{[valid_pool_handle]() {
-        RMM_EXPECTS(nullptr != valid_pool_handle, "Unexpected null pool handle.");
-        return valid_pool_handle;
+  cuda_async_view_memory_resource(cudaMemPool_t pool_handle)
+    : cuda_pool_handle_{[pool_handle]() {
+        RMM_EXPECTS(nullptr != pool_handle, "Unexpected null pool handle.");
+        return pool_handle;
       }()}
   {
     // Check if cudaMallocAsync Memory pool supported
@@ -72,7 +72,8 @@ class cuda_async_view_memory_resource final : public device_memory_resource {
    */
   [[nodiscard]] cudaMemPool_t pool_handle() const noexcept { return cuda_pool_handle_; }
 
-  cuda_async_view_memory_resource() = default;
+  cuda_async_view_memory_resource()  = default;
+  ~cuda_async_view_memory_resource() = default;
   cuda_async_view_memory_resource(cuda_async_view_memory_resource const&) =
     default;  ///< @default_copy_constructor
   cuda_async_view_memory_resource(cuda_async_view_memory_resource&&) =
diff --git a/python/rmm/rmm/librmm/memory_resource.pxd b/python/rmm/rmm/librmm/memory_resource.pxd
index 9e7b70c4f..3ded7a9be 100644
--- a/python/rmm/rmm/librmm/memory_resource.pxd
+++ b/python/rmm/rmm/librmm/memory_resource.pxd
@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,6 +23,8 @@ from libcpp.optional cimport optional
 from libcpp.pair cimport pair
 from libcpp.string cimport string
 
+from cuda.bindings.cyruntime cimport cudaMemPool_t
+
 from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 from rmm.librmm.memory_resource cimport device_memory_resource
 
@@ -108,6 +110,14 @@ cdef extern from "rmm/mr/device/cuda_async_memory_resource.hpp" \
             optional[size_t] release_threshold,
             optional[allocation_handle_type] export_handle_type) except +
 
+cdef extern from "rmm/mr/device/cuda_async_view_memory_resource.hpp" \
+        namespace "rmm::mr" nogil:
+
+    cdef cppclass cuda_async_view_memory_resource(device_memory_resource):
+        cuda_async_view_memory_resource(
+            cudaMemPool_t pool_handle) except +
+        cudaMemPool_t pool_handle() const
+
 # TODO: when we adopt Cython 3.0 use enum class
 cdef extern from "rmm/mr/device/cuda_async_memory_resource.hpp" \
         namespace \
diff --git a/python/rmm/rmm/mr.py b/python/rmm/rmm/mr.py
index 673ffde82..eada9b21b 100644
--- a/python/rmm/rmm/mr.py
+++ b/python/rmm/rmm/mr.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
     BinningMemoryResource,
     CallbackMemoryResource,
     CudaAsyncMemoryResource,
+    CudaAsyncViewMemoryResource,
     CudaMemoryResource,
     DeviceMemoryResource,
     FailureCallbackResourceAdaptor,
@@ -50,6 +51,7 @@
     "BinningMemoryResource",
     "CallbackMemoryResource",
     "CudaAsyncMemoryResource",
+    "CudaAsyncViewMemoryResource",
     "CudaMemoryResource",
     "DeviceMemoryResource",
     "FailureCallbackResourceAdaptor",
diff --git a/python/rmm/rmm/pylibrmm/memory_resource.pxd b/python/rmm/rmm/pylibrmm/memory_resource.pxd
index d1e5610db..8ccd07a02 100644
--- a/python/rmm/rmm/pylibrmm/memory_resource.pxd
+++ b/python/rmm/rmm/pylibrmm/memory_resource.pxd
@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2024, NVIDIA CORPORATION.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -44,6 +44,9 @@ cdef class SamHeadroomMemoryResource(DeviceMemoryResource):
 cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
     pass
 
+cdef class CudaAsyncViewMemoryResource(DeviceMemoryResource):
+    pass
+
 cdef class PoolMemoryResource(UpstreamResourceAdaptor):
     pass
 
diff --git a/python/rmm/rmm/pylibrmm/memory_resource.pyx b/python/rmm/rmm/pylibrmm/memory_resource.pyx
index 0189a58b5..253352e7c 100644
--- a/python/rmm/rmm/pylibrmm/memory_resource.pyx
+++ b/python/rmm/rmm/pylibrmm/memory_resource.pyx
@@ -28,7 +28,8 @@ from libcpp.memory cimport make_unique, unique_ptr
 from libcpp.optional cimport optional
 from libcpp.pair cimport pair
 
-from cuda.bindings.runtime import cudaError_t
+from cuda.bindings cimport cyruntime
+from cuda.bindings import driver, runtime
 
 from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice
 
@@ -54,6 +55,7 @@ from rmm.librmm.memory_resource cimport (
     binning_memory_resource,
     callback_memory_resource,
     cuda_async_memory_resource,
+    cuda_async_view_memory_resource,
     cuda_memory_resource,
     deallocate_callback_t,
     device_memory_resource,
@@ -203,6 +205,46 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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_handle : cudaMemPool_t or CUmemoryPool
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        pool_handle
+    ):
+        # Convert the pool_handle to a cyruntime.cudaMemPool_t
+        if not isinstance(pool_handle, (runtime.cudaMemPool_t, driver.CUmemoryPool)):
+            raise ValueError("pool_handle must be a cudaMemPool_t or CUmemoryPool")
+
+        cdef cyruntime.cudaMemPool_t c_pool_handle
+        c_pool_handle = <cyruntime.cudaMemPool_t><uintptr_t>int(pool_handle)
+
+        self.c_obj.reset(
+            new cuda_async_view_memory_resource(c_pool_handle)
+        )
+
+    def pool_handle(self):
+        cdef cuda_async_view_memory_resource* c_mr = \
+            <cuda_async_view_memory_resource*>self.c_obj.get()
+        return <uintptr_t>c_mr.pool_handle()
+
+
 cdef class ManagedMemoryResource(DeviceMemoryResource):
     def __cinit__(self):
         self.c_obj.reset(
@@ -991,7 +1033,7 @@ cpdef void _initialize(
     try:
         original_device = getDevice()
     except CUDARuntimeError as e:
-        if e.status == cudaError_t.cudaErrorNoDevice:
+        if e.status == runtime.cudaError_t.cudaErrorNoDevice:
             warnings.warn(e.msg)
         else:
             raise e
diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py
index ee02d5d0e..2fc917863 100644
--- a/python/rmm/rmm/tests/test_rmm.py
+++ b/python/rmm/rmm/tests/test_rmm.py
@@ -1078,3 +1078,42 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_default_pool(dtype, nelem, alloc):
+    # Get the default memory pool handle
+    current_device = rmm._cuda.gpu.getDevice()
+    err, pool = runtime.cudaDeviceGetDefaultMemPool(current_device)
+    assert err == runtime.cudaError_t.cudaSuccess
+
+    mr = rmm.mr.CudaAsyncViewMemoryResource(pool)
+    rmm.mr.set_current_device_resource(mr)
+    assert rmm.mr.get_current_device_resource_type() is type(mr)
+    array_tester(dtype, nelem, alloc)
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_custom_pool(dtype, nelem, alloc):
+    # Create a memory pool handle
+    props = runtime.cudaMemPoolProps()
+    props.allocType = runtime.cudaMemAllocationType.cudaMemAllocationTypePinned
+    props.location.id = rmm._cuda.gpu.getDevice()
+    props.location.type = runtime.cudaMemLocationType.cudaMemLocationTypeDevice
+    err, pool = runtime.cudaMemPoolCreate(props)
+    assert err == runtime.cudaError_t.cudaSuccess
+
+    mr = rmm.mr.CudaAsyncViewMemoryResource(pool)
+    rmm.mr.set_current_device_resource(mr)
+    assert rmm.mr.get_current_device_resource_type() is type(mr)
+    array_tester(dtype, nelem, alloc)
+
+    # After the pool is destroyed, new allocations should raise
+    (err,) = runtime.cudaMemPoolDestroy(pool)
+    assert err == runtime.cudaError_t.cudaSuccess
+    with pytest.raises(MemoryError):
+        array_tester(dtype, nelem, alloc)


Reviews:
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T18:11:48Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T18:12:19Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T18:52:07Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T18:53:41Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T20:28:03Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T21:21:30Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T21:31:10Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T21:32:08Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T21:32:43Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T21:33:34Z
user: login: wence-
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T21:34:10Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T21:35:01Z
user: login: wence-
body: I had only minor nitpicks, so approving as is.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T21:35:10Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T22:29:23Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T22:30:13Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T22:30:54Z
user: login: bdice
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T22:32:27Z

Review comments:
diff_hunk: @@ -203,6 +205,38 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: wence-
body: ```suggestion
        cyruntime.cudaMemPool_t valid_pool_handle not None
```
created_at: 2025-03-17T18:11:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -203,6 +205,38 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
+    ):
+        cdef cyruntime.cudaMemPool_t c_memory_pool_handle = \
+            <cyruntime.cudaMemPool_t>valid_pool_handle
+        self.c_obj.reset(
+            new cuda_async_view_memory_resource(c_memory_pool_handle)
+        )
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: wence-
body: ```suggestion
        self.c_obj.reset(
            new cuda_async_view_memory_resource(valid_pool_handle)
        )
```
created_at: 2025-03-17T18:12:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -203,6 +205,38 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: Build error:
```
/home/coder/rmm/python/rmm/rmm/pylibrmm/memory_resource.pyx:226:8: Only Python type arguments can have 'not None'
```

If I remove the `not None`, I get a different build error:

```
/home/coder/rmm/python/rmm/rmm/pylibrmm/memory_resource.pyx:226:8: Cannot convert Python object argument to type 'cudaMemPool_t'
```
created_at: 2025-03-17T18:52:07Z
author_association: COLLABORATOR
diff_hunk: @@ -203,6 +205,38 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: I also tried using `<cyruntime.cudaMempool_t ?>` for the cast, and got a segfault.
created_at: 2025-03-17T18:53:41Z
author_association: COLLABORATOR
diff_hunk: @@ -203,6 +205,38 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: Fixed in 88a2f71fb81d457e9e04c486162a116461ed5166.
created_at: 2025-03-17T20:28:03Z
author_association: COLLABORATOR
diff_hunk: @@ -72,7 +72,8 @@ class cuda_async_view_memory_resource final : public device_memory_resource {
    */
   [[nodiscard]] cudaMemPool_t pool_handle() const noexcept { return cuda_pool_handle_; }
 
-  cuda_async_view_memory_resource() = default;
+  cuda_async_view_memory_resource()  = default;
+  ~cuda_async_view_memory_resource() = default;
path: include/rmm/mr/device/cuda_async_view_memory_resource.hpp
user: login: bdice
body: Fixed a missing rule-of-6 destructor declaration here, which raised a warning in my IDE.
created_at: 2025-03-18T21:19:25Z
author_association: COLLABORATOR
diff_hunk: @@ -203,6 +205,45 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: I don't like the name `valid_pool_handle` but I copied it from C++. Should we consider renaming it?
created_at: 2025-03-18T21:20:33Z
author_association: COLLABORATOR
diff_hunk: @@ -203,6 +205,45 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: wence-
body: Sure, though please mention somewhere in this docstring that the user is responsible for keeping the mempool alive
created_at: 2025-03-18T21:31:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -203,6 +205,45 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
+    ):
+        # Convert the valid_pool_handle to a cyruntime.cudaMemPool_t
+        cdef cyruntime.cudaMemPool_t c_valid_pool_handle
+        if isinstance(valid_pool_handle, (runtime.cudaMemPool_t, driver.CUmemoryPool)):
+            raw_pool_handle = int(valid_pool_handle)
+            c_valid_pool_handle = <cyruntime.cudaMemPool_t><uintptr_t>raw_pool_handle
+        else:
+            raw_pool_handle = int(runtime.cudaMemPool_t(valid_pool_handle))
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: wence-
body: Is there any typing we can apply to the constructor argument that would prevent someone passing something bad in here (or does this constructor do that for hs?)
created_at: 2025-03-18T21:32:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -991,7 +1032,7 @@ cpdef void _initialize(
     try:
         original_device = getDevice()
     except CUDARuntimeError as e:
-        if e.status == cudaError_t.cudaErrorNoDevice:
+        if e.status == runtime.cudaError_t.cudaErrorNoDevice:
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: wence-
body: Why?
created_at: 2025-03-18T21:32:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1078,3 +1078,36 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_default_pool(dtype, nelem, alloc):
+    # Get the default memory pool handle
+    current_device = rmm._cuda.gpu.getDevice()
+    err, pool = runtime.cudaDeviceGetDefaultMemPool(current_device)
+    assert err == runtime.cudaError_t.cudaSuccess
+
+    mr = rmm.mr.CudaAsyncViewMemoryResource(pool)
+    rmm.mr.set_current_device_resource(mr)
+    assert rmm.mr.get_current_device_resource_type() is type(mr)
+    array_tester(dtype, nelem, alloc)
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_custom_pool(dtype, nelem, alloc):
+    # Create a memory pool handle
+    props = runtime.cudaMemPoolProps()
+    props.allocType = runtime.cudaMemAllocationType.cudaMemAllocationTypePinned
+    props.location.id = rmm._cuda.gpu.getDevice()
+    props.location.type = runtime.cudaMemLocationType.cudaMemLocationTypeDevice
+    err, pool = runtime.cudaMemPoolCreate(props)
path: python/rmm/rmm/tests/test_rmm.py
user: login: wence-
body: This test should also destroy the pool
created_at: 2025-03-18T21:33:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1078,3 +1078,36 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_default_pool(dtype, nelem, alloc):
+    # Get the default memory pool handle
+    current_device = rmm._cuda.gpu.getDevice()
+    err, pool = runtime.cudaDeviceGetDefaultMemPool(current_device)
path: python/rmm/rmm/tests/test_rmm.py
user: login: wence-
body: Does this pool need to be released? I don't know how the default mempool behaves
created_at: 2025-03-18T21:34:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -991,7 +1032,7 @@ cpdef void _initialize(
     try:
         original_device = getDevice()
     except CUDARuntimeError as e:
-        if e.status == cudaError_t.cudaErrorNoDevice:
+        if e.status == runtime.cudaError_t.cudaErrorNoDevice:
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: We adopted this naming pattern (`from cuda.bindings import runtime` and `runtime.*`, etc.) in several RAPIDS libraries when cuda-python changed its layout. This keeps us from having to import a bunch of specific names from `cuda.bindings.runtime`, and it makes it clear where they come from. In this particular file we have both `cyruntime` and `runtime` and I want it to be clear which names are from which modules.
created_at: 2025-03-18T21:35:01Z
author_association: COLLABORATOR
diff_hunk: @@ -203,6 +205,45 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
+        Handle to a CUDA memory pool which will be used to serve allocation
+        requests.
+    """
+    def __cinit__(
+        self,
+        valid_pool_handle
+    ):
+        # Convert the valid_pool_handle to a cyruntime.cudaMemPool_t
+        cdef cyruntime.cudaMemPool_t c_valid_pool_handle
+        if isinstance(valid_pool_handle, (runtime.cudaMemPool_t, driver.CUmemoryPool)):
+            raw_pool_handle = int(valid_pool_handle)
+            c_valid_pool_handle = <cyruntime.cudaMemPool_t><uintptr_t>raw_pool_handle
+        else:
+            raw_pool_handle = int(runtime.cudaMemPool_t(valid_pool_handle))
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: I narrowed the constructor a bit in b7e899d62127f06b6729ed1426869a255a4fa539. I am now forcing users to pass a `cuda.bindings.runtime.cudaMemPool_t` or `cuda.bindings.driver.CUmemoryPool` instance. The constructor no longer accepts raw ints as pointers to the corresponding types. I think that is much safer.
created_at: 2025-03-18T22:29:23Z
author_association: COLLABORATOR
diff_hunk: @@ -1078,3 +1078,36 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_default_pool(dtype, nelem, alloc):
+    # Get the default memory pool handle
+    current_device = rmm._cuda.gpu.getDevice()
+    err, pool = runtime.cudaDeviceGetDefaultMemPool(current_device)
+    assert err == runtime.cudaError_t.cudaSuccess
+
+    mr = rmm.mr.CudaAsyncViewMemoryResource(pool)
+    rmm.mr.set_current_device_resource(mr)
+    assert rmm.mr.get_current_device_resource_type() is type(mr)
+    array_tester(dtype, nelem, alloc)
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_custom_pool(dtype, nelem, alloc):
+    # Create a memory pool handle
+    props = runtime.cudaMemPoolProps()
+    props.allocType = runtime.cudaMemAllocationType.cudaMemAllocationTypePinned
+    props.location.id = rmm._cuda.gpu.getDevice()
+    props.location.type = runtime.cudaMemLocationType.cudaMemLocationTypeDevice
+    err, pool = runtime.cudaMemPoolCreate(props)
path: python/rmm/rmm/tests/test_rmm.py
user: login: bdice
body: Good idea! Tested in b7e899d62127f06b6729ed1426869a255a4fa539. It also makes sure that the memory resource raises a `MemoryError` when asked to allocate after the pool is destroyed.
created_at: 2025-03-18T22:30:13Z
author_association: COLLABORATOR
diff_hunk: @@ -1078,3 +1078,36 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+def test_cuda_async_view_memory_resource_default_pool(dtype, nelem, alloc):
+    # Get the default memory pool handle
+    current_device = rmm._cuda.gpu.getDevice()
+    err, pool = runtime.cudaDeviceGetDefaultMemPool(current_device)
path: python/rmm/rmm/tests/test_rmm.py
user: login: bdice
body: I don't think so. It cannot be destroyed, per the docs. Attempting to destroy it returns an error code. 😄

> **Note:**
> A device's default memory pool cannot be destroyed.

https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html#group__CUDART__MEMORY__POOLS_1g709113128c1c52c3bf170022dc7723dd
created_at: 2025-03-18T22:30:54Z
author_association: COLLABORATOR
diff_hunk: @@ -203,6 +205,45 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         )
 
 
+cdef class CudaAsyncViewMemoryResource(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.
+
+    Parameters
+    ----------
+    valid_pool_handle : cudaMemPool_t
path: python/rmm/rmm/pylibrmm/memory_resource.pyx
user: login: bdice
body: Done in b7e899d62127f06b6729ed1426869a255a4fa539. The new sentence says,

> The memory pool passed in must not be destroyed during the lifetime of this memory resource.
created_at: 2025-03-18T22:32:27Z
author_association: COLLABORATOR

PR comments:
user: login: copy-pr-bot[bot]
created_at: 2025-03-17T17:01:14Z
body: Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message [here](https://docs.gha-runners.nvidia.com/cpr/auto-sync).

author_association: NONE
user: login: bdice
created_at: 2025-03-17T17:01:25Z
body: /ok to test
author_association: COLLABORATOR
user: login: bdice
created_at: 2025-03-17T17:03:07Z
body: Currently tests are failing locally and I am struggling to identify why. It seems like something is breaking when passing the pool from Python to C++?
author_association: COLLABORATOR
user: login: bdice
created_at: 2025-03-17T17:43:18Z
body: /ok to test
author_association: COLLABORATOR
user: login: bdice
created_at: 2025-03-18T12:48:38Z
body: This PR is out of date with the commits I've pushed. I will close and reopen to retrigger it.
author_association: COLLABORATOR
user: login: bdice
created_at: 2025-03-18T22:32:45Z
body: /merge
author_association: COLLABORATOR