Expose devices according to ZE_AFFINITY_MASK

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

level_zero/core/source/device/device.h

@@ -101,7 +101,7 @@ struct Device : _ze_device_handle_t {
 
     inline ze_device_handle_t toHandle() { return this; }
 
-    static Device *create(DriverHandle *driverHandle, NEO::Device *neoDevice);
+    static Device *create(DriverHandle *driverHandle, NEO::Device *neoDevice, uint32_t currentDeviceMask);
 
     virtual NEO::PreemptionMode getDevicePreemptionMode() const = 0;
     virtual const NEO::DeviceInfo &getDeviceInfo() const = 0;

level_zero/core/source/device/device_imp.cpp

@@ -556,7 +556,7 @@ bool DeviceImp::isMultiDeviceCapable() const {
     return neoDevice->getNumAvailableDevices() > 1u;
 }
 
-Device *Device::create(DriverHandle *driverHandle, NEO::Device *neoDevice) {
+Device *Device::create(DriverHandle *driverHandle, NEO::Device *neoDevice, uint32_t currentDeviceMask) {
     auto device = new DeviceImp;
     UNRECOVERABLE_IF(device == nullptr);
 
@@ -575,16 +575,22 @@ Device *Device::create(DriverHandle *driverHandle, NEO::Device *neoDevice) {
     if (device->neoDevice->getNumAvailableDevices() == 1) {
         device->numSubDevices = 0;
     } else {
-        device->numSubDevices = device->neoDevice->getNumAvailableDevices();
+        for (uint32_t i = 0; i < device->neoDevice->getNumAvailableDevices(); i++) {
-        for (uint32_t i = 0; i < device->numSubDevices; i++) {
+
+            if (!((1UL << i) & currentDeviceMask)) {
+                continue;
+            }
+
             ze_device_handle_t subDevice = Device::create(driverHandle,
-                                                          device->neoDevice->getDeviceById(i));
+                                                          device->neoDevice->getDeviceById(i),
+                                                          0);
             if (subDevice == nullptr) {
                 return nullptr;
             }
             reinterpret_cast<DeviceImp *>(subDevice)->isSubdevice = true;
             device->subDevices.push_back(static_cast<Device *>(subDevice));
         }
+        device->numSubDevices = static_cast<uint32_t>(device->subDevices.size());
     }
 
     if (neoDevice->getCompilerInterface()) {

level_zero/core/source/driver/driver.cpp

@@ -30,10 +30,10 @@ void DriverImp::initialize(bool *result) {
     UNRECOVERABLE_IF(nullptr == executionEnvironment);
 
     executionEnvironment->incRefInternal();
-    auto devices = NEO::DeviceFactory::createDevices(*executionEnvironment);
+    auto neoDevices = NEO::DeviceFactory::createDevices(*executionEnvironment);
     executionEnvironment->decRefInternal();
-    if (!devices.empty()) {
+    if (!neoDevices.empty()) {
-        GlobalDriver.reset(DriverHandle::create(std::move(devices)));
+        GlobalDriver.reset(DriverHandle::create(std::move(neoDevices)));
         if (GlobalDriver != nullptr) {
             *result = true;
         }

level_zero/core/source/driver/driver_handle_imp.cpp

@@ -112,12 +112,22 @@ DriverHandleImp::~DriverHandleImp() {
     }
 }
 
-ze_result_t DriverHandleImp::initialize(std::vector<std::unique_ptr<NEO::Device>> devices) {
+ze_result_t DriverHandleImp::initialize(std::vector<std::unique_ptr<NEO::Device>> neoDevices) {
-    for (auto &neoDevice : devices) {
+
+    uint32_t currentMaskOffset = 0;
+    for (auto &neoDevice : neoDevices) {
         if (!neoDevice->getHardwareInfo().capabilityTable.levelZeroSupported) {
             continue;
         }
 
+        uint32_t currentDeviceMask = (affinityMask >> currentMaskOffset) & ((1UL << neoDevice->getNumAvailableDevices()) - 1);
+        bool isDeviceExposed = currentDeviceMask ? true : false;
+
+        currentMaskOffset += neoDevice->getNumAvailableDevices();
+        if (!isDeviceExposed) {
+            continue;
+        }
+
         if (this->memoryManager == nullptr) {
             this->memoryManager = neoDevice->getMemoryManager();
             if (this->memoryManager == nullptr) {
@@ -130,7 +140,7 @@ ze_result_t DriverHandleImp::initialize(std::vector<std::unique_ptr<NEO::Device>
             }
         }
 
-        auto device = Device::create(this, neoDevice.release());
+        auto device = Device::create(this, neoDevice.release(), currentDeviceMask);
         this->devices.push_back(device);
     }
 
@@ -149,6 +159,8 @@ DriverHandle *DriverHandle::create(std::vector<std::unique_ptr<NEO::Device>> dev
     DriverHandleImp *driverHandle = new DriverHandleImp;
     UNRECOVERABLE_IF(nullptr == driverHandle);
 
+    driverHandle->getEnv("ZE_AFFINITY_MASK", driverHandle->affinityMask);
+
     ze_result_t res = driverHandle->initialize(std::move(devices));
     if (res != ZE_RESULT_SUCCESS) {
         delete driverHandle;

level_zero/core/source/driver/driver_handle_imp.h

@@ -50,7 +50,7 @@ struct DriverHandleImp : public DriverHandle {
     ze_result_t openEventPoolIpcHandle(ze_ipc_event_pool_handle_t hIpc, ze_event_pool_handle_t *phEventPool) override;
     ze_result_t checkMemoryAccessFromDevice(Device *device, const void *ptr) override;
     NEO::SVMAllocsManager *getSvmAllocsManager() override;
-    ze_result_t initialize(std::vector<std::unique_ptr<NEO::Device>> devices);
+    ze_result_t initialize(std::vector<std::unique_ptr<NEO::Device>> neoDevices);
     bool findAllocationDataForRange(const void *buffer,
                                     size_t size,
                                     NEO::SvmAllocationData **allocData) override;
@@ -58,11 +58,23 @@ struct DriverHandleImp : public DriverHandle {
                                                                      size_t size,
                                                                      bool *allocationRangeCovered) override;
 
+    template <typename T>
+    bool getEnv(const char *varName, T &varValue) {
+        char *varChar = getenv(varName);
+        if (varChar) {
+            varValue = static_cast<T>(atoi(varChar));
+            return true;
+        }
+        return false;
+    }
+
     uint32_t numDevices = 0;
     std::unordered_map<std::string, void *> extensionFunctionsLookupMap;
     std::vector<Device *> devices;
     NEO::MemoryManager *memoryManager = nullptr;
     NEO::SVMAllocsManager *svmAllocsManager = nullptr;
+
+    uint32_t affinityMask = std::numeric_limits<uint32_t>::max();
 };
 
 } // namespace L0

level_zero/core/test/unit_tests/mocks/mock_driver_handle.cpp

@@ -96,7 +96,7 @@ ze_result_t Mock<DriverHandle>::doFreeMem(const void *ptr) {
 void Mock<DriverHandle>::setupDevices(std::vector<std::unique_ptr<NEO::Device>> neoDevices) {
     this->numDevices = static_cast<uint32_t>(neoDevices.size());
     for (auto &neoDevice : neoDevices) {
-        auto device = Device::create(this, neoDevice.release());
+        auto device = Device::create(this, neoDevice.release(), std::numeric_limits<uint32_t>::max());
         this->devices.push_back(device);
     }
 }

level_zero/core/test/unit_tests/sources/builtin/builtin_functions_tests.cpp

@@ -76,21 +76,21 @@ HWTEST_F(TestBuiltinFunctionsLibImpl, givenInitFunctionWhenBultinsTableConstainN
 
 HWTEST_F(TestBuiltinFunctionsLibImpl, givenCompilerInterfaceWhenCreateDeviceAndImageSupportedThenBuiltinsImageFunctionsAreLoaded) {
     neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[neoDevice->getRootDeviceIndex()]->compilerInterface.reset(new NEO::MockCompilerInterface());
-    std::unique_ptr<L0::Device> testDvice(Device::create(device->getDriverHandle(), neoDevice));
+    std::unique_ptr<L0::Device> testDevice(Device::create(device->getDriverHandle(), neoDevice, std::numeric_limits<uint32_t>::max()));
 
     if (device->getHwInfo().capabilityTable.supportsImages) {
         for (uint32_t builtId = 0; builtId < static_cast<uint32_t>(ImageBuiltin::COUNT); builtId++) {
-            EXPECT_NE(nullptr, testDvice->getBuiltinFunctionsLib()->getImageFunction(static_cast<L0::ImageBuiltin>(builtId)));
+            EXPECT_NE(nullptr, testDevice->getBuiltinFunctionsLib()->getImageFunction(static_cast<L0::ImageBuiltin>(builtId)));
         }
     }
 }
 
 HWTEST_F(TestBuiltinFunctionsLibImpl, givenCompilerInterfaceWhenCreateDeviceThenBuiltinsFunctionsAreLoaded) {
     neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[neoDevice->getRootDeviceIndex()]->compilerInterface.reset(new NEO::MockCompilerInterface());
-    std::unique_ptr<L0::Device> testDvice(Device::create(device->getDriverHandle(), neoDevice));
+    std::unique_ptr<L0::Device> testDevice(Device::create(device->getDriverHandle(), neoDevice, std::numeric_limits<uint32_t>::max()));
 
     for (uint32_t builtId = 0; builtId < static_cast<uint32_t>(Builtin::COUNT); builtId++) {
-        EXPECT_NE(nullptr, testDvice->getBuiltinFunctionsLib()->getFunction(static_cast<L0::Builtin>(builtId)));
+        EXPECT_NE(nullptr, testDevice->getBuiltinFunctionsLib()->getFunction(static_cast<L0::Builtin>(builtId)));
     }
 }
 } // namespace ult

level_zero/core/test/unit_tests/sources/device/test_device.cpp

@@ -191,6 +191,32 @@ struct MultipleDevicesTest : public ::testing::Test {
     const uint32_t numSubDevices = 2u;
 };
 
+TEST_F(MultipleDevicesTest, whenDeviceContainsSubDevicesThenItIsMultiDeviceCapable) {
+    L0::Device *device0 = driverHandle->devices[0];
+    EXPECT_TRUE(device0->isMultiDeviceCapable());
+
+    L0::Device *device1 = driverHandle->devices[1];
+    EXPECT_TRUE(device1->isMultiDeviceCapable());
+}
+
+TEST_F(MultipleDevicesTest, whenRetrievingNumberOfSubdevicesThenCorrectNumberIsReturned) {
+    L0::Device *device0 = driverHandle->devices[0];
+
+    uint32_t count = 0;
+    auto result = device0->getSubDevices(&count, nullptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_EQ(numSubDevices, count);
+
+    std::vector<ze_device_handle_t> subDevices(count);
+    count++;
+    result = device0->getSubDevices(&count, subDevices.data());
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_EQ(numSubDevices, count);
+    for (auto subDevice : subDevices) {
+        EXPECT_NE(nullptr, subDevice);
+    }
+}
+
 TEST_F(MultipleDevicesTest, givenTheSameDeviceThenCanAccessPeerReturnsTrue) {
     L0::Device *device0 = driverHandle->devices[0];
 

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

@@ -14,6 +14,8 @@
 #include "level_zero/core/source/driver/driver_handle_imp.h"
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
 
+#include <bitset>
+
 namespace L0 {
 namespace ult {
 
@@ -126,5 +128,119 @@ TEST_F(DriverTestMultipleFamilyNoSupport, whenInitializingDriverWithArrayOfNotSu
     EXPECT_EQ(nullptr, driverHandle);
 }
 
+struct DriverTestMultipleDeviceWithAffinityMask : public ::testing::WithParamInterface<std::tuple<int, int>>,
+                                                  public ::testing::Test {
+    void SetUp() override {
+        DebugManager.flags.CreateMultipleSubDevices.set(numSubDevices);
+        VariableBackup<bool> mockDeviceFlagBackup(&MockDevice::createSingleDevice, false);
+
+        NEO::ExecutionEnvironment *executionEnvironment = new NEO::ExecutionEnvironment();
+        executionEnvironment->prepareRootDeviceEnvironments(numRootDevices);
+        for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
+            executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(NEO::defaultHwInfo.get());
+        }
+
+        for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
+            devices.push_back(std::unique_ptr<NEO::MockDevice>(NEO::MockDevice::createWithExecutionEnvironment<NEO::MockDevice>(
+                NEO::defaultHwInfo.get(),
+                executionEnvironment, i)));
+        }
+    }
+
+    void getNumOfExposedDevices(uint32_t mask, uint32_t &rootDeviceExposed, uint32_t &numOfSubDevicesExposed) {
+        rootDeviceExposed = (((1UL << numSubDevices) - 1) & mask) ? 1 : 0;
+        numOfSubDevicesExposed = static_cast<uint32_t>(static_cast<std::bitset<sizeof(uint32_t) * 8>>(mask).count());
+    }
+
+    DebugManagerStateRestore restorer;
+    std::vector<std::unique_ptr<NEO::Device>> devices;
+
+    const uint32_t numRootDevices = 2u;
+    const uint32_t numSubDevices = 4u;
+};
+
+TEST_F(DriverTestMultipleDeviceWithAffinityMask, whenNotSettingAffinityThenAllRootDevicesAndSubDevicesAreExposed) {
+    L0::DriverHandleImp *driverHandle = new DriverHandleImp;
+
+    ze_result_t res = driverHandle->initialize(std::move(devices));
+    uint32_t deviceCount = 0;
+    res = zeDeviceGet(driverHandle->toHandle(), &deviceCount, nullptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_EQ(deviceCount, numRootDevices);
+
+    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 < numRootDevices; i++) {
+        ze_device_handle_t hDevice = phDevices[i];
+        EXPECT_NE(nullptr, hDevice);
+
+        uint32_t subDeviceCount = 0;
+        res = zeDeviceGetSubDevices(hDevice, &subDeviceCount, nullptr);
+        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+        EXPECT_EQ(numSubDevices, subDeviceCount);
+    }
+
+    delete driverHandle;
+}
+
+TEST_P(DriverTestMultipleDeviceWithAffinityMask, whenSettingAffinityMaskToDifferentValuesThenCorrectNumberOfDevicesIsExposed) {
+    L0::DriverHandleImp *driverHandle = new DriverHandleImp;
+
+    uint32_t device0Mask = std::get<0>(GetParam());
+    uint32_t rootDevice0Exposed = 0;
+    uint32_t numOfSubDevicesExposedInDevice0 = 0;
+    getNumOfExposedDevices(device0Mask, rootDevice0Exposed, numOfSubDevicesExposedInDevice0);
+    uint32_t device1Mask = std::get<1>(GetParam());
+    uint32_t rootDevice1Exposed = 0;
+    uint32_t numOfSubDevicesExposedInDevice1 = 0;
+    getNumOfExposedDevices(device1Mask, rootDevice1Exposed, numOfSubDevicesExposedInDevice1);
+
+    driverHandle->affinityMask = device0Mask | (device1Mask << numSubDevices);
+
+    ze_result_t res = driverHandle->initialize(std::move(devices));
+    uint32_t deviceCount = 0;
+    res = zeDeviceGet(driverHandle->toHandle(), &deviceCount, nullptr);
+
+    uint32_t totalRootDevices = rootDevice0Exposed + rootDevice1Exposed;
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_EQ(deviceCount, totalRootDevices);
+
+    if (deviceCount) {
+        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);
+
+            uint32_t subDeviceCount = 0;
+            res = zeDeviceGetSubDevices(hDevice, &subDeviceCount, nullptr);
+            EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+            if (rootDevice0Exposed && !rootDevice1Exposed) {
+                EXPECT_EQ(numOfSubDevicesExposedInDevice0, subDeviceCount);
+            } else if (!rootDevice0Exposed && rootDevice1Exposed) {
+                EXPECT_EQ(numOfSubDevicesExposedInDevice1, subDeviceCount);
+            } else {
+                if (i == 0) {
+                    EXPECT_EQ(numOfSubDevicesExposedInDevice0, subDeviceCount);
+                } else {
+                    EXPECT_EQ(numOfSubDevicesExposedInDevice1, subDeviceCount);
+                }
+            }
+        }
+    }
+
+    delete driverHandle;
+}
+
+INSTANTIATE_TEST_SUITE_P(DriverTestMultipleDeviceWithAffinityMaskTests,
+                         DriverTestMultipleDeviceWithAffinityMask,
+                         ::testing::Combine(
+                             ::testing::Range(0, 15), // Masks for 1 root device with 4 sub devices
+                             ::testing::Range(0, 15)));
+
 } // namespace ult
 } // namespace L0