From d6c6c39419f6889df237dfe4b94e858117e5d4a9 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Fri, 21 Aug 2026 08:27:24 -0700 Subject: [PATCH 1/3] Isolate CUmemLocation construction in a versioned helper Build CUmemLocation via field assignment in to_cumemlocation() so cuda.core compiles against both the 13.3 and 13.4 layouts. The localized arm is an optional helper argument that exists only when CUDA_VERSION >= 13040. --- cuda_core/build_hooks.py | 30 +++++++++- .../core/_memory/_device_memory_resource.pyx | 6 +- cuda_core/cuda/core/_memory/_location.pxd | 58 ++++++++++++++----- .../cuda/core/_memory/_managed_memory_ops.pyx | 4 +- .../cuda/core/_memory/_peer_access_utils.pyx | 6 +- cuda_core/cuda/core/graph/_graph_node.pyx | 8 +-- cuda_core/tests/test_build_hooks.py | 54 +++++++++++++++++ 7 files changed, 135 insertions(+), 31 deletions(-) diff --git a/cuda_core/build_hooks.py b/cuda_core/build_hooks.py index 626d50355ab..0334e6ed6c5 100644 --- a/cuda_core/build_hooks.py +++ b/cuda_core/build_hooks.py @@ -121,6 +121,31 @@ def _determine_cuda_major_version() -> str: ) +@functools.cache +def _cuda_core_has_localized_location() -> bool: + """Whether CUmemLocation exposes the ``localized`` union arm (CUDA 13.4+). + + CUDA 13.4 adds ``CUmemLocation.localized`` and + ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN``. This flag is independent of + ``CUDA_CORE_BUILD_MAJOR`` so 13.3 and 13.4 can share a major version while + still compiling different ``to_cumemlocation`` signatures. + """ + override = os.environ.get("CUDA_CORE_HAS_LOCALIZED_LOCATION") + if override is not None: + return bool(int(override)) + cuda_path = _get_cuda_path() + cuda_h = os.path.join(cuda_path, "include", "cuda.h") + try: + with open(cuda_h, encoding="utf-8") as f: + for line in f: + m = re.match(r"^#\s*define\s+CUDA_VERSION\s+(\d+)\s*$", line) + if m: + return int(m.group(1)) >= 13040 + except OSError: + pass + return False + + # used later by setup() _extensions = None @@ -220,7 +245,10 @@ def get_sources(mod_name): ) nthreads = int(os.environ.get("CUDA_PYTHON_PARALLEL_LEVEL", os.cpu_count() // 2)) - compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_determine_cuda_major_version())} + compile_time_env = { + "CUDA_CORE_BUILD_MAJOR": int(_determine_cuda_major_version()), + "CUDA_CORE_HAS_LOCALIZED_LOCATION": int(_cuda_core_has_localized_location()), + } compiler_directives = {"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True} _CythonOptions.warning_errors = True if COMPILE_FOR_COVERAGE: diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index 1ee492edd77..36203a317f6 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 to_cumemlocation from cuda.core._memory._memory_pool cimport ( _MemPool, MP_init_create_pool, MP_raise_release_threshold, ) @@ -321,10 +322,7 @@ 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 = to_cumemlocation("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..b46d366f38f 100644 --- a/cuda_core/cuda/core/_memory/_location.pxd +++ b/cuda_core/cuda/core/_memory/_location.pxd @@ -9,32 +9,62 @@ # 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. +# +# Construction uses field assignment rather than Cython struct literals so +# the same source compiles against both the CUDA 13.3 two-member declaration +# and the CUDA 13.4 declaration that adds the ``localized`` union arm. +# ``CUDA_CORE_HAS_LOCALIZED_LOCATION`` selects a helper signature with an +# optional ``localized`` argument (13.4+) or without it (13.3). Passing +# ``localized=...`` is therefore a Cython compile-time error on 13.3. from cuda.bindings cimport cydriver IF CUDA_CORE_BUILD_MAJOR >= 13: - cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id): + cdef inline void _fill_id_location( + cydriver.CUmemLocation* cu_loc, str kind, int loc_id + ) except *: 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}") + + IF CUDA_CORE_HAS_LOCALIZED_LOCATION: + cdef inline cydriver.CUmemLocation to_cumemlocation( + str kind, int loc_id=0, tuple localized=None + ): + cdef cydriver.CUmemLocation cu_loc + if kind == "device_locality_domain": + if localized is None: + raise ValueError( + "kind='device_locality_domain' requires " + "localized=(device_id, locality_domain_id)" + ) + cu_loc.type = ( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN + ) + cu_loc.localized.deviceId = localized[0] + cu_loc.localized.localityDomainId = localized[1] + return cu_loc + _fill_id_location(&cu_loc, kind, loc_id) + return cu_loc + ELSE: + cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id=0): + cdef cydriver.CUmemLocation cu_loc + _fill_id_location(&cu_loc, kind, loc_id) + return cu_loc ELSE: - cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id): + cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id=0): raise NotImplementedError( "CUmemLocation requires cuda.core built against CUDA 13 headers" ) diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index dcda07aab06..fd793774bd5 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -193,9 +193,7 @@ 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 = to_cumemlocation("host", 0) else: cu_loc = to_cumemlocation(loc.kind, loc.id) with nogil: diff --git a/cuda_core/cuda/core/_memory/_peer_access_utils.pyx b/cuda_core/cuda/core/_memory/_peer_access_utils.pyx index 69d59f9e005..c960a218da1 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 to_cumemlocation 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,7 @@ 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 = to_cumemlocation("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..62ef354d191 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 to_cumemlocation from cuda.core._module cimport Kernel from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition from cuda.core.graph._subclasses cimport ( @@ -835,11 +836,8 @@ 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 + to_cumemlocation("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) diff --git a/cuda_core/tests/test_build_hooks.py b/cuda_core/tests/test_build_hooks.py index c08ad4cd3c5..37afb7b53bb 100644 --- a/cuda_core/tests/test_build_hooks.py +++ b/cuda_core/tests/test_build_hooks.py @@ -100,6 +100,7 @@ def _check_version_detection( build_hooks._get_cuda_path.cache_clear() build_hooks._determine_cuda_major_version.cache_clear() + build_hooks._cuda_core_has_localized_location.cache_clear() get_cuda_path_or_home.cache_clear() mock_env = { @@ -125,6 +126,7 @@ def test_env_var_override(self, version): """CUDA_CORE_BUILD_MAJOR env var override works with various versions.""" build_hooks._get_cuda_path.cache_clear() build_hooks._determine_cuda_major_version.cache_clear() + build_hooks._cuda_core_has_localized_location.cache_clear() get_cuda_path_or_home.cache_clear() with mock.patch.dict(os.environ, {"CUDA_CORE_BUILD_MAJOR": version}, clear=False): result = build_hooks._determine_cuda_major_version() @@ -159,9 +161,61 @@ def test_missing_cuda_path_raises_error(self): """RuntimeError is raised when CUDA_PATH/CUDA_HOME not set and no env var override.""" build_hooks._get_cuda_path.cache_clear() build_hooks._determine_cuda_major_version.cache_clear() + build_hooks._cuda_core_has_localized_location.cache_clear() get_cuda_path_or_home.cache_clear() with ( mock.patch.dict(os.environ, {}, clear=True), pytest.raises(RuntimeError, match="CUDA_PATH or CUDA_HOME"), ): build_hooks._determine_cuda_major_version() + + +def _check_localized_location_detection(cuda_version, expected, *, env_override=None): + """Test localized-arm detection with a mock cuda.h.""" + with tempfile.TemporaryDirectory() as tmpdir: + include_dir = Path(tmpdir) / "include" + include_dir.mkdir() + (include_dir / "cuda.h").write_text(f"#define CUDA_VERSION {cuda_version}\n") + + build_hooks._get_cuda_path.cache_clear() + build_hooks._cuda_core_has_localized_location.cache_clear() + get_cuda_path_or_home.cache_clear() + + mock_env = {"CUDA_PATH": tmpdir} + if env_override is not None: + mock_env["CUDA_CORE_HAS_LOCALIZED_LOCATION"] = env_override + + with mock.patch.dict(os.environ, mock_env, clear=True): + assert build_hooks._cuda_core_has_localized_location() is expected + + +class TestHasLocalizedLocation: + """Tests for _cuda_core_has_localized_location().""" + + @pytest.mark.agent_authored(model="grok-4.6") + @pytest.mark.parametrize( + ("cuda_version", "expected"), + [ + (12080, False), + (13000, False), + (13030, False), + (13040, True), + (14000, True), + ], + ids=["12.8", "13.0", "13.3", "13.4", "14.0"], + ) + def test_cuda_headers_parsing(self, cuda_version, expected): + """CUDA_VERSION 13040+ enables the localized CUmemLocation arm.""" + _check_localized_location_detection(cuda_version, expected) + + @pytest.mark.agent_authored(model="grok-4.6") + @pytest.mark.parametrize( + ("override", "expected"), + [ + ("0", False), + ("1", True), + ], + ) + def test_env_var_override(self, override, expected): + """CUDA_CORE_HAS_LOCALIZED_LOCATION overrides the header-derived value.""" + _check_localized_location_detection(13030, expected, env_override=override) From f9c452992bad3d92dcf9dd7e9ea509b7774841e8 Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Fri, 21 Aug 2026 11:50:46 -0700 Subject: [PATCH 2/3] address review feedback --- cuda_core/build_hooks.py | 30 +------ .../core/_memory/_device_memory_resource.pyx | 5 +- cuda_core/cuda/core/_memory/_location.pxd | 85 +++++++++---------- .../cuda/core/_memory/_managed_memory_ops.pyx | 9 +- cuda_core/cuda/core/_memory/_memory_pool.pyx | 6 +- .../cuda/core/_memory/_peer_access_utils.pyx | 5 +- cuda_core/cuda/core/graph/_graph_node.pyx | 5 +- cuda_core/tests/test_build_hooks.py | 54 ------------ 8 files changed, 60 insertions(+), 139 deletions(-) diff --git a/cuda_core/build_hooks.py b/cuda_core/build_hooks.py index 0334e6ed6c5..626d50355ab 100644 --- a/cuda_core/build_hooks.py +++ b/cuda_core/build_hooks.py @@ -121,31 +121,6 @@ def _determine_cuda_major_version() -> str: ) -@functools.cache -def _cuda_core_has_localized_location() -> bool: - """Whether CUmemLocation exposes the ``localized`` union arm (CUDA 13.4+). - - CUDA 13.4 adds ``CUmemLocation.localized`` and - ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN``. This flag is independent of - ``CUDA_CORE_BUILD_MAJOR`` so 13.3 and 13.4 can share a major version while - still compiling different ``to_cumemlocation`` signatures. - """ - override = os.environ.get("CUDA_CORE_HAS_LOCALIZED_LOCATION") - if override is not None: - return bool(int(override)) - cuda_path = _get_cuda_path() - cuda_h = os.path.join(cuda_path, "include", "cuda.h") - try: - with open(cuda_h, encoding="utf-8") as f: - for line in f: - m = re.match(r"^#\s*define\s+CUDA_VERSION\s+(\d+)\s*$", line) - if m: - return int(m.group(1)) >= 13040 - except OSError: - pass - return False - - # used later by setup() _extensions = None @@ -245,10 +220,7 @@ def get_sources(mod_name): ) nthreads = int(os.environ.get("CUDA_PYTHON_PARALLEL_LEVEL", os.cpu_count() // 2)) - compile_time_env = { - "CUDA_CORE_BUILD_MAJOR": int(_determine_cuda_major_version()), - "CUDA_CORE_HAS_LOCALIZED_LOCATION": int(_cuda_core_has_localized_location()), - } + compile_time_env = {"CUDA_CORE_BUILD_MAJOR": int(_determine_cuda_major_version())} compiler_directives = {"embedsignature": True, "warn.deprecated.IF": False, "freethreading_compatible": True} _CythonOptions.warning_errors = True if COMPILE_FOR_COVERAGE: diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index 36203a317f6..7f909f71d90 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -5,7 +5,7 @@ from __future__ import annotations from cuda.bindings cimport cydriver -from cuda.core._memory._location cimport to_cumemlocation +from cuda.core._memory._location cimport cumemlocation_from_type from cuda.core._memory._memory_pool cimport ( _MemPool, MP_init_create_pool, MP_raise_release_threshold, ) @@ -322,7 +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 = to_cumemlocation("device", dev_id) + cdef cydriver.CUmemLocation location = cumemlocation_from_type( + 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 b46d366f38f..0982b472c75 100644 --- a/cuda_core/cuda/core/_memory/_location.pxd +++ b/cuda_core/cuda/core/_memory/_location.pxd @@ -2,69 +2,64 @@ # # 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. # -# Construction uses field assignment rather than Cython struct literals so -# the same source compiles against both the CUDA 13.3 two-member declaration -# and the CUDA 13.4 declaration that adds the ``localized`` union arm. -# ``CUDA_CORE_HAS_LOCALIZED_LOCATION`` selects a helper signature with an -# optional ``localized`` argument (13.4+) or without it (13.3). Passing -# ``localized=...`` is therefore a Cython compile-time error on 13.3. +# 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 IF CUDA_CORE_BUILD_MAJOR >= 13: - cdef inline void _fill_id_location( - cydriver.CUmemLocation* cu_loc, str kind, int loc_id - ) except *: + cdef inline cydriver.CUmemLocation cumemlocation_from_type( + cydriver.CUmemLocationType loc_type, int loc_id + ): + """Build a ``CUmemLocation`` from an already-known ``CUmemLocationType``. + + 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 + + cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id): if kind == "device": - cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - cu_loc.id = loc_id + return cumemlocation_from_type( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, loc_id) elif kind == "host": - cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST - cu_loc.id = 0 + return cumemlocation_from_type( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) elif kind == "host_numa": - cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA - cu_loc.id = loc_id + return cumemlocation_from_type( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, loc_id) elif kind == "host_numa_current": - cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT - cu_loc.id = 0 + return cumemlocation_from_type( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, 0) else: raise ValueError(f"unknown location kind: {kind!r}") - - IF CUDA_CORE_HAS_LOCALIZED_LOCATION: - cdef inline cydriver.CUmemLocation to_cumemlocation( - str kind, int loc_id=0, tuple localized=None - ): - cdef cydriver.CUmemLocation cu_loc - if kind == "device_locality_domain": - if localized is None: - raise ValueError( - "kind='device_locality_domain' requires " - "localized=(device_id, locality_domain_id)" - ) - cu_loc.type = ( - cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN - ) - cu_loc.localized.deviceId = localized[0] - cu_loc.localized.localityDomainId = localized[1] - return cu_loc - _fill_id_location(&cu_loc, kind, loc_id) - return cu_loc - ELSE: - cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id=0): - cdef cydriver.CUmemLocation cu_loc - _fill_id_location(&cu_loc, kind, loc_id) - return cu_loc ELSE: - cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id=0): + cdef inline cydriver.CUmemLocation cumemlocation_from_type( + cydriver.CUmemLocationType loc_type, int loc_id + ): + cdef cydriver.CUmemLocation cu_loc + cu_loc.type = loc_type + cu_loc.id = loc_id + return cu_loc + + cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id): raise NotImplementedError( "CUmemLocation requires cuda.core built against CUDA 13 headers" ) diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index fd793774bd5..c48fc8ebf6d 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_type 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_type # 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,7 +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 = to_cumemlocation("host", 0) + cu_loc = cumemlocation_from_type( + 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..c7173c2b8f3 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_type 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. +from cuda.core._memory._location cimport cumemlocation_from_type # 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_type(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 c960a218da1..1a7b1c6ae6a 100644 --- a/cuda_core/cuda/core/_memory/_peer_access_utils.pyx +++ b/cuda_core/cuda/core/_memory/_peer_access_utils.pyx @@ -10,7 +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 to_cumemlocation +from cuda.core._memory._location cimport cumemlocation_from_type 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 @@ -114,7 +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 = to_cumemlocation("device", dev_id) + cdef cydriver.CUmemLocation location = cumemlocation_from_type( + 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 62ef354d191..36a0215da15 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -21,7 +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 to_cumemlocation +from cuda.core._memory._location cimport cumemlocation_from_type from cuda.core._module cimport Kernel from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition from cuda.core.graph._subclasses cimport ( @@ -836,7 +836,8 @@ 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( - to_cumemlocation("device", peer_id), + cumemlocation_from_type( + cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, peer_id), cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, )) diff --git a/cuda_core/tests/test_build_hooks.py b/cuda_core/tests/test_build_hooks.py index 37afb7b53bb..c08ad4cd3c5 100644 --- a/cuda_core/tests/test_build_hooks.py +++ b/cuda_core/tests/test_build_hooks.py @@ -100,7 +100,6 @@ def _check_version_detection( build_hooks._get_cuda_path.cache_clear() build_hooks._determine_cuda_major_version.cache_clear() - build_hooks._cuda_core_has_localized_location.cache_clear() get_cuda_path_or_home.cache_clear() mock_env = { @@ -126,7 +125,6 @@ def test_env_var_override(self, version): """CUDA_CORE_BUILD_MAJOR env var override works with various versions.""" build_hooks._get_cuda_path.cache_clear() build_hooks._determine_cuda_major_version.cache_clear() - build_hooks._cuda_core_has_localized_location.cache_clear() get_cuda_path_or_home.cache_clear() with mock.patch.dict(os.environ, {"CUDA_CORE_BUILD_MAJOR": version}, clear=False): result = build_hooks._determine_cuda_major_version() @@ -161,61 +159,9 @@ def test_missing_cuda_path_raises_error(self): """RuntimeError is raised when CUDA_PATH/CUDA_HOME not set and no env var override.""" build_hooks._get_cuda_path.cache_clear() build_hooks._determine_cuda_major_version.cache_clear() - build_hooks._cuda_core_has_localized_location.cache_clear() get_cuda_path_or_home.cache_clear() with ( mock.patch.dict(os.environ, {}, clear=True), pytest.raises(RuntimeError, match="CUDA_PATH or CUDA_HOME"), ): build_hooks._determine_cuda_major_version() - - -def _check_localized_location_detection(cuda_version, expected, *, env_override=None): - """Test localized-arm detection with a mock cuda.h.""" - with tempfile.TemporaryDirectory() as tmpdir: - include_dir = Path(tmpdir) / "include" - include_dir.mkdir() - (include_dir / "cuda.h").write_text(f"#define CUDA_VERSION {cuda_version}\n") - - build_hooks._get_cuda_path.cache_clear() - build_hooks._cuda_core_has_localized_location.cache_clear() - get_cuda_path_or_home.cache_clear() - - mock_env = {"CUDA_PATH": tmpdir} - if env_override is not None: - mock_env["CUDA_CORE_HAS_LOCALIZED_LOCATION"] = env_override - - with mock.patch.dict(os.environ, mock_env, clear=True): - assert build_hooks._cuda_core_has_localized_location() is expected - - -class TestHasLocalizedLocation: - """Tests for _cuda_core_has_localized_location().""" - - @pytest.mark.agent_authored(model="grok-4.6") - @pytest.mark.parametrize( - ("cuda_version", "expected"), - [ - (12080, False), - (13000, False), - (13030, False), - (13040, True), - (14000, True), - ], - ids=["12.8", "13.0", "13.3", "13.4", "14.0"], - ) - def test_cuda_headers_parsing(self, cuda_version, expected): - """CUDA_VERSION 13040+ enables the localized CUmemLocation arm.""" - _check_localized_location_detection(cuda_version, expected) - - @pytest.mark.agent_authored(model="grok-4.6") - @pytest.mark.parametrize( - ("override", "expected"), - [ - ("0", False), - ("1", True), - ], - ) - def test_env_var_override(self, override, expected): - """CUDA_CORE_HAS_LOCALIZED_LOCATION overrides the header-derived value.""" - _check_localized_location_detection(13030, expected, env_override=override) From 6589041d1e717d16fb2fbc0f37248d1a8d6b8adf Mon Sep 17 00:00:00 2001 From: Ralf Juengling Date: Fri, 21 Aug 2026 17:56:21 -0700 Subject: [PATCH 3/3] address review feedback --- .../core/_memory/_device_memory_resource.pyx | 4 +- cuda_core/cuda/core/_memory/_location.pxd | 53 +++++++++---------- .../cuda/core/_memory/_managed_memory_ops.pyx | 6 +-- cuda_core/cuda/core/_memory/_memory_pool.pyx | 6 +-- .../cuda/core/_memory/_peer_access_utils.pyx | 4 +- cuda_core/cuda/core/graph/_graph_node.pyx | 4 +- 6 files changed, 37 insertions(+), 40 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index 7f909f71d90..d91ae88c949 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -5,7 +5,7 @@ from __future__ import annotations from cuda.bindings cimport cydriver -from cuda.core._memory._location cimport cumemlocation_from_type +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, ) @@ -322,7 +322,7 @@ 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 = cumemlocation_from_type( + cdef cydriver.CUmemLocation location = cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, dev_id) with nogil: diff --git a/cuda_core/cuda/core/_memory/_location.pxd b/cuda_core/cuda/core/_memory/_location.pxd index 0982b472c75..a9cc525827b 100644 --- a/cuda_core/cuda/core/_memory/_location.pxd +++ b/cuda_core/cuda/core/_memory/_location.pxd @@ -5,9 +5,7 @@ # 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 @@ -20,45 +18,44 @@ from cuda.bindings cimport cydriver -IF CUDA_CORE_BUILD_MAJOR >= 13: - cdef inline cydriver.CUmemLocation cumemlocation_from_type( - cydriver.CUmemLocationType loc_type, int loc_id - ): - """Build a ``CUmemLocation`` from an already-known ``CUmemLocationType``. +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 - 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 cumemlocation_from_type( + return cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, loc_id) elif kind == "host": - return cumemlocation_from_type( + return cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) elif kind == "host_numa": - return cumemlocation_from_type( + return cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA, loc_id) elif kind == "host_numa_current": - return cumemlocation_from_type( + return cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT, 0) else: raise ValueError(f"unknown location kind: {kind!r}") ELSE: - cdef inline cydriver.CUmemLocation cumemlocation_from_type( - cydriver.CUmemLocationType loc_type, int loc_id - ): - cdef cydriver.CUmemLocation cu_loc - cu_loc.type = loc_type - cu_loc.id = loc_id - return cu_loc - cdef inline cydriver.CUmemLocation to_cumemlocation(str kind, int loc_id): raise NotImplementedError( "CUmemLocation requires cuda.core built against CUDA 13 headers" diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index c48fc8ebf6d..2fa1ab81487 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -13,10 +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 / cumemlocation_from_type are referenced only from CUDA 13 +# 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_type # no-cython-lint +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 @@ -195,7 +195,7 @@ 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 = cumemlocation_from_type( + cu_loc = cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST, 0) else: cu_loc = to_cumemlocation(loc.kind, loc.id) diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index c7173c2b8f3..67630bec7c3 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -12,9 +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_type is referenced only from CUDA 13 branches. cython-lint +# 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_type # no-cython-lint +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, @@ -282,7 +282,7 @@ cdef int MP_init_current_pool( """ IF CUDA_CORE_BUILD_MAJOR >= 13: cdef cydriver.CUmemoryPool pool - cdef cydriver.CUmemLocation loc = cumemlocation_from_type(loc_type, 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 1a7b1c6ae6a..048e7f180db 100644 --- a/cuda_core/cuda/core/_memory/_peer_access_utils.pyx +++ b/cuda_core/cuda/core/_memory/_peer_access_utils.pyx @@ -10,7 +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_type +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 @@ -114,7 +114,7 @@ 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 = cumemlocation_from_type( + 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)) diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 36a0215da15..4ecce8bb1f9 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -21,7 +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_type +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 ( @@ -836,7 +836,7 @@ 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( - cumemlocation_from_type( + cumemlocation_from_id( cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, peer_id), cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, ))