Copy pinned memory tests to test_launcher.py

Author: shwina

cuda_core/examples/memory_ops.py

@@ -72,7 +72,7 @@
 pinned_original = pinned_array.copy()
 
 # Sync before kernel launch
-dev.sync()
+stream.sync()
 
 # Launch kernel
 block = 256

cuda_core/tests/test_launcher.py

@@ -9,7 +9,15 @@
 import pytest
 from conftest import skipif_need_cuda_headers
 
-from cuda.core.experimental import Device, LaunchConfig, LegacyPinnedMemoryResource, Program, ProgramOptions, launch
+from cuda.core.experimental import (
+    Device,
+    DeviceMemoryResource,
+    LaunchConfig,
+    LegacyPinnedMemoryResource,
+    Program,
+    ProgramOptions,
+    launch,
+)
 
 
 def test_launch_config_init(init_cuda):
@@ -197,3 +205,99 @@ def test_cooperative_launch():
     config = LaunchConfig(grid=1, block=1, cooperative_launch=True)
     launch(s, config, ker)
     s.sync()
+
+
+@pytest.mark.parametrize(
+    "memory_resource_class",
+    [
+        DeviceMemoryResource,
+        LegacyPinnedMemoryResource,
+    ],
+)
+def test_launch_with_buffers_allocated_by_memory_resource(init_cuda, memory_resource_class):
+    """Test that kernels can access memory allocated by memory resources."""
+    dev = Device()
+    dev.set_current()
+    stream = dev.create_stream()
+
+    # Kernel that operates on memory
+    code = """
+    extern "C"
+    __global__ void memory_ops(float* data, size_t N) {
+        const unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
+        if (tid < N) {
+            // Access memory (device or pinned)
+            data[tid] = data[tid] * 3.0f;
+        }
+    }
+    """
+
+    # Compile kernel
+    arch = "".join(f"{i}" for i in dev.compute_capability)
+    program_options = ProgramOptions(std="c++17", arch=f"sm_{arch}")
+    prog = Program(code, code_type="c++", options=program_options)
+    mod = prog.compile("cubin")
+    kernel = mod.get_kernel("memory_ops")
+
+    # Create memory resource
+    if memory_resource_class == DeviceMemoryResource:
+        mr = memory_resource_class(dev.device_id)
+    else:  # LegacyPinnedMemoryResource
+        mr = memory_resource_class()
+
+    # Allocate memory
+    size = 1024
+    dtype = np.float32
+    element_size = dtype().itemsize
+    total_size = size * element_size
+
+    buffer = mr.allocate(total_size, stream=stream)
+
+    # Create array view based on memory type
+    if mr.is_host_accessible:
+        # For pinned memory, use numpy
+        array = np.from_dlpack(buffer).view(dtype=dtype)
+    else:
+        # For device memory, use cupy
+        import cupy as cp
+
+        array = cp.from_dlpack(buffer).view(dtype=dtype)
+
+    # Initialize data with random values
+    if mr.is_host_accessible:
+        rng = np.random.default_rng()
+        array[:] = rng.random(size, dtype=dtype)
+    else:
+        import cupy as cp
+
+        rng = cp.random.default_rng()
+        array[:] = rng.random(size, dtype=dtype)
+
+    # Store original values for verification
+    original = array.copy()
+
+    # Sync before kernel launch
+    stream.sync()
+
+    # Launch kernel
+    block = 256
+    grid = (size + block - 1) // block
+    config = LaunchConfig(grid=grid, block=block)
+
+    launch(stream, config, kernel, buffer, np.uint64(size))
+    stream.sync()
+
+    # Verify kernel operations
+    if mr.is_host_accessible:
+        assert np.allclose(array, original * 3.0), f"{memory_resource_class.__name__} operation failed"
+    else:
+        import cupy as cp
+
+        assert cp.allclose(array, original * 3.0), f"{memory_resource_class.__name__} operation failed"
+
+    # Clean up
+    buffer.close(stream)
+    stream.close()
+
+    # Verify buffer is properly closed
+    assert buffer.handle == 0, f"{memory_resource_class.__name__} buffer should be closed"