From ab850a64229bcd3f9794652940197f208161b833 Mon Sep 17 00:00:00 2001 From: Phillip Cloud <417981+cpcloud@users.noreply.github.com> Date: Tue, 6 Jan 2026 08:12:52 -0500 Subject: [PATCH] feat: improve structured dtype array support in `StridedMemoryView` (#1425) * test: add test for structured dtype properties * feat: improve support for structured dtype arrays held by `StridedMemoryView` * chore: simplify using an lru cache instead of manual caching (cherry picked from commit 36742ddc5b91ed388f55d5ec7ade8d245835fb35) --- cuda_core/cuda/core/_memoryview.pyx | 32 +++++++++-------------------- cuda_core/tests/test_utils.py | 17 +++++++++++++++ 2 files changed, 27 insertions(+), 22 deletions(-) diff --git a/cuda_core/cuda/core/_memoryview.pyx b/cuda_core/cuda/core/_memoryview.pyx index c12cbbaa8a..1dbebe7789 100644 --- a/cuda_core/cuda/core/_memoryview.pyx +++ b/cuda_core/cuda/core/_memoryview.pyx @@ -365,8 +365,7 @@ cdef class StridedMemoryView: if self.dl_tensor != NULL: self._dtype = dtype_dlpack_to_numpy(&self.dl_tensor.dtype) elif self.metadata is not None: - # TODO: this only works for built-in numeric types - self._dtype = _typestr2dtype[self.metadata["typestr"]] + self._dtype = _typestr2dtype(self.metadata["typestr"]) return self._dtype @@ -486,25 +485,14 @@ cdef StridedMemoryView view_as_dlpack(obj, stream_ptr, view=None): return buf -_builtin_numeric_dtypes = [ - numpy.dtype("uint8"), - numpy.dtype("uint16"), - numpy.dtype("uint32"), - numpy.dtype("uint64"), - numpy.dtype("int8"), - numpy.dtype("int16"), - numpy.dtype("int32"), - numpy.dtype("int64"), - numpy.dtype("float16"), - numpy.dtype("float32"), - numpy.dtype("float64"), - numpy.dtype("complex64"), - numpy.dtype("complex128"), - numpy.dtype("bool"), -] -# Doing it once to avoid repeated overhead -_typestr2dtype = {dtype.str: dtype for dtype in _builtin_numeric_dtypes} -_typestr2itemsize = {dtype.str: dtype.itemsize for dtype in _builtin_numeric_dtypes} +@functools.lru_cache +def _typestr2dtype(str typestr): + return numpy.dtype(typestr) + + +@functools.lru_cache +def _typestr2itemsize(str typestr): + return _typestr2dtype(typestr).itemsize cdef object dtype_dlpack_to_numpy(DLDataType* dtype): @@ -664,7 +652,7 @@ cdef _StridedLayout layout_from_cai(object metadata): cdef _StridedLayout layout = _StridedLayout.__new__(_StridedLayout) cdef object shape = metadata["shape"] cdef object strides = metadata.get("strides") - cdef int itemsize = _typestr2itemsize[metadata["typestr"]] + cdef int itemsize = _typestr2itemsize(metadata["typestr"]) layout.init_from_tuple(shape, strides, itemsize, True) return layout diff --git a/cuda_core/tests/test_utils.py b/cuda_core/tests/test_utils.py index 06ee3520e2..eb883cd3f3 100644 --- a/cuda_core/tests/test_utils.py +++ b/cuda_core/tests/test_utils.py @@ -446,3 +446,20 @@ def test_from_buffer_with_non_power_of_two_itemsize(): buffer = dev.memory_resource.allocate(required_size) view = StridedMemoryView.from_buffer(buffer, shape=shape, strides=layout.strides, dtype=dtype, is_readonly=True) assert view.dtype == dtype + + +def test_struct_array(): + cp = pytest.importorskip("cupy") + + x = np.array([(1.0, 2), (2.0, 3)], dtype=[("array1", np.float64), ("array2", np.int64)]) + + y = cp.empty(2, dtype=x.dtype) + y.set(x) + + smv = StridedMemoryView.from_cuda_array_interface(y, stream_ptr=0) + assert smv.size * smv.dtype.itemsize == x.nbytes + assert smv.size == x.size + assert smv.shape == x.shape + # full dtype information doesn't seem to be preserved due to use of type strings, + # which are lossy, e.g., dtype([("a", "int")]).str == "V8" + assert smv.dtype == np.dtype(f"V{x.itemsize}")