Simplify pinned memory example

Author: shwina

cuda_core/examples/memory_ops.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import cupy as cp
+import numpy as np
 
 from cuda.core.experimental import (
     Device,
@@ -55,28 +56,16 @@
 
 # 1. Device Memory (GPU-only)
 device_buffer = device_mr.allocate(total_size, stream=stream)
-device_array = cp.ndarray(
-    size,
-    dtype=dtype,
-    memptr=cp.cuda.MemoryPointer(
-        cp.cuda.UnownedMemory(int(device_buffer.handle), device_buffer.size, device_buffer), 0
-    ),
-)
+device_array = cp.from_dlpack(device_buffer).view(dtype=dtype)
 
 # 2. Pinned Memory (CPU memory, GPU accessible)
 pinned_buffer = pinned_mr.allocate(total_size, stream=stream)
-pinned_array = cp.ndarray(
-    size,
-    dtype=dtype,
-    memptr=cp.cuda.MemoryPointer(
-        cp.cuda.UnownedMemory(int(pinned_buffer.handle), pinned_buffer.size, pinned_buffer), 0
-    ),
-)
+pinned_array = np.from_dlpack(pinned_buffer).view(dtype=dtype)
 
 # Initialize data
 rng = cp.random.default_rng()
 device_array[:] = rng.random(size, dtype=dtype)
-pinned_array[:] = rng.random(size, dtype=dtype)
+pinned_array[:] = rng.random(size, dtype=dtype).get()
 
 # Store original values for verification
 device_original = device_array.copy()
@@ -97,17 +86,6 @@
 assert cp.allclose(device_array, device_original + 1.0), "Device memory operation failed"
 assert cp.allclose(pinned_array, pinned_original * 3.0), "Pinned memory operation failed"
 
-# Demonstrate buffer copying operations
-print("Memory buffer properties:")
-print(f"Device buffer - Device accessible: {device_buffer.is_device_accessible}")
-print(f"Pinned buffer - Device accessible: {pinned_buffer.is_device_accessible}")
-
-# Assert memory properties
-assert device_buffer.is_device_accessible, "Device buffer should be device accessible"
-assert not device_buffer.is_host_accessible, "Device buffer should not be host accessible"
-assert pinned_buffer.is_device_accessible, "Pinned buffer should be device accessible"
-assert pinned_buffer.is_host_accessible, "Pinned buffer should be host accessible"
-
 # Copy data between different memory types
 print("\nCopying data between memory types...")
 
@@ -120,13 +98,7 @@
 
 # Create a new device buffer and copy from pinned
 new_device_buffer = device_mr.allocate(total_size, stream=stream)
-new_device_array = cp.ndarray(
-    size,
-    dtype=dtype,
-    memptr=cp.cuda.MemoryPointer(
-        cp.cuda.UnownedMemory(int(new_device_buffer.handle), new_device_buffer.size, new_device_buffer), 0
-    ),
-)
+new_device_array = cp.from_dlpack(new_device_buffer).view(dtype=dtype)
 
 pinned_buffer.copy_to(new_device_buffer, stream=stream)
 stream.sync()
@@ -146,18 +118,6 @@
 assert device_dlpack[0] == DLDeviceType.kDLCUDA, "Device buffer should have CUDA device type"
 assert pinned_dlpack[0] == DLDeviceType.kDLCUDAHost, "Pinned buffer should have CUDA host device type"
 
-# Test buffer size properties
-assert device_buffer.size == total_size, f"Device buffer size mismatch: expected {total_size}, got {device_buffer.size}"
-assert pinned_buffer.size == total_size, f"Pinned buffer size mismatch: expected {total_size}, got {pinned_buffer.size}"
-assert new_device_buffer.size == total_size, (
-    f"New device buffer size mismatch: expected {total_size}, got {new_device_buffer.size}"
-)
-
-# Test memory resource properties
-assert device_buffer.memory_resource == device_mr, "Device buffer should use device memory resource"
-assert pinned_buffer.memory_resource == pinned_mr, "Pinned buffer should use pinned memory resource"
-assert new_device_buffer.memory_resource == device_mr, "New device buffer should use device memory resource"
-
 # Clean up
 device_buffer.close(stream)
 pinned_buffer.close(stream)