Affinity mask helper

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>

level_zero/core/source/driver/driver_handle_imp.cpp

@@ -158,20 +158,25 @@ ze_result_t DriverHandleImp::initialize(std::vector<std::unique_ptr<NEO::Device>
             }
         }
 
+        const auto rootDeviceIndex = neoDevice->getRootDeviceIndex();
+        auto rootDeviceEnvironment = neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[rootDeviceIndex].get();
+
         if (enableProgramDebugging) {
             if (neoDevice->getDebugger() != nullptr) {
                 NEO::printDebugString(NEO::DebugManager.flags.PrintDebugMessages.get(), stderr,
                                       "%s", "Source Level Debugger cannot be used with Environment Variable enabling program debugging.\n");
                 UNRECOVERABLE_IF(neoDevice->getDebugger() != nullptr && enableProgramDebugging);
             }
-            neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[neoDevice->getRootDeviceIndex()]->debugger = DebuggerL0::create(neoDevice.get());
+            rootDeviceEnvironment->debugger = DebuggerL0::create(neoDevice.get());
         }
 
-        this->rootDeviceIndices.insert(neoDevice->getRootDeviceIndex());
-        this->deviceBitfields.insert({neoDevice->getRootDeviceIndex(), neoDevice->getDeviceBitfield()});
+        this->rootDeviceIndices.insert(rootDeviceIndex);
+        this->deviceBitfields.insert({rootDeviceIndex, neoDevice->getDeviceBitfield()});
 
         auto pNeoDevice = neoDevice.release();
-        auto device = Device::create(this, pNeoDevice, pNeoDevice->getExecutionEnvironment()->rootDeviceEnvironments[pNeoDevice->getRootDeviceIndex()]->deviceAffinityMask, false, &returnValue);
+
+        auto subDevicesMask = static_cast<uint32_t>(rootDeviceEnvironment->deviceAffinityMask.getGenericSubDevicesMask().to_ulong());
+        auto device = Device::create(this, pNeoDevice, subDevicesMask, false, &returnValue);
         this->devices.push_back(device);
 
         multiOsContextDriver |= device->isMultiDeviceCapable();

shared/source/device/device.cpp

