Add python wrapper for system memory resource

python/rmm/rmm/_lib/memory_resource.pxd

@@ -52,6 +52,12 @@ cdef class CudaMemoryResource(DeviceMemoryResource):
 cdef class ManagedMemoryResource(DeviceMemoryResource):
     pass
 
+cdef class SystemMemoryResource(DeviceMemoryResource):
+    pass
+
+cdef class SamHeadroomResourceAdaptor(DeviceMemoryResource):
+    cdef readonly DeviceMemoryResource system_mr
+
 cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
     pass
 

python/rmm/rmm/_lib/memory_resource.pyx

@@ -94,6 +94,11 @@ cdef extern from "rmm/mr/device/managed_memory_resource.hpp" \
     cdef cppclass managed_memory_resource(device_memory_resource):
         managed_memory_resource() except +
 
+cdef extern from "rmm/mr/device/system_memory_resource.hpp" \
+        namespace "rmm::mr" nogil:
+    cdef cppclass system_memory_resource(device_memory_resource):
+        system_memory_resource() except +
+
 cdef extern from "rmm/mr/device/cuda_async_memory_resource.hpp" \
         namespace "rmm::mr" nogil:
 
@@ -170,6 +175,13 @@ cdef extern from "rmm/mr/device/limiting_resource_adaptor.hpp" \
         size_t get_allocated_bytes() except +
         size_t get_allocation_limit() except +
 
+cdef extern from "rmm/mr/device/sam_headroom_resource_adaptor.hpp" \
+        namespace "rmm::mr" nogil:
+    cdef cppclass sam_headroom_resource_adaptor[Upstream](device_memory_resource):
+        sam_headroom_resource_adaptor(
+            Upstream* upstream_mr,
+            size_t headroom) except +
+
 cdef extern from "rmm/mr/device/logging_resource_adaptor.hpp" \
         namespace "rmm::mr" nogil:
     cdef cppclass logging_resource_adaptor[Upstream](device_memory_resource):
@@ -366,6 +378,49 @@ cdef class ManagedMemoryResource(DeviceMemoryResource):
         pass
 
 
+cdef class SystemMemoryResource(DeviceMemoryResource):
+    def __cinit__(self):
+        self.c_obj.reset(
+            new system_memory_resource()
+        )
+
+    def __init__(self):
+        """
+        Memory resource that uses ``malloc``/``free`` for
+        allocation/deallocation.
+        """
+        pass
+
+
+cdef class SamHeadroomResourceAdaptor(DeviceMemoryResource):
+    def __cinit__(
+        self,
+        size_t headroom
+    ):
+        self.system_mr = SystemMemoryResource()
+        self.c_obj.reset(
+            new sam_headroom_resource_adaptor[system_memory_resource](
+                <system_memory_resource*> self.system_mr.get_mr(),
+                headroom
+            )
+        )
+
+    def __init__(
+        self,
+        size_t headroom
+    ):
+        """
+        Memory resource that adapts system memory resource to allocate memory
+        with a headroom.
+
+        Parameters
+        ----------
+        headroom : size_t
+            Size of the reserved GPU memory as headroom
+        """
+        pass
+
+
 cdef class PoolMemoryResource(UpstreamResourceAdaptor):
 
     def __cinit__(

python/rmm/rmm/mr.py

@@ -24,7 +24,9 @@
     ManagedMemoryResource,
     PoolMemoryResource,
     PrefetchResourceAdaptor,
+    SamHeadroomResourceAdaptor,
     StatisticsResourceAdaptor,
+    SystemMemoryResource,
     TrackingResourceAdaptor,
     UpstreamResourceAdaptor,
     _flush_logs,
@@ -54,7 +56,9 @@
     "ManagedMemoryResource",
     "PoolMemoryResource",
     "PrefetchResourceAdaptor",
+    "SamHeadroomResourceAdaptor",
     "StatisticsResourceAdaptor",
+    "SystemMemoryResource",
     "TrackingResourceAdaptor",
     "FailureCallbackResourceAdaptor",
     "UpstreamResourceAdaptor",

python/rmm/rmm/rmm.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019, NVIDIA CORPORATION.
+# Copyright (c) 2019-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

python/rmm/rmm/tests/test_rmm.py

@@ -38,6 +38,11 @@
     _runtime_version >= 11020
 )
 
+_SYSTEM_MEMORY_SUPPORTED = rmm._cuda.gpu.getDeviceAttribute(
+    cudart.cudaDeviceAttr.cudaDevAttrPageableMemoryAccess,
+    rmm._cuda.gpu.getDevice(),
+)
+
 
 def array_tester(dtype, nelem, alloc):
     # data
@@ -91,6 +96,37 @@ def test_rmm_modes(dtype, nelem, alloc, managed, pool):
     array_tester(dtype, nelem, alloc)
 
 
+@pytest.mark.skipif(
+    not _SYSTEM_MEMORY_SUPPORTED,
+    reason="System memory not supported",
+)
+@pytest.mark.parametrize("dtype", _dtypes)
+@pytest.mark.parametrize("nelem", _nelems)
+@pytest.mark.parametrize("alloc", _allocs)
+@pytest.mark.parametrize(
+    "system, pool, headroom", list(product([False, True], [False, True], [False, True]))
+)
+def test_rmm_modes_system_memory(dtype, nelem, alloc, system, pool, headroom):
+    assert rmm.is_initialized()
+    array_tester(dtype, nelem, alloc)
+
+    if system and headroom:
+        base_mr = rmm.mr.SamHeadroomResourceAdaptor(headroom=1 << 20)
+    elif system:
+        base_mr = rmm.mr.SystemMemoryResource()
+    else:
+        base_mr = rmm.mr.CudaMemoryResource()
+    if pool:
+        mr = rmm.mr.PoolMemoryResource(base_mr)
+    else:
+        mr = base_mr
+    rmm.mr.set_current_device_resource(mr)
+
+    assert rmm.is_initialized()
+
+    array_tester(dtype, nelem, alloc)
+
+
 @pytest.mark.parametrize("dtype", _dtypes)
 @pytest.mark.parametrize("nelem", _nelems)
 @pytest.mark.parametrize("alloc", _allocs)
@@ -410,7 +446,13 @@ def test_pool_memory_resource(dtype, nelem, alloc):
     [
         lambda: rmm.mr.CudaMemoryResource(),
         lambda: rmm.mr.ManagedMemoryResource(),
-    ],
+    ]
+    + (
+        [lambda: rmm.mr.SystemMemoryResource(),
+         lambda: rmm.mr.SamHeadroomResourceAdaptor(headroom=1 << 20)]
+        if _SYSTEM_MEMORY_SUPPORTED
+        else []
+    ),
 )
 def test_fixed_size_memory_resource(dtype, nelem, alloc, upstream):
     mr = rmm.mr.FixedSizeMemoryResource(
@@ -432,7 +474,13 @@ def test_fixed_size_memory_resource(dtype, nelem, alloc, upstream):
         lambda: rmm.mr.PoolMemoryResource(
             rmm.mr.CudaMemoryResource(), 1 << 20
         ),
-    ],
+    ]
+    + (
+        [lambda: rmm.mr.SystemMemoryResource(),
+         lambda: rmm.mr.SamHeadroomResourceAdaptor(headroom=1 << 20)]
+        if _SYSTEM_MEMORY_SUPPORTED
+        else []
+    ),
 )
 def test_binning_memory_resource(dtype, nelem, alloc, upstream_mr):
     upstream = upstream_mr()