Avoid aggregate CUmemLocation initialization

Author: rwgk

cuda_core/cuda/core/_memory/_device_memory_resource.pyx

@@ -321,10 +321,9 @@ cpdef str DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id):
 
     cdef int dev_id = Device(device_id).device_id
     cdef cydriver.CUmemAccess_flags flags
-    cdef cydriver.CUmemLocation location = cydriver.CUmemLocation(
-        type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
-        id=dev_id,
-    )
+    cdef cydriver.CUmemLocation location
+    location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+    location.id = dev_id
 
     with nogil:
         HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location))

cuda_core/cuda/core/_memory/_location.pxd

@@ -15,24 +15,22 @@ from cuda.bindings cimport cydriver
 
 IF CUDA_CORE_BUILD_MAJOR >= 13:
     cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id):
+        cdef cydriver.CUmemLocation cu_loc
         if kind == "device":
-            return cydriver.CUmemLocation(
-                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
-                id=loc_id)
+            cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+            cu_loc.id = loc_id
         elif kind == "host":
-            return cydriver.CUmemLocation(
-                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
-                id=0)
+            cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
+            cu_loc.id = 0
         elif kind == "host_numa":
-            return cydriver.CUmemLocation(
-                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA,
-                id=loc_id)
+            cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA
+            cu_loc.id = loc_id
         elif kind == "host_numa_current":
-            return cydriver.CUmemLocation(
-                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT,
-                id=0)
+            cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT
+            cu_loc.id = 0
         else:
             raise ValueError(f"unknown location kind: {kind!r}")
+        return cu_loc
 ELSE:
     cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id):
         raise NotImplementedError(

cuda_core/cuda/core/_memory/_managed_memory_ops.pyx

@@ -193,9 +193,8 @@ cdef void _do_single_advise(Buffer buf, object advice_value, object loc, bint al
             # Driver ignores location for read_mostly / unset_preferred_location
             # advice values but still validates the CUmemLocation; pass a
             # host placeholder.
-            cu_loc = cydriver.CUmemLocation(
-                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
-                id=0)
+            cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
+            cu_loc.id = 0
         else:
             cu_loc = to_cumemlocation(loc.kind, loc.id)
         with nogil:

cuda_core/cuda/core/_memory/_memory_pool.pyx

@@ -279,8 +279,9 @@ cdef int MP_init_current_pool(
     """
     IF CUDA_CORE_BUILD_MAJOR >= 13:
         cdef cydriver.CUmemoryPool pool
-        cdef cydriver.CUmemLocation loc = cydriver.CUmemLocation(
-            type=loc_type, id=loc_id)
+        cdef cydriver.CUmemLocation loc
+        loc.type = loc_type
+        loc.id = loc_id
         with nogil:
             HANDLE_RETURN(cydriver.cuMemGetMemPool(&pool, &loc, alloc_type))
         self._h_pool = create_mempool_handle_ref(pool)

cuda_core/cuda/core/_memory/_peer_access_utils.pyx

@@ -113,10 +113,9 @@ cdef inline tuple _query_peer_access_ids(DeviceMemoryResource mr):
 cdef inline bint _peer_access_includes(DeviceMemoryResource mr, int dev_id):
     """Return True if peer access from ``dev_id`` is currently granted."""
     cdef cydriver.CUmemAccess_flags flags
-    cdef cydriver.CUmemLocation location = cydriver.CUmemLocation(
-        type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
-        id=dev_id,
-    )
+    cdef cydriver.CUmemLocation location
+    location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+    location.id = dev_id
     with nogil:
         HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(mr._h_pool), &location))
     return flags == cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE

cuda_core/cuda/core/graph/_graph_node.pyx

@@ -827,20 +827,18 @@ cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object device,
         num_deps = 1
 
     cdef vector[cydriver.CUmemAccessDesc] access_descs
+    cdef cydriver.CUmemAccessDesc access_desc
     cdef int peer_id
     cdef list peer_ids = []
 
     if peer_access is not None:
         for peer_dev in peer_access:
             peer_id = getattr(peer_dev, 'device_id', peer_dev)
             peer_ids.append(peer_id)
-            access_descs.push_back(cydriver.CUmemAccessDesc_st(
-                cydriver.CUmemLocation_st(
-                    cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
-                    peer_id
-                ),
-                cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
-            ))
+            access_desc.location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+            access_desc.location.id = peer_id
+            access_desc.flags = cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+            access_descs.push_back(access_desc)
 
     cdef str memory_type_str = "device" if memory_type is None else str(memory_type)