Add support for new affinity mask format (2)

Change-Id: I8a1d277ff86b0fbdcde96d4d02d6a235a56e197b
Signed-off: Jaime Arteaga <jaime.a.arteaga.molina@intel.com>

Author: Jaime Arteaga

level_zero/core/source/driver/driver_handle_imp.cpp

@@ -153,17 +153,21 @@ uint32_t DriverHandleImp::parseAffinityMask(std::vector<std::unique_ptr<NEO::Dev
         size_t posNextComma = this->affinityMaskString.find_first_of(",", pos);
         std::string rootDeviceString = this->affinityMaskString.substr(pos, std::min(posNextDot, posNextComma) - pos);
         uint32_t rootDeviceIndex = static_cast<uint32_t>(std::stoul(rootDeviceString, nullptr, 0));
-        pos += rootDeviceString.size();
-        if (posNextDot != std::string::npos &&
-            this->affinityMaskString.at(pos) == '.' && posNextDot < posNextComma) {
-            pos++;
-            std::string subDeviceString = this->affinityMaskString.substr(pos, posNextComma - pos);
-            uint32_t subDeviceIndex = static_cast<uint32_t>(std::stoul(subDeviceString, nullptr, 0));
-            affinityMaskBitSet[rootDeviceIndex][subDeviceIndex] = true;
-        } else {
-            std::fill(affinityMaskBitSet[rootDeviceIndex].begin(),
-                      affinityMaskBitSet[rootDeviceIndex].end(),
-                      true);
+        if (rootDeviceIndex < neoDevices.size()) {
+            pos += rootDeviceString.size();
+            if (posNextDot != std::string::npos &&
+                this->affinityMaskString.at(pos) == '.' && posNextDot < posNextComma) {
+                pos++;
+                std::string subDeviceString = this->affinityMaskString.substr(pos, posNextComma - pos);
+                uint32_t subDeviceIndex = static_cast<uint32_t>(std::stoul(subDeviceString, nullptr, 0));
+                if (subDeviceIndex < neoDevices[rootDeviceIndex]->getNumAvailableDevices()) {
+                    affinityMaskBitSet[rootDeviceIndex][subDeviceIndex] = true;
+                }
+            } else {
+                std::fill(affinityMaskBitSet[rootDeviceIndex].begin(),
+                          affinityMaskBitSet[rootDeviceIndex].end(),
+                          true);
+            }
         }
         if (posNextComma == std::string::npos) {
             break;

level_zero/core/test/unit_tests/sources/driver/test_driver.cpp

@@ -527,5 +527,85 @@ INSTANTIATE_TEST_SUITE_P(DriverTestMultipleDeviceWithAffinityMaskTests,
                          ::testing::Combine(
                              ::testing::ValuesIn(maskArray.masks),
                              ::testing::ValuesIn(maskArray.masks)));
+
+TEST_F(DriverTestMultipleDeviceWithAffinityMask,
+       whenSettingAffinityMaskWithDeviceLargerThanAvailableDevicesThenRootDeviceValueIsIgnored) {
+    L0::DriverHandleImp *driverHandle = new DriverHandleImp;
+
+    constexpr uint32_t totalRootDevices = 2;
+    uint32_t subDevice1Index = 0;
+    driverHandle->affinityMaskString = "0,23,1." + std::to_string(subDevice1Index);
+
+    ze_result_t res = driverHandle->initialize(std::move(devices));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    uint32_t deviceCount = 0;
+    res = zeDeviceGet(driverHandle->toHandle(), &deviceCount, nullptr);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_EQ(deviceCount, totalRootDevices);
+
+    std::vector<ze_device_handle_t> phDevices(deviceCount);
+    res = zeDeviceGet(driverHandle->toHandle(), &deviceCount, phDevices.data());
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    for (uint32_t i = 0; i < deviceCount; i++) {
+        ze_device_handle_t hDevice = phDevices[i];
+        EXPECT_NE(nullptr, hDevice);
+
+        DeviceImp *device = reinterpret_cast<DeviceImp *>(L0::Device::fromHandle(hDevice));
+
+        uint32_t subDeviceCount = 0;
+        res = zeDeviceGetSubDevices(hDevice, &subDeviceCount, nullptr);
+        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+        if (i == 0) {
+            EXPECT_EQ(4u, subDeviceCount);
+        } else {
+            ze_device_handle_t hSubDevice;
+            res = zeDeviceGetSubDevices(hDevice, &subDeviceCount, &hSubDevice);
+            EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+            DeviceImp *subDevice = reinterpret_cast<DeviceImp *>(L0::Device::fromHandle(hSubDevice));
+
+            EXPECT_EQ(1u, subDeviceCount);
+            EXPECT_EQ(subDevice->neoDevice, device->neoDevice->getDeviceById(subDevice1Index));
+        }
+    }
+
+    delete driverHandle;
+    L0::GlobalDriver = nullptr;
+}
+
+TEST_F(DriverTestMultipleDeviceWithAffinityMask,
+       whenSettingAffinityMaskWithSubDeviceLargerThanAvailableSubDevicesThenSubDeviceValueIsIgnored) {
+    L0::DriverHandleImp *driverHandle = new DriverHandleImp;
+
+    driverHandle->affinityMaskString = "0,1.77";
+
+    ze_result_t res = driverHandle->initialize(std::move(devices));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    uint32_t deviceCount = 0;
+    res = zeDeviceGet(driverHandle->toHandle(), &deviceCount, nullptr);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_EQ(deviceCount, 1u);
+
+    std::vector<ze_device_handle_t> phDevices(deviceCount);
+    res = zeDeviceGet(driverHandle->toHandle(), &deviceCount, phDevices.data());
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    ze_device_handle_t hDevice = phDevices[0];
+    EXPECT_NE(nullptr, hDevice);
+
+    uint32_t subDeviceCount = 0;
+    res = zeDeviceGetSubDevices(hDevice, &subDeviceCount, nullptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_EQ(4u, subDeviceCount);
+
+    delete driverHandle;
+    L0::GlobalDriver = nullptr;
+}
+
 } // namespace ult
 } // namespace L0