fix

Author: leofang

cuda_core/cuda/core/experimental/_kernel_arg_handler.pyx

@@ -3,21 +3,23 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from cpython.mem cimport PyMem_Malloc, PyMem_Free
-from libc.stdint cimport (intptr_t,
+from libc.stdint cimport (intptr_t, uintptr_t,
                           int8_t, int16_t, int32_t, int64_t,
                           uint8_t, uint16_t, uint32_t, uint64_t,)
 from libcpp cimport bool as cpp_bool
 from libcpp.complex cimport complex as cpp_complex
 from libcpp cimport nullptr
 from libcpp cimport vector
 
+from cuda.bindings cimport cydriver
+from cuda.core.experimental._memoryview cimport _MDSPAN
+
 import ctypes
 
 import numpy
 
 from cuda.core.experimental._memory import Buffer
 from cuda.core.experimental._utils.cuda_utils import driver
-from cuda.bindings cimport cydriver
 
 
 ctypedef cpp_complex.complex[float] cpp_single_complex
@@ -265,6 +267,8 @@ cdef class ParamHolder:
         cdef size_t i
         cdef int not_prepared
         cdef object arg_type
+        cdef _MDSPAN mdspan_obj
+        cdef uintptr_t mdspan_ptr
         self.data = vector.vector[voidptr](n_args, nullptr)
         self.data_addresses = vector.vector[voidptr](n_args)
         for i, arg in enumerate(kernel_args):
@@ -296,6 +300,14 @@ cdef class ParamHolder:
             elif arg_type is complex:
                 prepare_arg[cpp_double_complex](self.data, self.data_addresses, arg, i)
                 continue
+            elif arg_type is _MDSPAN:
+                # The mdspan struct is allocated on the host and owned by the CuPy mdspan object.
+                # We pass a pointer to the struct so the driver can copy it by value to the kernel.
+                # Access _ptr at C level to avoid creating a temporary Python object.
+                mdspan_obj = <_MDSPAN>arg
+                mdspan_ptr = mdspan_obj._ptr
+                self.data_addresses[i] = <void*>mdspan_ptr
+                continue
 
             not_prepared = prepare_numpy_arg(self.data, self.data_addresses, arg, i)
             if not_prepared:

cuda_core/cuda/core/experimental/_memoryview.pxd

@@ -0,0 +1,10 @@
+from libc.stdint cimport uintptr_t
+
+
+cdef class _MDSPAN:
+    cdef:
+        # this must be a pointer to a host mdspan object
+        readonly uintptr_t _ptr
+        # if the host mdspan is exported from any Python object,
+        # we need to keep a reference to that object alive
+        readonly object _exporting_obj

cuda_core/cuda/core/experimental/_memoryview.pyx

@@ -2,7 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from libc.stdint cimport uintptr_t
+
 from ._dlpack cimport *
+from cuda.core.experimental._utils cimport cuda_utils
 
 import functools
 import warnings
@@ -11,12 +14,26 @@ from typing import Optional
 import numpy
 
 from cuda.core.experimental._utils.cuda_utils import handle_return, driver
-from cuda.core.experimental._utils cimport cuda_utils
 
 
 # TODO(leofang): support NumPy structured dtypes
 
 
+cdef class _MDSPAN:
+
+    def __cinit__(self):
+        self._ptr = 0
+
+    def __init__(self, uintptr_t ptr, object obj=None):
+        self._ptr = ptr
+        self._exporting_obj = obj
+
+    def __dealloc__(self):
+        self._ptr = 0
+        self._exporting_obj = None
+
+
+
 cdef class StridedMemoryView:
     """A dataclass holding metadata of a strided dense array/tensor.
 
@@ -98,6 +115,7 @@ cdef class StridedMemoryView:
         # this flag helps prevent unnecessary recompuation of _strides
         bint _strides_init
         object _dtype
+        _MDSPAN _mdspan
 
     def __init__(self, obj: object = None, stream_ptr: int | None = None) -> None:
         cdef str clsname = self.__class__.__name__
@@ -224,6 +242,27 @@ cdef class StridedMemoryView:
                     self._dtype = numpy.dtype(self.metadata["typestr"])
         return self._dtype
 
+    @property
+    def as_mdspan(self) -> _MDSPAN:
+        """A C++ mdspan view of the tensor.
+
+        Returns
+        -------
+        mdspan : _MDSPAN
+        """
+        if self._mdspan is None:
+            arr = self.exporting_obj
+            module = self.exporting_obj.__class__.__module__.split(".")[0]
+            if module == "cupy":
+                mdspan = arr.mdspan
+                #mdspan = arr.cstruct
+                self._mdspan = _MDSPAN(<uintptr_t>(mdspan.ptr), mdspan)
+            else:
+                raise NotImplementedError(
+                    f"as_mdspan is not implemented for objects from module '{module}'"
+                )
+        return self._mdspan
+
     def __repr__(self):
         return (f"StridedMemoryView(ptr={self.ptr},\n"
               + f"                  shape={self.shape},\n"

cuda_core/examples/debug_cupy_arr.py

@@ -0,0 +1,33 @@
+import cupy as cp
+from cuda.core.experimental import Program, Device, LaunchConfig, launch
+from cuda.core.experimental.utils import StridedMemoryView
+
+
+code = r"""
+extern "C"
+__global__ void debug_cupy_arr(const float* arr) {
+    printf("ptr: %p\n", arr);
+}
+
+"""
+
+
+ker = cp.RawKernel(code, 'debug_cupy_arr')
+arr = cp.array([1, 2, 3], dtype=cp.float32)
+print(f"arr device ptr: {arr.data.ptr:#x}")
+ker((1,), (1,), (arr,))
+cp.cuda.Device().synchronize()
+
+
+
+mod = Program(code, code_type='c++')
+obj = mod.compile(target_type='cubin')
+ker = obj.get_kernel('debug_cupy_arr')
+dev = Device()
+dev.set_current()
+
+#launch(dev.default_stream, LaunchConfig(grid=1, block=1), ker, arr.data.ptr)
+launch(dev.default_stream, LaunchConfig(grid=1, block=1), ker, StridedMemoryView(arr, stream_ptr=-1).as_mdspan)
+dev.default_stream.sync()
+
+print("Done.")

cuda_core/examples/mdspan_verify_args.py

@@ -21,6 +21,7 @@
 import os, sys
 import cupy as cp
 from cuda.core.experimental import Device, LaunchConfig, Program, ProgramOptions, launch
+from cuda.core.experimental.utils import StridedMemoryView
 
 # prepare include
 cuda_path = os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME"))
@@ -42,37 +43,81 @@
 code_verify = """
 #include <cuda/std/mdspan>
 
+typedef struct {
+    void* ptr;
+    size_t ext1;
+    size_t ext2;
+} mdspan_view_t;
+
+
 // Kernel to verify layout_right (C-order) mdspan arguments
 template<typename T>
 __global__ void verify_mdspan_layout_right(
-    cuda::std::mdspan<T, cuda::std::extents<size_t, cuda::std::dynamic_extent, cuda::std::dynamic_extent>, cuda::std::layout_right> arr
+    mdspan_view_t arr
 ) {
     // Only thread 0 prints to avoid cluttered output
     if (threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
         printf("=== layout_right (C-order) mdspan ===\\n");
-        printf("Data pointer: %p\\n", arr.data_handle());
-        printf("Extent 0 (rows): %zu\\n", arr.extent(0));
-        printf("Extent 1 (cols): %zu\\n", arr.extent(1));
-        printf("Size: %zu\\n", arr.size());
-        
-        // For layout_right, strides are implicit but we can query them
-        printf("Stride 0: %zu\\n", arr.stride(0));
-        printf("Stride 1: %zu\\n", arr.stride(1));
-        
-        // Verify memory layout: for layout_right (C-order)
-        // stride(0) should equal extent(1), stride(1) should be 1
-        printf("Expected stride(0) = extent(1): %s\\n", 
-               (arr.stride(0) == arr.extent(1)) ? "PASS" : "FAIL");
-        printf("Expected stride(1) = 1: %s\\n", 
-               (arr.stride(1) == 1) ? "PASS" : "FAIL");
-        
-        // Test element access
-        if (arr.extent(0) > 0 && arr.extent(1) > 0) {
-            printf("First element arr(0,0): %f\\n", static_cast<float>(arr(0, 0)));
-        }
+        printf("sizeof(mdspan_view_t): %llu\\n", sizeof(arr));
+        printf("view - ptr: %p\\n", reinterpret_cast<mdspan_view_t*>(&arr)->ptr);
+        printf("view2 : %p\\n", *(void**)((char*)(&arr) + 0));
+        printf("view - ext1: %p\\n", reinterpret_cast<mdspan_view_t*>(&arr)->ext1);
+        printf("view - ext2: %p\\n", reinterpret_cast<mdspan_view_t*>(&arr)->ext2);
     }
 }
 
+// // Kernel to verify layout_right (C-order) mdspan arguments
+//
+// typedef struct {
+//     void* ptr;
+//     void* ext1;
+//     void* ext2;
+// } mdspan_view_t;
+//
+// template<typename T>
+// __global__ void verify_mdspan_layout_right(
+//     cuda::std::mdspan<T, cuda::std::extents<size_t, cuda::std::dynamic_extent, cuda::std::dynamic_extent>, cuda::std::layout_right> arr
+// ) {
+//     // Only thread 0 prints to avoid cluttered output
+//     if (threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
+//         printf("=== layout_right (C-order) mdspan ===\\n");
+//         printf("sizeof(mdspan): %llu\\n", sizeof(arr));
+//         printf("view - ptr: %p\\n", reinterpret_cast<mdspan_view_t*>(&arr)->ptr);
+//         printf("view2 : %p\\n", (void**)((char*)(&arr) + 0));
+//         //printf("view - ext1: %llu\\n", *((size_t*)(reinterpret_cast<mdspan_view_t*>(&arr)->ext1)));
+//         //printf("view - ext2: %llu\\n", *((size_t*)(reinterpret_cast<mdspan_view_t*>(&arr)->ext2)));
+//         printf("view - ext1: %p\\n", reinterpret_cast<mdspan_view_t*>(&arr)->ext1);
+//         printf("view - ext2: %p\\n", reinterpret_cast<mdspan_view_t*>(&arr)->ext2);
+// 
+//         printf("Data pointer: %p\\n", arr.data_handle());
+//         printf("Data pointer (actual): %p\\n", (void*)((char*)(&arr) + 0));
+//         printf("Data pointer (actual): %p\\n", addressof(arr));
+//         printf("Extent 0 (rows): %llu\\n", arr.extent(0));
+//         printf("Extent 1 (cols): %llu\\n", arr.extent(1));
+//         printf("Extent 0 (rows) (actual): %llu\\n", (size_t)(*((char*)(&arr) + 8)));
+//         printf("Extent 1 (cols) (actual): %llu\\n", (size_t)(*((char*)(&arr) + 16)));
+//         printf("Size: %zu\\n", arr.size());
+//         
+//         // For layout_right, strides are implicit but we can query them
+//         printf("Stride 0: %llu\\n", arr.stride(0));
+//         printf("Stride 1: %llu\\n", arr.stride(1));
+//         printf("Stride 0 (actual): %llu\\n", (size_t)((char*)(&arr) + 24));
+//         printf("Stride 1 (actual): %llu\\n", (size_t)((char*)(&arr) + 32));
+//         
+//         // Verify memory layout: for layout_right (C-order)
+//         // stride(0) should equal extent(1), stride(1) should be 1
+//         printf("Expected stride(0) = extent(1): %s\\n", 
+//                (arr.stride(0) == arr.extent(1)) ? "PASS" : "FAIL");
+//         printf("Expected stride(1) = 1: %s\\n", 
+//                (arr.stride(1) == 1) ? "PASS" : "FAIL");
+//         
+//         // Test element access
+//         if (arr.extent(0) > 0 && arr.extent(1) > 0) {
+//             printf("First element arr(0,0): %f\\n", static_cast<float>(arr(0, 0)));
+//         }
+//     }
+// }
+
 // Kernel to verify layout_left (F-order) mdspan arguments
 template<typename T>
 __global__ void verify_mdspan_layout_left(
@@ -162,10 +207,13 @@ def prepare_mdspan_args_layout_right(arr, dtype, shape):
     tuple
         Arguments to pass to the kernel (needs investigation)
     """
-    data_ptr = arr.data.ptr
-    rows, cols = shape
-    # TODO: Determine exact argument structure
-    return (data_ptr, rows, cols)
+    #obj = arr.mdspan
+    #print(f"{hex(obj.ptr)=}, {obj.ptr=}")
+    #return (obj.ptr,)
+
+    obj = StridedMemoryView(arr, stream_ptr=-1).as_mdspan
+    print(f"{hex(obj._ptr)=}, {obj._ptr=}, type={type(obj)}")
+    return (obj,)
 
 
 def prepare_mdspan_args_layout_left(arr, dtype, shape):
@@ -266,7 +314,7 @@ def verify_layout_right():
 
     # Verify array is in C-order
     assert arr.flags['C_CONTIGUOUS']
-    
+    print(f"Array pointer: {hex(arr.data.ptr)}")
     print(f"Array shape: {arr.shape}")
     print(f"Array strides (bytes): {arr.strides}")
     print(f"Array strides (elements): ({arr.strides[0]//arr.itemsize}, {arr.strides[1]//arr.itemsize})")
@@ -282,8 +330,8 @@ def verify_layout_right():
     config = LaunchConfig(grid=1, block=1)
 
     # TODO: Launch kernel with proper mdspan arguments
-    # launch(s, config, ker, *args)
-    # s.sync()
+    launch(s, config, ker, *args)
+    s.sync()
 
     print("Verification kernel prepared (not executed)")
     print()