Skip to content

Commit 08b4cf4

Browse files
committed
cuda.core: expose last prefetch location on ManagedBuffer
1 parent fb57dbf commit 08b4cf4

7 files changed

Lines changed: 102 additions & 48 deletions

File tree

‎cuda_core/cuda/core/_host.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ class Host:
2323
``Host`` is the symmetric counterpart of :class:`~cuda.core.Device`
2424
for managed-memory `prefetch`, `advise`, and `discard_prefetch`
2525
targets. Pass either a ``Device`` or a ``Host`` to those operations
26-
and to ``ManagedBuffer.preferred_location`` / ``accessed_by``.
26+
and to ``ManagedBuffer.preferred_location`` / ``accessed_by``. A
27+
``Host`` may also be returned by ``ManagedBuffer.last_prefetch_location``.
2728
2829
``Host`` is a singleton class, mirroring :class:`~cuda.core.Device`:
2930
constructor calls with the same arguments return the same instance,

‎cuda_core/cuda/core/_memory/_managed_buffer.py‎

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
_do_single_discard_prefetch_py,
1616
_do_single_discard_py,
1717
_do_single_prefetch_py,
18+
_read_last_prefetch_location_v2,
1819
_read_preferred_location_v2,
1920
)
2021
from cuda.core._utils.cuda_utils import driver, handle_return
@@ -43,6 +44,7 @@
4344
_ATTR_READ_MOSTLY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_READ_MOSTLY
4445
_ATTR_PREFERRED = _RANGE.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION
4546
_ATTR_ACCESSED_BY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY
47+
_ATTR_LAST_PREFETCH = _RANGE.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION
4648

4749

