diff --git a/cuda_core/cuda/core/_host.py b/cuda_core/cuda/core/_host.py index 30464409871..9b703039865 100644 --- a/cuda_core/cuda/core/_host.py +++ b/cuda_core/cuda/core/_host.py @@ -23,7 +23,8 @@ class Host: ``Host`` is the symmetric counterpart of :class:`~cuda.core.Device` for managed-memory `prefetch`, `advise`, and `discard_prefetch` targets. Pass either a ``Device`` or a ``Host`` to those operations - and to ``ManagedBuffer.preferred_location`` / ``accessed_by``. + and to ``ManagedBuffer.preferred_location`` / ``accessed_by``. A + ``Host`` may also be returned by ``ManagedBuffer.last_prefetch_location``. ``Host`` is a singleton class, mirroring :class:`~cuda.core.Device`: constructor calls with the same arguments return the same instance, diff --git a/cuda_core/cuda/core/_memory/_managed_buffer.py b/cuda_core/cuda/core/_memory/_managed_buffer.py index 646f83221e8..44b84ff241e 100644 --- a/cuda_core/cuda/core/_memory/_managed_buffer.py +++ b/cuda_core/cuda/core/_memory/_managed_buffer.py @@ -15,6 +15,7 @@ _do_single_discard_prefetch_py, _do_single_discard_py, _do_single_prefetch_py, + _read_last_prefetch_location_v2, _read_preferred_location_v2, ) from cuda.core._utils.cuda_utils import driver, handle_return @@ -43,6 +44,7 @@ _ATTR_READ_MOSTLY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_READ_MOSTLY _ATTR_PREFERRED = _RANGE.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION _ATTR_ACCESSED_BY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY +_ATTR_LAST_PREFETCH = _RANGE.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION def _check_open(buf: Buffer) -> None: @@ -147,11 +149,12 @@ class ManagedBuffer(Buffer): Note ---- - On CUDA 13 builds, ``preferred_location`` round-trips full NUMA - information. On CUDA 12 builds, ``Host(numa_id=...)`` and - ``Host.numa_current()`` are rejected with ``TypeError`` at the call - boundary — only ``Device(...)`` and the generic ``Host()`` are - accepted. Use ``Host()`` to target the host on CUDA 12. + On CUDA 13 builds, ``preferred_location`` and + ``last_prefetch_location`` round-trip full NUMA information. On CUDA 12 + builds, ``Host(numa_id=...)`` and ``Host.numa_current()`` are rejected + with ``TypeError`` at the call boundary — only ``Device(...)`` and the + generic ``Host()`` are accepted. Use ``Host()`` to target the host on + CUDA 12. """ @classmethod @@ -168,7 +171,8 @@ def from_handle( Use this when you have an externally-allocated managed pointer and want the property-style advice API (:attr:`read_mostly`, - :attr:`preferred_location`, :attr:`accessed_by`). + :attr:`preferred_location`, :attr:`last_prefetch_location`, + :attr:`accessed_by`). Parameters ---------- @@ -232,6 +236,28 @@ def preferred_location(self, value: Device | Host | None) -> None: else: _advise_one(self, _SET_PREFERRED, value) + @property + def last_prefetch_location(self) -> Device | Host | None: + """Location targeted by the most recent explicit prefetch. + + Returns ``None`` if any page in the range has never been prefetched, + or if the pages do not share one last-prefetch location. This reports + the application's requested destination, not current residency or + completion of the asynchronous prefetch operation. + + On CUDA 13 builds, fully round-trips ``Host(numa_id=N)``. On CUDA 12, + the legacy attribute carries only a device ordinal (or ``-1`` for + host), so host NUMA details are unavailable. + """ + if binding_version() >= (13, 0, 0) and driver_version() >= (13, 0, 0): + return _read_last_prefetch_location_v2(self) + loc_id = _get_int_attr(self, _ATTR_LAST_PREFETCH) + if loc_id == -2: + return None + if loc_id == -1: + return Host() + return Device(loc_id) + @property def accessed_by(self) -> AccessedBySetProxy: """Live set-like view of ``set_accessed_by`` locations.""" diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi index 4523a39d728..2409d729f0a 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyi @@ -69,6 +69,7 @@ def _do_single_prefetch_py(buf: Buffer, location: Device | Host | None, stream: Uses cuMemPrefetchAsync (works on CUDA 12 and 13). """ def _read_preferred_location_v2(buf: Buffer) -> Device | Host | None: ... +def _read_last_prefetch_location_v2(buf: Buffer) -> Device | Host | None: ... def discard_prefetch_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer], locations: Device | Host | Sequence[Device | Host]) -> None: """Discard a batch of managed-memory ranges and prefetch them to target locations. diff --git a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx index 649468ad020..5ae01887e2c 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx @@ -281,15 +281,13 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: ) except ?cydriver.CUDA_ERROR_NOT_FOUND nogil - def _read_preferred_location_v2(Buffer buf) -> Device | Host | None: - """Internal: read preferred_location with full NUMA detail. - - Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose - attribute allowlist doesn't yet include the cu13 _TYPE / _ID - attributes) by calling cydriver directly. - - Returns Device | Host | None. - """ + cdef object _read_location_v2( + Buffer buf, + cydriver.CUmem_range_attribute type_attribute, + cydriver.CUmem_range_attribute id_attribute, + ): + # cuda.bindings.driver.cuMemRangeGetAttribute does not yet accept the + # CUDA 13 _TYPE / _ID attributes, so query them through cydriver. Buffer_check_open(buf) cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr) cdef size_t nbytes = buf._size @@ -298,12 +296,12 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: with nogil: HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( &loc_type, sizeof(int), - cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE, + type_attribute, cu_ptr, nbytes, )) HANDLE_RETURN(cydriver.cuMemRangeGetAttribute( &loc_id, sizeof(int), - cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID, + id_attribute, cu_ptr, nbytes, )) if loc_type == cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE: @@ -315,7 +313,25 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: return Host(numa_id=loc_id) if loc_type == cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT: return Host.numa_current() - return None # CU_MEM_LOCATION_TYPE_INVALID — no preferred location + return None # CU_MEM_LOCATION_TYPE_INVALID + + + def _read_preferred_location_v2(Buffer buf) -> Device | Host | None: + """Internal: read preferred_location with full NUMA detail.""" + return _read_location_v2( + buf, + cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE, + cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID, + ) + + + def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None: + """Internal: read last_prefetch_location with full NUMA detail.""" + return _read_location_v2( + buf, + cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_TYPE, + cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_ID, + ) cdef void _do_batch_prefetch_op(tuple bufs, tuple locs, Stream s, _BatchPrefetchFn fn): @@ -348,16 +364,21 @@ IF CUDA_CORE_BUILD_MAJOR >= 13: )) ELSE: def _read_preferred_location_v2(Buffer buf) -> Device | Host | None: - # Symbol exists so _managed_buffer.py can `from ... import - # _read_preferred_location_v2` unconditionally at module top. - # `ManagedBuffer.preferred_location` gates on both - # binding_version() and driver_version() >= (13, 0, 0) before - # calling, so this path is unreachable on a cu12 build. + # Symbols exist so _managed_buffer.py can import the v2 readers + # unconditionally. Their properties gate on both binding_version() + # and driver_version() >= (13, 0, 0), so these paths are unreachable + # on a CUDA 12 build. raise NotImplementedError( "_read_preferred_location_v2 requires a CUDA 13 build of cuda.core" ) + def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None: + raise NotImplementedError( + "_read_last_prefetch_location_v2 requires a CUDA 13 build of cuda.core" + ) + + cdef void _do_batch_prefetch(tuple bufs, tuple locs, Stream s): IF CUDA_CORE_BUILD_MAJOR >= 13: _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync) diff --git a/cuda_core/cuda/core/_memory/_managed_memory_resource.pyi b/cuda_core/cuda/core/_memory/_managed_memory_resource.pyi index 44523ae0f53..67a9443189d 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_managed_memory_resource.pyi @@ -93,9 +93,9 @@ class ManagedMemoryResource(_MemPool): ManagedBuffer A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that exposes the property-style advice API - (``read_mostly``, ``preferred_location``, ``accessed_by``) - and instance methods (``prefetch``, ``discard``, - ``discard_prefetch``). + (``read_mostly``, ``preferred_location``, + ``last_prefetch_location``, ``accessed_by``) and instance methods + (``prefetch``, ``discard``, ``discard_prefetch``). """ @property def device_id(self) -> int: diff --git a/cuda_core/cuda/core/_memory/_managed_memory_resource.pyx b/cuda_core/cuda/core/_memory/_managed_memory_resource.pyx index 738a2ff4af7..8c47ff487ab 100644 --- a/cuda_core/cuda/core/_memory/_managed_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_managed_memory_resource.pyx @@ -119,9 +119,9 @@ cdef class ManagedMemoryResource(_MemPool): ManagedBuffer A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that exposes the property-style advice API - (``read_mostly``, ``preferred_location``, ``accessed_by``) - and instance methods (``prefetch``, ``discard``, - ``discard_prefetch``). + (``read_mostly``, ``preferred_location``, + ``last_prefetch_location``, ``accessed_by``) and instance methods + (``prefetch``, ``discard``, ``discard_prefetch``). """ MP_check_open(self) assert isinstance(stream, Stream), "Only Stream is supported for managed memory allocations" diff --git a/cuda_core/tests/memory/test_managed_ops.py b/cuda_core/tests/memory/test_managed_ops.py index 7b51b85dd2b..c4a2acd4189 100644 --- a/cuda_core/tests/memory/test_managed_ops.py +++ b/cuda_core/tests/memory/test_managed_ops.py @@ -9,7 +9,6 @@ from cuda.bindings import driver from cuda.core import Device, Host, ManagedBuffer -from cuda.core._memory._managed_buffer import _get_int_attr # Managed-memory prefetch and CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION # operate at physical-page granularity. Test buffers must each occupy a full @@ -20,16 +19,9 @@ _PAGE_SIZE = mmap.PAGESIZE _MANAGED_TEST_ALLOCATION_SIZE = _PAGE_SIZE _READ_MOSTLY_ENABLED = 1 -_HOST_LOCATION_ID = -1 _INVALID_HOST_DEVICE_ORDINAL = 0 -# TODO(#2109): replace with ``buf.last_prefetch_location`` once -# ``ManagedBuffer`` exposes mem-range attributes directly. -def _last_prefetch_location(buf): - return _get_int_attr(buf, driver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION) - - def _page_base(buf): # Page-aligned base of the buffer's start address; two buffers sharing a # page cannot be prefetched to different locations independently. @@ -224,8 +216,7 @@ def test_same_location(self, location_ops_device, location_ops_mr): stream.sync() for buf in bufs: - last = _last_prefetch_location(buf) - assert last == device.device_id + assert buf.last_prefetch_location == device buf.close() def test_per_buffer_location(self, location_ops_device, location_ops_mr): @@ -242,10 +233,8 @@ def test_per_buffer_location(self, location_ops_device, location_ops_mr): prefetch_batch(stream, bufs, [Host(), device]) stream.sync() - last0 = _last_prefetch_location(bufs[0]) - last1 = _last_prefetch_location(bufs[1]) - assert last0 == _HOST_LOCATION_ID - assert last1 == device.device_id + assert bufs[0].last_prefetch_location == Host() + assert bufs[1].last_prefetch_location == device for buf in bufs: buf.close() @@ -285,8 +274,7 @@ def test_same_location(self, location_ops_device, location_ops_mr): discard_prefetch_batch(stream, bufs, device) stream.sync() for buf in bufs: - last = _last_prefetch_location(buf) - assert last == device.device_id + assert buf.last_prefetch_location == device buf.close() @@ -366,6 +354,23 @@ def test_from_handle(self, init_cuda): finally: plain.close() + @pytest.mark.agent_authored(model="gpt-5") + def test_last_prefetch_location_initially_none(self, external_managed_buffer): + assert external_managed_buffer.last_prefetch_location is None + + @pytest.mark.agent_authored(model="gpt-5") + def test_last_prefetch_location_roundtrip_host_numa(self, location_ops_device, managed_buffer): + from cuda.core._utils.version import binding_version, driver_version + + if binding_version() < (13, 0, 0) or driver_version() < (13, 0, 0): + pytest.skip("Host NUMA last-prefetch location requires CUDA 13") + + stream = location_ops_device.create_stream() + location = Host(numa_id=0) + managed_buffer.prefetch(location, stream=stream) + stream.sync() + assert managed_buffer.last_prefetch_location == location + @pytest.mark.thread_unsafe(reason="external_managed_buffer is shared between threads") def test_read_mostly_roundtrip(self, external_managed_buffer): buf = external_managed_buffer @@ -478,7 +483,7 @@ def test_instance_prefetch(self, location_ops_device, managed_buffer): stream = device.create_stream() buf.prefetch(device, stream=stream) stream.sync() - assert _last_prefetch_location(buf) == device.device_id + assert buf.last_prefetch_location == device def test_instance_discard(self, location_ops_device, managed_buffer): if not hasattr(driver, "cuMemDiscardBatchAsync"): @@ -501,7 +506,7 @@ def test_instance_discard_prefetch(self, discard_prefetch_device): stream.sync() buf.discard_prefetch(device, stream=stream) stream.sync() - assert _last_prefetch_location(buf) == device.device_id + assert buf.last_prefetch_location == device finally: buf.close()