Skip to content

Commit 27b309f

Browse files
cpcloudclaude
andcommitted
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>
1 parent 384c5c5 commit 27b309f

File tree

1 file changed

+45
-0
lines changed

1 file changed

+45
-0
lines changed

cuda_core/tests/test_memory.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -605,6 +605,51 @@ def test_managed_buffer_dlpack_roundtrip_device_type():
605605
assert view.__dlpack_device__() == (int(DLDeviceType.kDLCUDAManaged), 0)
606606

607607

608+
def test_managed_memory_resource_buffer_dlpack_device_type():
609+
"""Verify that pool-allocated managed memory reports kDLCUDAManaged.
610+
611+
Allocations from ManagedMemoryResource go through cuMemAllocFromPoolAsync
612+
on a CU_MEM_ALLOCATION_TYPE_MANAGED pool, which does not set
613+
CU_POINTER_ATTRIBUTE_IS_MANAGED. The Buffer must still report itself as
614+
managed via its memory resource so that DLPack consumers (e.g. CCCL's
615+
make_tma_descriptor) accept the buffer.
616+
"""
617+
device = Device()
618+
device.set_current()
619+
skip_if_managed_memory_unsupported(device)
620+
mr = create_managed_memory_resource_or_skip(
621+
ManagedMemoryResourceOptions(preferred_location=device.device_id)
622+
)
623+
buf = mr.allocate(1024)
624+
625+
assert mr.is_managed
626+
assert buf.is_managed
627+
assert buf.__dlpack_device__() == (DLDeviceType.kDLCUDAManaged, 0)
628+
629+
view = StridedMemoryView.from_any_interface(buf, stream_ptr=-1)
630+
assert view.__dlpack_device__() == (int(DLDeviceType.kDLCUDAManaged), 0)
631+
632+
633+
@pytest.mark.parametrize(
634+
("mr_factory", "expected_managed"),
635+
[
636+
(lambda dev: DeviceMemoryResource(dev), False),
637+
(lambda dev: PinnedMemoryResource(), False),
638+
],
639+
)
640+
def test_non_managed_resources_report_not_managed(mr_factory, expected_managed):
641+
"""Non-managed memory resources must report is_managed=False."""
642+
device = Device()
643+
device.set_current()
644+
if not device.properties.memory_pools_supported:
645+
pytest.skip("Device does not support mempool operations")
646+
mr = mr_factory(device)
647+
assert mr.is_managed is expected_managed
648+
buf = mr.allocate(1024)
649+
assert buf.is_managed is expected_managed
650+
buf.close()
651+
652+
608653
@pytest.mark.parametrize("use_device_object", [True, False])
609654
def test_device_memory_resource_initialization(use_device_object):
610655
"""Test that DeviceMemoryResource can be initialized successfully.

0 commit comments

Comments
 (0)