Implement release threshold configuration for DeviceMemoryResource performance optimization (#875)

* Initial plan

* Implement release threshold configuration for DeviceMemoryResource performance optimization

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

* Add performance demo for DeviceMemoryResource release threshold optimization

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

* Fix linting issues and format code with ruff

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

* Remove try-except wrapper and performance demo per code review feedback

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

* Add release note for DeviceMemoryResource performance optimization

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

* Add skip decorator for mempool support check in device memory test

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

* Address code review feedback: move skip logic, add docstring note, update release note

Co-authored-by: kkraus14 <3665167+kkraus14@users.noreply.github.com>

* Remove verbose docstring Notes section per code review feedback

Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: leofang <5534781+leofang@users.noreply.github.com>
Co-authored-by: kkraus14 <3665167+kkraus14@users.noreply.github.com>

Author: 3 people

cuda_core/cuda/core/experimental/_memory.py

@@ -326,6 +326,23 @@ def __init__(self, device_id: int):
         self._handle = handle_return(driver.cuDeviceGetMemPool(device_id))
         self._dev_id = device_id
 
+        # Set a higher release threshold to improve performance when there are no active allocations.
+        # By default, the release threshold is 0, which means memory is immediately released back
+        # to the OS when there are no active suballocations, causing performance issues.
+        # Check current release threshold
+        current_threshold = handle_return(
+            driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD)
+        )
+        # If threshold is 0 (default), set it to maximum to retain memory in the pool
+        if int(current_threshold) == 0:
+            handle_return(
+                driver.cuMemPoolSetAttribute(
+                    self._handle,
+                    driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
+                    driver.cuuint64_t(0xFFFFFFFFFFFFFFFF),
+                )
+            )
+
     def allocate(self, size: int, stream: Stream = None) -> Buffer:
         """Allocate a buffer of the requested size.
 

cuda_core/docs/source/release/0.X.Y-notes.rst

@@ -36,4 +36,5 @@ None.
 Fixes and enhancements
 ----------------------
 
+- Improved :class:`DeviceMemoryResource` allocation performance when there are no active allocations by setting a higher release threshold (addresses issue #771).
 - Fix :class:`LaunchConfig` grid unit conversion when cluster is set (addresses issue #867).

cuda_core/tests/test_memory.py

@@ -10,7 +10,7 @@
 
 import pytest
 
-from cuda.core.experimental import Buffer, Device, MemoryResource
+from cuda.core.experimental import Buffer, Device, DeviceMemoryResource, MemoryResource
 from cuda.core.experimental._memory import DLDeviceType
 from cuda.core.experimental._utils.cuda_utils import handle_return
 
@@ -257,3 +257,29 @@ def test_buffer_dunder_dlpack_device_failure():
     buffer = dummy_mr.allocate(size=1024)
     with pytest.raises(BufferError, match=r"^buffer is neither device-accessible nor host-accessible$"):
         buffer.__dlpack_device__()
+
+
+def test_device_memory_resource_initialization():
+    """Test that DeviceMemoryResource can be initialized successfully.
+
+    This test verifies that the DeviceMemoryResource initializes properly,
+    including the release threshold configuration for performance optimization.
+    """
+    device = Device()
+    if not device.properties.memory_pools_supported:
+        pytest.skip("memory pools not supported")
+    device.set_current()
+
+    # This should succeed and configure the memory pool release threshold
+    mr = DeviceMemoryResource(device.device_id)
+
+    # Verify basic properties
+    assert mr.device_id == device.device_id
+    assert mr.is_device_accessible is True
+    assert mr.is_host_accessible is False
+
+    # Test allocation/deallocation works
+    buffer = mr.allocate(1024)
+    assert buffer.size == 1024
+    assert buffer.device_id == device.device_id
+    buffer.close()