@@ -66,7 +66,7 @@ SubDevice *Device::createEngineInstancedSubDevice(uint32_t subDeviceIndex, aub_s
 }
 
 bool Device::genericSubDevicesAllowed() {
-    auto deviceMask = executionEnvironment->rootDeviceEnvironments[getRootDeviceIndex()]->deviceAffinityMask;
+    auto deviceMask = executionEnvironment->rootDeviceEnvironments[getRootDeviceIndex()]->deviceAffinityMask.getGenericSubDevicesMask();
     uint32_t subDeviceCount = HwHelper::getSubDevicesCount(&getHardwareInfo());
     deviceBitfield = maxNBitValue(subDeviceCount);
     deviceBitfield &= deviceMask;

shared/source/execution_environment/execution_environment.cpp

@@ -10,6 +10,7 @@
 #include "shared/source/built_ins/built_ins.h"
 #include "shared/source/built_ins/sip.h"
 #include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/helpers/affinity_mask.h"
 #include "shared/source/helpers/hw_helper.h"
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/memory_manager/os_agnostic_memory_manager.h"
@@ -91,11 +92,9 @@ void ExecutionEnvironment::parseAffinityMask() {
         return;
     }
 
-    std::vector<std::vector<bool>> affinityMaskBitSet(rootDeviceEnvironments.size());
-    for (uint32_t i = 0; i < affinityMaskBitSet.size(); i++) {
-        auto hwInfo = rootDeviceEnvironments[i]->getHardwareInfo();
-        affinityMaskBitSet[i].resize(HwHelper::getSubDevicesCount(hwInfo));
-    }
+    const uint32_t numRootDevices = static_cast<uint32_t>(rootDeviceEnvironments.size());
+
+    std::vector<AffinityMaskHelper> affinityMaskHelper(numRootDevices);
 
     size_t pos = 0;
     while (pos < affinityMaskString.size()) {
@@ -103,21 +102,21 @@ void ExecutionEnvironment::parseAffinityMask() {
         size_t posNextComma = affinityMaskString.find_first_of(",", pos);
         std::string rootDeviceString = affinityMaskString.substr(pos, std::min(posNextDot, posNextComma) - pos);
         uint32_t rootDeviceIndex = static_cast<uint32_t>(std::stoul(rootDeviceString, nullptr, 0));
-        if (rootDeviceIndex < rootDeviceEnvironments.size()) {
+        if (rootDeviceIndex < numRootDevices) {
+            auto hwInfo = rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
+            auto subDevicesCount = HwHelper::getSubDevicesCount(hwInfo);
+
             pos += rootDeviceString.size();
             if (posNextDot != std::string::npos &&
                 affinityMaskString.at(pos) == '.' && posNextDot < posNextComma) {
                 pos++;
                 std::string subDeviceString = affinityMaskString.substr(pos, posNextComma - pos);
                 uint32_t subDeviceIndex = static_cast<uint32_t>(std::stoul(subDeviceString, nullptr, 0));
-                auto hwInfo = rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
-                if (subDeviceIndex < HwHelper::getSubDevicesCount(hwInfo)) {
-                    affinityMaskBitSet[rootDeviceIndex][subDeviceIndex] = true;
+                if (subDeviceIndex < subDevicesCount) {
+                    affinityMaskHelper[rootDeviceIndex].enableGenericSubDevice(subDeviceIndex);
                 }
             } else {
-                std::fill(affinityMaskBitSet[rootDeviceIndex].begin(),
-                          affinityMaskBitSet[rootDeviceIndex].end(),
-                          true);
+                affinityMaskHelper[rootDeviceIndex].enableAllGenericSubDevices(subDevicesCount);
             }
         }
         if (posNextComma == std::string::npos) {
@@ -126,31 +125,13 @@ void ExecutionEnvironment::parseAffinityMask() {
         pos = posNextComma + 1;
     }
 
-    uint32_t offset = 0;
-    uint32_t affinityMask = 0;
-    for (uint32_t i = 0; i < affinityMaskBitSet.size(); i++) {
-        for (uint32_t j = 0; j < affinityMaskBitSet[i].size(); j++) {
-            if (affinityMaskBitSet[i][j] == true) {
-                affinityMask |= (1UL << offset);
-            }
-            offset++;
-        }
-    }
-
-    uint32_t currentMaskOffset = 0;
     std::vector<std::unique_ptr<RootDeviceEnvironment>> filteredEnvironments;
-    for (size_t i = 0u; i < this->rootDeviceEnvironments.size(); i++) {
-        auto hwInfo = rootDeviceEnvironments[i]->getHardwareInfo();
-
-        uint32_t currentDeviceMask = (affinityMask >> currentMaskOffset) & ((1UL << HwHelper::getSubDevicesCount(hwInfo)) - 1);
-        bool isDeviceExposed = currentDeviceMask > 0;
-
-        currentMaskOffset += HwHelper::getSubDevicesCount(hwInfo);
-        if (!isDeviceExposed) {
+    for (uint32_t i = 0u; i < numRootDevices; i++) {
+        if (!affinityMaskHelper[i].isDeviceEnabled()) {
             continue;
         }
 
-        rootDeviceEnvironments[i]->deviceAffinityMask = currentDeviceMask;
+        rootDeviceEnvironments[i]->deviceAffinityMask = affinityMaskHelper[i];
         filteredEnvironments.emplace_back(rootDeviceEnvironments[i].release());
     }
 

shared/source/execution_environment/root_device_environment.h

@@ -7,6 +7,7 @@
 
 #pragma once
 #include "shared/source/built_ins/sip_kernel_type.h"
+#include "shared/source/helpers/affinity_mask.h"
 #include "shared/source/helpers/options.h"
 
 #include <cstdint>
@@ -33,8 +34,6 @@ class SipKernel;
 class SWTagsManager;
 struct HardwareInfo;
 
-constexpr uint32_t allSubDevicesActive = std::numeric_limits<uint32_t>::max();
-
 struct RootDeviceEnvironment {
   protected:
     std::unique_ptr<HardwareInfo> hwInfo;
@@ -74,7 +73,7 @@ struct RootDeviceEnvironment {
     std::unique_ptr<SWTagsManager> tagsManager;
     ExecutionEnvironment &executionEnvironment;
 
-    uint32_t deviceAffinityMask = allSubDevicesActive;
+    AffinityMaskHelper deviceAffinityMask{true};
 
   private:
     std::mutex mtx;

shared/source/helpers/CMakeLists.txt

@@ -8,6 +8,7 @@ set(NEO_CORE_HELPERS
     ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
     ${CMAKE_CURRENT_SOURCE_DIR}/abort.h
     ${CMAKE_CURRENT_SOURCE_DIR}/address_patch.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/affinity_mask.h
     ${CMAKE_CURRENT_SOURCE_DIR}/aligned_memory.h
     ${CMAKE_CURRENT_SOURCE_DIR}/api_specific_config.h
     ${CMAKE_CURRENT_SOURCE_DIR}/array_count.h

shared/source/helpers/affinity_mask.h

@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2021 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/helpers/constants.h"
+#include "shared/source/helpers/debug_helpers.h"
+
+#include <array>
+#include <bitset>
+#include <vector>
+
+namespace NEO {
+
+class AffinityMaskHelper {
+  public:
+    using AffinityMaskContainer = std::vector<std::bitset<32>>;
+
+    AffinityMaskHelper(bool allSubdevicesActive) {
+        if (!allSubdevicesActive) {
+            return;
+        }
+
+        constexpr size_t maxInitialSubdeviceCount = 4;
+
+        enableAllGenericSubDevices(maxInitialSubdeviceCount);
+    }
+
+    AffinityMaskHelper() : AffinityMaskHelper(false) {}
+
+    void enableGenericSubDevice(uint32_t subDeviceIndex) {
+        subDevicesWithEnginesMasks.resize(subDeviceIndex + 1);
+
+        genericSubDevicesMask.set(subDeviceIndex);
+        subDevicesWithEnginesMasks[subDeviceIndex] = std::numeric_limits<uint32_t>::max();
+    }
+
+    void enableAllGenericSubDevices(uint32_t subDeviceCount) {
+        for (uint32_t i = 0; i < subDeviceCount; i++) {
+            enableGenericSubDevice(i);
+        }
+    }
+
+    DeviceBitfield getGenericSubDevicesMask() const {
+        return genericSubDevicesMask;
+    }
+
+    bool isDeviceEnabled() const {
+        return genericSubDevicesMask.any();
+    }
+
+  protected:
+    AffinityMaskContainer subDevicesWithEnginesMasks;
+    DeviceBitfield genericSubDevicesMask = 0;
+};
+} // namespace NEO