Merge pull request #1005 from IntelPython/add-partition-max-sub-devices

Author: oleksandr-pavlyk

dpctl/_backend.pxd

@@ -206,6 +206,7 @@ cdef extern from "syclinterface/dpctl_sycl_device_interface.h":
         const DPCTLSyclDeviceRef DRef)
     cdef size_t *DPCTLDevice_GetSubGroupSizes(const DPCTLSyclDeviceRef DRef,
         size_t *res_len)
+    cdef uint32_t DPCTLDevice_GetPartitionMaxSubDevices(const DPCTLSyclDeviceRef DRef)
 
 
 cdef extern from "syclinterface/dpctl_sycl_device_manager.h":

dpctl/_sycl_device.pyx

@@ -55,6 +55,7 @@ from ._backend cimport (  # noqa: E211
     DPCTLDevice_GetMaxWriteImageArgs,
     DPCTLDevice_GetName,
     DPCTLDevice_GetParentDevice,
+    DPCTLDevice_GetPartitionMaxSubDevices,
     DPCTLDevice_GetPlatform,
     DPCTLDevice_GetPreferredVectorWidthChar,
     DPCTLDevice_GetPreferredVectorWidthDouble,
@@ -1292,9 +1293,25 @@ cdef class SyclDevice(_SyclDevice):
         """
         cdef uint64_t cache_line_sz = DPCTLDevice_GetGlobalMemCacheLineSize(
             self._device_ref
-	)
+	    )
         return cache_line_sz
 
+    @property
+    def partition_max_sub_devices(self):
+        """ The maximum number of sub-devices this :class:`dpctl.SyclDevice`
+        instance can be partitioned into. The value returned cannot exceed the
+        value returned by :attr:`dpctl.SyclDevice.max_compute_units`.
+
+        Returns:
+            int: The maximum number of sub-devices that can be created when this
+                device is partitioned. Zero value indicates that device can not
+                be partitioned.
+        """
+        cdef uint32_t max_part = DPCTLDevice_GetPartitionMaxSubDevices(
+            self._device_ref
+        )
+        return max_part
+
     cdef cpp_bool equals(self, SyclDevice other):
         """ Returns ``True`` if the :class:`dpctl.SyclDevice` argument has the
         same _device_ref as this SyclDevice.

dpctl/tests/_device_attributes_checks.py

@@ -540,6 +540,13 @@ def check_parent_device(device):
     assert pd is None or isinstance(pd, dpctl.SyclDevice)
 
 
+def check_partition_max_sub_devices(device):
+    max_part = device.partition_max_sub_devices
+    assert isinstance(max_part, int)
+    assert max_part >= 0
+    assert max_part <= device.max_compute_units
+
+
 def check_filter_string(device):
     try:
         fs = device.filter_string
@@ -670,6 +677,7 @@ def check_global_mem_cache_line_size(device):
     check_profiling_timer_resolution,
     check_platform,
     check_parent_device,
+    check_partition_max_sub_devices,
     check_filter_string,
     check_vendor,
     check_driver_version,

examples/pybind11/use_dpctl_syclqueue/tests/test_queue_device.py

@@ -63,3 +63,11 @@ def test_get_sub_group_sizes():
     assert type(szs) is list
     assert all(type(el) is int for el in szs)
     szs == d.sub_group_sizes
+
+
+def test_get_partition_max_sub_devices():
+    d = dpctl.SyclDevice()
+    mt = uqd.get_partition_max_sub_devices(d)
+    assert type(mt) is int
+    assert mt >= 0
+    assert mt <= d.max_compute_units

examples/pybind11/use_dpctl_syclqueue/use_queue_device/__init__.py

@@ -20,6 +20,7 @@
     get_device_global_mem_size,
     get_device_local_mem_size,
     get_max_compute_units,
+    get_partition_max_sub_devices,
     get_sub_group_sizes,
     offloaded_array_mod,
 )
@@ -30,6 +31,7 @@
     "get_device_local_mem_size",
     "offloaded_array_mod",
     "get_sub_group_sizes",
