99
1010from cuda .bindings import driver
1111from 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
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-
3325def _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