Added support for consuming DLPack allocated on a sub-device

USM allocation must be bound to the default context. When producing
DLPack, device_id is populated with id of the ancestor root device.

Code remains functional on systems where default_context extension
support is not enabled (e.g. Windows), but DLPack sharing is limited
to allocation made on root devices only.

Author: oleksandr-pavlyk

dpctl/tensor/_dlpack.pyx

@@ -140,6 +140,7 @@ cpdef to_dlpack_capsule(usm_ndarray usm_ary) except+:
     cdef c_dpctl.SyclQueue ary_sycl_queue
     cdef c_dpctl.SyclDevice ary_sycl_device
     cdef DPCTLSyclDeviceRef pDRef = NULL
+    cdef DPCTLSyclDeviceRef tDRef = NULL
     cdef DLManagedTensor *dlm_tensor = NULL
     cdef DLTensor *dl_tensor = NULL
     cdef int nd = usm_ary.get_ndim()
@@ -157,19 +158,42 @@ cpdef to_dlpack_capsule(usm_ndarray usm_ary) except+:
     ary_sycl_queue = usm_ary.get_sycl_queue()
     ary_sycl_device = ary_sycl_queue.get_sycl_device()
 
-    # check that ary_sycl_device is a non-partitioned device
-    pDRef = DPCTLDevice_GetParentDevice(ary_sycl_device.get_device_ref())
-    if pDRef is not NULL:
-        DPCTLDevice_Delete(pDRef)
-        raise DLPackCreationError(
-            "to_dlpack_capsule: DLPack can only export arrays allocated on "
-            "non-partitioned SYCL devices."
-        )
-    default_context = dpctl.SyclQueue(ary_sycl_device).sycl_context
-    if not usm_ary.sycl_context == default_context:
+    try:
+        default_context = ary_sycl_device.sycl_platform.default_context
+    except RuntimeError:
+        # RT does not support default_context, e.g. Windows
+        default_context = None
+    if default_context is None:
+        # check that ary_sycl_device is a non-partitioned device
+        pDRef = DPCTLDevice_GetParentDevice(ary_sycl_device.get_device_ref())
+        if pDRef is not NULL:
+            DPCTLDevice_Delete(pDRef)
+            raise DLPackCreationError(
+                "to_dlpack_capsule: DLPack can only export arrays allocated "
+                "on non-partitioned SYCL devices on platforms where "
+                "default_context oneAPI extension is not supported."
+            )
+    else:
+        if not usm_ary.sycl_context == default_context:
+            raise DLPackCreationError(
+                "to_dlpack_capsule: DLPack can only export arrays based on USM "
+                "allocations bound to a default platform SYCL context"
+            )
+        # Find the unpartitioned parent of the allocation device
+        pDRef = DPCTLDevice_GetParentDevice(ary_sycl_device.get_device_ref())
+        if pDRef is not NULL:
+            tDRef = DPCTLDevice_GetParentDevice(pDRef)
+            while tDRef is not NULL:
+                DPCTLDevice_Delete(pDRef)
+                pDRef = tDRef
+                tDRef = DPCTLDevice_GetParentDevice(pDRef)
+            ary_sycl_device = c_dpctl.SyclDevice._create(pDRef)
+
+    # Find ordinal number of the parent device
+    device_id = ary_sycl_device.get_overall_ordinal()
+    if device_id < 0:
         raise DLPackCreationError(
-            "to_dlpack_capsule: DLPack can only export arrays based on USM "
-            "allocations bound to a default platform SYCL context"
+            "to_dlpack_capsule: failed to determine device_id"
         )
 
     dlm_tensor = <DLManagedTensor *> stdlib.malloc(
@@ -192,14 +216,6 @@ cpdef to_dlpack_capsule(usm_ndarray usm_ary) except+:
         for i in range(nd):
             shape_strides_ptr[nd + i] = strides_ptr[i]
 
-    device_id = ary_sycl_device.get_overall_ordinal()
-    if device_id < 0:
-        stdlib.free(shape_strides_ptr)
-        stdlib.free(dlm_tensor)
-        raise DLPackCreationError(
-            "to_dlpack_capsule: failed to determine device_id"
-        )
-
     ary_dt = usm_ary.dtype
     ary_dtk = ary_dt.kind
     element_offset = usm_ary.get_offset()
@@ -278,15 +294,16 @@ cpdef usm_ndarray from_dlpack_capsule(object py_caps) except +:
         success.
     Raises:
         TypeError: if argument is not a "dltensor" capsule.
-        ValueError: if argument is "used_dltensor" capsule,
-             if the USM pointer is not bound to the reconstructed
+        ValueError: if argument is "used_dltensor" capsule
+        BufferError:  if the USM pointer is not bound to the reconstructed
              sycl context, or the DLPack's device_type is not supported
              by dpctl.
     """
     cdef DLManagedTensor *dlm_tensor = NULL
     cdef bytes usm_type
     cdef size_t sz = 1
     cdef int i
+    cdef int device_id = -1
     cdef int element_bytesize = 0
     cdef Py_ssize_t offset_min = 0
     cdef Py_ssize_t offset_max = 0
@@ -308,26 +325,37 @@ cpdef usm_ndarray from_dlpack_capsule(object py_caps) except +:
             py_caps, "dltensor")
     # Verify that we can work with this device
     if dlm_tensor.dl_tensor.device.device_type == kDLOneAPI:
-        q = dpctl.SyclQueue(str(<int>dlm_tensor.dl_tensor.device.device_id))
+        device_id = dlm_tensor.dl_tensor.device.device_id
+        root_device = dpctl.SyclDevice(str(<int>device_id))
+        try:
+            default_context = root_device.sycl_platform.default_context
+        except RuntimeError:
+            default_context = dpctl.SyclQueue(root_device).sycl_context
         if dlm_tensor.dl_tensor.data is NULL:
             usm_type = b"device"
+            q = dpctl.SyclQueue(default_context, root_device)
         else:
             usm_type = c_dpmem._Memory.get_pointer_type(
                 <DPCTLSyclUSMRef> dlm_tensor.dl_tensor.data,
-                <c_dpctl.SyclContext>q.sycl_context)
-        if usm_type == b"unknown":
-            raise ValueError(
-                f"Data pointer in DLPack is not bound to default sycl "
-                "context of device '{device_id}', translated to "
-                "{q.sycl_device.filter_string}"
+                <c_dpctl.SyclContext>default_context)
+            if usm_type == b"unknown":
+                raise BufferError(
+                    "Data pointer in DLPack is not bound to default sycl "
+                    f"context of device '{device_id}', translated to "
+                    f"{root_device.filter_string}"
+                )
+            alloc_device = c_dpmem._Memory.get_pointer_device(
+                <DPCTLSyclUSMRef> dlm_tensor.dl_tensor.data,
+                <c_dpctl.SyclContext>default_context
             )
+            q = dpctl.SyclQueue(default_context, alloc_device)
         if dlm_tensor.dl_tensor.dtype.bits % 8:
-            raise ValueError(
+            raise BufferError(
                 "Can not import DLPack tensor whose element's "
                 "bitsize is not a multiple of 8"
             )
         if dlm_tensor.dl_tensor.dtype.lanes != 1:
-            raise ValueError(
+            raise BufferError(
                 "Can not import DLPack tensor with lanes != 1"
             )
         if dlm_tensor.dl_tensor.strides is NULL: