Skip to content

Commit 0650f4b

Browse files
rparolinclaude
andcommitted
Fix managed memory incorrectly classified as kDLCUDAHost in DLPack device mapping
_smv_get_dl_device() treated all buffers that are both device- and host-accessible as kDLCUDAHost. Managed (unified) memory is also both- accessible, so it was misclassified. CCCL's make_tma_descriptor then rejected the descriptor with "Device type must be kDLCUDA or kDLCUDAManaged". Preserve the is_managed flag already queried via CU_POINTER_ATTRIBUTE_IS_MANAGED in _query_memory_attrs(), expose it on Buffer, and use it in _smv_get_dl_device() to return kDLCUDAManaged for managed memory. Fixes: https://nvbugspro.nvidia.com/bug/6044342 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent c2f79a1 commit 0650f4b

File tree

3 files changed

+14
-2
lines changed

3 files changed

+14
-2
lines changed

cuda_core/cuda/core/_memory/_buffer.pxd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ cdef struct _MemAttrs:
1212
int device_id
1313
bint is_device_accessible
1414
bint is_host_accessible
15+
bint is_managed
1516

1617

1718
cdef class Buffer:

cuda_core/cuda/core/_memory/_buffer.pyx

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,12 @@ cdef class Buffer:
396396
_init_mem_attrs(self)
397397
return self._mem_attrs.is_host_accessible
398398

399+
@property
400+
def is_managed(self) -> bool:
401+
"""Return True if this buffer is CUDA managed (unified) memory, otherwise False."""
402+
_init_mem_attrs(self)
403+
return self._mem_attrs.is_managed
404+
399405
@property
400406
def is_mapped(self) -> bool:
401407
"""Return True if this buffer is mapped into the process via IPC."""
@@ -459,6 +465,7 @@ cdef inline int _query_memory_attrs(
459465
out.is_host_accessible = True
460466
out.is_device_accessible = False
461467
out.device_id = -1
468+
out.is_managed = False
462469
elif (
463470
is_managed
464471
or memory_type == cydriver.CUmemorytype.CU_MEMORYTYPE_HOST
@@ -467,10 +474,12 @@ cdef inline int _query_memory_attrs(
467474
out.is_host_accessible = True
468475
out.is_device_accessible = True
469476
out.device_id = device_id
477+
out.is_managed = is_managed
470478
elif memory_type == cydriver.CUmemorytype.CU_MEMORYTYPE_DEVICE:
471479
out.is_host_accessible = False
472480
out.is_device_accessible = True
473481
out.device_id = device_id
482+
out.is_managed = False
474483
else:
475484
with cython.gil:
476485
raise ValueError(f"Unsupported memory type: {memory_type}")

cuda_core/cuda/core/_memoryview.pyx

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -607,8 +607,10 @@ cdef inline int _smv_get_dl_device(
607607
device_type = _kDLCUDA
608608
device_id = buf.device_id
609609
elif d and h:
610-
# We do not currently differentiate pinned vs managed here.
611-
device_type = _kDLCUDAHost
610+
if buf.is_managed:
611+
device_type = _kDLCUDAManaged
612+
else:
613+
device_type = _kDLCUDAHost
612614
device_id = 0
613615
elif (not d) and h:
614616
device_type = _kDLCPU

0 commit comments

Comments
 (0)