diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index 1ee492edd77..d91ae88c949 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -5,6 +5,7 @@ from __future__ import annotations from cuda.bindings cimport cydriver +from cuda.core._memory._location cimport cumemlocation_from_id from cuda.core._memory._memory_pool cimport ( _MemPool, MP_init_create_pool, MP_raise_release_threshold, ) @@ -321,10 +322,8 @@ 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 = cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev_id) with nogil: HANDLE_RETURN(cydriver.cuMemPoolGetAccess(&flags, as_cu(dmr._h_pool), &location)) diff --git a/cuda_core/cuda/core/_memory/_location.pxd b/cuda_core/cuda/core/_memory/_location.pxd index e46850ca886..a9cc525827b 100644 --- a/cuda_core/cuda/core/_memory/_location.pxd +++ b/cuda_core/cuda/core/_memory/_location.pxd @@ -2,35 +2,57 @@ # # SPDX-License-Identifier: Apache-2.0 -# Conversion from the internal ``_LocSpec`` record produced by -# ``_managed_location._coerce_location`` to the driver's ``CUmemLocation``. +# Conversion helpers for the driver's ``CUmemLocation`` struct. # # Header-only so both the managed-memory ops and the batched copy path can -# cimport it without either module depending on the other. ``CUmemLocation`` -# is only populated on a CUDA 13 build; the CUDA 12 stub exists so callers -# compiled there still resolve the symbol. +# cimport it without either module depending on the other. +# +# Both helpers use field assignment rather than Cython struct literals +# (``CUmemLocation(type=..., id=...)``) so this source keeps compiling if a +# future generated ``cydriver.pxd`` adds a sibling member to the struct's +# anonymous union (e.g. CUDA 13.4's ``localized`` arm): Cython's struct-literal +# coercion warns "Not all members given for struct" whenever a call site does +# not name every declared member, and cuda_core promotes that warning to a +# build error. from cuda.bindings cimport cydriver +cdef inline cydriver.CUmemLocation cumemlocation_from_id( + cydriver.CUmemLocationType loc_type, int loc_id +): + """Build a ``CUmemLocation`` whose active payload is the ``id`` field. + + ``loc_type`` must be one of the kinds whose payload is ``id`` + (``CU_MEM_LOCATION_TYPE_DEVICE``, ``HOST``, ``HOST_NUMA``, or + ``HOST_NUMA_CURRENT``); it must not be used for + ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN``, whose payload is a + separate ``localized`` union member. + + For call sites that already carry a ``CUmemLocationType`` value (e.g. + from a pool-configuration parameter), rather than the ``kind`` string + used by :func:`to_cumemlocation`. + """ + cdef cydriver.CUmemLocation cu_loc + cu_loc.type = loc_type + cu_loc.id = loc_id + return cu_loc + + IF CUDA_CORE_BUILD_MAJOR >= 13: cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id): if kind == "device": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, - id=loc_id) + return cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, loc_id) elif kind == "host": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, - id=0) + return cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) elif kind == "host_numa": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, - id=loc_id) + return cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, loc_id) elif kind == "host_numa_current": - return cydriver.CUmemLocation( - type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, - id=0) + return cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, 0) else: raise ValueError(f"unknown location kind: {kind!r}") ELSE: diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index dcda07aab06..2fa1ab81487 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -13,8 +13,10 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: from cuda.bindings cimport cydriver from cuda.core._memory._buffer cimport Buffer, Buffer_coerce_batch -# to_cumemlocation is referenced only from CUDA 13 branches. cython-lint does -# not evaluate compile-time IF blocks, so it needs a pragma to be seen as used. +# to_cumemlocation / cumemlocation_from_id are referenced only from CUDA 13 +# branches. cython-lint does not evaluate compile-time IF blocks, so they +# need a pragma to be seen as used. +from cuda.core._memory._location cimport cumemlocation_from_id # no-cython-lint from cuda.core._memory._location cimport to_cumemlocation # no-cython-lint from cuda.core._resource_handles cimport as_cu from cuda.core._stream cimport Stream, Stream_accept @@ -193,9 +195,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 = cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) else: cu_loc = to_cumemlocation(loc.kind, loc.id) with nogil: diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index cccc95a01a2..67630bec7c3 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -12,6 +12,9 @@ from libc.string cimport memset from cuda.bindings cimport cydriver from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle, MemoryResource from cuda.core._memory cimport _ipc +# cumemlocation_from_id is referenced only from a CUDA 13 branch. cython-lint +# does not evaluate compile-time IF blocks, so it needs a pragma to be seen as used. +from cuda.core._memory._location cimport cumemlocation_from_id # no-cython-lint from cuda.core._stream cimport Stream_accept, Stream from cuda.core._resource_handles cimport ( MemoryPoolHandle, @@ -279,8 +282,7 @@ 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 = cumemlocation_from_id(loc_type, loc_id) with nogil: HANDLE_RETURN(cydriver.cuMemGetMemPool(&pool, &loc, alloc_type)) self._h_pool = create_mempool_handle_ref(pool) diff --git a/cuda_core/cuda/core/_memory/_peer_access_utils.pyx b/cuda_core/cuda/core/_memory/_peer_access_utils.pyx index 69d59f9e005..048e7f180db 100644 --- a/cuda_core/cuda/core/_memory/_peer_access_utils.pyx +++ b/cuda_core/cuda/core/_memory/_peer_access_utils.pyx @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any from cuda.bindings cimport cydriver from cuda.core._memory._device_memory_resource cimport DeviceMemoryResource +from cuda.core._memory._location cimport cumemlocation_from_id from cuda.core._resource_handles cimport as_cu from cuda.core._utils.cuda_utils cimport HANDLE_RETURN from cpython.mem cimport PyMem_Malloc, PyMem_Free @@ -113,10 +114,8 @@ 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 = cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, 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 diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 2c9c07e6b3a..4ecce8bb1f9 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -21,6 +21,7 @@ from cuda.core._event cimport Event from cuda.core._kernel_arg_handler cimport ParamHolder from cuda.core._launch_config cimport LaunchConfig from cuda.core._memory._buffer cimport Buffer +from cuda.core._memory._location cimport cumemlocation_from_id from cuda.core._module cimport Kernel from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition from cuda.core.graph._subclasses cimport ( @@ -835,11 +836,9 @@ cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object device, 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 + cumemlocation_from_id( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, peer_id), + cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, )) cdef str memory_type_str = "device" if memory_type is None else str(memory_type)