Support for composite devices

dpctl/__init__.py

@@ -35,6 +35,7 @@
     SyclSubDeviceCreationError,
 )
 from ._sycl_device_factory import (
+    get_composite_devices,
     get_devices,
     get_num_devices,
     has_accelerator_devices,
@@ -87,6 +88,7 @@
     "has_gpu_devices",
     "has_accelerator_devices",
     "has_host_device",
+    "get_composite_devices",
 ]
 __all__ += [
     "SyclEvent",

dpctl/_backend.pxd

@@ -98,6 +98,8 @@ cdef extern from "syclinterface/dpctl_sycl_enum_types.h":
         _usm_atomic_shared_allocations      'usm_atomic_shared_allocations',
         _host_debuggable                    'host_debuggable',
         _emulated                           'emulated',
+        _is_component                       'is_component',
+        _is_composite                       'is_composite',
 
     ctypedef enum _partition_affinity_domain_type 'DPCTLPartitionAffinityDomainType':
         _not_applicable                     'not_applicable',
@@ -217,6 +219,8 @@ cdef extern from "syclinterface/dpctl_sycl_device_interface.h":
     cdef uint32_t DPCTLDevice_GetPartitionMaxSubDevices(const DPCTLSyclDeviceRef DRef)
     cdef uint32_t DPCTLDevice_GetMaxClockFrequency(const DPCTLSyclDeviceRef DRef)
     cdef uint64_t DPCTLDevice_GetMaxMemAllocSize(const DPCTLSyclDeviceRef DRef)
+    cdef DPCTLSyclDeviceRef DPCTLDevice_GetCompositeDevice(const DPCTLSyclDeviceRef DRef)
+    cdef DPCTLDeviceVectorRef DPCTLDevice_GetComponentDevices(const DPCTLSyclDeviceRef DRef)
 
 
 cdef extern from "syclinterface/dpctl_sycl_device_manager.h":
@@ -238,6 +242,7 @@ cdef extern from "syclinterface/dpctl_sycl_device_manager.h":
     cdef DPCTLSyclContextRef DPCTLDeviceMgr_GetCachedContext(
         const DPCTLSyclDeviceRef DRef)
     cdef int64_t DPCTLDeviceMgr_GetRelativeId(const DPCTLSyclDeviceRef DRef)
+    cdef DPCTLDeviceVectorRef DPCTLDeviceMgr_GetCompositeDevices()
 
 
 cdef extern from "syclinterface/dpctl_sycl_device_selector_interface.h":
@@ -316,6 +321,8 @@ cdef extern from "syclinterface/dpctl_sycl_platform_interface.h":
         const DPCTLSyclPlatformRef)
     cdef DPCTLDeviceVectorRef DPCTLPlatform_GetDevices(
         const DPCTLSyclPlatformRef PRef, _device_type DTy)
+    cdef DPCTLDeviceVectorRef DPCTLPlatform_GetCompositeDevices(
+        const DPCTLSyclPlatformRef PRef)
 
 
 cdef extern from "syclinterface/dpctl_sycl_context_interface.h":

dpctl/_sycl_device.pyx

