Add managed-memory advise, prefetch, and discard-prefetch free functions

cuda_core/cuda/core/_memory/_buffer.pxd

@@ -4,6 +4,7 @@
 
 from libc.stdint cimport uintptr_t
 
+from cuda.bindings cimport cydriver
 from cuda.core._resource_handles cimport DevicePtrHandle
 from cuda.core._stream cimport Stream
 
@@ -38,3 +39,10 @@ cdef Buffer Buffer_from_deviceptr_handle(
     MemoryResource mr,
     object ipc_descriptor = *
 )
+
+# Memory attribute query helpers (used by _managed_memory_ops)
+cdef void _init_mem_attrs(Buffer self)
+cdef int _query_memory_attrs(
+    _MemAttrs& out,
+    cydriver.CUdeviceptr ptr,
+) except -1 nogil

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -71,6 +71,7 @@ A type union of :obj:`~driver.CUdeviceptr`, `int` and `None` for hinting
 :attr:`Buffer.handle`.
 """
 
+
 cdef class Buffer:
     """Represent a handle to allocated memory.
 
@@ -455,12 +456,15 @@ cdef inline int _query_memory_attrs(
         ret = cydriver.cuPointerGetAttributes(3, attrs, <void**>vals, ptr)
     HANDLE_RETURN(ret)
 
+    # TODO: HMM/ATS-enabled sysmem should also report is_managed=True; the
+    # CU_POINTER_ATTRIBUTE_IS_MANAGED query does not capture that yet.
+    out.is_managed = is_managed != 0
+
     if memory_type == 0:
         # unregistered host pointer
         out.is_host_accessible = True
         out.is_device_accessible = False
         out.device_id = -1
-        out.is_managed = False
     elif (
         is_managed
         or memory_type == cydriver.CUmemorytype.CU_MEMORYTYPE_HOST
@@ -469,12 +473,10 @@ cdef inline int _query_memory_attrs(
         out.is_host_accessible = True
         out.is_device_accessible = True
         out.device_id = device_id
-        out.is_managed = is_managed
     elif memory_type == cydriver.CUmemorytype.CU_MEMORYTYPE_DEVICE:
         out.is_host_accessible = False
         out.is_device_accessible = True
         out.device_id = device_id
-        out.is_managed = False
     else:
         with cython.gil:
             raise ValueError(f"Unsupported memory type: {memory_type}")

cuda_core/cuda/core/_memory/_managed_location.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+_VALID_KINDS = ("device", "host", "host_numa", "host_numa_current")
+LocationKind = Literal["device", "host", "host_numa", "host_numa_current"]
+
+
+@dataclass(frozen=True)
+class Location:
+    """Typed managed-memory location.
+
+    Use the classmethod constructors (``device``, ``host``, ``host_numa``,
+    ``host_numa_current``) rather than constructing directly.
+    """
+
+    kind: LocationKind
+    id: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.kind not in _VALID_KINDS:
+            raise ValueError(f"kind must be one of {_VALID_KINDS!r}, got {self.kind!r}")
+        if self.kind == "device":
+            if not isinstance(self.id, int) or self.id < 0:
+                raise ValueError("device id must be >= 0")
+        elif self.kind == "host_numa":
+            if not isinstance(self.id, int) or self.id < 0:
+                raise ValueError("host_numa id must be >= 0")
+        elif self.kind in ("host", "host_numa_current"):
+            if self.id is not None:
+                raise ValueError(f"{self.kind} location must have id=None")
+
+    @classmethod
+    def device(cls, device_id: int) -> Location:
+        return cls(kind="device", id=device_id)
+
+    @classmethod
+    def host(cls) -> Location:
+        return cls(kind="host", id=None)
+
+    @classmethod
+    def host_numa(cls, numa_id: int) -> Location:
+        return cls(kind="host_numa", id=numa_id)
+
+    @classmethod
+    def host_numa_current(cls) -> Location:
+        return cls(kind="host_numa_current", id=None)
+
+
+def _coerce_location(value, *, allow_none: bool = False) -> Location | None:
+    """Coerce ``Location`` / ``Device`` / int / ``None`` to ``Location``.
+
+    Maps int ``-1`` to host and other non-negative ints to that device ordinal.
+    """
+    from cuda.core._device import Device  # avoid import cycle at module load
+
+    if isinstance(value, Location):
+        return value
+    if isinstance(value, Device):
+        return Location.device(value.device_id)
+    if value is None:
+        if allow_none:
+            return None
+        raise ValueError("location is required")
+    if isinstance(value, int):
+        if value == -1:
+            return Location.host()
+        if value >= 0:
+            return Location.device(value)
+        raise ValueError(f"device ordinal must be >= 0 (or -1 for host), got {value}")
+    raise TypeError(f"location must be a Location, Device, int, or None; got {type(value).__name__}")