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
@@ -296,6 +298,12 @@ 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.
+                self.data_addresses[i] = <void*>((<_MDSPAN>arg)._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/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,22 +43,66 @@
 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(
+//     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("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("Extent 0 (rows): %zu\\n", arr.extent(0));
-        printf("Extent 1 (cols): %zu\\n", arr.extent(1));
+        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: %zu\\n", arr.stride(0));
-        printf("Stride 1: %zu\\n", arr.stride(1));
+        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
@@ -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()