Skip to content

Commit cb3c132

Browse files
cpcloudclaude
andauthored
Fix is_managed reporting for pool-allocated managed memory (#1924)
* Fix is_managed reporting for pool-allocated managed memory Pool-allocated managed memory via cuMemAllocFromPoolAsync (from a pool created with CU_MEM_ALLOCATION_TYPE_MANAGED) does not set CU_POINTER_ATTRIBUTE_IS_MANAGED=1. _query_memory_attrs therefore classified the allocation as pinned host memory, causing classify_dl_device to return kDLCUDAHost instead of kDLCUDAManaged. CCCL's make_tma_descriptor only accepts kDLCUDA or kDLCUDAManaged, so as_tensor_map() failed with "Device type must be kDLCUDA or kDLCUDAManaged" on managed buffers. Buffer.is_device_accessible / is_host_accessible already delegate to the memory resource when one is attached. Apply the same pattern to is_managed, and expose is_managed on the MemoryResource base (defaulting to False) with ManagedMemoryResource overriding it to True. Also ignore .claude/settings.local.json in .gitignore. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test(memory): add direct is_managed coverage for ManagedMemoryResource The existing test_managed_buffer_dlpack_roundtrip_device_type uses a DummyUnifiedMemoryResource backed by cuMemAllocManaged, which sets CU_POINTER_ATTRIBUTE_IS_MANAGED and so never exercised the pool-allocated path that surfaced the bug. Add two targeted tests: - test_managed_memory_resource_buffer_dlpack_device_type: allocates from ManagedMemoryResource (cuMemAllocFromPoolAsync on a managed pool) and asserts is_managed and kDLCUDAManaged through Buffer and view. - test_non_managed_resources_report_not_managed: parametrized smoke test ensuring DeviceMemoryResource and PinnedMemoryResource still report is_managed=False so the new MemoryResource.is_managed default does not silently misclassify non-managed resources. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix Buffer.is_managed to merge driver and resource signals Previous fix unconditionally delegated Buffer.is_managed to _memory_resource.is_managed, which returns False for any MemoryResource subclass that does not opt in. That broke DummyUnifiedMemoryResource (and any user-defined MR wrapping cuMemAllocManaged) where the driver pointer attribute correctly reports IS_MANAGED=1 but the resource does not override is_managed. Query the driver first; only fall back to the memory resource when the driver does not report IS_MANAGED (the pool-allocated managed memory path). This keeps both old-style cuMemAllocManaged buffers and ManagedMemoryResource pool allocations correctly classified. Also rework the regression test parametrization to skip the pinned case when PinnedMemoryResource is unavailable (CUDA < 13.0), and pick up the ruff-format reflow of the helper call site. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: refresh pixi.lock for upstream package updates Pick up cuda-nvrtc 13.2.78, libcufile 1.17.1.22, and other transitive package updates from conda-forge. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 62df918 commit cb3c132

File tree

6 files changed

+331
-175
lines changed

6 files changed

+331
-175
lines changed

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,3 +196,4 @@ cython_debug/
196196

197197
# Cursor
198198
.cursorrules
199+
.claude/settings.local.json

cuda_bindings/pixi.lock

Lines changed: 71 additions & 11 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

cuda_core/cuda/core/_memory/_buffer.pyx

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,11 @@ cdef class Buffer:
391391
def is_managed(self) -> bool:
392392
"""Return True if this buffer is CUDA managed (unified) memory, otherwise False."""
393393
_init_mem_attrs(self)
394-
return self._mem_attrs.is_managed
394+
if self._mem_attrs.is_managed:
395+
return True
396+
# Pool-allocated managed memory does not set CU_POINTER_ATTRIBUTE_IS_MANAGED,
397+
# so fall back to the memory resource when it advertises managed allocations.
398+
return self._memory_resource is not None and self._memory_resource.is_managed
395399

396400
@property
397401
def is_mapped(self) -> bool:
@@ -535,6 +539,11 @@ cdef class MemoryResource:
535539
"""Whether buffers allocated by this resource are host-accessible."""
536540
raise TypeError("MemoryResource.is_host_accessible must be implemented by subclasses.")
537541

542+
@property
543+
def is_managed(self) -> bool:
544+
"""Whether buffers allocated by this resource are CUDA managed (unified) memory."""
545+
return False
546+
538547
@property
539548
def device_id(self) -> int:
540549
"""Device ID associated with this memory resource, or -1 if not applicable."""

cuda_core/cuda/core/_memory/_managed_memory_resource.pyx

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,11 @@ cdef class ManagedMemoryResource(_MemPool):
121121
"""Return True. This memory resource provides host-accessible buffers."""
122122
return True
123123

124+
@property
125+
def is_managed(self) -> bool:
126+
"""Return True. This memory resource provides managed (unified) memory buffers."""
127+
return True
128+
124129

125130
IF CUDA_CORE_BUILD_MAJOR >= 13:
126131
cdef tuple _VALID_LOCATION_TYPES = ("device", "host", "host_numa")

0 commit comments

Comments
 (0)