@@ -32,6 +32,8 @@ from ._backend cimport (  # noqa: E211
     DPCTLDevice_CreateSubDevicesEqually,
     DPCTLDevice_Delete,
     DPCTLDevice_GetBackend,
+    DPCTLDevice_GetComponentDevices,
+    DPCTLDevice_GetCompositeDevice,
     DPCTLDevice_GetDeviceType,
     DPCTLDevice_GetDriverVersion,
     DPCTLDevice_GetGlobalMemCacheLineSize,
@@ -795,6 +797,32 @@ cdef class SyclDevice(_SyclDevice):
         cdef _aspect_type AT = _aspect_type._emulated
         return DPCTLDevice_HasAspect(self._device_ref, AT)
 
+    @property
+    def has_aspect_is_component(self):
+        """ Returns ``True`` if this device is a component device, ``False``
+        otherwise. A device with this aspect will have a composite device
+        from which it is descended.
+
+        Returns:
+            bool:
+                Indicates if device is a component device.
+        """
+        cdef _aspect_type AT = _aspect_type._is_component
+        return DPCTLDevice_HasAspect(self._device_ref, AT)
+
+
+    @property
+    def has_aspect_is_composite(self):
+        """ Returns ``True`` if this device is a composite device, ``False``
+        otherwise. A device with this aspect contains component devices.
+
+        Returns:
+            bool:
+                Indicates if device is a composite device.
+        """
+        cdef _aspect_type AT = _aspect_type._is_composite
+        return DPCTLDevice_HasAspect(self._device_ref, AT)
+
     @property
     def image_2d_max_width(self):
         """ Returns the maximum width of a 2D image or 1D image in pixels.
@@ -1728,6 +1756,41 @@ cdef class SyclDevice(_SyclDevice):
             return None
         return SyclDevice._create(pDRef)
 
+    @property
+    def composite_device(self):
+        """ The composite device for a component device, or None for a non-component device.
+
+        Returns:
+            dpctl.SyclDevice:
+                The composite :class:`dpctl.SyclDevice` instance for a
+                component device, or ``None`` for a non-component device.
+        """
+        cdef DPCTLSyclDeviceRef CDRef = NULL
+        CDRef = DPCTLDevice_GetCompositeDevice(self._device_ref)
+        if (CDRef is NULL):
+            return None
+        return SyclDevice._create(CDRef)
+
+    def component_devices(self):
+        """ Returns a list of component devices contained in this SYCL device.
+
+        The returned list will be empty if this SYCL device is not a composite
+        device, i.e., if `is_composite` is ``False``.
+
+        Returns:
+            List[:class:`dpctl.SyclDevice`]:
+                List of component devices.
+
+        Raises:
+            dpctl.SyclSubdeviceCreationError:
+                if sub-devices can not be created.
+        """
+        cdef DPCTLDeviceVectorRef cDVRef = NULL
+        cDVRef = DPCTLDevice_GetComponentDevices(self._device_ref)
+        if cDVRef is NULL:
+            raise ValueError("Internal error: NULL device vector encountered")
+        return _get_devices(cDVRef)
+
     @property
     def profiling_timer_resolution(self):
         """ Profiling timer resolution.

dpctl/_sycl_device_factory.pxd

@@ -32,6 +32,7 @@ cpdef SyclDevice select_cpu_device()
 cpdef SyclDevice select_default_device()
 cpdef SyclDevice select_gpu_device()
 cpdef list get_devices(backend=*, device_type=*)
+cpdef list get_composite_devices()
 cpdef int get_num_devices(backend=*, device_type=*)
 cpdef cpp_bool has_gpu_devices()
 cpdef cpp_bool has_cpu_devices()

dpctl/_sycl_device_factory.pyx

@@ -31,6 +31,7 @@ from ._backend cimport (  # noqa: E211
     DPCTLCPUSelector_Create,
     DPCTLDefaultSelector_Create,
     DPCTLDevice_CreateFromSelector,
+    DPCTLDeviceMgr_GetCompositeDevices,
     DPCTLDeviceMgr_GetDevices,
     DPCTLDeviceMgr_GetNumDevices,
     DPCTLDeviceSelector_Delete,
@@ -62,6 +63,7 @@ __all__ = [
     "has_gpu_devices",
     "has_accelerator_devices",
     "_cached_default_device",
+    "get_composite_devices",
 ]
 
 
@@ -202,6 +204,32 @@ cpdef list get_devices(backend=backend_type.all, device_type=device_type_t.all):
     return devices
 
 
+cpdef list get_composite_devices():
+    """
+    Returns a list of the available composite :class:`dpctl.SyclDevice`
+    instances.
+
+    Only available when `ZE_FLAT_DEVICE_HIERARCHY=COMBINED` is set in
+    the environment, and only for specific Level Zero devices
+    (i.e., those which expose multiple tiles as root devices).
+
+    For more information, see:
+    https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/experimental/sycl_ext_oneapi_composite_device.asciidoc
+
+    Returns:
+        list:
+            A list of available composite :class:`dpctl.SyclDevice` instances.
+    """
+    cdef DPCTLDeviceVectorRef DVRef = NULL
+    cdef list composite_devices
+
+    DVRef = DPCTLDeviceMgr_GetCompositeDevices()
+    composite_devices = _get_devices(DVRef)
+    DPCTLDeviceVector_Delete(DVRef)
+
+    return composite_devices
+
+
 cpdef int get_num_devices(
     backend=backend_type.all, device_type=device_type_t.all
 ):

dpctl/_sycl_platform.pyx

@@ -37,6 +37,7 @@ from ._backend cimport (  # noqa: E211
     DPCTLPlatform_CreateFromSelector,
     DPCTLPlatform_Delete,
     DPCTLPlatform_GetBackend,
+    DPCTLPlatform_GetCompositeDevices,
     DPCTLPlatform_GetDefaultContext,
     DPCTLPlatform_GetDevices,
     DPCTLPlatform_GetName,
@@ -448,6 +449,41 @@ cdef class SyclPlatform(_SyclPlatform):
 
         return devices
 
+    def get_composite_devices(self):
+        """
+        Returns the list of composite :class:`dpctl.SyclDevice` objects
+        associated with :class:`dpctl.SyclPlatform` instance.
+
+        Returns:
+            list:
+                A :obj:`list` of composite :class:`dpctl.SyclDevice` objects
+                that belong to this platform.
+
+        Raises:
+            TypeError:
+                If `device_type` is not a str or :class:`dpctl.device_type`
+                enum.
+            ValueError:
+                If the ``DPCTLPlatform_GetCompositeDevices`` call returned
+                ``NULL`` instead of a ``DPCTLDeviceVectorRef`` object.
+        """
+        cdef DPCTLDeviceVectorRef DVRef = NULL
+        cdef size_t num_devs
+        cdef size_t i
+        cdef DPCTLSyclDeviceRef DRef
+
+        DVRef = DPCTLPlatform_GetCompositeDevices(self.get_platform_ref())
+        if (DVRef is NULL):
+            raise ValueError("Internal error: NULL device vector encountered")
+        num_devs = DPCTLDeviceVector_Size(DVRef)
+        composite_devices = []
+        for i in range(num_devs):
+            DRef = DPCTLDeviceVector_GetAt(DVRef, i)
+            composite_devices.append(SyclDevice._create(DRef))
+        DPCTLDeviceVector_Delete(DVRef)
+
+        return composite_devices
+
 
 def lsplatform(verbosity=0):
     """

dpctl/tests/test_sycl_device.py

@@ -165,6 +165,8 @@ def test_equal():
     "usm_atomic_host_allocations",
     "usm_atomic_shared_allocations",
     "emulated",
+    "is_component",
+    "is_composite",
 ]
 
 # SYCL 2020 spec aspects not presently
@@ -310,3 +312,30 @@ def test_get_unpartitioned_parent_device_from_sub_device():
         pytest.skip("Default device can not be partitioned")
     assert isinstance(sdevs, list) and len(sdevs) > 0
     assert dev == sdevs[0].get_unpartitioned_parent_device()
+
+
+def test_composite_device_method():
+    """
+    Test that the composite_device method returns a composite
+    device found in ``dpctl.get_composite_devices()``
+    """
+    devices = dpctl.get_devices()
+    composite_devices = dpctl.get_composite_devices()
+    for d in devices:
+        if d.has_aspect_is_component:
+            Cd = d.composite_device
+            assert Cd in composite_devices
+
+
+def test_get_component_devices_from_composite():
+    """
+    Test that the component_devices method returns component
+    root devices.
+    """
+    devices = dpctl.get_devices()
+    composite_devices = dpctl.get_composite_devices()
+    for Cd in composite_devices:
+        component_devices = Cd.component_devices()
+        for d in component_devices:
+            # component devices are root devices
+            assert d in devices

dpctl/tests/test_sycl_device_factory.py

@@ -185,3 +185,13 @@ def test_get_devices_with_device_type_str(device_type_str):
             assert dev == devices[i]
     else:
         pytest.skip()
+
+
+def test_get_composite_devices():
+    devices = dpctl.get_composite_devices()
+    if devices:
+        num_devices = len(devices)
+        for i in range(num_devices):
+            assert devices[i].has_aspect_is_composite
+    else:
+        pytest.skip()

dpctl/tests/test_sycl_platform.py

@@ -259,3 +259,20 @@ def test_platform_get_devices_enum_device_type():
                 devices = p.get_devices(device_type=dty)
                 if len(devices):
                     assert (d.device_type == dty for d in devices)
+
+
+def test_platform_get_composite_devices():
+    platforms = dpctl.get_platforms()
+    if platforms:
+        for p in platforms:
+            composite_devices = p.get_composite_devices()
+            if not composite_devices:
+                pass
+            devices = p.get_devices()
+            for Cd in composite_devices:
+                assert Cd.has_aspect_is_composite
+                component_devices = Cd.component_devices()
+                # all component devices are root devices
+                assert all(d in devices for d in component_devices)
+    else:
+        pytest.skip("No platforms available")

libsyclinterface/helper/source/dpctl_utils_helper.cpp

@@ -218,6 +218,12 @@ std::string DPCTL_AspectToStr(aspect aspectTy)
     case aspect::emulated:
         ss << "emulated";
         break;
+    case aspect::ext_oneapi_is_component:
+        ss << "is_component";
+        break;
+    case aspect::ext_oneapi_is_composite:
+        ss << "is_composite";
+        break;
     default:
         throw std::runtime_error("Unsupported aspect type");
     }
@@ -287,6 +293,12 @@ aspect DPCTL_StrToAspectType(const std::string &aspectTyStr)
     else if (aspectTyStr == "emulated") {
         aspectTy = aspect::emulated;
     }
+    else if (aspectTyStr == "is_component") {
+        aspectTy = aspect::ext_oneapi_is_component;
+    }
+    else if (aspectTyStr == "is_composite") {
+        aspectTy = aspect::ext_oneapi_is_composite;
+    }
     else {
         // \todo handle the error
         throw std::runtime_error("Unsupported aspect type");
@@ -335,6 +347,10 @@ aspect DPCTL_DPCTLAspectTypeToSyclAspect(DPCTLSyclAspectType AspectTy)
         return aspect::host_debuggable;
     case DPCTLSyclAspectType::emulated:
         return aspect::emulated;
+    case DPCTLSyclAspectType::is_component:
+        return aspect::ext_oneapi_is_component;
+    case DPCTLSyclAspectType::is_composite:
+        return aspect::ext_oneapi_is_composite;
     default:
         throw std::runtime_error("Unsupported aspect type");
     }
@@ -381,6 +397,10 @@ DPCTLSyclAspectType DPCTL_SyclAspectToDPCTLAspectType(aspect Aspect)
         return DPCTLSyclAspectType::host_debuggable;
     case aspect::emulated:
         return DPCTLSyclAspectType::emulated;
+    case aspect::ext_oneapi_is_composite:
+        return DPCTLSyclAspectType::is_composite;
+    case aspect::ext_oneapi_is_component:
+        return DPCTLSyclAspectType::is_component;
     default:
         throw std::runtime_error("Unsupported aspect type");
     }

libsyclinterface/include/syclinterface/dpctl_sycl_device_interface.h

@@ -766,4 +766,30 @@ __dpctl_keep size_t *
 DPCTLDevice_GetSubGroupSizes(__dpctl_keep const DPCTLSyclDeviceRef DRef,
                              size_t *res_len);
 
+/*!
+ * @brief Wrapper over
+ * device.get_info<info::device::parent_device>
+ *
+ * @param    DRef           Opaque pointer to a sycl::device
+ * @return   Returns an opaque pointer to the composite device for a
+ * component device, or nullptr if the device is not a component device.
+ */
+DPCTL_API
+__dpctl_give DPCTLSyclDeviceRef
+DPCTLDevice_GetCompositeDevice(__dpctl_keep const DPCTLSyclDeviceRef DRef);
+
+/*!
+ * @brief Returns a vector of component devices that are contained by the
+ * provided composite device. If the device is not a composite device,
+ * returns an empty vector.
+ *
+ * @param    DRef         Opaque pointer to a ``sycl::device``
+ * @return   A #DPCTLDeviceVectorRef containing component
+ * #DPCTLSyclDeviceRef objects
+ * @ingroup DeviceInterface
+ */
+DPCTL_API
+__dpctl_give DPCTLDeviceVectorRef
+DPCTLDevice_GetComponentDevices(__dpctl_keep const DPCTLSyclDeviceRef DRef);
+
 DPCTL_C_EXTERN_C_END