feat: Introduce StridedLayout, support wrapping external allocations in Buffer, add StridedMemoryView.from_buffer

cuda_core/cuda/core/experimental/_memory/_buffer.pxd

@@ -7,14 +7,23 @@ from libc.stdint cimport uintptr_t
 from cuda.core.experimental._stream cimport Stream
 
 
+cdef struct _MemAttrs:
+ int device_id
+ bint is_device_accessible
+ bint is_host_accessible
+
+
 cdef class Buffer:
  cdef:
  uintptr_t _ptr
  size_t _size
  MemoryResource _memory_resource
  object _ipc_data
+ object _owner
  object _ptr_obj
  Stream _alloc_stream
+ _MemAttrs _mem_attrs
+ bint _mem_attrs_inited
 
 
 cdef class MemoryResource:

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+cimport cython
 from libc.stdint cimport uintptr_t, int64_t, uint64_t
 
 from cuda.bindings cimport cydriver
@@ -18,6 +19,7 @@ from typing import TypeVar, Union
 
 from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
 from cuda.core.experimental._utils.cuda_utils import driver
+from cuda.core.experimental._device import Device
 
 __all__ = ['Buffer', 'MemoryResource']
 
@@ -47,6 +49,8 @@ cdef class Buffer:
  self._ipc_data = None
  self._ptr_obj = None
  self._alloc_stream = None
