Support H2D/D2D/D2H SMV copy

cuda_core/cuda/core/experimental/_layout.pxd

@@ -170,6 +170,12 @@ cdef class StridedLayout:
  _init_base_layout_from_tuple(base, shape, None)
  return self._init_dense(base, itemsize, order_flag, &stride_order_vec)
 
+ cdef inline StridedLayout copy(StridedLayout self):
+ cdef StridedLayout new_layout = StridedLayout.__new__(StridedLayout)
+ new_layout.init_from_ptr(self.base.ndim, self.base.shape, self.base.strides, self.itemsize)
+ new_layout.slice_offset = self.slice_offset
+ return new_layout
+
  # ==============================
  # Properties
  # ==============================
@@ -318,6 +324,18 @@ cdef class StridedLayout:
  # Layout manipulation
  # ==============================
 
+ cdef inline int make_dense(StridedLayout self, OrderFlag order_flag, axis_vec_t *stride_order_vec) except -1 nogil:
+ if order_flag == ORDER_C:
+ _dense_strides_c(self.base)
+ elif order_flag == ORDER_F:
+ _dense_strides_f(self.base)
+ elif order_flag == ORDER_PERM:
+ _dense_strides_in_order(self.base, deref(stride_order_vec))
+ else:
+ raise ValueError("The stride_order must be 'C', 'F', or a permutation.")
+ self.slice_offset = 0
+ self._prop_mask = 0
+ return 0
 
  cdef int reshape_into(StridedLayout self, StridedLayout out_layout, BaseLayout& new_shape) except -1 nogil
  cdef int permute_into(StridedLayout self, StridedLayout out_layout, axis_vec_t& axis_order) except -1 nogil
@@ -375,21 +393,32 @@ cdef inline stride_t *get_strides_ptr(BaseLayout& base) except? NULL nogil:
  return tmp_strides
 
 
-cdef inline bint _base_layout_equal(BaseLayout& a, BaseLayout& b) noexcept nogil:
+cdef inline bint base_equal_shapes(BaseLayout& a, BaseLayout& b) noexcept nogil:
  if a.ndim != b.ndim:
  return False
  for i in range(a.ndim):
  if a.shape[i] != b.shape[i]:
  return False
- if a.strides != NULL or b.strides != NULL:
- if a.strides == NULL or b.strides == NULL:
+ return True
+
+
+cdef inline bint _base_equal_strides(BaseLayout& a, BaseLayout& b) noexcept nogil:
+ if a.strides == NULL or b.strides == NULL:
+ return a.strides == b.strides
+ for i in range(a.ndim):
+ if a.strides[i] != b.strides[i]:
  return False
- for i in range(a.ndim):
- if a.strides[i] != b.strides[i]:
- return False
  return True
 
 
+cdef inline bint base_equal_strides(BaseLayout& a, BaseLayout& b) noexcept nogil:
+ return a.ndim == b.ndim and _base_equal_strides(a, b)
+
+
+cdef inline bint base_layout_equal(BaseLayout& a, BaseLayout& b) noexcept nogil:
+ return base_equal_shapes(a, b) and _base_equal_strides(a, b)
+
+
 @cython.overflowcheck(True)
 cdef inline int64_t _volume(BaseLayout& base) except? -1 nogil:
  cdef int64_t vol = 1

cuda_core/cuda/core/experimental/_layout.pyx

@@ -177,7 +177,7 @@ cdef class StridedLayout:
  )
 
  def __eq__(self : StridedLayout, other : StridedLayout) -> bool:
- return self.itemsize == other.itemsize and self.slice_offset == other.slice_offset and _base_layout_equal(self.base, other.base)
+ return self.itemsize == other.itemsize and self.slice_offset == other.slice_offset and base_layout_equal(self.base, other.base)
 
  @property
  def ndim(self : StridedLayout):

cuda_core/cuda/core/experimental/_memoryview.pyx

@@ -5,7 +5,11 @@
 from ._dlpack cimport *
 from libc.stdint cimport intptr_t
 from cuda.core.experimental._layout cimport StridedLayout