4850
def _check_open(buf: Buffer) -> None:
@@ -147,11 +149,12 @@ class ManagedBuffer(Buffer):
147149
148150
Note
149151
----
150-
On CUDA 13 builds, ``preferred_location`` round-trips full NUMA
151-
information. On CUDA 12 builds, ``Host(numa_id=...)`` and
152-
``Host.numa_current()`` are rejected with ``TypeError`` at the call
153-
boundary — only ``Device(...)`` and the generic ``Host()`` are
154-
accepted. Use ``Host()`` to target the host on CUDA 12.
152+
On CUDA 13 builds, ``preferred_location`` and
153+
``last_prefetch_location`` round-trip full NUMA information. On CUDA 12
154+
builds, ``Host(numa_id=...)`` and ``Host.numa_current()`` are rejected
155+
with ``TypeError`` at the call boundary — only ``Device(...)`` and the
156+
generic ``Host()`` are accepted. Use ``Host()`` to target the host on
157+
CUDA 12.
155158
"""
156159

157160
@classmethod
@@ -168,7 +171,8 @@ def from_handle(
168171
169172
Use this when you have an externally-allocated managed pointer
170173
and want the property-style advice API (:attr:`read_mostly`,
171-
:attr:`preferred_location`, :attr:`accessed_by`).
174+
:attr:`preferred_location`, :attr:`last_prefetch_location`,
175+
:attr:`accessed_by`).
172176
173177
Parameters
174178
----------
@@ -232,6 +236,28 @@ def preferred_location(self, value: Device | Host | None) -> None:
232236
else:
233237
_advise_one(self, _SET_PREFERRED, value)
234238

239+
@property
240+
def last_prefetch_location(self) -> Device | Host | None:
241+
"""Location targeted by the most recent explicit prefetch.
242+
243+
Returns ``None`` if any page in the range has never been prefetched,
244+
or if the pages do not share one last-prefetch location. This reports
245+
the application's requested destination, not current residency or
246+
completion of the asynchronous prefetch operation.
247+
248+
On CUDA 13 builds, fully round-trips ``Host(numa_id=N)``. On CUDA 12,
249+
the legacy attribute carries only a device ordinal (or ``-1`` for
250+
host), so host NUMA details are unavailable.
251+
"""
252+
if binding_version() >= (13, 0, 0) and driver_version() >= (13, 0, 0):
253+
return _read_last_prefetch_location_v2(self)
254+
loc_id = _get_int_attr(self, _ATTR_LAST_PREFETCH)
255+
if loc_id == -2:
256+
return None
257+
if loc_id == -1:
258+
return Host()
259+
return Device(loc_id)
260+
235261
@property
236262
def accessed_by(self) -> AccessedBySetProxy:
237263
"""Live set-like view of ``set_accessed_by`` locations."""

‎cuda_core/cuda/core/_memory/_managed_memory_ops.pyi‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ def _do_single_prefetch_py(buf: Buffer, location: Device | Host | None, stream:
6969
Uses cuMemPrefetchAsync (works on CUDA 12 and 13).
7070
"""
7171
def _read_preferred_location_v2(buf: Buffer) -> Device | Host | None: ...
72+
def _read_last_prefetch_location_v2(buf: Buffer) -> Device | Host | None: ...
7273
def discard_prefetch_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer], locations: Device | Host | Sequence[Device | Host]) -> None:
7374
"""Discard a batch of managed-memory ranges and prefetch them to target locations.
7475

‎cuda_core/cuda/core/_memory/_managed_memory_ops.pyx‎

Lines changed: 38 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -281,15 +281,13 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
281281
) except ?cydriver.CUDA_ERROR_NOT_FOUND nogil
282282

283283

284-
def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
285-
"""Internal: read preferred_location with full NUMA detail.
286-
287-
Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose
288-
attribute allowlist doesn't yet include the cu13 _TYPE / _ID
289-
attributes) by calling cydriver directly.
290-
291-
Returns Device | Host | None.
292-
"""
284+
cdef object _read_location_v2(
285+
Buffer buf,
286+
cydriver.CUmem_range_attribute type_attribute,
287+
cydriver.CUmem_range_attribute id_attribute,
288+
):
289+
# cuda.bindings.driver.cuMemRangeGetAttribute does not yet accept the
290+
# CUDA 13 _TYPE / _ID attributes, so query them through cydriver.
293291
Buffer_check_open(buf)
294292
cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr)
295293
cdef size_t nbytes = buf._size
@@ -298,12 +296,12 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
298296
with nogil:
299297
HANDLE_RETURN(cydriver.cuMemRangeGetAttribute(
300298
<void*>&loc_type, sizeof(int),
301-
cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE,
299+
type_attribute,
302300
cu_ptr, nbytes,
303301
))
304302
HANDLE_RETURN(cydriver.cuMemRangeGetAttribute(
305303
<void*>&loc_id, sizeof(int),
306-
cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID,
304+
id_attribute,
307305
cu_ptr, nbytes,
308306
))
309307
if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE:
@@ -315,7 +313,25 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
315313
return Host(numa_id=loc_id)
316314
if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT:
317315
return Host.numa_current()
318-
return None # CU_MEM_LOCATION_TYPE_INVALID — no preferred location
316+
return None # CU_MEM_LOCATION_TYPE_INVALID
317+
318+
319+
def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
320+
"""Internal: read preferred_location with full NUMA detail."""
321+
return _read_location_v2(
322+
buf,
323+
cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE,
324+
cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID,
325+
)
326+
327+
328+
def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None:
329+
"""Internal: read last_prefetch_location with full NUMA detail."""
330+
return _read_location_v2(
331+
buf,
332+
cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_TYPE,
333+
cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_ID,
334+
)
319335

320336

321337
cdef void _do_batch_prefetch_op(tuple bufs, tuple locs, Stream s, _BatchPrefetchFn fn):
@@ -348,16 +364,21 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
348364
))
349365
ELSE:
350366
def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
351-
# Symbol exists so _managed_buffer.py can `from ... import
352-
# _read_preferred_location_v2` unconditionally at module top.
353-
# `ManagedBuffer.preferred_location` gates on both
354-
# binding_version() and driver_version() >= (13, 0, 0) before
355-
# calling, so this path is unreachable on a cu12 build.
367+
# Symbols exist so _managed_buffer.py can import the v2 readers
368+
# unconditionally. Their properties gate on both binding_version()
369+
# and driver_version() >= (13, 0, 0), so these paths are unreachable
370+
# on a CUDA 12 build.
356371
raise NotImplementedError(
357372
"_read_preferred_location_v2 requires a CUDA 13 build of cuda.core"
358373
)
359374

360375

376+
def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None:
377+
raise NotImplementedError(
378+
"_read_last_prefetch_location_v2 requires a CUDA 13 build of cuda.core"
379+
)
380+
381+
361382
cdef void _do_batch_prefetch(tuple bufs, tuple locs, Stream s):
362383
IF CUDA_CORE_BUILD_MAJOR >= 13:
363384
_do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync)

‎cuda_core/cuda/core/_memory/_managed_memory_resource.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,9 +93,9 @@ class ManagedMemoryResource(_MemPool):
9393
ManagedBuffer
9494
A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that
9595
exposes the property-style advice API
96-
(``read_mostly``, ``preferred_location``, ``accessed_by``)
97-
and instance methods (``prefetch``, ``discard``,
98-
``discard_prefetch``).
96+
(``read_mostly``, ``preferred_location``,
97+
``last_prefetch_location``, ``accessed_by``) and instance methods
98+
(``prefetch``, ``discard``, ``discard_prefetch``).
9999
"""
100100
@property
101101
def device_id(self) -> int:

