Add method to get device from device

Related-To: NEO-3691

Change-Id: I710d740d82803e2a844b30a5cd3e4b017192f6f1
Signed-off-by: Mateusz Jablonski <[email protected]>

Author: JablonskiMateusz

runtime/device/device.h

@@ -98,6 +98,7 @@ class Device : public BaseObject<_cl_device_id> {
         return this->deviceInfo.sharedSystemMemCapabilities != 0u;
     }
     virtual uint32_t getNumAvailableDevices() const = 0;
+    virtual Device *getDeviceById(uint32_t deviceId) const = 0;
 
   protected:
     Device() = delete;

runtime/device/root_device.cpp

@@ -25,6 +25,14 @@ uint32_t RootDevice::getNumAvailableDevices() const {
     return getNumSubDevices();
 }
 
+Device *RootDevice::getDeviceById(uint32_t deviceId) const {
+    UNRECOVERABLE_IF(deviceId >= getNumAvailableDevices());
+    if (subdevices.empty()) {
+        return const_cast<RootDevice *>(this);
+    }
+    return subdevices[deviceId].get();
+};
+
 RootDevice::RootDevice(ExecutionEnvironment *executionEnvironment, uint32_t deviceIndex) : Device(executionEnvironment, deviceIndex) {}
 bool RootDevice::createDeviceImpl() {
     auto status = Device::createDeviceImpl();

runtime/device/root_device.h

@@ -19,6 +19,7 @@ class RootDevice : public Device {
     bool createDeviceImpl() override;
     uint32_t getNumAvailableDevices() const override;
     uint32_t getNumSubDevices() const;
+    Device *getDeviceById(uint32_t deviceId) const override;
 
     /* We hide the retain and release function of BaseObject. */
     void retain() override;

runtime/device/sub_device.cpp

@@ -35,5 +35,9 @@ DeviceBitfield SubDevice::getDeviceBitfieldForOsContext() const {
 uint32_t SubDevice::getNumAvailableDevices() const {
     return 1u;
 }
+Device *SubDevice::getDeviceById(uint32_t deviceId) const {
+    UNRECOVERABLE_IF(deviceId >= getNumAvailableDevices());
+    return const_cast<SubDevice *>(this);
+}
 
 } // namespace NEO

runtime/device/sub_device.h

@@ -18,6 +18,7 @@ class SubDevice : public Device {
     void retainInternal();
     void releaseInternal();
     uint32_t getNumAvailableDevices() const override;
+    Device *getDeviceById(uint32_t deviceId) const override;
 
   protected:
     DeviceBitfield getDeviceBitfieldForOsContext() const override;

unit_tests/device/sub_device_tests.cpp

@@ -100,3 +100,32 @@ TEST(SubDevicesTest, givenSubDeviceWhenOsContextIsCreatedThenItsBitfieldBasesOnS
     EXPECT_EQ(firstSubDeviceMask, static_cast<uint32_t>(firstSubDevice->getDefaultEngine().osContext->getDeviceBitfield().to_ulong()));
     EXPECT_EQ(secondSubDeviceMask, static_cast<uint32_t>(secondSubDevice->getDefaultEngine().osContext->getDeviceBitfield().to_ulong()));
 }
+
+TEST(SubDevicesTest, givenDeviceWithoutSubDevicesWhenGettingDeviceByIdZeroThenGetThisDevice) {
+    auto device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(*platformDevices));
+
+    EXPECT_EQ(1u, device->getNumAvailableDevices());
+    EXPECT_EQ(device.get(), device->getDeviceById(0u));
+    EXPECT_THROW(device->getDeviceById(1), std::exception);
+}
+
+TEST(SubDevicesTest, givenDeviceWithSubDevicesWhenGettingDeviceByIdThenGetCorrectSubDevice) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.CreateMultipleSubDevices.set(2);
+    auto device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(*platformDevices));
+    EXPECT_EQ(2u, device->getNumSubDevices());
+    EXPECT_EQ(device->subdevices.at(0).get(), device->getDeviceById(0));
+    EXPECT_EQ(device->subdevices.at(1).get(), device->getDeviceById(1));
+    EXPECT_THROW(device->getDeviceById(2), std::exception);
+}
+
+TEST(SubDevicesTest, givenSubDevicesWhenGettingDeviceByIdZeroThenGetThisSubDevice) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.CreateMultipleSubDevices.set(2);
+    auto device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(*platformDevices));
+    EXPECT_EQ(2u, device->getNumSubDevices());
+    auto subDevice = device->subdevices.at(0).get();
+
+    EXPECT_EQ(subDevice, subDevice->getDeviceById(0));
+    EXPECT_THROW(subDevice->getDeviceById(1), std::exception);
+}