-from cuda.core.experimental._stream import Stream
+from cuda.core.experimental._memory._buffer cimport Buffer
+from cuda.core.experimental._strided_copy._copy cimport copy_into_d2d, copy_into_d2h, copy_into_h2d
+from cuda.core.experimental._strided_copy._copy_utils cimport get_data_ptr
+from cuda.core.experimental._stream cimport Stream
+from cuda.core.experimental._strided_copy._copy import CopyAllocatorOptions
 
 import functools
 import warnings
@@ -15,9 +19,6 @@ import numpy
 
 from cuda.core.experimental._utils.cuda_utils import handle_return, driver
 
-
-from cuda.core.experimental._memory import Buffer
-
 # TODO(leofang): support NumPy structured dtypes
 
 
@@ -160,7 +161,7 @@ cdef class StridedMemoryView:
 
  @classmethod
  def from_buffer(
- cls, buffer : Buffer, layout : StridedLayout,
+ cls, Buffer buffer, StridedLayout layout,
  dtype : numpy.dtype | None = None,
  is_readonly : bool = False
  ) -> StridedMemoryView:
@@ -218,7 +219,7 @@ cdef class StridedMemoryView:
  dlm_tensor.deleter(dlm_tensor)
 
  def view(
- self, layout : StridedLayout | None = None, dtype : numpy.dtype | None = None
+ self, StridedLayout layout = None, object dtype = None
  ) -> StridedMemoryView:
  """
  Creates a new view with adjusted layout and dtype.
@@ -236,7 +237,7 @@ cdef class StridedMemoryView:
 
  def copy_from(
  self, other : StridedMemoryView, stream : Stream,
- allocator = None,
+ allocator : CopyAllocatorOptions | dict[str, MemoryResource] | None = None,
  blocking : bool | None = None,
  ):
  """
@@ -253,16 +254,19 @@ cdef class StridedMemoryView:
  if :attr:`dtype` is not specified).
  * The destination's layout must be unique (see :meth:`StridedLayout.is_unique`).
 
+ It is the user's responsibility to ensure proper current device is set
+ when calling this method.
+
  Parameters
  ----------
  other : StridedMemoryView
  The view to copy data from.
  stream : Stream | None, optional
  The stream to schedule the copy on.
- allocator : MemoryResource | None, optional
- If temporary buffers are needed, the specifed memory resources
- will be used to allocate the memory. If not specified, default
-resources will be used.
+ allocator : :obj:`~utils.CopyAllocatorOptions` | dict[str, :obj:`~_memory.MemoryResource`] | None, optional
+ If temporary buffers are needed, the specifed ``allocator.host`` and
+``allocator.device`` memory resources will be used to allocate the memory
+on the host and device respectively.
  blocking : bool | None, optional
  Whether the call should block until the copy is complete.
  * ``True``: the ``stream`` is synchronized with the host at the end of the call,
@@ -274,19 +278,19 @@ cdef class StridedMemoryView:
  * for device-to-device, it defaults to ``False`` (non-blocking),
  * for host-to-device or device-to-host, it defaults to ``True`` (blocking).
  """
-raise NotImplementedError("Sorry, not supported: copy_from")
+copy_into(self, other, stream, allocator, blocking)
 
  def copy_to(
- self, other : StridedMemoryView, stream : Stream | None = None,
- allocator = None,
+ self, other : StridedMemoryView, Stream stream,
+ allocator : CopyAllocatorOptions | dict[str, MemoryResource] | None = None,
  blocking : bool | None = None,
  ):
  """
  Copies the data from this view into the ``other`` view.
 
  For details, see :meth:`copy_from`.
  """
-raise NotImplementedError("Sorry, not supported: copy_to")
+copy_into(other, self, stream, allocator, blocking)
 
  @property
  def layout(self) -> StridedLayout:
@@ -342,7 +346,7 @@ cdef class StridedMemoryView:
  raise ValueError("Cannot infer layout from the exporting object")
  return self._layout
 
- cdef inline object get_buffer(self):
+ cdef inline Buffer get_buffer(self):
  """
  Returns Buffer instance with the underlying data.
  If the SMV was created from a Buffer, it will return the same Buffer instance.
@@ -661,13 +665,9 @@ cdef StridedLayout layout_from_cai(object metadata):
  return layout
 
 
