make attributes properties

Author: keenan-simpson

cuda_core/cuda/core/experimental/_memory.py

@@ -302,6 +302,25 @@ class Mempool(MemoryResource):
     Use the static methods `create` or `from_shared_handle` to instantiate a Mempool.
     Direct instantiation is not supported.
 
+    Properties
+    ----------
+    reuse_follow_event_dependencies : bool
+        Allow memory to be reused when there are event dependencies between streams.
+    reuse_allow_opportunistic : bool
+        Allow reuse of completed frees without dependencies.
+    reuse_allow_internal_dependencies : bool
+        Allow insertion of new stream dependencies for memory reuse.
+    release_threshold : int
+        Amount of reserved memory to hold before OS release.
+    reserved_mem_current : int
+        Current amount of backing memory allocated.
+    reserved_mem_high : int
+        High watermark of backing memory allocated.
+    used_mem_current : int
+        Current amount of memory in use.
+    used_mem_high : int
+        High watermark of memory in use.
+
     See Also
     --------
     create : Create a new memory pool
@@ -590,38 +609,85 @@ def device_id(self) -> int:
         """
         return self._dev_id
 
-    def get_attribute(self, attr: driver.CUmemPool_attribute) -> int:
-        """Get a memory pool attribute.
+    @property
+    def reuse_follow_event_dependencies(self) -> bool:
+        """Allow memory to be reused when there are event dependencies between streams."""
+        return bool(
+            handle_return(
+                driver.cuMemPoolGetAttribute(
+                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES
+                )
+            )
+        )
 
-        Parameters
-        ----------
-        attr : CUmemPool_attribute
-            The attribute to query. Supported attributes are:
-            - CU_MEMPOOL_ATTR_RELEASE_THRESHOLD: Amount of reserved memory to hold before releasing
-            - CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES: Allow reuse with event dependencies
-            - CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC: Allow reuse without dependencies
-            - CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES: Allow reuse with internal dependencies
-            - CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT: Current reserved memory
-            - CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH: High watermark of reserved memory
-            - CU_MEMPOOL_ATTR_USED_MEM_CURRENT: Current used memory
-            - CU_MEMPOOL_ATTR_USED_MEM_HIGH: High watermark of used memory
+    @property
+    def reuse_allow_opportunistic(self) -> bool:
+        """Allow reuse of completed frees without dependencies."""
+        return bool(
+            handle_return(
+                driver.cuMemPoolGetAttribute(
+                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC
+                )
+            )
+        )
 
-        Returns
-        -------
-        int
-            The value of the requested attribute
+    @property
+    def reuse_allow_internal_dependencies(self) -> bool:
+        """Allow insertion of new stream dependencies for memory reuse."""
+        return bool(
+            handle_return(
+                driver.cuMemPoolGetAttribute(
+                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES
+                )
+            )
+        )
 
-        Raises
-        ------
-        CUDAError
-            If the attribute query fails
+    @property
+    def release_threshold(self) -> int:
+        """Amount of reserved memory to hold before OS release."""
+        return int(
+            handle_return(
+                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD)
+            )
+        )
 
-        Notes
-        -----
-        The meaning and units of the returned value depend on the specific attribute
-        being queried. See the CUDA documentation for details on each attribute.
-        """
-        return handle_return(driver.cuMemPoolGetAttribute(self._handle, attr))
+    @property
+    def reserved_mem_current(self) -> int:
+        """Current amount of backing memory allocated."""
+        return int(
+            handle_return(
+                driver.cuMemPoolGetAttribute(
+                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT
+                )
+            )
+        )
+
+    @property
+    def reserved_mem_high(self) -> int:
+        """High watermark of backing memory allocated."""
+        return int(
+            handle_return(
+                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH)
+            )
+        )
+
+    @property
+    def used_mem_current(self) -> int:
+        """Current amount of memory in use."""
+        return int(
+            handle_return(
+                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_CURRENT)
+            )
+        )
+
+    @property
+    def used_mem_high(self) -> int:
+        """High watermark of memory in use."""
+        return int(
+            handle_return(
+                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH)
+            )
+        )
 
 
 class _DefaultAsyncMempool(MemoryResource):

cuda_core/tests/test_memory.py

@@ -260,9 +260,9 @@ def test_mempool():
     src_buffer.close()
 
     # Test pool attributes
-    used_mem = int(mr.get_attribute(driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_CURRENT))
+    used_mem = mr.used_mem_current
     assert used_mem >= 0
-    reserved_mem = int(mr.get_attribute(driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT))
+    reserved_mem = mr.reserved_mem_current
     assert reserved_mem >= 0
 
     # Test error cases