dpctl/ Python API change to add support for subdevices from #326

Author: etotmeni

dpctl/_backend.pxd

@@ -104,6 +104,18 @@ cdef extern from "dpctl_sycl_enum_types.h":
     ctypedef _aspect_type DPCTLSyclAspectType
 
 
+    cdef enum _partition_affinity_domain_type 'DPCTLPartitionAffinityDomainType':
+        _not_applicable                     'not_applicable',
+        _numa                               'numa',
+        _L4_cache                           'L4_cache',
+        _L3_cache                           'L3_cache',
+        _L2_cache                           'L2_cache',
+        _L1_cache                           'L1_cache',
+        _next_partitionable                 'next_partitionable',
+
+    ctypedef _partition_affinity_domain_type DPCTLPartitionAffinityDomainType
+
+
 cdef extern from "dpctl_sycl_types.h":
     cdef struct DPCTLOpaqueSyclContext
     cdef struct DPCTLOpaqueSyclDevice
@@ -126,6 +138,11 @@ cdef extern from "dpctl_sycl_types.h":
     ctypedef DPCTLOpaqueSyclUSM            *DPCTLSyclUSMRef
 
 
+cdef extern from "dpctl_sycl_device_manager.h":
+    cdef struct DPCTLDeviceVector
+    ctypedef DPCTLDeviceVector *DPCTLDeviceVectorRef
+
+
 cdef extern from "dpctl_sycl_device_interface.h":
     cdef bool DPCTLDevice_AreEq(const DPCTLSyclDeviceRef DRef1,
                                 const DPCTLSyclDeviceRef DRef2)
@@ -163,11 +180,16 @@ cdef extern from "dpctl_sycl_device_interface.h":
     cdef uint32_t DPCTLDevice_GetPreferredVectorWidthHalf(const DPCTLSyclDeviceRef DRef)
     cpdef bool DPCTLDevice_HasAspect(
         const DPCTLSyclDeviceRef DRef, DPCTLSyclAspectType AT)
+    cdef DPCTLDeviceVectorRef DPCTLDevice_CreateSubDevicesEqually(
+        const DPCTLSyclDeviceRef DRef, size_t count)
+    cdef DPCTLDeviceVectorRef DPCTLDevice_CreateSubDevicesByCounts(
+        const DPCTLSyclDeviceRef DRef, size_t *counts, size_t ncounts)
+    cdef DPCTLDeviceVectorRef DPCTLDevice_CreateSubDevicesByAffinity(
+        const DPCTLSyclDeviceRef DRef,
+        DPCTLPartitionAffinityDomainType PartitionAffinityDomainTy)
 
 
 cdef extern from "dpctl_sycl_device_manager.h":
-    cdef struct DPCTLDeviceVector
-    ctypedef DPCTLDeviceVector *DPCTLDeviceVectorRef
     cdef DPCTLDeviceVectorRef DPCTLDeviceVector_CreateFromArray(
         size_t nelems,
         DPCTLSyclDeviceRef *elems)

dpctl/_sycl_device.pxd

@@ -23,6 +23,7 @@
 from ._backend cimport (
     DPCTLSyclDeviceRef,
     DPCTLSyclDeviceSelectorRef,
+    _partition_affinity_domain_type
 )
 
 
@@ -44,3 +45,7 @@ cdef class SyclDevice(_SyclDevice):
     cdef int _init_from__SyclDevice(self, _SyclDevice other)
     cdef int _init_from_selector(self, DPCTLSyclDeviceSelectorRef DSRef)
     cdef DPCTLSyclDeviceRef get_device_ref(self)
+    cdef _raise_sub_devices_creation_error(self, fname, errcode)
+    cdef list create_sub_devices_equally(self, size_t count)
+    cdef list create_sub_devices_by_counts(self, list counts)
+    cdef list create_sub_devices_by_affinity(self, _partition_affinity_domain_type domain)

dpctl/_sycl_device.pyx