-cdef inline intptr_t get_data_ptr(object buffer, StridedLayout layout) except? 0:
- return <intptr_t>(int(buffer.handle)) + layout.get_slice_offset_in_bytes()
-
-
 cdef inline int view_buffer_strided(
  StridedMemoryView view,
-object buffer,
+Buffer buffer,
  StridedLayout layout,
  object dtype,
  bint is_readonly,
@@ -706,3 +706,54 @@ cdef inline int view_buffer_strided(
  view._layout = layout
  view._dtype = dtype
  return 0
+
+
+cdef int copy_into(
+ StridedMemoryView dst,
+ StridedMemoryView src,
+ Stream stream,
+ allocator : CopyAllocatorOptions | dict[str, MemoryResource] | None = None,
+ blocking : bool | None = None,
+) except -1:
+ cdef object dst_dtype = dst.get_dtype()
+ cdef object src_dtype = src.get_dtype()
+ if dst_dtype is not None and src_dtype is not None and dst_dtype != src_dtype:
+ raise ValueError(
+ f"The destination and source dtypes must be the same, "
+ f"got {dst_dtype} and {src_dtype}."
+ )
+ if dst.readonly:
+ raise ValueError("The destination view is readonly.")
+ cdef bint is_src_device_accessible = src.is_device_accessible
+ cdef bint is_dst_device_accessible = dst.is_device_accessible
+ cdef bint is_blocking
+ cdef int device_id
+ cdef Buffer dst_buffer = dst.get_buffer()
+ cdef Buffer src_buffer = src.get_buffer()
+ cdef StridedLayout dst_layout = dst.get_layout()
+ cdef StridedLayout src_layout = src.get_layout()
+ if is_src_device_accessible and is_dst_device_accessible:
+ device_id = dst.device_id
+ if src.device_id != device_id:
+ raise ValueError(
+ f"The destination and source views must be on the "
+ f"same device, got {device_id} and {src.device_id}."
+ )
+ is_blocking = blocking if blocking is not None else False
+ copy_into_d2d(dst_buffer, dst_layout, src_buffer, src_layout, device_id, stream, is_blocking)
+ return 0
+ elif is_src_device_accessible:
+ device_id = src.device_id
+ is_blocking = blocking if blocking is not None else True
+ copy_into_d2h(dst_buffer, dst_layout, src_buffer, src_layout, device_id, stream, allocator, is_blocking)
+ return 0
+ elif is_dst_device_accessible:
+ device_id = dst.device_id
+ is_blocking = blocking if blocking is not None else True
+ copy_into_h2d(dst_buffer, dst_layout, src_buffer, src_layout, device_id, stream, allocator, is_blocking)
+ return 0
+ else:
+ raise ValueError(
+ "The host-to-host copy is not supported, "
+ "at least one of the views must be device-accessible."
+ )

cuda_core/cuda/core/experimental/_strided_copy/_copy.pxd

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cuda.core.experimental._stream cimport Stream
+from cuda.core.experimental._layout cimport StridedLayout
+from cuda.core.experimental._memory._buffer cimport Buffer
+
+from cuda.core.experimental._memory import MemoryResource
+
+
+cdef int copy_into_d2d(
+ Buffer dst_buffer,
+ StridedLayout dst_layout,
+ Buffer src_buffer,
+ StridedLayout src_layout,
+ int device_id,
+ Stream stream,
+ bint blocking,
+) except -1
+
+
+cdef int copy_into_d2h(
+ Buffer dst_buffer,
+ StridedLayout dst_layout,
+ Buffer src_buffer,
+ StridedLayout src_layout,
+ int device_id,
+ Stream stream,
+ allocator : CopyAllocatorOptions | dict[str, MemoryResource] | None,
+ bint blocking,
+) except -1
+
+
+cdef int copy_into_h2d(
+ Buffer dst_buffer,
+ StridedLayout dst_layout,
+ Buffer src_buffer,
+ StridedLayout src_layout,
+ int device_id,
+ Stream stream,
+ allocator : CopyAllocatorOptions | dict[str, MemoryResource] | None,
+ bint blocking,
+) except -1