Simplify resource handle patterns and clean up tests

- Remove Kernel._module (ObjectCode reference no longer needed since
  KernelHandle keeps library alive via LibraryHandle dependency)
- Simplify Kernel._from_obj signature (remove unused ObjectCode param)
- Replace weakref patterns with direct handle storage:
  - KernelAttributes: store KernelHandle instead of weakref to Kernel
  - _MemPoolAttributes: store MemoryPoolHandle instead of weakref to _MemPool
- Rename get_kernel_from_library to create_kernel_handle for consistency
- Remove fragile annotation introspection from test_saxpy_arguments
- Update test_mempool_attributes_ownership to reflect new ownership semantics

Author: Andy-Jost

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -749,7 +749,7 @@ struct KernelBox {
 };
 }  // namespace
 
-KernelHandle get_kernel_from_library(LibraryHandle h_library, const char* name) {
+KernelHandle create_kernel_handle(LibraryHandle h_library, const char* name) {
     GILReleaseGuard gil;
     CUkernel kernel;
     if (CUDA_SUCCESS != (err = p_cuLibraryGetKernel(&kernel, *h_library, name))) {

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -254,7 +254,7 @@ LibraryHandle create_library_handle_ref(CUlibrary library);
 // The kernel structurally depends on the provided library handle.
 // Kernels have no explicit destroy - their lifetime is tied to the library.
 // Returns empty handle on error (caller must check).
-KernelHandle get_kernel_from_library(LibraryHandle h_library, const char* name);
+KernelHandle create_kernel_handle(LibraryHandle h_library, const char* name);
 
 // Create a non-owning kernel handle with library dependency.
 // Use for borrowed kernels. The library handle keeps the library alive.

cuda_core/cuda/core/_memory/_memory_pool.pxd

@@ -16,7 +16,16 @@ cdef class _MemPool(MemoryResource):
         IPCDataForMR          _ipc_data
         object                _attributes
         object                _peer_accessible_by
-        object                __weakref__
+
+
+cdef class _MemPoolAttributes:
+    cdef:
+        MemoryPoolHandle _h_pool
+
+    @staticmethod
+    cdef _MemPoolAttributes _init(MemoryPoolHandle h_pool)
+
+    cdef int _getattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except? -1
 
 
 cdef class _MemPoolOptions:

cuda_core/cuda/core/_memory/_memory_pool.pyx

@@ -30,7 +30,6 @@ from cuda.core._utils.cuda_utils cimport (
 
 from typing import TYPE_CHECKING
 import platform  # no-cython-lint
-import weakref
 
 from cuda.core._utils.cuda_utils import driver
 
@@ -50,16 +49,15 @@ cdef class _MemPoolOptions:
 
 
 cdef class _MemPoolAttributes:
-    cdef:
-        object _mr_weakref
+    """Provides access to memory pool attributes."""
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError("_MemPoolAttributes cannot be instantiated directly. Please use MemoryResource APIs.")
 
-    @classmethod
-    def _init(cls, mr):
-        cdef _MemPoolAttributes self = _MemPoolAttributes.__new__(cls)
-        self._mr_weakref = mr
+    @staticmethod
+    cdef _MemPoolAttributes _init(MemoryPoolHandle h_pool):
+        cdef _MemPoolAttributes self = _MemPoolAttributes.__new__(_MemPoolAttributes)
+        self._h_pool = h_pool
         return self
 
     def __repr__(self):
@@ -69,12 +67,8 @@ cdef class _MemPoolAttributes:
         )
 
     cdef int _getattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1:
-        cdef _MemPool mr = <_MemPool>(self._mr_weakref())
-        if mr is None:
-            raise RuntimeError("_MemPool is expired")
-        cdef cydriver.CUmemoryPool pool_handle = as_cu(mr._h_pool)
         with nogil:
-            HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(pool_handle, attr_enum, value))
+            HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(as_cu(self._h_pool), attr_enum, value))
         return 0
 
     @property
@@ -202,8 +196,7 @@ cdef class _MemPool(MemoryResource):
     def attributes(self) -> _MemPoolAttributes:
         """Memory pool attributes."""
         if self._attributes is None:
-            ref = weakref.ref(self)
-            self._attributes = _MemPoolAttributes._init(ref)
+            self._attributes = _MemPoolAttributes._init(self._h_pool)
         return self._attributes
 
     @property

cuda_core/cuda/core/_module.pxd

@@ -14,13 +14,11 @@ cdef class KernelAttributes
 cdef class Kernel:
     cdef:
         KernelHandle _h_kernel
-        ObjectCode _module  # ObjectCode reference
         KernelAttributes _attributes  # lazy
         KernelOccupancy _occupancy  # lazy
-        object __weakref__  # Enable weak references
 
     @staticmethod
-    cdef Kernel _from_obj(KernelHandle h_kernel, ObjectCode mod)
+    cdef Kernel _from_obj(KernelHandle h_kernel)
 
     cdef tuple _get_arguments_info(self, bint param_info=*)
 
@@ -46,8 +44,11 @@ cdef class KernelOccupancy:
 
 cdef class KernelAttributes:
     cdef:
-        object _kernel_weakref
+        KernelHandle _h_kernel
         dict _cache
 
+    @staticmethod
+    cdef KernelAttributes _init(KernelHandle h_kernel)
+
     cdef int _get_cached_attribute(self, int device_id, cydriver.CUfunction_attribute attribute) except? -1
     cdef int _resolve_device_id(self, device_id) except? -1

cuda_core/cuda/core/_module.pyx

@@ -8,7 +8,6 @@ from libc.stddef cimport size_t
 
 import functools
 import threading
-import weakref
 from collections import namedtuple
 
 from cuda.core._device import Device
@@ -21,7 +20,7 @@ from cuda.core._resource_handles cimport (
     create_library_handle_from_file,
     create_library_handle_from_data,
     create_library_handle_ref,
-    get_kernel_from_library,
+    create_kernel_handle,
     create_kernel_handle_ref,
     get_last_error,
     as_cu,
@@ -139,15 +138,15 @@ cdef inline LibraryHandle _make_empty_library_handle():
 
 
 cdef class KernelAttributes:
-    """Provides access to kernel attributes. Uses weakref to avoid preventing Kernel GC."""
+    """Provides access to kernel attributes."""
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError("KernelAttributes cannot be instantiated directly. Please use Kernel APIs.")
 
-    @classmethod
-    def _init(cls, kernel):
-        cdef KernelAttributes self = KernelAttributes.__new__(cls)
-        self._kernel_weakref = weakref.ref(kernel)
+    @staticmethod
+    cdef KernelAttributes _init(KernelHandle h_kernel):
+        cdef KernelAttributes self = KernelAttributes.__new__(KernelAttributes)
+        self._h_kernel = h_kernel
         self._cache = {}
         _lazy_init()
         return self
@@ -158,12 +157,9 @@ cdef class KernelAttributes:
         cached = self._cache.get(cache_key, cache_key)
         if cached is not cache_key:
             return cached
-        cdef Kernel kernel = <Kernel>(self._kernel_weakref())
-        if kernel is None:
-            raise RuntimeError("Cannot access kernel attributes for expired Kernel object")
         cdef int result
         with nogil:
-            HANDLE_RETURN(cydriver.cuKernelGetAttribute(&result, attribute, as_cu(kernel._h_kernel), device_id))
+            HANDLE_RETURN(cydriver.cuKernelGetAttribute(&result, attribute, as_cu(self._h_kernel), device_id))
         self._cache[cache_key] = result
         return result
 
@@ -496,10 +492,9 @@ cdef class Kernel:
         raise RuntimeError("Kernel objects cannot be instantiated directly. Please use ObjectCode APIs.")
 
     @staticmethod
-    cdef Kernel _from_obj(KernelHandle h_kernel, ObjectCode mod):
+    cdef Kernel _from_obj(KernelHandle h_kernel):
         cdef Kernel ker = Kernel.__new__(Kernel)
         ker._h_kernel = h_kernel
-        ker._module = mod
         ker._attributes = None
         ker._occupancy = None
         return ker
@@ -508,7 +503,7 @@ cdef class Kernel:
     def attributes(self) -> KernelAttributes:
         """Get the read-only attributes of this kernel."""
         if self._attributes is None:
-            self._attributes = KernelAttributes._init(self)
+            self._attributes = KernelAttributes._init(self._h_kernel)
         return self._attributes
 
     cdef tuple _get_arguments_info(self, bint param_info=False):
@@ -607,7 +602,7 @@ cdef class Kernel:
         if not h_kernel:
             HANDLE_RETURN(get_last_error())
 
-        return Kernel._from_obj(h_kernel, mod)
+        return Kernel._from_obj(h_kernel)
 
 
 CodeTypeT = bytes | bytearray | str
@@ -812,10 +807,10 @@ cdef class ObjectCode:
         except KeyError:
             name = name.encode()
 
-        cdef KernelHandle h_kernel = get_kernel_from_library(self._h_library, <const char*>name)
+        cdef KernelHandle h_kernel = create_kernel_handle(self._h_library, <const char*>name)
         if not h_kernel:
             HANDLE_RETURN(get_last_error())
-        return Kernel._from_obj(h_kernel, self)
+        return Kernel._from_obj(h_kernel)
 
     @property
     def code(self) -> CodeTypeT:

cuda_core/cuda/core/_resource_handles.pxd

@@ -109,6 +109,6 @@ cdef LibraryHandle create_library_handle_from_data(const void* data) nogil excep
 cdef LibraryHandle create_library_handle_ref(cydriver.CUlibrary library) nogil except+
 
 # Kernel handles
-cdef KernelHandle get_kernel_from_library(LibraryHandle h_library, const char* name) nogil except+
+cdef KernelHandle create_kernel_handle(LibraryHandle h_library, const char* name) nogil except+
 cdef KernelHandle create_kernel_handle_ref(
     cydriver.CUkernel kernel, LibraryHandle h_library) nogil except+

cuda_core/cuda/core/_resource_handles.pyx

@@ -102,7 +102,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
         cydriver.CUlibrary library) nogil except+
 
     # Kernel handles
-    KernelHandle get_kernel_from_library "cuda_core::get_kernel_from_library" (
+    KernelHandle create_kernel_handle "cuda_core::create_kernel_handle" (
         LibraryHandle h_library, const char* name) nogil except+
     KernelHandle create_kernel_handle_ref "cuda_core::create_kernel_handle_ref" (
         cydriver.CUkernel kernel, LibraryHandle h_library) nogil except+

cuda_core/tests/test_memory.py

@@ -1162,7 +1162,7 @@ def test_mempool_attributes_repr(memory_resource_factory):
 
 
 def test_mempool_attributes_ownership(memory_resource_factory):
-    """Ensure the attributes bundle handles references correctly for all memory resource types."""
+    """Ensure the attributes bundle keeps the pool alive via the handle."""
     MR, MRops = memory_resource_factory
     device = Device()
 
@@ -1190,21 +1190,9 @@ def test_mempool_attributes_ownership(memory_resource_factory):
     mr.close()
     del mr
 
-    # After deleting the memory resource, the attributes suite is disconnected.
-    with pytest.raises(RuntimeError, match="is expired"):
-        _ = attributes.used_mem_high
-
-    # Even when a new object is created (we found a case where the same
-    # mempool handle was really reused).
-    if MR is DeviceMemoryResource:
-        mr = MR(device, dict(max_size=POOL_SIZE))  # noqa: F841
-    elif MR is PinnedMemoryResource:
-        mr = MR(dict(max_size=POOL_SIZE))  # noqa: F841
-    elif MR is ManagedMemoryResource:
-        mr = create_managed_memory_resource_or_skip(dict())  # noqa: F841
-
-    with pytest.raises(RuntimeError, match="is expired"):
-        _ = attributes.used_mem_high
+    # The attributes bundle keeps the pool alive via MemoryPoolHandle,
+    # so accessing attributes still works even after the MR is deleted.
+    _ = attributes.used_mem_high  # Should not raise
 
 
 # Ensure that memory views dellocate their reference to dlpack tensors

cuda_core/tests/test_module.py

@@ -229,12 +229,6 @@ def test_saxpy_arguments(get_saxpy_kernel_cubin, cuda12_4_prerequisite_check):
             _ = krn.num_arguments
         return
 
-    # Check that arguments_info returns ParamInfo objects (works for both Python and Cython classes)
-    # For Python classes: type(krn).arguments_info.fget.__annotations__ contains ParamInfo
-    # For Cython cdef classes: property descriptors don't have .fget, so we check the actual values
-    prop = type(krn).arguments_info
-    if hasattr(prop, "fget") and hasattr(prop.fget, "__annotations__"):
-        assert "ParamInfo" in str(prop.fget.__annotations__)
     arg_info = krn.arguments_info
     n_args = len(arg_info)
     assert n_args == krn.num_arguments