Merge pull request #623 from diptorupd/fix/issue_601

Correct the logic to return device count.

Author: oleksandr-pavlyk

dpctl-capi/source/dpctl_sycl_device_manager.cpp

@@ -66,6 +66,39 @@ std::string get_device_info_str(const device &Device)
     return ss.str();
 }
 
+/*!
+ * @brief Canonicalizes a device identifier bit flag to have a valid (i.e., not
+ * UNKNOWN) backend and device type bits.
+ *
+ * The device id is bit flag that indicates the backend and device type, both
+ * of which are optional, that are to be queried. The function makes sure if a
+ * device identifier only provides a device type value the backend is set to
+ * DPCTL_ALL_BACKENDS. Similarly, if only backend is provided the device type
+ * is set to DPCTL_ALL.
+ *
+ * @param    device_id     A bit flag storing a backend and a device type value.
+ * @return   Canonicalized bit flag that makes sure neither backend nor device
+ * type is UNKNOWN (0). For cases where the input device id does not provide
+ * either one of the values, we set the value to ALL.
+ */
+int to_canonical_device_id(int device_id)
+{ // If the identifier is 0 (UNKNOWN_DEVICE) return 0.
+    if (!device_id)
+        return 0;
+
+    // Check if the device identifier has a backend specified. If not, then
+    // toggle all backend specifier bits, i.e. set the backend to
+    // DPCTL_ALL_BACKENDS.
+    if (!(device_id & DPCTL_ALL_BACKENDS))
+        device_id |= DPCTL_ALL_BACKENDS;
+
+    // Check if a device type was specified. If not, set device type to ALL.
+    if (!(device_id & ~DPCTL_ALL_BACKENDS))
+        device_id |= DPCTL_ALL;
+
+    return device_id;
+}
+
 struct DeviceCacheBuilder
 {
     using DeviceCache = std::unordered_map<device, context>;
@@ -146,12 +179,18 @@ DPCTLDeviceMgr_GetDevices(int device_identifier)
 {
     std::vector<DPCTLSyclDeviceRef> *Devices = nullptr;
 
+    device_identifier = to_canonical_device_id(device_identifier);
+
     try {
         Devices = new std::vector<DPCTLSyclDeviceRef>();
     } catch (std::bad_alloc const &ba) {
         delete Devices;
         return nullptr;
     }
+
+    if (!device_identifier)
+        return wrap(Devices);
+
     const auto &root_devices = device::get_devices();
     default_selector mRanker;
 
@@ -195,6 +234,10 @@ int DPCTLDeviceMgr_GetPositionInDevices(__dpctl_keep DPCTLSyclDeviceRef DRef,
         return not_found;
     }
 
+    device_identifier = to_canonical_device_id(device_identifier);
+    if (!device_identifier)
+        return not_found;
+
     const auto &root_devices = device::get_devices();
     default_selector mRanker;
     int index = not_found;
@@ -224,6 +267,11 @@ size_t DPCTLDeviceMgr_GetNumDevices(int device_identifier)
 {
     size_t nDevices = 0;
     auto &cache = DeviceCacheBuilder::getDeviceCache();
+
+    device_identifier = to_canonical_device_id(device_identifier);
+    if (!device_identifier)
+        return 0;
+
     for (const auto &entry : cache) {
         auto Bty(DPCTL_SyclBackendToDPCTLBackendType(
             entry.first.get_platform().get_backend()));

dpctl-capi/tests/test_sycl_device_manager.cpp

@@ -137,6 +137,57 @@ INSTANTIATE_TEST_SUITE_P(
                       DPCTLSyclBackendType::DPCTL_OPENCL |
                           DPCTLSyclDeviceType::DPCTL_GPU));
 
+struct TestGetNumDevicesForDTy : public ::testing::TestWithParam<int>
+{
+    size_t nDevices = 0;
+    TestGetNumDevicesForDTy()
+    {
+        cl::sycl::info::device_type sycl_dty =
+            DPCTL_DPCTLDeviceTypeToSyclDeviceType(
+                DPCTLSyclDeviceType(GetParam()));
+
+        auto devices = cl::sycl::device::get_devices(sycl_dty);
+        EXPECT_TRUE(devices.size() == DPCTLDeviceMgr_GetNumDevices(GetParam()));
+    }
+};
+
+INSTANTIATE_TEST_SUITE_P(
+    GetDevices,
+    TestGetNumDevicesForDTy,
+    ::testing::Values(DPCTLSyclDeviceType::DPCTL_ACCELERATOR,
+                      DPCTLSyclDeviceType::DPCTL_ALL,
+                      DPCTLSyclDeviceType::DPCTL_CPU,
+                      DPCTLSyclDeviceType::DPCTL_GPU,
+                      DPCTLSyclDeviceType::DPCTL_HOST_DEVICE));
+
+struct TestGetNumDevicesForBTy : public ::testing::TestWithParam<int>
+{
+    size_t nDevices = 0;
+    TestGetNumDevicesForBTy()
+    {
+        cl::sycl::backend sycl_bty = DPCTL_DPCTLBackendTypeToSyclBackend(
+            DPCTLSyclBackendType(GetParam()));
+
+        auto platforms = cl::sycl::platform::get_platforms();
+        for (const auto &P : platforms) {
+            if (P.get_backend() == sycl_bty) {
+                auto devices = P.get_devices();
+                EXPECT_TRUE(devices.size() ==
+                            DPCTLDeviceMgr_GetNumDevices(GetParam()));
+            }
+        }
+    }
+};
+
+INSTANTIATE_TEST_SUITE_P(
+    GetDevices,
+    TestGetNumDevicesForBTy,
+    ::testing::Values(DPCTLSyclBackendType::DPCTL_CUDA,
+                      DPCTLSyclBackendType::DPCTL_ALL_BACKENDS,
+                      DPCTLSyclBackendType::DPCTL_HOST,
+                      DPCTLSyclBackendType::DPCTL_LEVEL_ZERO,
+                      DPCTLSyclBackendType::DPCTL_OPENCL));
+
 struct TestDPCTLDeviceVector : public ::testing::Test
 {
 };

dpctl/_sycl_device.pyx

@@ -351,7 +351,7 @@ cdef class SyclDevice(_SyclDevice):
         elif DTy == _device_type._GPU:
             return device_type.gpu
         elif DTy == _device_type._HOST_DEVICE:
-            return device_type.host_device
+            return device_type.host
         else:
             raise ValueError("Unknown device type.")
 
@@ -1014,8 +1014,7 @@ cdef class SyclDevice(_SyclDevice):
         cdef int64_t relId = -1
 
         DTy = DPCTLDevice_GetDeviceType(self._device_ref)
-        relId = DPCTLDeviceMgr_GetPositionInDevices(
-            self._device_ref, _backend_type._ALL_BACKENDS | DTy)
+        relId = DPCTLDeviceMgr_GetPositionInDevices(self._device_ref, DTy)
         return relId
 
     cdef int get_backend_ordinal(self):
@@ -1032,8 +1031,7 @@ cdef class SyclDevice(_SyclDevice):
         cdef int64_t relId = -1
 
         BTy = DPCTLDevice_GetBackend(self._device_ref)
-        relId = DPCTLDeviceMgr_GetPositionInDevices(
-            self._device_ref, BTy | _device_type._ALL_DEVICES)
+        relId = DPCTLDeviceMgr_GetPositionInDevices(self._device_ref, BTy)
         return relId
 
     cdef int get_overall_ordinal(self):

dpctl/_sycl_device_factory.pyx

@@ -94,7 +94,7 @@ cdef _device_type _string_to_dpctl_sycl_device_ty(str dty_str):
         return _device_type._CUSTOM
     elif dty_str == "gpu":
         return _device_type._GPU
-    elif dty_str == "host_device":
+    elif dty_str == "host":
         return _device_type._HOST_DEVICE
     else:
         return _device_type._UNKNOWN_DEVICE
@@ -128,7 +128,7 @@ cdef _device_type _enum_to_dpctl_sycl_device_ty(DTy):
         return _device_type._CUSTOM
     elif DTy == device_type_t.gpu:
         return _device_type._GPU
-    elif DTy == device_type_t.host_device:
+    elif DTy == device_type_t.host:
         return _device_type._HOST_DEVICE
     else:
         return _device_type._UNKNOWN_DEVICE
@@ -164,7 +164,7 @@ cpdef list get_devices(backend=backend_type.all, device_type=device_type_t.all):
         device_type (optional): Defaults to ``dpctl.device_type.all``.
             A :class:`dpctl.device_type` enum value or a string that
             specifies a SYCL device type. Currently, accepted values are:
-            "gpu", "cpu", "accelerator", "host_device", or "all".
+            "gpu", "cpu", "accelerator", "host", or "all".
     Returns:
         list: A list of available :class:`dpctl.SyclDevice` instances that
         satisfy the provided :class:`dpctl.backend_type` and
@@ -217,7 +217,7 @@ cpdef int get_num_devices(
         device_type (optional): Defaults to ``dpctl.device_type.all``.
             A :class:`dpctl.device_type` enum value or a string that
             specifies a SYCL device type. Currently, accepted values are:
-            "gpu", "cpu", "accelerator", "host_device", or "all".
+            "gpu", "cpu", "accelerator", "host", or "all".
     Returns:
         int: The number of available SYCL devices that satisfy the provided
         :class:`dpctl.backend_type` and :class:`dpctl.device_type` values.

dpctl/enum_types.py

@@ -39,7 +39,7 @@ class device_type(Enum):
     gpu                  1
     cpu                  2
     accelerator          3
-    host_device          4
+    host                 4
     ==================   ============
     """
 
@@ -49,7 +49,7 @@ class device_type(Enum):
     cpu = auto()
     custom = auto()
     gpu = auto()
-    host_device = auto()
+    host = auto()
 
 
 class backend_type(Enum):

dpctl/tests/test_sycl_device_factory.py

@@ -27,7 +27,7 @@
     (bty.level_zero, dty.gpu),
     (bty.opencl, dty.gpu),
     (bty.opencl, dty.cpu),
-    (bty.host, dty.host_device),
+    (bty.host, dty.host),
 ]
 
 argument_list_2 = [