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.

Author: juenglin

cuda_core/build_hooks.py

@@ -75,6 +75,23 @@ def _get_cuda_path() -> str:
     return cuda_path
 
 
+@functools.cache
+def _read_cuda_version_int() -> int | None:
+    """Return the numeric CUDA_VERSION macro from cuda.h, or None if unavailable."""
+    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:
+                    # CUDA_VERSION is e.g. 12020 for 12.2, 13040 for 13.4.
+                    return int(m.group(1))
+    except OSError:
+        pass
+    return None
+
+
 @functools.cache
 def _determine_cuda_major_version() -> str:
     """Determine the CUDA major version for building cuda.core.
@@ -97,20 +114,11 @@ def _determine_cuda_major_version() -> str:
         return cuda_major
 
     # Derive from the CUDA headers (the authoritative source for what we compile against).
-    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:
-                    v = int(m.group(1))
-                    # CUDA_VERSION is e.g. 12020 for 12.2.
-                    cuda_major = str(v // 1000)
-                    print("CUDA MAJOR VERSION:", cuda_major)
-                    return cuda_major
-    except OSError:
-        pass
+    version = _read_cuda_version_int()
+    if version is not None:
+        cuda_major = str(version // 1000)
+        print("CUDA MAJOR VERSION:", cuda_major)
+        return cuda_major
 
     # CUDA_PATH or CUDA_HOME is required for the build, so we should not reach here
     # in normal circumstances. Raise an error to make the issue clear.
@@ -121,6 +129,22 @@ 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))
+    version = _read_cuda_version_int()
+    return version is not None and version >= 13040
+
+
 # used later by setup()
 _extensions = None
 
@@ -220,7 +244,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:

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))

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"
         )

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:

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

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)

cuda_core/tests/test_build_hooks.py

@@ -99,7 +99,9 @@ def _check_version_detection(
         cuda_h.write_text(f"#define CUDA_VERSION {cuda_version}\n")
 
         build_hooks._get_cuda_path.cache_clear()
+        build_hooks._read_cuda_version_int.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 = {
@@ -124,7 +126,9 @@ class TestGetCudaMajorVersion:
     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._read_cuda_version_int.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()
@@ -158,10 +162,64 @@ def test_env_var_takes_priority_over_headers(self):
     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._read_cuda_version_int.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._read_cuda_version_int.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)