@@ -24,11 +24,16 @@ from ._backend cimport (
     _aspect_type,
     _backend_type,
     _device_type,
+    _partition_affinity_domain_type,
     DPCTLCString_Delete,
     DPCTLDefaultSelector_Create,
     DPCTLDevice_Copy,
     DPCTLDevice_CreateFromSelector,
     DPCTLDevice_Delete,
+    DPCTLDeviceVectorRef,
+    DPCTLDeviceVector_Delete,
+    DPCTLDeviceVector_GetAt,
+    DPCTLDeviceVector_Size,
     DPCTLDevice_GetBackend,
     DPCTLDevice_GetDeviceType,
     DPCTLDevice_GetDriverInfo,
@@ -61,16 +66,29 @@ from ._backend cimport (
     DPCTLDevice_GetPreferredVectorWidthFloat,
     DPCTLDevice_GetPreferredVectorWidthDouble,
     DPCTLDevice_GetPreferredVectorWidthHalf,
+    DPCTLDevice_CreateSubDevicesEqually,
+    DPCTLDevice_CreateSubDevicesByCounts,
+    DPCTLDevice_CreateSubDevicesByAffinity,
 )
 from . import backend_type, device_type
 from libc.stdint cimport uint32_t
 import warnings
+from libc.stdlib cimport malloc, free
 
 __all__ = [
     "SyclDevice",
 ]
 
 
+cdef class SubDeviceCreationError(Exception):
+    """
+    A SubDeviceCreationError exception is raised when
+    sub-devices were not created.
+
+    """
+    pass
+
+
 cdef class _SyclDevice:
     """ A helper metaclass to abstract a cl::sycl::device instance.
     """
@@ -83,6 +101,19 @@ cdef class _SyclDevice:
         DPCTLSize_t_Array_Delete(self._max_work_item_sizes)
 
 
+cdef list _get_devices(DPCTLDeviceVectorRef DVRef):
+    cdef list devices = []
+    cdef size_t nelems = 0
+    if DVRef:
+        nelems = DPCTLDeviceVector_Size(DVRef)
+        for i in range(0, nelems):
+            DRef = DPCTLDeviceVector_GetAt(DVRef, i)
+            D = SyclDevice._create(DRef)
+            devices.append(D)
+
+    return devices
+
+
 cdef class SyclDevice(_SyclDevice):
     """ Python equivalent for cl::sycl::device class.
 
@@ -167,6 +198,87 @@ cdef class SyclDevice(_SyclDevice):
             SyclDevice._init_helper(self, DRef)
             return 0
 
+    cdef _raise_sub_devices_creation_error(self, fname, errcode):
+        e = SubDeviceCreationError("Sub-devices were not created.")
+        e.fname = fname
+        e.code = errcode
+        raise e
+
+    cdef list create_sub_devices_equally(self, size_t count):
+        """ Returns a vector of sub devices partitioned from this SYCL device
+            based on the count parameter. The returned
+            vector contains as many sub devices as can be created such that each sub
+            device contains count compute units. If the device’s total number of compute
+            units is not evenly divided by count, then the remaining compute units are
+            not included in any of the sub devices.
+        """
+        cdef DPCTLDeviceVectorRef DVRef = NULL
+        DVRef = DPCTLDevice_CreateSubDevicesEqually(self._device_ref, count)
+        if DVRef is NULL:
+            self._raise_sub_devices_creation_error("DPCTLSubDeviceCreationError", -1)
+        cdef list devices = _get_devices(DVRef)
+        DPCTLDeviceVector_Delete(DVRef)
+        return devices
+
+    cdef list create_sub_devices_by_counts(self, list counts):
+        """ Returns a vector of sub devices
+            partitioned from this SYCL device based on the counts parameter. For each
+            non-zero value M in the counts vector, a sub device with M compute units
+            is created.
+        """
+        cdef size_t ncounts = len(counts)
+        cdef size_t *counts_buff = <size_t *> malloc(ncounts * sizeof(size_t))
+        cdef DPCTLDeviceVectorRef DVRef = NULL
+        cdef int i
+        for i in range(ncounts):
+            counts_buff[i] = counts[i]
+        DVRef = DPCTLDevice_CreateSubDevicesByCounts(self._device_ref, counts_buff, ncounts)
+        if DVRef is NULL:
+            self._raise_sub_devices_creation_error("DPCTLSubDeviceCreationError", -1)
+        cdef list devices = _get_devices(DVRef)
+        free(counts_buff)
+        DPCTLDeviceVector_Delete(DVRef)
+        return devices
+
+    cdef list create_sub_devices_by_affinity(self, _partition_affinity_domain_type domain):
+        """ Returns a vector of sub devices
+            partitioned from this SYCL device by affinity domain based on the domain
+            parameter.
+        """
+        cdef DPCTLDeviceVectorRef DVRef = NULL
+        DVRef = DPCTLDevice_CreateSubDevicesByAffinity(self._device_ref, domain)
+        if DVRef is NULL:
+            self._raise_sub_devices_creation_error("DPCTLSubDeviceCreationError", -1)
+        cdef list devices = _get_devices(DVRef)
+        DPCTLDeviceVector_Delete(DVRef)
+        return devices
+
+    def create_sub_devices(self, partition):
+        if isinstance(partition, int) and partition > 0:
+            return self.create_sub_devices_equally(partition)
+        elif isinstance(partition, list) and all([i > 0 for i in partition]):
+            return self.create_sub_devices_by_counts(partition)
+        elif isinstance(partition, str):
+            if partition == "not_applicable":
+                domain_type = _partition_affinity_domain_type._not_applicable
+            elif partition == "numa":
+                domain_type = _partition_affinity_domain_type._numa
+            elif partition == "L4_cache":
+                domain_type = _partition_affinity_domain_type._L4_cache
+            elif partition == "L3_cache":
+                domain_type = _partition_affinity_domain_type._L3_cache
+            elif partition == "L2_cache":
+                domain_type = _partition_affinity_domain_type._L2_cache
+            elif partition == "L1_cache":
+                domain_type = _partition_affinity_domain_type._L1_cache
+            elif partition == "next_partitionable":
+                domain_type = _partition_affinity_domain_type._next_partitionable
+            else:
+                raise Exception('Unsupported type of domain')
+            return self.create_sub_devices_by_affinity(domain_type)
+        else:
+            raise Exception('Unsupported type of sub-device argument')
+
     def __cinit__(self, arg=None):
         cdef DPCTLSyclDeviceSelectorRef DSRef = NULL
         cdef const char *filter_c_str = NULL

dpctl/tests/test_sycl_device.py

@@ -19,6 +19,7 @@
 
 import dpctl
 import pytest
+from dpctl._sycl_device import SubDeviceCreationError
 
 list_of_standard_selectors = [
     dpctl.select_accelerator_device,
@@ -296,6 +297,89 @@ def check_get_preferred_vector_width_half(device):
         pytest.fail("preferred_vector_width_half call failed")
 
 
+def check_create_sub_devices_equally(device):
+    try:
+        n = int(device.max_compute_units / 2)
+        device.create_sub_devices(n)
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_counts(device):
+    try:
+        n = device.max_compute_units / 2
+        device.create_sub_devices([n, n])
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_not_applicable(device):
+    try:
+        device.create_sub_devices("not_applicable")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_numa(device):
+    try:
+        device.create_sub_devices("numa")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_L4_cache(device):
+    try:
+        device.create_sub_devices("L4_cache")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_L3_cache(device):
+    try:
+        device.create_sub_devices("L3_cache")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_L2_cache(device):
+    try:
+        device.create_sub_devices("L2_cache")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_L1_cache(device):
+    try:
+        device.create_sub_devices("L1_cache")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
+def check_create_sub_devices_by_affinity_next_partitionable(device):
+    try:
+        device.create_sub_devices("next_partitionable")
+    except SubDeviceCreationError:
+        pytest.skip("create_sub_devices can't create sub-devices on this device")
+    except Exception:
+        pytest.fail("create_sub_devices failed")
+
+
 def check_print_device_info(device):
     try:
         device.print_device_info()
@@ -339,6 +423,15 @@ def check_print_device_info(device):
     check_has_aspect_usm_shared_allocations,
     check_has_aspect_usm_restricted_shared_allocations,
     check_has_aspect_usm_system_allocator,
+    check_create_sub_devices_equally,
+    check_create_sub_devices_by_counts,
+    check_create_sub_devices_by_affinity_not_applicable,
+    check_create_sub_devices_by_affinity_numa,
+    check_create_sub_devices_by_affinity_L4_cache,
+    check_create_sub_devices_by_affinity_L3_cache,
+    check_create_sub_devices_by_affinity_L2_cache,
+    check_create_sub_devices_by_affinity_L1_cache,
+    check_create_sub_devices_by_affinity_next_partitionable,
     check_print_device_info,
 ]