Add dpctl.get_composite_devices

This leverages oneAPI extension for composite devices to add the free function `ext_oneapi_get_composite_devices` to the main dpctl namespace

Author: ndgrigorian

dpctl/__init__.py

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

dpctl/_backend.pxd

@@ -237,6 +237,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":

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
 ):

libsyclinterface/include/syclinterface/dpctl_sycl_device_manager.h

@@ -173,6 +173,17 @@ DPCTL_API
 int64_t
 DPCTLDeviceMgr_GetRelativeId(__dpctl_keep const DPCTLSyclDeviceRef DRef);
 
+/*!
+ * @brief Returns a pointer to a std::vector<sycl::DPCTLSyclDeviceRef>
+ * containing the set of ::DPCTLSyclDeviceRef pointers to the
+ * available composite devices.
+ *
+ * @return   A #DPCTLDeviceVectorRef containing #DPCTLSyclDeviceRef objects
+ * that are composite devices.
+ */
+DPCTL_API
+__dpctl_give DPCTLDeviceVectorRef DPCTLDeviceMgr_GetCompositeDevices();
+
 /*! @} */
 
 DPCTL_C_EXTERN_C_END

libsyclinterface/source/dpctl_sycl_device_manager.cpp

@@ -112,7 +112,7 @@ struct DeviceCacheBuilder
      * avoid the performance overhead of context creation for every queue.
      *
      * The singleton pattern implemented here ensures that the map is
-     * created once in a thread-safe manner. Since, the map is ony read
+     * created once in a thread-safe manner. Since, the map is only read
      * post-creation we do not need any further protection to ensure
      * thread-safety.
      */
@@ -136,7 +136,7 @@ struct DeviceCacheBuilder
 
                     try {
                         // Per https://github.com/intel/llvm/blob/sycl/sycl/doc/
-                        // extensions/PlatformContext/PlatformContext.adoc
+                        // extensions/supported/sycl_ext_oneapi_default_context.asciidoc
                         // sycl::queue(D) would create default platform context
                         // for capable compiler, sycl::context(D) otherwise
                         auto Q = queue(D);
@@ -357,3 +357,35 @@ int64_t DPCTLDeviceMgr_GetRelativeId(__dpctl_keep const DPCTLSyclDeviceRef DRef)
 
     return -1;
 }
+
+/*!
+ * Returns a list of the available composite devices, or an empty list if
+ * there are none.
+ */
+__dpctl_give DPCTLDeviceVectorRef DPCTLDeviceMgr_GetCompositeDevices()
+{
+    using vecTy = std::vector<DPCTLSyclDeviceRef>;
+    vecTy *Devices = nullptr;
+
+    try {
+        Devices = new std::vector<DPCTLSyclDeviceRef>();
+    } catch (std::exception const &e) {
+        delete Devices;
+        error_handler(e, __FILE__, __func__, __LINE__);
+        return nullptr;
+    }
+
+    try {
+        auto composite_devices =
+            ext::oneapi::experimental::get_composite_devices();
+        Devices->reserve(composite_devices.size());
+        for (const auto &CDev : composite_devices) {
+            Devices->emplace_back(wrap<device>(new device(std::move(CDev))));
+        }
+        return wrap<vecTy>(Devices);
+    } catch (std::exception const &e) {
+        delete Devices;
+        error_handler(e, __FILE__, __func__, __LINE__);
+        return nullptr;
+    }
+}