+    "get_partition_max_sub_devices",
 ]
 
 __doc__ = """

examples/pybind11/use_dpctl_syclqueue/use_queue_device/_example.cpp

@@ -90,6 +90,11 @@ std::vector<std::size_t> get_sub_group_sizes(const sycl::device &d)
     return d.get_info<sycl::info::device::sub_group_sizes>();
 }
 
+std::uint32_t get_partition_max_sub_devices(const sycl::device &d)
+{
+    return d.get_info<sycl::info::device::partition_max_sub_devices>();
+}
+
 PYBIND11_MODULE(_use_queue_device, m)
 {
     m.def(
@@ -108,4 +113,7 @@ PYBIND11_MODULE(_use_queue_device, m)
           "Compute offloaded modular reduction of integer-valued NumPy array");
     m.def("get_sub_group_sizes", &get_sub_group_sizes,
           "Gets info::device::sub_group_sizes property of given device");
+    m.def("get_partition_max_sub_devices", &get_partition_max_sub_devices,
+          "Gets info::device::partition_max_sub_devices for given "
+          "dpctl.SyclDevice");
 }

libsyclinterface/include/dpctl_sycl_device_interface.h

@@ -597,6 +597,18 @@ DPCTL_API
 __dpctl_give DPCTLSyclDeviceRef
 DPCTLDevice_GetParentDevice(__dpctl_keep const DPCTLSyclDeviceRef DRef);
 
+/*!
+ * @brief Wrapper over
+ * device.get_info<info::device::partition_max_sub_devices>
+ *
+ * @param    DRef           Opaque pointer to a sycl::device
+ * @return   Returns the maximum number of sub-devices that can be created
+ * when this device is partitioned.
+ */
+DPCTL_API
+uint32_t DPCTLDevice_GetPartitionMaxSubDevices(
+    __dpctl_keep const DPCTLSyclDeviceRef DRef);
+
 /*!
  * @brief Wrapper over
  * std::hash<sycl::device>'s operator()

libsyclinterface/source/dpctl_sycl_device_interface.cpp

@@ -573,6 +573,24 @@ DPCTLDevice_GetParentDevice(__dpctl_keep const DPCTLSyclDeviceRef DRef)
         return nullptr;
 }
 
+uint32_t DPCTLDevice_GetPartitionMaxSubDevices(
+    __dpctl_keep const DPCTLSyclDeviceRef DRef)
+{
+    auto D = unwrap<device>(DRef);
+    if (D) {
+        try {
+            uint32_t part_max_sub_devs =
+                D->get_info<info::device::partition_max_sub_devices>();
+            return part_max_sub_devs;
+        } catch (std::exception const &e) {
+            error_handler(e, __FILE__, __func__, __LINE__);
+            return 0;
+        }
+    }
+    else
+        return 0;
+}
+
 __dpctl_give DPCTLDeviceVectorRef
 DPCTLDevice_CreateSubDevicesEqually(__dpctl_keep const DPCTLSyclDeviceRef DRef,
                                     size_t count)

libsyclinterface/tests/test_sycl_device_interface.cpp

@@ -415,6 +415,17 @@ TEST_P(TestDPCTLSyclDeviceInterface, ChkGetParentDevice)
     EXPECT_TRUE(pDRef == nullptr);
 }
 
+TEST_P(TestDPCTLSyclDeviceInterface, ChkGetPartitionMaxSubDevices)
+{
+    uint32_t max_part = 0;
+    size_t max_cu = 0;
+    EXPECT_NO_FATAL_FAILURE(max_part =
+                                DPCTLDevice_GetPartitionMaxSubDevices(DRef));
+    EXPECT_TRUE(max_part >= 0);
+    EXPECT_NO_FATAL_FAILURE(max_cu = DPCTLDevice_GetMaxComputeUnits(DRef));
+    EXPECT_TRUE(max_part <= max_cu);
+}
+
 TEST_P(TestDPCTLSyclDeviceInterface, ChkGetProfilingTimerResolution)
 {
     size_t res = 0;
@@ -705,6 +716,14 @@ TEST_F(TestDPCTLSyclDeviceNullArgs, ChkGetParentDevice)
     ASSERT_TRUE(pDRef == nullptr);
 }
 
+TEST_F(TestDPCTLSyclDeviceNullArgs, ChkGetPartitionMaxSubDevices)
+{
+    uint32_t max_part = 0;
+    EXPECT_NO_FATAL_FAILURE(
+        max_part = DPCTLDevice_GetPartitionMaxSubDevices(Null_DRef));
+    ASSERT_TRUE(max_part == 0);
+}
+
 TEST_F(TestDPCTLSyclDeviceNullArgs, ChkCreateSubDevicesEqually)
 {
     DPCTLDeviceVectorRef DVRef = nullptr;