From ea64a3fbc355d5382087bf59083c8c64130f88fa Mon Sep 17 00:00:00 2001
From: Mateusz Jablonski <mateusz.jablonski@intel.com>
Date: Mon, 23 Sep 2019 14:04:18 +0200
Subject: [PATCH] Add method for getting number of available devices

Related-To: NEO-3691

Change-Id: Ida2454adbfc9b177acbbb320f3f9577849a6fd74
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>
---
 runtime/device/device.h                | 1 +
 runtime/device/root_device.cpp         | 7 +++++++
 runtime/device/root_device.h           | 1 +
 runtime/device/sub_device.cpp          | 3 +++
 runtime/device/sub_device.h            | 1 +
 unit_tests/device/sub_device_tests.cpp | 5 +++++
 6 files changed, 18 insertions(+)

diff --git a/runtime/device/device.h b/runtime/device/device.h
index e1502e7255..556cfba91f 100644
--- a/runtime/device/device.h
+++ b/runtime/device/device.h
@@ -97,6 +97,7 @@ class Device : public BaseObject<_cl_device_id> {
     bool areSharedSystemAllocationsAllowed() const {
         return this->deviceInfo.sharedSystemMemCapabilities != 0u;
     }
+    virtual uint32_t getNumAvailableDevices() const = 0;
 
   protected:
     Device() = delete;
diff --git a/runtime/device/root_device.cpp b/runtime/device/root_device.cpp
index 69d8ddf4fe..a25c41a8f3 100644
--- a/runtime/device/root_device.cpp
+++ b/runtime/device/root_device.cpp
@@ -18,6 +18,13 @@ uint32_t RootDevice::getNumSubDevices() const {
     return static_cast<uint32_t>(subdevices.size());
 }
 
+uint32_t RootDevice::getNumAvailableDevices() const {
+    if (subdevices.empty()) {
+        return 1u;
+    }
+    return getNumSubDevices();
+}
+
 RootDevice::RootDevice(ExecutionEnvironment *executionEnvironment, uint32_t deviceIndex) : Device(executionEnvironment, deviceIndex) {}
 bool RootDevice::createDeviceImpl() {
     auto status = Device::createDeviceImpl();
diff --git a/runtime/device/root_device.h b/runtime/device/root_device.h
index f21b6f9806..93d2c9c0c9 100644
--- a/runtime/device/root_device.h
+++ b/runtime/device/root_device.h
@@ -17,6 +17,7 @@ class RootDevice : public Device {
     RootDevice(ExecutionEnvironment *executionEnvironment, uint32_t deviceIndex);
     ~RootDevice() override;
     bool createDeviceImpl() override;
+    uint32_t getNumAvailableDevices() const override;
     uint32_t getNumSubDevices() const;
 
     /* We hide the retain and release function of BaseObject. */
diff --git a/runtime/device/sub_device.cpp b/runtime/device/sub_device.cpp
index 6042346ebd..282d678c40 100644
--- a/runtime/device/sub_device.cpp
+++ b/runtime/device/sub_device.cpp
@@ -32,5 +32,8 @@ DeviceBitfield SubDevice::getDeviceBitfieldForOsContext() const {
     deviceBitfield.set(subDeviceIndex);
     return deviceBitfield;
 }
+uint32_t SubDevice::getNumAvailableDevices() const {
+    return 1u;
+}
 
 } // namespace NEO
diff --git a/runtime/device/sub_device.h b/runtime/device/sub_device.h
index 17ce25e9d9..81b78b86ce 100644
--- a/runtime/device/sub_device.h
+++ b/runtime/device/sub_device.h
@@ -17,6 +17,7 @@ class SubDevice : public Device {
     unique_ptr_if_unused<Device> release() override;
     void retainInternal();
     void releaseInternal();
+    uint32_t getNumAvailableDevices() const override;
 
   protected:
     DeviceBitfield getDeviceBitfieldForOsContext() const override;
diff --git a/unit_tests/device/sub_device_tests.cpp b/unit_tests/device/sub_device_tests.cpp
index b018ca3590..f2cfdeee88 100644
--- a/unit_tests/device/sub_device_tests.cpp
+++ b/unit_tests/device/sub_device_tests.cpp
@@ -19,6 +19,7 @@ TEST(SubDevicesTest, givenDefaultConfigWhenCreateRootDeviceThenItDoesntContainSu
     auto device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(*platformDevices));
 
     EXPECT_EQ(0u, device->getNumSubDevices());
+    EXPECT_EQ(1u, device->getNumAvailableDevices());
 }
 
 TEST(SubDevicesTest, givenCreateMultipleSubDevicesFlagSetWhenCreateRootDeviceThenItContainsSubDevices) {
@@ -31,6 +32,10 @@ TEST(SubDevicesTest, givenCreateMultipleSubDevicesFlagSetWhenCreateRootDeviceThe
     EXPECT_EQ(0u, device->getDeviceIndex());
     EXPECT_EQ(1u, device->subdevices.at(0)->getDeviceIndex());
     EXPECT_EQ(2u, device->subdevices.at(1)->getDeviceIndex());
+
+    EXPECT_EQ(2u, device->getNumAvailableDevices());
+    EXPECT_EQ(1u, device->subdevices.at(0)->getNumAvailableDevices());
+    EXPECT_EQ(1u, device->subdevices.at(1)->getNumAvailableDevices());
 }
 
 TEST(SubDevicesTest, givenDeviceWithSubDevicesWhenSubDeviceRefcountsAreChangedThenChangeIsPropagatedToRootDevice) {