‎cuda_core/cuda/core/_memory/_managed_memory_resource.pyx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,9 +119,9 @@ cdef class ManagedMemoryResource(_MemPool):
119119
ManagedBuffer
120120
A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that
121121
exposes the property-style advice API
122-
(``read_mostly``, ``preferred_location``, ``accessed_by``)
123-
and instance methods (``prefetch``, ``discard``,
124-
``discard_prefetch``).
122+
(``read_mostly``, ``preferred_location``,
123+
``last_prefetch_location``, ``accessed_by``) and instance methods
124+
(``prefetch``, ``discard``, ``discard_prefetch``).
125125
"""
126126
MP_check_open(self)
127127
assert isinstance(stream, Stream), "Only Stream is supported for managed memory allocations"

‎cuda_core/tests/memory/test_managed_ops.py‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99

1010
from cuda.bindings import driver
1111
from cuda.core import Device, Host, ManagedBuffer
12-
from cuda.core._memory._managed_buffer import _get_int_attr
1312

1413
# Managed-memory prefetch and CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION
1514
# operate at physical-page granularity. Test buffers must each occupy a full
@@ -20,16 +19,9 @@
2019
_PAGE_SIZE = mmap.PAGESIZE
2120
_MANAGED_TEST_ALLOCATION_SIZE = _PAGE_SIZE
2221
_READ_MOSTLY_ENABLED = 1
23-
_HOST_LOCATION_ID = -1
2422
_INVALID_HOST_DEVICE_ORDINAL = 0
2523

2624

27-
# TODO(#2109): replace with ``buf.last_prefetch_location`` once
28-
# ``ManagedBuffer`` exposes mem-range attributes directly.
29-
def _last_prefetch_location(buf):
30-
return _get_int_attr(buf, driver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION)
31-
32-
3325
def _page_base(buf):
3426
# Page-aligned base of the buffer's start address; two buffers sharing a
3527
# page cannot be prefetched to different locations independently.
@@ -224,8 +216,7 @@ def test_same_location(self, location_ops_device, location_ops_mr):
224216
stream.sync()
225217

226218
for buf in bufs:
227-
last = _last_prefetch_location(buf)
228-
assert last == device.device_id
219+
assert buf.last_prefetch_location == device
229220
buf.close()
230221

231222
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):
242233
prefetch_batch(stream, bufs, [Host(), device])
243234
stream.sync()
244235

245-
last0 = _last_prefetch_location(bufs[0])
246-
last1 = _last_prefetch_location(bufs[1])
247-
assert last0 == _HOST_LOCATION_ID
248-
assert last1 == device.device_id
236+
assert bufs[0].last_prefetch_location == Host()
237+
assert bufs[1].last_prefetch_location == device
249238
for buf in bufs:
250239
buf.close()
251240

@@ -285,8 +274,7 @@ def test_same_location(self, location_ops_device, location_ops_mr):
285274
discard_prefetch_batch(stream, bufs, device)
286275
stream.sync()
287276
for buf in bufs:
288-
last = _last_prefetch_location(buf)
289-
assert last == device.device_id
277+
assert buf.last_prefetch_location == device
290278
buf.close()
291279

292280

@@ -366,6 +354,23 @@ def test_from_handle(self, init_cuda):
366354
finally:
367355
plain.close()
368356

357+
@pytest.mark.agent_authored(model="gpt-5")
358+
def test_last_prefetch_location_initially_none(self, external_managed_buffer):
359+
assert external_managed_buffer.last_prefetch_location is None
360+
361+
@pytest.mark.agent_authored(model="gpt-5")
362+
def test_last_prefetch_location_roundtrip_host_numa(self, location_ops_device, managed_buffer):
363+
from cuda.core._utils.version import binding_version, driver_version
364+
365+
if binding_version() < (13, 0, 0) or driver_version() < (13, 0, 0):
366+
pytest.skip("Host NUMA last-prefetch location requires CUDA 13")
367+
368+
stream = location_ops_device.create_stream()
369+
location = Host(numa_id=0)
370+
managed_buffer.prefetch(location, stream=stream)
371+
stream.sync()
372+
assert managed_buffer.last_prefetch_location == location
373+
369374
@pytest.mark.thread_unsafe(reason="external_managed_buffer is shared between threads")
370375
def test_read_mostly_roundtrip(self, external_managed_buffer):
371376
buf = external_managed_buffer
@@ -478,7 +483,7 @@ def test_instance_prefetch(self, location_ops_device, managed_buffer):
478483
stream = device.create_stream()
479484
buf.prefetch(device, stream=stream)
480485
stream.sync()
481-
assert _last_prefetch_location(buf) == device.device_id
486+
assert buf.last_prefetch_location == device
482487

483488
def test_instance_discard(self, location_ops_device, managed_buffer):
484489
if not hasattr(driver, "cuMemDiscardBatchAsync"):
@@ -501,7 +506,7 @@ def test_instance_discard_prefetch(self, discard_prefetch_device):
501506
stream.sync()
502507
buf.discard_prefetch(device, stream=stream)
503508
stream.sync()
504-
assert _last_prefetch_location(buf) == device.device_id
509+
assert buf.last_prefetch_location == device
505510
finally:
506511
buf.close()
507512

0 commit comments

Comments
 (0)