+ self._owner = None
+ self._mem_attrs_inited = False
 
  def __init__(self, *args, **kwargs):
  raise RuntimeError("Buffer objects cannot be instantiated directly. "
@@ -55,15 +59,19 @@ cdef class Buffer:
  @classmethod
  def _init(
  cls, ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None,
- stream: Stream | None = None, ipc_descriptor: IPCBufferDescriptor | None = None
+ stream: Stream | None = None, ipc_descriptor: IPCBufferDescriptor | None = None,
+ owner : object | None = None
  ):
  cdef Buffer self = Buffer.__new__(cls)
  self._ptr = <uintptr_t>(int(ptr))
  self._ptr_obj = ptr
  self._size = size
+ if mr is not None and owner is not None:
+ raise ValueError("owner and memory resource cannot be both specified together")
  self._memory_resource = mr
  self._ipc_data = IPCDataForBuffer(ipc_descriptor, True) if ipc_descriptor is not None else None
  self._alloc_stream = <Stream>(stream) if stream is not None else None
+ self._owner = owner
  return self
 
  def __dealloc__(self):
@@ -75,7 +83,8 @@ cdef class Buffer:
 
  @staticmethod
  def from_handle(
- ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None
+ ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None,
+ owner: object | None = None,
  ) -> Buffer:
  """Create a new :class:`Buffer` object from a pointer.
 
@@ -87,9 +96,13 @@ cdef class Buffer:
  Memory size of the buffer
  mr : :obj:`~_memory.MemoryResource`, optional
  Memory resource associated with the buffer
+ owner : object, optional
+ An object holding external allocation that the ``ptr`` points to.
+ The reference is kept as long as the buffer is alive.
+ The ``owner`` and ``mr`` cannot be specified together.
  """
  # TODO: It is better to take a stream for latter deallocation
- return Buffer._init(ptr, size, mr=mr)
+ return Buffer._init(ptr, size, mr=mr, owner=owner)
 
  @classmethod
  def from_ipc_descriptor(
@@ -297,7 +310,9 @@ cdef class Buffer:
  """Return the device ordinal of this buffer."""
  if self._memory_resource is not None:
  return self._memory_resource.device_id
- raise NotImplementedError("WIP: Currently this property only supports buffers with associated MemoryResource")
+ else:
+ Buffer_init_mem_attrs(self)
+ return self._mem_attrs.device_id
 
  @property
  def handle(self) -> DevicePointerT:
@@ -321,14 +336,18 @@ cdef class Buffer:
  """Return True if this buffer can be accessed by the GPU, otherwise False."""
  if self._memory_resource is not None:
  return self._memory_resource.is_device_accessible
- raise NotImplementedError("WIP: Currently this property only supports buffers with associated MemoryResource")
+ else:
+ Buffer_init_mem_attrs(self)
+ return self._mem_attrs.is_device_accessible
 
  @property
  def is_host_accessible(self) -> bool:
  """Return True if this buffer can be accessed by the CPU, otherwise False."""
  if self._memory_resource is not None:
  return self._memory_resource.is_host_accessible
- raise NotImplementedError("WIP: Currently this property only supports buffers with associated MemoryResource")
+ else:
+ Buffer_init_mem_attrs(self)
+ return self._mem_attrs.is_host_accessible
 
  @property
  def is_mapped(self) -> bool:
@@ -346,20 +365,92 @@ cdef class Buffer:
  """Return the memory size of this buffer."""
  return self._size
 
+ @property
+ def owner(self) -> object:
+ """Return the object holding external allocation."""
+ return self._owner
+
 
 # Buffer Implementation
 # ---------------------
 cdef inline void Buffer_close(Buffer self, stream):
  cdef Stream s
- if self._ptr and self._memory_resource is not None:
- s = Stream_accept(stream) if stream is not None else self._alloc_stream
- self._memory_resource.deallocate(self._ptr, self._size, s)
+ if self._ptr:
+ if self._memory_resource is not None:
+ s = Stream_accept(stream) if stream is not None else self._alloc_stream
+ self._memory_resource.deallocate(self._ptr, self._size, s)
  self._ptr = 0
  self._memory_resource = None
+ self._owner = None
  self._ptr_obj = None
  self._alloc_stream = None
 
 
+cdef Buffer_init_mem_attrs(Buffer self):
+ if not self._mem_attrs_inited:
+ query_memory_attrs(self._mem_attrs, self._ptr)
+ self._mem_attrs_inited = True
+
+
+cdef int query_memory_attrs(_MemAttrs &out, uintptr_t ptr) except -1 nogil:
+ cdef unsigned int memory_type = 0
+ cdef int is_managed = 0
+ cdef int device_id = 0
+ _query_memory_attrs(memory_type, is_managed, device_id, <cydriver.CUdeviceptr>ptr)
+
+ if memory_type == 0:
+ # unregistered host pointer
+ out.is_host_accessible = True
+ out.is_device_accessible = False
+ out.device_id = -1
+ # for managed memory, the memory type can be CU_MEMORYTYPE_DEVICE,
+ # so we need to check it first not to falsely claim it is not
+ # host accessible.
+ elif (
+ is_managed
+ or memory_type == cydriver.CUmemorytype.CU_MEMORYTYPE_HOST
+ ):
+ # For pinned memory allocated with cudaMallocHost or paged-locked
+ # with cudaHostRegister, the memory_type is
+ # cydriver.CUmemorytype.CU_MEMORYTYPE_HOST.
+ # TODO(ktokarski): In some cases, the registered memory requires
+ # using different ptr for device and host, we could check
+ # cuMemHostGetDevicePointer and
+ # CU_DEVICE_ATTRIBUTE_CAN_USE_HOST_POINTER_FOR_REGISTERED_MEM
+ # to double check the device accessibility.
+ out.is_host_accessible = True
+ out.is_device_accessible = True
+ out.device_id = device_id
+ elif memory_type == cydriver.CUmemorytype.CU_MEMORYTYPE_DEVICE:
+ out.is_host_accessible = False
+ out.is_device_accessible = True
+ out.device_id = device_id
+ else:
+ raise ValueError(f"Unsupported memory type: {memory_type}")
+ return 0
+
+
+cdef inline int _query_memory_attrs(unsigned int& memory_type, int & is_managed, int& device_id, cydriver.CUdeviceptr ptr) except -1 nogil:
+ cdef cydriver.CUpointer_attribute attrs[3]
+ cdef uintptr_t vals[3]
+ attrs[0] = cydriver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_MEMORY_TYPE
+ attrs[1] = cydriver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_IS_MANAGED
+ attrs[2] = cydriver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL
+ vals[0] = <uintptr_t><void*>&memory_type
+ vals[1] = <uintptr_t><void*>&is_managed
+ vals[2] = <uintptr_t><void*>&device_id
+
+ cdef cydriver.CUresult ret
+ ret = cydriver.cuPointerGetAttributes(3, attrs, <void**>vals, ptr)
+ if ret == cydriver.CUresult.CUDA_ERROR_NOT_INITIALIZED:
+ with cython.gil:
+ # Device class handles the cuInit call internally
+ Device()
+ ret = cydriver.cuPointerGetAttributes(3, attrs, <void**>vals, ptr)
+ HANDLE_RETURN(ret)
+ return 0
+
+
 cdef class MemoryResource:
  """Abstract base class for memory resources that manage allocation and
  deallocation of buffers.