fix: store registered engines per root device

in most cases we need to iterate over engines associated to single root device

Related-To: NEO-7925
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>

level_zero/core/source/cmdlist/cmdlist_imp.cpp

@@ -58,7 +58,8 @@ ze_result_t CommandListImp::destroy() {
             auto heapAllocation = surfaceStateHeap->getGraphicsAllocation();
             NEO::WaitParams defaultWaitParams{false, false, NEO::TimeoutControls::maxTimeout};
 
-            auto &deviceEngines = device->getNEODevice()->getMemoryManager()->getRegisteredEngines();
+            auto rootDeviceIndex = device->getRootDeviceIndex();
+            auto &deviceEngines = device->getNEODevice()->getMemoryManager()->getRegisteredEngines(rootDeviceIndex);
             for (auto &engine : deviceEngines) {
                 if (NEO::EngineHelpers::isComputeEngine(engine.getEngineType())) {
                     auto contextId = engine.osContext->getContextId();

level_zero/core/test/unit_tests/fixtures/device_fixture.cpp

@@ -30,7 +30,7 @@ void DeviceFixture::setUp() {
 }
 void DeviceFixture::setupWithExecutionEnvironment(NEO::ExecutionEnvironment &executionEnvironment) {
     execEnv = &executionEnvironment;
-    neoDevice = NEO::MockDevice::createWithExecutionEnvironment<NEO::MockDevice>(NEO::defaultHwInfo.get(), &executionEnvironment, 0u);
+    neoDevice = NEO::MockDevice::createWithExecutionEnvironment<NEO::MockDevice>(NEO::defaultHwInfo.get(), &executionEnvironment, rootDeviceIndex);
     mockBuiltIns = new MockBuiltins();
     neoDevice->executionEnvironment->rootDeviceEnvironments[0]->builtins.reset(mockBuiltIns);
     NEO::DeviceVector devices;

level_zero/core/test/unit_tests/fixtures/device_fixture.h

@@ -44,6 +44,7 @@ struct DeviceFixture {
     MockBuiltins *mockBuiltIns = nullptr;
     NEO::ExecutionEnvironment *execEnv = nullptr;
 
+    const uint32_t rootDeviceIndex = 0u;
     template <typename HelperType>
     HelperType &getHelper() const;
 };

level_zero/core/test/unit_tests/sources/debugger/linux/test_l0_debugger_linux.cpp

@@ -45,7 +45,7 @@ struct L0DebuggerLinuxFixture {
         executionEnvironment->rootDeviceEnvironments[0]->osInterface.reset(osInterface);
         executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<Drm>(drmMock));
 
-        neoDevice = NEO::MockDevice::create<NEO::MockDevice>(executionEnvironment, 0u);
+        neoDevice = NEO::MockDevice::create<NEO::MockDevice>(executionEnvironment, rootDeviceIndex);
 
         NEO::DeviceVector devices;
         devices.push_back(std::unique_ptr<NEO::Device>(neoDevice));
@@ -61,7 +61,7 @@ struct L0DebuggerLinuxFixture {
         device = nullptr;
         neoDevice = nullptr;
     };
-
+    const uint32_t rootDeviceIndex = 0u;
     std::unique_ptr<Mock<L0::DriverHandleImp>> driverHandle;
     NEO::MockDevice *neoDevice = nullptr;
     L0::Device *device = nullptr;
@@ -180,8 +180,8 @@ TEST_F(L0DebuggerLinuxTest, whenRegisterElfAndLinkWithAllocationIsCalledThenItRe
     NEO::DebugData debugData;
     debugData.vIsa = "01234567890";
     debugData.vIsaSize = 10;
-    MockDrmAllocation isaAllocation(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockDrmAllocation isaAllocation(rootDeviceIndex, AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
-    MockBufferObject bo(drmMock, 3, 0, 0, 1);
+    MockBufferObject bo(rootDeviceIndex, drmMock, 3, 0, 0, 1);
     isaAllocation.bufferObjects[0] = &bo;
     device->getL0Debugger()->registerElfAndLinkWithAllocation(&debugData, &isaAllocation);
 
@@ -200,7 +200,7 @@ TEST_F(L0DebuggerLinuxTest, whenRegisterElfAndLinkWithAllocationIsCalledInAlloca
     NEO::DebugData debugData;
     debugData.vIsa = "01234567890";
     debugData.vIsaSize = 10;
-    MockDrmAllocation isaAllocation(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockDrmAllocation isaAllocation(rootDeviceIndex, AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
     device->getL0Debugger()->registerElfAndLinkWithAllocation(&debugData, &isaAllocation);
 
     EXPECT_EQ(static_cast<size_t>(10u), drmMock->registeredDataSize);
@@ -210,8 +210,8 @@ HWTEST_F(L0DebuggerLinuxTest, givenFailureToRegisterElfWhenRegisterElfAndLinkWit
     NEO::DebugData debugData;
     debugData.vIsa = "01234567890";
     debugData.vIsaSize = 10;
-    MockDrmAllocation isaAllocation(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockDrmAllocation isaAllocation(rootDeviceIndex, AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
-    MockBufferObject bo(drmMock, 3, 0, 0, 1);
+    MockBufferObject bo(rootDeviceIndex, drmMock, 3, 0, 0, 1);
     isaAllocation.bufferObjects[0] = &bo;
 
     auto debuggerL0Hw = static_cast<MockDebuggerL0Hw<FamilyType> *>(device->getL0Debugger());
@@ -228,12 +228,12 @@ HWTEST_F(L0DebuggerLinuxTest, givenFailureToRegisterElfWhenRegisterElfAndLinkWit
 }
 
 TEST_F(L0DebuggerLinuxTest, givenAllocationsWhenAttachingZebinModuleThenAllAllocationsHaveRegisteredHandles) {
-    MockDrmAllocation isaAllocation(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockDrmAllocation isaAllocation(rootDeviceIndex, AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
-    MockBufferObject bo(drmMock, 3, 0, 0, 1);
+    MockBufferObject bo(rootDeviceIndex, drmMock, 3, 0, 0, 1);
     isaAllocation.bufferObjects[0] = &bo;
 
-    MockDrmAllocation isaAllocation2(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockDrmAllocation isaAllocation2(rootDeviceIndex, AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
-    MockBufferObject bo2(drmMock, 3, 0, 0, 1);
+    MockBufferObject bo2(rootDeviceIndex, drmMock, 3, 0, 0, 1);
     isaAllocation2.bufferObjects[0] = &bo2;
 
     uint32_t handle = 0;

level_zero/core/test/unit_tests/sources/debugger/test_l0_debugger_sba_tracking.cpp

@@ -48,7 +48,7 @@ HWTEST_P(L0DebuggerParameterizedTests, givenL0DebuggerWhenCreatedThenPerContextS
     }
     std::vector<NEO::GraphicsAllocation *> allocations;
 
-    auto &allEngines = neoDevice->getMemoryManager()->getRegisteredEngines();
+    auto &allEngines = neoDevice->getMemoryManager()->getRegisteredEngines(neoDevice->getRootDeviceIndex());
 
     for (auto &engine : allEngines) {
         auto sbaAllocation = debugger->getSbaTrackingBuffer(engine.osContext->getContextId());

level_zero/core/test/unit_tests/sources/memory/linux/test_memory_linux.cpp

@@ -103,28 +103,8 @@ class MemoryManagerIpcImplicitScalingObtainFdMock : public NEO::DrmMemoryManager
                                                                           MemoryPool::System4KBPages,
                                                                           canonizedGpuAddress);
         auto &drm = this->getDrm(0u);
-        alloc->bufferObjects[0] = mockBos.emplace_back(new MockBufferObject{&drm}).get();
+        alloc->bufferObjects[0] = mockBos.emplace_back(new MockBufferObject{properties.rootDeviceIndex, &drm}).get();
-        alloc->bufferObjects[1] = mockBos.emplace_back(new MockBufferObject{&drm}).get();
+        alloc->bufferObjects[1] = mockBos.emplace_back(new MockBufferObject{properties.rootDeviceIndex, &drm}).get();
-        alloc->setGpuBaseAddress(0xabcd);
-        return alloc;
-    }
-
-    NEO::GraphicsAllocation *allocateGraphicsMemoryWithProperties(const AllocationProperties &properties) override {
-        auto ptr = reinterpret_cast<void *>(sharedHandleAddress++);
-        auto gmmHelper = getGmmHelper(0);
-        auto canonizedGpuAddress = gmmHelper->canonize(castToUint64(ptr));
-        size_t size = 0x1000;
-        auto alloc = new IpcImplicitScalingObtainFdMockGraphicsAllocation(0u,
-                                                                          NEO::AllocationType::BUFFER,
-                                                                          nullptr,
-                                                                          ptr,
-                                                                          size,
-                                                                          0u,
-                                                                          MemoryPool::System4KBPages,
-                                                                          canonizedGpuAddress);
-        auto &drm = this->getDrm(0u);
-        alloc->bufferObjects[0] = mockBos.emplace_back(new MockBufferObject{&drm}).get();
-        alloc->bufferObjects[1] = mockBos.emplace_back(new MockBufferObject{&drm}).get();
         alloc->setGpuBaseAddress(0xabcd);
         return alloc;
     }
@@ -150,8 +130,8 @@ class MemoryManagerIpcImplicitScalingObtainFdMock : public NEO::DrmMemoryManager
                                                                           MemoryPool::System4KBPages,
                                                                           canonizedGpuAddress);
         auto &drm = this->getDrm(0u);
-        alloc->bufferObjects[0] = mockBos.emplace_back(new MockBufferObject{&drm}).get();
+        alloc->bufferObjects[0] = mockBos.emplace_back(new MockBufferObject{properties.rootDeviceIndex, &drm}).get();
-        alloc->bufferObjects[1] = mockBos.emplace_back(new MockBufferObject{&drm}).get();
+        alloc->bufferObjects[1] = mockBos.emplace_back(new MockBufferObject{properties.rootDeviceIndex, &drm}).get();
         alloc->setGpuBaseAddress(0xabcd);
         return alloc;
     }
@@ -531,27 +511,7 @@ class MemoryManagerIpcObtainFdMock : public NEO::DrmMemoryManager {
                                                            MemoryPool::System4KBPages,
                                                            canonizedGpuAddress);
         auto &drm = this->getDrm(0u);
-        MockBufferObject bo0(&drm);
+        MockBufferObject bo0(properties.rootDeviceIndex, &drm);
-        alloc->bufferObjects[0] = &bo0;
-        alloc->setGpuBaseAddress(0xabcd);
-        return alloc;
-    }
-
-    NEO::GraphicsAllocation *allocateGraphicsMemoryWithProperties(const AllocationProperties &properties) override {
-        auto ptr = reinterpret_cast<void *>(sharedHandleAddress++);
-        auto gmmHelper = getGmmHelper(0);
-        auto canonizedGpuAddress = gmmHelper->canonize(castToUint64(ptr));
-        size_t size = 0x1000;
-        auto alloc = new IpcObtainFdMockGraphicsAllocation(0u,
-                                                           NEO::AllocationType::BUFFER,
-                                                           nullptr,
-                                                           ptr,
-                                                           size,
-                                                           0u,
-                                                           MemoryPool::System4KBPages,
-                                                           canonizedGpuAddress);
-        auto &drm = this->getDrm(0u);
-        MockBufferObject bo0(&drm);
         alloc->bufferObjects[0] = &bo0;
         alloc->setGpuBaseAddress(0xabcd);
         return alloc;
@@ -578,7 +538,7 @@ class MemoryManagerIpcObtainFdMock : public NEO::DrmMemoryManager {
                                                            MemoryPool::System4KBPages,
                                                            canonizedGpuAddress);
         auto &drm = this->getDrm(0u);
-        MockBufferObject bo0(&drm);
+        MockBufferObject bo0(properties.rootDeviceIndex, &drm);
         alloc->bufferObjects[0] = &bo0;
         alloc->setGpuBaseAddress(0xabcd);
         return alloc;

opencl/source/program/program.cpp

@@ -310,7 +310,7 @@ void Program::cleanCurrentKernelInfo(uint32_t rootDeviceIndex) {
     for (auto &kernelInfo : buildInfo.kernelInfoArray) {
         if (kernelInfo->kernelAllocation) {
             // register cache flush in all csrs where kernel allocation was used
-            for (auto &engine : this->executionEnvironment.memoryManager->getRegisteredEngines()) {
+            for (auto &engine : this->executionEnvironment.memoryManager->getRegisteredEngines(rootDeviceIndex)) {
                 auto contextId = engine.osContext->getContextId();
                 if (kernelInfo->kernelAllocation->isUsedByOsContext(contextId)) {
                     engine.commandStreamReceiver->registerInstructionCacheFlush();

opencl/test/unit_test/device/device_tests.cpp

@@ -216,7 +216,7 @@ HWTEST_F(DeviceTest, givenDeviceWithoutSubDevicesWhenCreatingContextsThenMemoryM
     auto rootDeviceIndex = device.getRootDeviceIndex();
     MockMemoryManager *memoryManager = static_cast<MockMemoryManager *>(device.getMemoryManager());
 
-    OsContext *defaultOsContextMemoryManager = memoryManager->registeredEngines[memoryManager->defaultEngineIndex[rootDeviceIndex]].osContext;
+    OsContext *defaultOsContextMemoryManager = memoryManager->getRegisteredEngines(rootDeviceIndex)[memoryManager->defaultEngineIndex[rootDeviceIndex]].osContext;
     OsContext *defaultOsContextRootDevice = device.getDefaultEngine().osContext;
     EXPECT_EQ(defaultOsContextRootDevice, defaultOsContextMemoryManager);
 }
@@ -227,7 +227,7 @@ HWTEST_F(DeviceTest, givenDeviceWithSubDevicesWhenCreatingContextsThenMemoryMana
     auto rootDeviceIndex = device.getRootDeviceIndex();
     MockMemoryManager *memoryManager = static_cast<MockMemoryManager *>(device.getMemoryManager());
 
-    OsContext *defaultOsContextMemoryManager = memoryManager->registeredEngines[memoryManager->defaultEngineIndex[rootDeviceIndex]].osContext;
+    OsContext *defaultOsContextMemoryManager = memoryManager->getRegisteredEngines(rootDeviceIndex)[memoryManager->defaultEngineIndex[rootDeviceIndex]].osContext;
     OsContext *defaultOsContextRootDevice = device.getDefaultEngine().osContext;
     EXPECT_EQ(defaultOsContextRootDevice, defaultOsContextMemoryManager);
 }
@@ -238,7 +238,8 @@ HWTEST_F(DeviceTest, givenMultiDeviceWhenCreatingContextsThenMemoryManagerDefaul
     MockMemoryManager *memoryManager = static_cast<MockMemoryManager *>(device.getMemoryManager());
 
     for (auto &pRootDevice : factory.rootDevices) {
-        OsContext *defaultOsContextMemoryManager = memoryManager->registeredEngines[memoryManager->defaultEngineIndex[pRootDevice->getRootDeviceIndex()]].osContext;
+        auto rootDeviceIndex = pRootDevice->getRootDeviceIndex();
+        OsContext *defaultOsContextMemoryManager = memoryManager->getRegisteredEngines(rootDeviceIndex)[memoryManager->defaultEngineIndex[rootDeviceIndex]].osContext;
         OsContext *defaultOsContextRootDevice = pRootDevice->getDefaultEngine().osContext;
         EXPECT_EQ(defaultOsContextRootDevice, defaultOsContextMemoryManager);
     }
@@ -283,7 +284,7 @@ TEST(DeviceCreation, givenDeviceWhenItIsCreatedThenOsContextIsRegistredInMemoryM
     } else if (device->getNumSubDevices() > 0) {
         numEnginesForDevice += device->getNumSubDevices();
     }
-    EXPECT_EQ(numEnginesForDevice, memoryManager->getRegisteredEnginesCount());
+    EXPECT_EQ(numEnginesForDevice, memoryManager->getRegisteredEngines(device->getRootDeviceIndex()).size());
 }
 
 TEST(DeviceCreation, givenMultiRootDeviceWhenTheyAreCreatedThenEachOsContextHasUniqueId) {
@@ -301,7 +302,6 @@ TEST(DeviceCreation, givenMultiRootDeviceWhenTheyAreCreatedThenEachOsContextHasU
 
     MockDevice *devices[] = {device1.get(), device2.get()};
 
-    auto &registeredEngines = executionEnvironment->memoryManager->getRegisteredEngines();
     auto &gfxCoreHelper = device1->getGfxCoreHelper();
     const auto &numGpgpuEngines = static_cast<uint32_t>(gfxCoreHelper.getGpgpuEngineInstances(device1->getRootDeviceEnvironment()).size());
 
@@ -311,13 +311,14 @@ TEST(DeviceCreation, givenMultiRootDeviceWhenTheyAreCreatedThenEachOsContextHasU
         numExpectedEngineInstancedEnginesPerDevice = device1->getNumSubDevices();
     }
 
-    auto expectedTotalRegisteredEngines = (numExpectedGenericEnginesPerDevice + numExpectedEngineInstancedEnginesPerDevice) * numDevices;
+    auto expectedTotalRegisteredEnginesPerRootDevice = numExpectedGenericEnginesPerDevice + numExpectedEngineInstancedEnginesPerDevice;
-
-    EXPECT_EQ(expectedTotalRegisteredEngines, registeredEngines.size());
 
     uint32_t contextId = 0;
-
     for (uint32_t i = 0; i < numDevices; i++) {
+        uint32_t contextWithinRootDevice = 0;
+        auto &registeredEngines = executionEnvironment->memoryManager->getRegisteredEngines(i);
+
+        EXPECT_EQ(expectedTotalRegisteredEnginesPerRootDevice, registeredEngines.size());
         auto device = devices[i];
 
         for (uint32_t j = 0; j < numExpectedEngineInstancedEnginesPerDevice; j++) {
@@ -326,10 +327,11 @@ TEST(DeviceCreation, givenMultiRootDeviceWhenTheyAreCreatedThenEachOsContextHasU
             EXPECT_EQ(contextId, engine.osContext->getContextId());
             EXPECT_EQ(1u, engine.osContext->getDeviceBitfield().to_ulong());
 
-            EXPECT_EQ(registeredEngines[contextId].commandStreamReceiver,
+            EXPECT_EQ(registeredEngines[contextWithinRootDevice].commandStreamReceiver,
                       engine.commandStreamReceiver);
 
             contextId++;
+            contextWithinRootDevice++;
         }
 
         for (uint32_t j = 0; j < numExpectedGenericEnginesPerDevice; j++) {
@@ -338,14 +340,13 @@ TEST(DeviceCreation, givenMultiRootDeviceWhenTheyAreCreatedThenEachOsContextHasU
             EXPECT_EQ(contextId, engine.osContext->getContextId());
             EXPECT_EQ(1u, engine.osContext->getDeviceBitfield().to_ulong());
 
-            EXPECT_EQ(registeredEngines[contextId].commandStreamReceiver,
+            EXPECT_EQ(registeredEngines[contextWithinRootDevice].commandStreamReceiver,
                       engine.commandStreamReceiver);
 
             contextId++;
+            contextWithinRootDevice++;
         }
     }
-
-    EXPECT_EQ(expectedTotalRegisteredEngines, executionEnvironment->memoryManager->getRegisteredEnginesCount());
 }
 
 TEST(DeviceCreation, givenMultiRootDeviceWhenTheyAreCreatedThenEachDeviceHasSeperateDeviceIndex) {

opencl/test/unit_test/memory_manager/memory_manager_tests.cpp

@@ -2122,19 +2122,6 @@ TEST(GraphicsAllocation, givenSharedHandleBasedConstructorWhenGraphicsAllocation
     EXPECT_EQ(expectedGpuAddress, graphicsAllocation.getGpuAddress());
 }
 
-TEST(ResidencyDataTest, givenOsContextWhenItIsRegisteredToMemoryManagerThenRefCountIncreases) {
-    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
-    auto &gfxCoreHelper = executionEnvironment.rootDeviceEnvironments[0]->getHelper<GfxCoreHelper>();
-    auto memoryManager = new MockMemoryManager(false, false, executionEnvironment);
-    executionEnvironment.memoryManager.reset(memoryManager);
-    DeviceBitfield deviceBitfield(1);
-    std::unique_ptr<CommandStreamReceiver> csr(createCommandStream(executionEnvironment, 0u, deviceBitfield));
-    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[0])[0],
-                                                                                                      PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo)));
-    EXPECT_EQ(1u, memoryManager->getRegisteredEnginesCount());
-    EXPECT_EQ(1, memoryManager->registeredEngines[0].osContext->getRefInternalCount());
-}
-
 TEST(MemoryManagerRegisteredEnginesTest, givenOsContextWhenItIsUnregisteredFromMemoryManagerThenRefCountDecreases) {
     auto device = std::unique_ptr<Device>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(defaultHwInfo.get()));
     auto memoryManager = device->getMemoryManager();
@@ -2160,9 +2147,10 @@ TEST(ResidencyDataTest, givenDeviceBitfieldWhenCreatingOsContextThenSetValidValu
     memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[0])[0],
                                                                                                       preemptionMode,
                                                                                                       deviceBitfield));
-    EXPECT_EQ(2u, memoryManager->registeredEngines[0].osContext->getNumSupportedDevices());
+    auto &engine = memoryManager->getRegisteredEngines(0u)[0];
-    EXPECT_EQ(deviceBitfield, memoryManager->registeredEngines[0].osContext->getDeviceBitfield());
+    EXPECT_EQ(2u, engine.osContext->getNumSupportedDevices());
-    EXPECT_EQ(preemptionMode, memoryManager->registeredEngines[0].osContext->getPreemptionMode());
+    EXPECT_EQ(deviceBitfield, engine.osContext->getDeviceBitfield());
+    EXPECT_EQ(preemptionMode, engine.osContext->getPreemptionMode());
 }
 
 TEST(ResidencyDataTest, givenTwoOsContextsWhenTheyAreRegisteredFromHigherToLowerThenProperSizeIsReturned) {
@@ -2177,9 +2165,11 @@ TEST(ResidencyDataTest, givenTwoOsContextsWhenTheyAreRegisteredFromHigherToLower
                                                                                                       PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo)));
     memoryManager->createAndRegisterOsContext(csr1.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[0])[0],
                                                                                                        PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo)));
-    EXPECT_EQ(2u, memoryManager->getRegisteredEnginesCount());
+    EXPECT_EQ(2u, memoryManager->allRegisteredEngines.size());
-    EXPECT_EQ(1, memoryManager->registeredEngines[0].osContext->getRefInternalCount());
+    EXPECT_EQ(1u, memoryManager->allRegisteredEngines[0].size());
-    EXPECT_EQ(1, memoryManager->registeredEngines[1].osContext->getRefInternalCount());
+    EXPECT_EQ(1u, memoryManager->allRegisteredEngines[1].size());
+    EXPECT_EQ(1, memoryManager->allRegisteredEngines[0][0].osContext->getRefInternalCount());
+    EXPECT_EQ(1, memoryManager->allRegisteredEngines[1][0].osContext->getRefInternalCount());
 }
 
 TEST(ResidencyDataTest, givenGpgpuEnginesWhenAskedForMaxOsContextCountThenValueIsGreaterOrEqual) {
@@ -2587,8 +2577,8 @@ TEST_F(MemoryAllocatorTest, whenCommandStreamerIsRegisteredThenReturnAssociatedE
 }
 
 TEST_F(MemoryAllocatorTest, whenCommandStreamerIsNotRegisteredThenReturnNullEngineControl) {
-    CommandStreamReceiver *dummyCsr = reinterpret_cast<CommandStreamReceiver *>(0x1);
+    MockCommandStreamReceiver csr{*executionEnvironment, 0u, {}};
-    auto engineControl = memoryManager->getRegisteredEngineForCsr(dummyCsr);
+    auto engineControl = memoryManager->getRegisteredEngineForCsr(&csr);
     EXPECT_EQ(nullptr, engineControl);
 }
 
@@ -2808,8 +2798,8 @@ HWTEST_F(PageTableManagerTest, givenPageTableManagerWhenMapAuxGpuVaThenForAllEng
     auto mockMngr = new MockGmmPageTableMngr();
     auto mockMngr2 = new MockGmmPageTableMngr();
 
-    memoryManager->getRegisteredEngines()[0].commandStreamReceiver->pageTableManager.reset(mockMngr);
+    memoryManager->getRegisteredEngines(1)[0].commandStreamReceiver->pageTableManager.reset(mockMngr);
-    memoryManager->getRegisteredEngines()[1].commandStreamReceiver->pageTableManager.reset(mockMngr2);
+    memoryManager->getRegisteredEngines(1)[1].commandStreamReceiver->pageTableManager.reset(mockMngr2);
 
     MockGraphicsAllocation allocation(1u, AllocationType::UNKNOWN, nullptr, 0, 0, MemoryPool::MemoryNull, MemoryManager::maxOsContextCount, 0llu);
     MockGmm gmm(executionEnvironment->rootDeviceEnvironments[allocation.getRootDeviceIndex()]->getGmmHelper());
@@ -2853,7 +2843,7 @@ HWTEST_F(PageTableManagerTest, givenPageTableManagerWhenUpdateAuxTableGmmErrorTh
     auto mockMngr = new MockGmmPageTableMngr();
     mockMngr->updateAuxTableResult = GMM_ERROR;
 
-    memoryManager->getRegisteredEngines()[0].commandStreamReceiver->pageTableManager.reset(mockMngr);
+    memoryManager->getRegisteredEngines(1)[0].commandStreamReceiver->pageTableManager.reset(mockMngr);
 
     MockGraphicsAllocation allocation(1u, AllocationType::UNKNOWN, nullptr, 0, 0, MemoryPool::MemoryNull, MemoryManager::maxOsContextCount, 0llu);
     MockGmm gmm(executionEnvironment->rootDeviceEnvironments[allocation.getRootDeviceIndex()]->getGmmHelper());
@@ -2885,7 +2875,7 @@ HWTEST_F(PageTableManagerTest, givenNullPageTableManagerWhenMapAuxGpuVaThenNoThr
     memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
                                                                                                       PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
-    memoryManager->getRegisteredEngines()[0].commandStreamReceiver->pageTableManager.reset(nullptr);
+    memoryManager->getRegisteredEngines(1)[0].commandStreamReceiver->pageTableManager.reset(nullptr);
 
     MockGraphicsAllocation allocation(1u, AllocationType::UNKNOWN, nullptr, 0, 0, MemoryPool::MemoryNull, MemoryManager::maxOsContextCount, 0llu);
     MockGmm gmm(executionEnvironment->rootDeviceEnvironments[allocation.getRootDeviceIndex()]->getGmmHelper());
@@ -2912,7 +2902,7 @@ HWTEST_F(PageTableManagerTest, givenNullPageTableManagerWhenMapAuxGpuVaThenRetur
     memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment->rootDeviceEnvironments[0])[0],
                                                                                                       PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
-    for (auto engine : memoryManager->getRegisteredEngines()) {
+    for (auto engine : memoryManager->getRegisteredEngines(1)) {
         engine.commandStreamReceiver->pageTableManager.reset(nullptr);
     }
 

opencl/test/unit_test/os_interface/windows/wddm_memory_manager_allocate_in_device_pool_tests.cpp

@@ -19,7 +19,7 @@ using namespace NEO;
 using namespace ::testing;
 
 TEST_F(WddmMemoryManagerSimpleTest, givenUseSystemMemorySetToTrueWhenAllocateInDevicePoolIsCalledThenNullptrIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;
     allocData.size = MemoryConstants::pageSize;
@@ -34,7 +34,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenUseSystemMemorySetToTrueWhenAllocateInD
 TEST_F(WddmMemoryManagerSimpleTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenLocalMemoryAllocationIsReturned) {
     const bool localMemoryEnabled = true;
 
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -55,10 +55,10 @@ TEST_F(WddmMemoryManagerSimpleTest, givenShareableAllocationWhenAllocateInDevice
 
     NEO::HardwareInfo hwInfo = *NEO::defaultHwInfo.get();
     hwInfo.featureTable.flags.ftrLocalMemory = true;
-    executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hwInfo);
+    executionEnvironment.rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hwInfo);
-    executionEnvironment->rootDeviceEnvironments[0]->initGmm();
+    executionEnvironment.rootDeviceEnvironments[0]->initGmm();
 
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -90,10 +90,10 @@ TEST_F(WddmMemoryManagerSimpleTest, givenShareableAllocationWhenAllocateGraphics
 
     NEO::HardwareInfo hwInfo = *NEO::defaultHwInfo.get();
     hwInfo.featureTable.flags.ftrLocalMemory = true;
-    executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hwInfo);
+    executionEnvironment.rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hwInfo);
-    executionEnvironment->rootDeviceEnvironments[0]->initGmm();
+    executionEnvironment.rootDeviceEnvironments[0]->initGmm();
 
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
     AllocationProperties properties{mockRootDeviceIndex, MemoryConstants::pageSize, AllocationType::SVM_GPU, mockDeviceBitfield};
     properties.allFlags = 0;
     properties.size = MemoryConstants::pageSize;
@@ -118,7 +118,7 @@ struct WddmMemoryManagerDevicePoolAlignmentTests : WddmMemoryManagerSimpleTest {
     void testAlignment(uint32_t allocationSize, uint32_t expectedAlignment) {
         const bool enable64kbPages = false;
         const bool localMemoryEnabled = true;
-        memoryManager = std::make_unique<MockWddmMemoryManager>(enable64kbPages, localMemoryEnabled, *executionEnvironment);
+        memoryManager = std::make_unique<MockWddmMemoryManager>(enable64kbPages, localMemoryEnabled, executionEnvironment);
 
         MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
         AllocationData allocData;
@@ -229,7 +229,7 @@ TEST_F(WddmMemoryManagerDevicePoolAlignmentTests, givenAtLeast2MbAllocationWhenA
 }
 
 HWTEST_F(WddmMemoryManagerSimpleTest, givenLinearStreamWhenItIsAllocatedThenItIsInLocalMemoryHasCpuPointerAndHasStandardHeap64kbAsGpuAddress) {
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
 
     auto graphicsAllocation = memoryManager->allocateGraphicsMemoryWithProperties({mockRootDeviceIndex, 4096u, AllocationType::LINEAR_STREAM, mockDeviceBitfield});
 
@@ -243,7 +243,7 @@ HWTEST_F(WddmMemoryManagerSimpleTest, givenLinearStreamWhenItIsAllocatedThenItIs
 
     if (is64bit) {
         auto gmmHelper = memoryManager->getGmmHelper(graphicsAllocation->getRootDeviceIndex());
-        if (executionEnvironment->rootDeviceEnvironments[graphicsAllocation->getRootDeviceIndex()]->isFullRangeSvm()) {
+        if (executionEnvironment.rootDeviceEnvironments[graphicsAllocation->getRootDeviceIndex()]->isFullRangeSvm()) {
             EXPECT_GE(gpuAddress, gmmHelper->canonize(partition.Standard64KB.Base));
             EXPECT_LE(gpuAddressEnd, gmmHelper->canonize(partition.Standard64KB.Limit));
         } else {
@@ -251,7 +251,7 @@ HWTEST_F(WddmMemoryManagerSimpleTest, givenLinearStreamWhenItIsAllocatedThenItIs
             EXPECT_LE(gpuAddressEnd, gmmHelper->canonize(partition.Standard.Limit));
         }
     } else {
-        if (executionEnvironment->rootDeviceEnvironments[graphicsAllocation->getRootDeviceIndex()]->isFullRangeSvm()) {
+        if (executionEnvironment.rootDeviceEnvironments[graphicsAllocation->getRootDeviceIndex()]->isFullRangeSvm()) {
             EXPECT_GE(gpuAddress, 0ull);
             EXPECT_LE(gpuAddress, UINT32_MAX);
 
@@ -266,8 +266,7 @@ HWTEST_F(WddmMemoryManagerSimpleTest, givenLinearStreamWhenItIsAllocatedThenItIs
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenLocalMemoryAllocationHasCorrectStorageInfoAndFlushL3IsSet) {
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -290,8 +289,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenNotSetUseSystemMemoryWhenGraphicsAlloca
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryAndUseSytemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -305,8 +303,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryAndUseSytemMemoryWhen
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryAndAllowed32BitAndForce32BitWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
     memoryManager->setForce32BitAllocations(true);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
@@ -324,8 +321,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryAndAllowed32BitAndFor
 TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryAndAllowed32BitWhen32BitIsNotForcedThenGraphicsAllocationInDevicePoolReturnsLocalMemoryAllocation) {
     const bool localMemoryEnabled = true;
 
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
     memoryManager->setForce32BitAllocations(false);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
@@ -346,8 +342,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryAndAllowed32BitWhen32
 
 TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryWhenAllocateFailsThenGraphicsAllocationInDevicePoolReturnsError) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;
@@ -367,8 +362,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryWhenAllocateFailsThen
 
 TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryWhenAllocateFailsThenGraphicsAllocationInPhysicalLocalDeviceMemoryReturnsError) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;
@@ -387,8 +381,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenEnabledLocalMemoryWhenAllocateFailsThen
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocatePhysicalLocalDeviceMemoryThenLocalMemoryAllocationHasCorrectStorageInfoAndNoGpuAddress) {
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -411,8 +404,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocatePhysicalLocalDeviceMemoryThenLo
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocatePhysicalLocalDeviceMemoryWithMultiStorageThenLocalMemoryAllocationHasCorrectStorageInfoAndNoGpuAddress) {
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -436,8 +428,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocatePhysicalLocalDeviceMemoryWithMu
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocatePhysicalLocalDeviceMemoryWithMultiBanksThenLocalMemoryAllocationHasCorrectStorageInfoAndNoGpuAddress) {
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
     allocData.allFlags = 0;
@@ -462,7 +453,6 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocatePhysicalLocalDeviceMemoryWithMu
 
 TEST_F(WddmMemoryManagerTest, givenLocalMemoryAllocationWhenCpuPointerNotMeetRestrictionsThenDontReserveMemRangeForMap) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
     memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
     void *cpuPtr = reinterpret_cast<void *>(memoryManager->getAlignedMallocRestrictions()->minAddress - 0x1000);
     size_t size = 0x1000;
@@ -483,8 +473,7 @@ TEST_F(WddmMemoryManagerTest, givenLocalMemoryAllocationWhenCpuPointerNotMeetRes
 
 TEST_F(WddmMemoryManagerSimpleTest, whenMemoryIsAllocatedInLocalMemoryThenTheAllocationNeedsMakeResidentBeforeLock) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
@@ -503,8 +492,7 @@ TEST_F(WddmMemoryManagerSimpleTest, whenMemoryIsAllocatedInLocalMemoryThenTheAll
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocationWithHighPriorityWhenMemoryIsAllocatedInLocalMemoryThenSetAllocationPriorityIsCalledWithHighPriority) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     AllocationType highPriorityTypes[] = {
         AllocationType::KERNEL_ISA,
@@ -536,8 +524,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocationWithHighPriorityWhenMemoryIsA
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocationWithoutHighPriorityWhenMemoryIsAllocatedInLocalMemoryThenSetAllocationPriorityIsCalledWithNormalPriority) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
@@ -557,8 +544,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocationWithoutHighPriorityWhenMemory
 
 TEST_F(WddmMemoryManagerSimpleTest, givenSetAllocationPriorityFailureWhenMemoryIsAllocatedInLocalMemoryThenNullptrIsReturned) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
@@ -576,8 +562,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenSetAllocationPriorityFailureWhenMemoryI
 
 TEST_F(WddmMemoryManagerSimpleTest, givenSetAllocationPriorityFailureWhenMemoryIsAllocatedInLocalPhysicalMemoryThenNullptrIsReturned) {
     const bool localMemoryEnabled = true;
-    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, executionEnvironment);
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, localMemoryEnabled, *executionEnvironment);
 
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
@@ -598,12 +583,7 @@ class WddmMemoryManagerSimpleTestWithLocalMemory : public MockWddmMemoryManagerF
     void SetUp() override {
         HardwareInfo localPlatformDevice = *defaultHwInfo;
         localPlatformDevice.featureTable.flags.ftrLocalMemory = true;
-
+        executionEnvironment.rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&localPlatformDevice);
-        platformsImpl->clear();
-        auto executionEnvironment = constructPlatform()->peekExecutionEnvironment();
-        executionEnvironment->prepareRootDeviceEnvironments(1u);
-        executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&localPlatformDevice);
-        executionEnvironment->rootDeviceEnvironments[0]->initGmm();
 
         MockWddmMemoryManagerFixture::SetUp();
         wddm->init();
@@ -616,7 +596,7 @@ class WddmMemoryManagerSimpleTestWithLocalMemory : public MockWddmMemoryManagerF
 };
 
 TEST_F(WddmMemoryManagerSimpleTestWithLocalMemory, givenLocalMemoryAndImageOrSharedResourceWhenAllocateInDevicePoolIsCalledThenLocalMemoryAllocationAndAndStatusSuccessIsReturned) {
-    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
+    memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
 
     ImageDescriptor imgDesc = {};
@@ -660,7 +640,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenSvmGpuAllocationWhenHostPtrProvidedThen
     EXPECT_EQ(nullptr, allocation->getUnderlyingBuffer());
     EXPECT_EQ(nullptr, allocation->getDriverAllocatedCpuPtr());
     // limited platforms will not use heap HeapIndex::HEAP_SVM
-    if (executionEnvironment->rootDeviceEnvironments[0]->isFullRangeSvm()) {
+    if (executionEnvironment.rootDeviceEnvironments[0]->isFullRangeSvm()) {
         EXPECT_EQ(svmPtr, reinterpret_cast<void *>(allocation->getGpuAddress()));
     }
     EXPECT_EQ(nullptr, allocation->getReservedAddressPtr());
@@ -767,7 +747,7 @@ struct WddmMemoryManagerSimple64BitTest : public WddmMemoryManagerSimpleTest {
             wddm->mapGpuVaStatus = true;
             VariableBackup<bool> restorer{&wddm->callBaseMapGpuVa, false};
 
-            memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, *executionEnvironment);
+            memoryManager = std::make_unique<MockWddmMemoryManager>(false, true, executionEnvironment);
             AllocationData allocData;
             allocData.allFlags = 0;
             allocData.size = static_cast<size_t>(MemoryConstants::gigaByte * 13);

opencl/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -64,7 +64,7 @@ void WddmMemoryManagerFixture::setUp() {
         memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
                                                                                                           PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
-        for (auto engine : memoryManager->getRegisteredEngines()) {
+        for (auto engine : memoryManager->getRegisteredEngines(rootDeviceIndex)) {
             if (engine.getEngineUsage() == EngineUsage::Regular) {
                 engine.commandStreamReceiver->pageTableManager.reset(GmmPageTableMngr::create(nullptr, 0, &dummyTTCallbacks));
             }
@@ -189,7 +189,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocateGraphicsMemoryWithPropertiesCal
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemoryIsCalledThenMemoryPoolIsSystem4KBPages) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     if (memoryManager->isLimitedGPU(0)) {
         GTEST_SKIP();
     }
@@ -197,7 +197,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemory
     auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), MemoryConstants::pageSize});
     EXPECT_NE(nullptr, allocation);
     EXPECT_EQ(MemoryPool::System4KBPages, allocation->getMemoryPool());
-    EXPECT_EQ(executionEnvironment->rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
+    EXPECT_EQ(executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
               allocation->getDefaultGmm()->resourceParams.Flags.Info.NonLocalOnly);
     memoryManager->freeGraphicsMemory(allocation);
 }
@@ -214,7 +214,7 @@ class MockCreateWddmAllocationMemoryManager : public MockWddmMemoryManager {
 };
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemoryFailedThenNullptrFromAllocateMemoryByKMDIsReturned) {
-    memoryManager.reset(new MockCreateWddmAllocationMemoryManager(*executionEnvironment));
+    memoryManager.reset(new MockCreateWddmAllocationMemoryManager(executionEnvironment));
     AllocationData allocationData;
     allocationData.size = MemoryConstants::pageSize;
     auto allocation = memoryManager->allocateMemoryByKMD(allocationData);
@@ -222,7 +222,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemory
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemoryFailedThenNullptrFromAllocatePhysicalDeviceMemoryIsReturned) {
-    memoryManager.reset(new MockCreateWddmAllocationMemoryManager(*executionEnvironment));
+    memoryManager.reset(new MockCreateWddmAllocationMemoryManager(executionEnvironment));
     AllocationData allocationData;
     allocationData.size = MemoryConstants::pageSize;
     MemoryManager::AllocationStatus status;
@@ -231,7 +231,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemory
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemoryFailedThenNullptrFromAllocatePhysicalLocalDeviceMemoryIsReturned) {
-    memoryManager.reset(new MockCreateWddmAllocationMemoryManager(*executionEnvironment));
+    memoryManager.reset(new MockCreateWddmAllocationMemoryManager(executionEnvironment));
     AllocationData allocationData;
     allocationData.size = MemoryConstants::pageSize;
     MemoryManager::AllocationStatus status;
@@ -240,21 +240,21 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemory
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWith64KBPagesEnabledWhenAllocateGraphicsMemory64kbIsCalledThenMemoryPoolIsSystem64KBPages) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     AllocationData allocationData;
     allocationData.size = 4096u;
     auto allocation = memoryManager->allocateGraphicsMemory64kb(allocationData);
     EXPECT_NE(nullptr, allocation);
     EXPECT_EQ(MemoryPool::System64KBPages, allocation->getMemoryPool());
 
-    EXPECT_EQ(executionEnvironment->rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
+    EXPECT_EQ(executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
               allocation->getDefaultGmm()->resourceParams.Flags.Info.NonLocalOnly);
 
     memoryManager->freeGraphicsMemory(allocation);
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocationDataWithStorageInfoWhenAllocateGraphicsMemory64kbThenStorageInfoInAllocationIsSetCorrectly) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     AllocationData allocationData;
     allocationData.storageInfo = {};
     auto allocation = memoryManager->allocateGraphicsMemory64kb(allocationData);
@@ -268,7 +268,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocationDataWithFlagsWhenAllocateGrap
       public:
         using GraphicsAllocation::allocationInfo;
     };
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     AllocationData allocationData;
     allocationData.flags.flushL3 = true;
     auto allocation = static_cast<MockGraphicsAllocation *>(memoryManager->allocateGraphicsMemory64kb(allocationData));
@@ -278,7 +278,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocationDataWithFlagsWhenAllocateGrap
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemoryWithPtrIsCalledThenMemoryPoolIsSystem4KBPages) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     if (memoryManager->isLimitedGPU(0)) {
         GTEST_SKIP();
     }
@@ -294,7 +294,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocateGraphicsMemory
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocate32BitGraphicsMemoryWithPtrIsCalledThenMemoryPoolIsSystem4KBPagesWith32BitGpuAddressing) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     void *ptr = reinterpret_cast<void *>(0x1001);
     auto size = MemoryConstants::pageSize;
 
@@ -302,14 +302,14 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenAllocate32BitGraphicsM
 
     ASSERT_NE(nullptr, allocation);
     EXPECT_EQ(MemoryPool::System4KBPagesWith32BitGpuAddressing, allocation->getMemoryPool());
-    EXPECT_EQ(executionEnvironment->rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
+    EXPECT_EQ(executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
               allocation->getDefaultGmm()->resourceParams.Flags.Info.NonLocalOnly);
 
     memoryManager->freeGraphicsMemory(allocation);
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWith64KBPagesDisabledWhenAllocateGraphicsMemoryForSVMThen4KBGraphicsAllocationIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     if (memoryManager->isLimitedGPU(0)) {
         GTEST_SKIP();
     }
@@ -318,13 +318,13 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWith64KBPagesDisabledWhenA
     auto svmAllocation = memoryManager->allocateGraphicsMemoryWithProperties({csr->getRootDeviceIndex(), size, AllocationType::SVM_ZERO_COPY, mockDeviceBitfield});
     EXPECT_NE(nullptr, svmAllocation);
     EXPECT_EQ(MemoryPool::System4KBPages, svmAllocation->getMemoryPool());
-    EXPECT_EQ(executionEnvironment->rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
+    EXPECT_EQ(executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo()->featureTable.flags.ftrLocalMemory,
               svmAllocation->getDefaultGmm()->resourceParams.Flags.Info.NonLocalOnly);
     memoryManager->freeGraphicsMemory(svmAllocation);
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWith64KBPagesEnabledWhenAllocateGraphicsMemoryForSVMThenMemoryPoolIsSystem64KBPages) {
-    memoryManager.reset(new MockWddmMemoryManager(true, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(true, false, executionEnvironment));
     if (memoryManager->isLimitedGPU(0)) {
         GTEST_SKIP();
     }
@@ -337,7 +337,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWith64KBPagesEnabledWhenAl
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenCreateAllocationFromHandleIsCalledThenMemoryPoolIsSystemCpuInaccessible) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     auto osHandle = 1u;
     gdi->getQueryResourceInfoArgOut().NumAllocations = 1;
     std::unique_ptr<Gmm> gmm(new Gmm(rootDeviceEnvironment->getGmmHelper(), nullptr, 0, 0, GMM_RESOURCE_USAGE_OCL_BUFFER, false, {}, true));
@@ -357,7 +357,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenMemoryManagerWhenCreateAllocationFromHa
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenSharedHandleWhenCreateGraphicsAllocationFromMultipleSharedHandlesIsCalledThenNullptrIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     auto handle = 1u;
     gdi->getQueryResourceInfoArgOut().NumAllocations = 1;
     std::unique_ptr<Gmm> gmm(new Gmm(rootDeviceEnvironment->getGmmHelper(), nullptr, 0, 0, GMM_RESOURCE_USAGE_OCL_BUFFER, false, {}, true));
@@ -376,7 +376,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenSharedHandleWhenCreateGraphicsAllocatio
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocationPropertiesWhenCreateAllocationFromHandleIsCalledThenCorrectAllocationTypeIsSet) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     auto osHandle = 1u;
     gdi->getQueryResourceInfoArgOut().NumAllocations = 1;
     std::unique_ptr<Gmm> gmm(new Gmm(rootDeviceEnvironment->getGmmHelper(), nullptr, 0, 0, GMM_RESOURCE_USAGE_OCL_BUFFER, false, {}, true));
@@ -402,7 +402,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocationPropertiesWhenCreateAllocatio
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, whenCreateAllocationFromHandleAndMapCallFailsThenFreeGraphicsMemoryIsCalled) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     auto osHandle = 1u;
     gdi->getQueryResourceInfoArgOut().NumAllocations = 1;
     auto gmm = std::make_unique<Gmm>(rootDeviceEnvironment->getGmmHelper(), nullptr, 0, 0, GMM_RESOURCE_USAGE_OCL_BUFFER, false, StorageInfo{}, true);
@@ -425,7 +425,7 @@ TEST_F(WddmMemoryManagerSimpleTest, whenCreateAllocationFromHandleAndMapCallFail
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocateGraphicsMemoryForNonSvmHostPtrIsCalledWhenNotAlignedPtrIsPassedThenAlignedGraphicsAllocationIsCreated) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     auto size = 13u;
     auto hostPtr = reinterpret_cast<const void *>(0x10001);
 
@@ -441,7 +441,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenAllocateGraphicsMemoryForNonSvmHostPtrI
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenAllocateGraphicsMemoryForNonSvmHostPtrIsCalledWhenNotAlignedPtrIsPassedAndImportedAllocationIsFalseThenAlignedGraphicsAllocationIsFreed) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     auto size = 13u;
     auto hostPtr = reinterpret_cast<const void *>(0x10001);
 
@@ -470,7 +470,7 @@ TEST_F(WddmMemoryManagerTest, givenAllocateGraphicsMemoryForNonSvmHostPtrIsCalle
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndSmallSizeWhenAskedToCreateGrahicsAllocationThenValidAllocationIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     memoryManager->hugeGfxMemoryChunkSize = MemoryConstants::pageSize64k;
     AllocationData allocationData;
     allocationData.size = 4096u;
@@ -482,7 +482,7 @@ TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndSmallSizeWhenAskedTo
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndHugeSizeWhenAskedToCreateGrahicsAllocationThenValidAllocationIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     memoryManager->hugeGfxMemoryChunkSize = MemoryConstants::pageSize64k;
     AllocationData allocationData;
     allocationData.size = 2ULL * MemoryConstants::pageSize64k;
@@ -494,7 +494,7 @@ TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndHugeSizeWhenAskedToC
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndSmallSizeWhenAskedToCreatePhysicalGraphicsAllocationThenValidAllocationIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     memoryManager->hugeGfxMemoryChunkSize = MemoryConstants::pageSize64k;
     AllocationData allocationData;
     allocationData.size = 4096u;
@@ -506,7 +506,7 @@ TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndSmallSizeWhenAskedTo
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndHugeSizeWhenAskedToCreatePhysicalLocalGraphicsAllocationThenValidAllocationIsReturned) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     memoryManager->hugeGfxMemoryChunkSize = MemoryConstants::pageSize64k;
     AllocationData allocationData;
     allocationData.size = 2ULL * MemoryConstants::pageSize64k;
@@ -518,7 +518,7 @@ TEST_F(WddmMemoryManagerSimpleTest, GivenShareableEnabledAndHugeSizeWhenAskedToC
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, GivenPhysicalMemoryAndVirtualMemoryThenMapSucceeds) {
-    memoryManager.reset(new MockWddmMemoryManager(false, false, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(false, false, executionEnvironment));
     memoryManager->hugeGfxMemoryChunkSize = MemoryConstants::pageSize64k;
     AllocationData allocationData;
     allocationData.size = 2ULL * MemoryConstants::pageSize64k;
@@ -533,7 +533,7 @@ TEST_F(WddmMemoryManagerSimpleTest, GivenPhysicalMemoryAndVirtualMemoryThenMapSu
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenZeroFenceValueOnSingleEngineRegisteredWhenHandleFenceCompletionIsCalledThenDoNotWaitOnCpu) {
-    ASSERT_EQ(1u, memoryManager->getRegisteredEnginesCount());
+    ASSERT_EQ(1u, memoryManager->getRegisteredEngines(0).size());
 
     auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({0, 32, AllocationType::BUFFER, mockDeviceBitfield}));
     allocation->getResidencyData().updateCompletionData(0u, 0u);
@@ -545,10 +545,10 @@ TEST_F(WddmMemoryManagerSimpleTest, givenZeroFenceValueOnSingleEngineRegisteredW
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenNonZeroFenceValueOnSingleEngineRegisteredWhenHandleFenceCompletionIsCalledThenWaitOnCpuOnce) {
-    ASSERT_EQ(1u, memoryManager->getRegisteredEnginesCount());
+    ASSERT_EQ(1u, memoryManager->getRegisteredEngines(0).size());
 
     auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({0, 32, AllocationType::BUFFER, mockDeviceBitfield}));
-    auto fence = &static_cast<OsContextWin *>(memoryManager->getRegisteredEngines()[0].osContext)->getResidencyController().getMonitoredFence();
+    auto fence = &static_cast<OsContextWin *>(memoryManager->getRegisteredEngines(0)[0].osContext)->getResidencyController().getMonitoredFence();
     allocation->getResidencyData().updateCompletionData(129u, 0u);
 
     memoryManager->handleFenceCompletion(allocation);
@@ -559,71 +559,64 @@ TEST_F(WddmMemoryManagerSimpleTest, givenNonZeroFenceValueOnSingleEngineRegister
     memoryManager->freeGraphicsMemory(allocation);
 }
 
-TEST_F(WddmMemoryManagerSimpleTest, givenNonZeroFenceValuesOnMultipleEnginesRegisteredWhenHandleFenceCompletionIsCalledThenWaitOnCpuForEachEngine) {
+TEST_F(WddmMemoryManagerSimpleTest, givenNonZeroFenceValuesOnMultipleEnginesRegisteredWhenHandleFenceCompletionIsCalledThenWaitOnCpuForEachEngineFromRootDevice) {
-    executionEnvironment->prepareRootDeviceEnvironments(2u);
-    for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
-        executionEnvironment->rootDeviceEnvironments[i]->setHwInfoAndInitHelpers(defaultHwInfo.get());
-        executionEnvironment->rootDeviceEnvironments[i]->initGmm();
-    }
     const uint32_t rootDeviceIndex = 1;
     DeviceBitfield deviceBitfield(2);
-    std::unique_ptr<CommandStreamReceiver> csr(createCommandStream(*executionEnvironment, rootDeviceIndex, deviceBitfield));
+    std::unique_ptr<CommandStreamReceiver> csr(createCommandStream(executionEnvironment, rootDeviceIndex, deviceBitfield));
 
-    auto wddm2 = static_cast<WddmMock *>(Wddm::createWddm(nullptr, *executionEnvironment->rootDeviceEnvironments[rootDeviceIndex].get()));
+    auto wddm2 = static_cast<WddmMock *>(executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->getDriverModel()->as<Wddm>());
-    wddm2->init();
+    wddm2->callBaseWaitFromCpu = false;
-    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm2);
+    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm2);
 
-    auto hwInfo = executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
+    auto hwInfo = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
-    auto &gfxCoreHelper = executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->getHelper<GfxCoreHelper>();
+    auto &gfxCoreHelper = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHelper<GfxCoreHelper>();
     OsContext *osContext = memoryManager->createAndRegisterOsContext(csr.get(),
-                                                                     EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex])[1],
+                                                                     EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[rootDeviceIndex])[1],
                                                                                                                   PreemptionHelper::getDefaultPreemptionMode(*hwInfo), deviceBitfield));
     osContext->ensureContextInitialized();
-    ASSERT_EQ(2u, memoryManager->getRegisteredEnginesCount());
+    ASSERT_EQ(1u, memoryManager->getRegisteredEngines(rootDeviceIndex).size());
 
     auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({0u, 32, AllocationType::BUFFER, mockDeviceBitfield}));
-    auto lastEngineFence = &static_cast<OsContextWin *>(osContext)->getResidencyController().getMonitoredFence();
+    auto lastEngineFence = &static_cast<OsContextWin *>(memoryManager->getRegisteredEngines(0)[0].osContext)->getResidencyController().getMonitoredFence();
     allocation->getResidencyData().updateCompletionData(129u, 0u);
     allocation->getResidencyData().updateCompletionData(152u, 1u);
 
     memoryManager->handleFenceCompletion(allocation);
     EXPECT_EQ(1u, wddm->waitFromCpuResult.called);
-    EXPECT_EQ(1u, wddm2->waitFromCpuResult.called);
+    EXPECT_EQ(0u, wddm2->waitFromCpuResult.called);
     EXPECT_EQ(129u, wddm->waitFromCpuResult.uint64ParamPassed);
-    EXPECT_EQ(152u, wddm2->waitFromCpuResult.uint64ParamPassed);
+    EXPECT_EQ(lastEngineFence, wddm->waitFromCpuResult.monitoredFence);
-    EXPECT_EQ(lastEngineFence, wddm2->waitFromCpuResult.monitoredFence);
 
     memoryManager->freeGraphicsMemory(allocation);
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenNonZeroFenceValueOnSomeOfMultipleEnginesRegisteredWhenHandleFenceCompletionIsCalledThenWaitOnCpuForTheseEngines) {
     const uint32_t rootDeviceIndex = 1;
-    executionEnvironment->prepareRootDeviceEnvironments(2u);
+
-    for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
-        executionEnvironment->rootDeviceEnvironments[i]->setHwInfoAndInitHelpers(defaultHwInfo.get());
-        executionEnvironment->rootDeviceEnvironments[i]->initGmm();
-    }
     DeviceBitfield deviceBitfield(2);
-    std::unique_ptr<CommandStreamReceiver> csr(createCommandStream(*executionEnvironment, rootDeviceIndex, deviceBitfield));
+    std::unique_ptr<CommandStreamReceiver> csr(createCommandStream(executionEnvironment, rootDeviceIndex, deviceBitfield));
+    std::unique_ptr<CommandStreamReceiver> csr2(createCommandStream(executionEnvironment, rootDeviceIndex, deviceBitfield));
 
-    auto wddm2 = static_cast<WddmMock *>(Wddm::createWddm(nullptr, *executionEnvironment->rootDeviceEnvironments[rootDeviceIndex].get()));
+    auto wddm2 = static_cast<WddmMock *>(executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->getDriverModel()->as<Wddm>());
-    wddm2->init();
+    wddm2->callBaseWaitFromCpu = false;
-    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm2);
+    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm2);
-    auto &gfxCoreHelper = executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->getHelper<GfxCoreHelper>();
+    auto &gfxCoreHelper = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHelper<GfxCoreHelper>();
-    auto hwInfo = executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
+    auto hwInfo = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
-    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex])[1],
+    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[rootDeviceIndex])[1],
                                                                                                       PreemptionHelper::getDefaultPreemptionMode(*hwInfo), deviceBitfield));
-    ASSERT_EQ(2u, memoryManager->getRegisteredEnginesCount());
+    memoryManager->createAndRegisterOsContext(csr2.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[rootDeviceIndex])[1],
+                                                                                                       PreemptionHelper::getDefaultPreemptionMode(*hwInfo), deviceBitfield));
+    ASSERT_EQ(2u, memoryManager->getRegisteredEngines(rootDeviceIndex).size());
 
-    auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({0u, 32, AllocationType::BUFFER, mockDeviceBitfield}));
+    auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({1u, 32, AllocationType::BUFFER, mockDeviceBitfield}));
-    auto lastEngineFence = &static_cast<OsContextWin *>(memoryManager->getRegisteredEngines()[0].osContext)->getResidencyController().getMonitoredFence();
+    auto lastEngineFence = &static_cast<OsContextWin *>(memoryManager->getRegisteredEngines(1)[0].osContext)->getResidencyController().getMonitoredFence();
-    allocation->getResidencyData().updateCompletionData(129u, 0u);
+    allocation->getResidencyData().updateCompletionData(129u, 1u);
-    allocation->getResidencyData().updateCompletionData(0, 1u);
+    allocation->getResidencyData().updateCompletionData(0, 2u);
 
     memoryManager->handleFenceCompletion(allocation);
-    EXPECT_EQ(1u, wddm->waitFromCpuResult.called);
+    EXPECT_EQ(1u, wddm2->waitFromCpuResult.called);
-    EXPECT_EQ(129u, wddm->waitFromCpuResult.uint64ParamPassed);
+    EXPECT_EQ(129u, wddm2->waitFromCpuResult.uint64ParamPassed);
-    EXPECT_EQ(lastEngineFence, wddm->waitFromCpuResult.monitoredFence);
+    EXPECT_EQ(lastEngineFence, wddm2->waitFromCpuResult.monitoredFence);
 
     memoryManager->freeGraphicsMemory(allocation);
 }
@@ -2027,7 +2020,7 @@ TEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerWithRegisteredOsContextWithE
                                                                                                        PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo), 2));
     memoryManager->createAndRegisterOsContext(csr2.get(), EngineDescriptorHelper::getDefaultDescriptor({aub_stream::ENGINE_RCS, EngineUsage::Regular},
                                                                                                        PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo), 3));
-    auto osContext = static_cast<OsContextWin *>(memoryManager->getRegisteredEngines()[1].osContext);
+    auto osContext = static_cast<OsContextWin *>(memoryManager->getRegisteredEngines(1)[0].osContext);
     osContext->getResidencyController().setMemoryBudgetExhausted();
     EXPECT_TRUE(memoryManager->isMemoryBudgetExhausted());
 }
@@ -2068,7 +2061,7 @@ TEST_F(MockWddmMemoryManagerTest, givenPageTableManagerWhenMapAuxGpuVaCalledThen
                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
     auto mockMngr = new MockGmmPageTableMngr();
-    for (auto engine : memoryManager.getRegisteredEngines()) {
+    for (auto engine : memoryManager.getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2109,7 +2102,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedAllocationWhenMappedGpuVaAndPag
                                                                                                                             PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
     auto mockMngr = new MockGmmPageTableMngr();
-    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines()) {
+    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2143,7 +2136,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedAllocationWhenMappedGpuVaAndPag
                                                                                                                             PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
     auto mockMngr = new MockGmmPageTableMngr();
-    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines()) {
+    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines(1)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2154,7 +2147,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedAllocationWhenMappedGpuVaAndPag
     expectedDdiUpdateAuxTable.DoNotWait = true;
     expectedDdiUpdateAuxTable.Map = true;
 
-    auto expectedCallCount = executionEnvironment->memoryManager->getRegisteredEnginesCount();
+    auto expectedCallCount = executionEnvironment->memoryManager->getRegisteredEngines(1).size();
 
     auto hwInfoMock = hardwareInfoTable[wddm.getGfxPlatform()->eProductFamily];
     ASSERT_NE(nullptr, hwInfoMock);
@@ -2183,7 +2176,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedAllocationAndPageTableSupported
                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
     auto mockMngr = new MockGmmPageTableMngr();
-    for (auto engine : memoryManager.getRegisteredEngines()) {
+    for (auto engine : memoryManager.getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2197,7 +2190,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedAllocationAndPageTableSupported
     expectedDdiUpdateAuxTable.DoNotWait = true;
     expectedDdiUpdateAuxTable.Map = false;
 
-    auto expectedCallCount = memoryManager.getRegisteredEnginesCount();
+    auto expectedCallCount = memoryManager.getRegisteredEngines(1u).size();
 
     memoryManager.freeGraphicsMemory(wddmAlloc);
 
@@ -2218,7 +2211,7 @@ TEST_F(MockWddmMemoryManagerTest, givenNonCompressedAllocationWhenReleaseingThen
                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
     auto mockMngr = new MockGmmPageTableMngr();
-    for (auto engine : memoryManager.getRegisteredEngines()) {
+    for (auto engine : memoryManager.getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2246,7 +2239,7 @@ TEST_F(MockWddmMemoryManagerTest, givenNonCompressedAllocationWhenMappedGpuVaThe
                                                                                                                             PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
 
     auto mockMngr = new MockGmmPageTableMngr();
-    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines()) {
+    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2282,7 +2275,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedFlagSetWhenInternalIsUnsetThenD
 
     auto mockMngr = new MockGmmPageTableMngr();
     auto rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[1].get();
-    for (auto engine : memoryManager.getRegisteredEngines()) {
+    for (auto engine : memoryManager.getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2320,7 +2313,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedFlagSetWhenInternalIsSetThenUpd
     auto mockMngr = new MockGmmPageTableMngr();
     auto rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[1].get();
     rootDeviceEnvironment->executionEnvironment.initializeMemoryManager();
-    for (auto engine : memoryManager.getRegisteredEngines()) {
+    for (auto engine : memoryManager.getRegisteredEngines(1u)) {
         engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
     }
 
@@ -2332,7 +2325,7 @@ TEST_F(MockWddmMemoryManagerTest, givenCompressedFlagSetWhenInternalIsSetThenUpd
     delete wddmAlloc->getDefaultGmm();
     wddmAlloc->setDefaultGmm(myGmm);
 
-    auto expectedCallCount = memoryManager.getRegisteredEnginesCount();
+    auto expectedCallCount = memoryManager.getRegisteredEngines(rootDeviceIndex).size();
 
     auto result = wddm->mapGpuVirtualAddress(myGmm, ALLOCATION_HANDLE, wddm->getGfxPartition().Standard.Base, wddm->getGfxPartition().Standard.Limit, 0u, gpuVa);
     EXPECT_TRUE(result);
@@ -2533,7 +2526,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenSvmCpuAllocationWhenSizeAndAlignmentPro
     EXPECT_NE(nullptr, allocation->getUnderlyingBuffer());
     EXPECT_EQ(allocation->getUnderlyingBuffer(), allocation->getDriverAllocatedCpuPtr());
     // limited platforms will not use heap HeapIndex::HEAP_SVM
-    if (executionEnvironment->rootDeviceEnvironments[allocation->getRootDeviceIndex()]->isFullRangeSvm()) {
+    if (executionEnvironment.rootDeviceEnvironments[allocation->getRootDeviceIndex()]->isFullRangeSvm()) {
         EXPECT_EQ(alignUp(allocation->getReservedAddressPtr(), size), reinterpret_cast<void *>(allocation->getGpuAddress()));
     }
     EXPECT_EQ((2 * size), allocation->getReservedAddressSize());
@@ -2545,7 +2538,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenWriteCombinedAllocationThenCpuAddressIs
     if (is32bit) {
         GTEST_SKIP();
     }
-    memoryManager.reset(new MockWddmMemoryManager(true, true, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(true, true, executionEnvironment));
     size_t size = 2 * MemoryConstants::megaByte;
     auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({0, size, AllocationType::WRITE_COMBINED, mockDeviceBitfield}));
     ASSERT_NE(nullptr, allocation);
@@ -2553,7 +2546,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenWriteCombinedAllocationThenCpuAddressIs
     EXPECT_NE(nullptr, allocation->getUnderlyingBuffer());
     EXPECT_NE(nullptr, reinterpret_cast<void *>(allocation->getGpuAddress()));
 
-    if (executionEnvironment->rootDeviceEnvironments[allocation->getRootDeviceIndex()]->isFullRangeSvm()) {
+    if (executionEnvironment.rootDeviceEnvironments[allocation->getRootDeviceIndex()]->isFullRangeSvm()) {
         EXPECT_EQ(allocation->getUnderlyingBuffer(), reinterpret_cast<void *>(allocation->getGpuAddress()));
     }
 
@@ -2566,23 +2559,23 @@ TEST_F(WddmMemoryManagerSimpleTest, givenDebugVariableWhenCreatingWddmMemoryMana
 
     {
         DebugManager.flags.EnableMultiStorageResources.set(0);
-        MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+        MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
         EXPECT_FALSE(memoryManager.supportsMultiStorageResources);
     }
 
     {
         DebugManager.flags.EnableMultiStorageResources.set(1);
-        MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+        MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
         EXPECT_TRUE(memoryManager.supportsMultiStorageResources);
     }
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenBufferHostMemoryAllocationAndLimitedRangeAnd32BitThenAllocationGoesToExternalHeap) {
-    if (executionEnvironment->rootDeviceEnvironments[0]->isFullRangeSvm() || !is32bit) {
+    if (executionEnvironment.rootDeviceEnvironments[0]->isFullRangeSvm() || !is32bit) {
         GTEST_SKIP();
     }
 
-    memoryManager.reset(new MockWddmMemoryManager(true, true, *executionEnvironment));
+    memoryManager.reset(new MockWddmMemoryManager(true, true, executionEnvironment));
     size_t size = 2 * MemoryConstants::megaByte;
     auto allocation = static_cast<WddmAllocation *>(memoryManager->allocateGraphicsMemoryWithProperties({0, size, AllocationType::BUFFER_HOST_MEMORY, mockDeviceBitfield}));
     ASSERT_NE(nullptr, allocation);
@@ -2689,7 +2682,7 @@ TEST_F(WddmMemoryManagerSimpleTest, whenWddmMemoryManagerIsCreatedThenAlignmentS
         {MemoryConstants::pageSize64k, true, AlignmentSelector::anyWastage, HeapIndex::TOTAL_HEAPS},
     };
 
-    MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+    MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
     EXPECT_EQ(expectedAlignments, memoryManager.alignmentSelector.peekCandidateAlignments());
 }
 
@@ -2701,7 +2694,7 @@ TEST_F(WddmMemoryManagerSimpleTest, given2MbPagesDisabledWhenWddmMemoryManagerIs
         {MemoryConstants::pageSize64k, true, AlignmentSelector::anyWastage, HeapIndex::TOTAL_HEAPS},
     };
 
-    MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+    MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
     EXPECT_EQ(expectedAlignments, memoryManager.alignmentSelector.peekCandidateAlignments());
 }
 
@@ -2715,7 +2708,7 @@ TEST_F(WddmMemoryManagerSimpleTest, givenCustomAlignmentWhenWddmMemoryManagerIsC
             {MemoryConstants::megaByte, false, AlignmentSelector::anyWastage, HeapIndex::TOTAL_HEAPS},
             {MemoryConstants::pageSize64k, true, AlignmentSelector::anyWastage, HeapIndex::TOTAL_HEAPS},
         };
-        MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+        MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
         EXPECT_EQ(expectedAlignments, memoryManager.alignmentSelector.peekCandidateAlignments());
     }
 
@@ -2726,13 +2719,13 @@ TEST_F(WddmMemoryManagerSimpleTest, givenCustomAlignmentWhenWddmMemoryManagerIsC
             {MemoryConstants::pageSize2Mb, false, 0.1f, HeapIndex::TOTAL_HEAPS},
             {MemoryConstants::pageSize64k, true, AlignmentSelector::anyWastage, HeapIndex::TOTAL_HEAPS},
         };
-        MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+        MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
         EXPECT_EQ(expectedAlignments, memoryManager.alignmentSelector.peekCandidateAlignments());
     }
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenGettingGlobalMemoryPercentThenCorrectValueIsReturned) {
-    MockWddmMemoryManager memoryManager(true, true, *executionEnvironment);
+    MockWddmMemoryManager memoryManager(true, true, executionEnvironment);
     uint32_t rootDeviceIndex = 0u;
     EXPECT_EQ(memoryManager.getPercentOfGlobalMemoryAvailable(rootDeviceIndex), 0.8);
 }
@@ -2756,7 +2749,7 @@ TEST_F(WddmMemoryManagerSimpleTest, whenAlignmentRequirementExceedsPageSizeThenA
         } callCount;
     };
 
-    MockWddmMemoryManagerAllocateWithAlignment memoryManager(true, true, *executionEnvironment);
+    MockWddmMemoryManagerAllocateWithAlignment memoryManager(true, true, executionEnvironment);
 
     AllocationData allocData = {};
     allocData.size = 1024;

opencl/test/unit_test/os_interface/windows/wddm_memory_manager_tests.h

@@ -21,7 +21,6 @@
 #include "shared/test/common/test_macros/hw_test.h"
 
 #include "opencl/test/unit_test/mocks/mock_context.h"
-#include "opencl/test/unit_test/mocks/mock_platform.h"
 
 #include "gtest/gtest.h"
 
@@ -50,22 +49,27 @@ typedef ::Test<WddmMemoryManagerFixture> WddmMemoryManagerTest;
 class MockWddmMemoryManagerFixture {
   public:
     void SetUp() {
-        executionEnvironment = platform()->peekExecutionEnvironment();
-        rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[0].get();
         auto osEnvironment = new OsEnvironmentWin();
         gdi = new MockGdi();
         osEnvironment->gdi.reset(gdi);
-        executionEnvironment->osEnvironment.reset(osEnvironment);
+        executionEnvironment.osEnvironment.reset(osEnvironment);
-        wddm = static_cast<WddmMock *>(Wddm::createWddm(nullptr, *rootDeviceEnvironment));
-        constexpr uint64_t heap32Base = (is32bit) ? 0x1000 : 0x800000000000;
-        wddm->setHeap32(heap32Base, 1000 * MemoryConstants::pageSize - 1);
-        wddm->init();
 
+        executionEnvironment.prepareRootDeviceEnvironments(2u);
+        for (auto i = 0u; i < executionEnvironment.rootDeviceEnvironments.size(); i++) {
+            executionEnvironment.rootDeviceEnvironments[i]->setHwInfoAndInitHelpers(defaultHwInfo.get());
+            executionEnvironment.rootDeviceEnvironments[i]->initGmm();
+            auto wddm = static_cast<WddmMock *>(Wddm::createWddm(nullptr, *executionEnvironment.rootDeviceEnvironments[i]));
+            constexpr uint64_t heap32Base = (is32bit) ? 0x1000 : 0x800000000000;
+            wddm->setHeap32(heap32Base, 1000 * MemoryConstants::pageSize - 1);
+            wddm->init();
+        }
+        rootDeviceEnvironment = executionEnvironment.rootDeviceEnvironments[0].get();
+        wddm = static_cast<WddmMock *>(rootDeviceEnvironment->osInterface->getDriverModel()->as<Wddm>());
         rootDeviceEnvironment->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm);
-        executionEnvironment->initializeMemoryManager();
+        executionEnvironment.initializeMemoryManager();
 
-        memoryManager = std::make_unique<MockWddmMemoryManager>(*executionEnvironment);
+        memoryManager = std::make_unique<MockWddmMemoryManager>(executionEnvironment);
-        csr.reset(createCommandStream(*executionEnvironment, 0u, 1));
+        csr.reset(createCommandStream(executionEnvironment, 0u, 1));
         auto hwInfo = rootDeviceEnvironment->getHardwareInfo();
         auto &gfxCoreHelper = rootDeviceEnvironment->getHelper<GfxCoreHelper>();
         osContext = memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*rootDeviceEnvironment)[0],
@@ -81,7 +85,7 @@ class MockWddmMemoryManagerFixture {
     }
 
     RootDeviceEnvironment *rootDeviceEnvironment = nullptr;
-    ExecutionEnvironment *executionEnvironment;
+    MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<MockWddmMemoryManager> memoryManager;
     std::unique_ptr<CommandStreamReceiver> csr;
 
@@ -95,11 +99,7 @@ typedef ::Test<MockWddmMemoryManagerFixture> WddmMemoryManagerResidencyTest;
 
 class ExecutionEnvironmentFixture : public ::testing::Test {
   public:
-    ExecutionEnvironmentFixture() {
+    MockExecutionEnvironment executionEnvironment;
-        executionEnvironment = platform()->peekExecutionEnvironment();
-    }
-
-    ExecutionEnvironment *executionEnvironment;
 };
 
 class WddmMemoryManagerFixtureWithGmockWddm : public ExecutionEnvironmentFixture {
@@ -109,19 +109,19 @@ class WddmMemoryManagerFixtureWithGmockWddm : public ExecutionEnvironmentFixture
     void SetUp() override {
         // wddm is deleted by memory manager
 
-        wddm = new WddmMock(*executionEnvironment->rootDeviceEnvironments[0].get());
+        wddm = new WddmMock(*executionEnvironment.rootDeviceEnvironments[0]);
         ASSERT_NE(nullptr, wddm);
         auto preemptionMode = PreemptionHelper::getDefaultPreemptionMode(*defaultHwInfo);
         wddm->init();
-        executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm);
+        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm);
-        osInterface = executionEnvironment->rootDeviceEnvironments[0]->osInterface.get();
+        osInterface = executionEnvironment.rootDeviceEnvironments[0]->osInterface.get();
-        memoryManager = new (std::nothrow) MockWddmMemoryManager(*executionEnvironment);
+        memoryManager = new (std::nothrow) MockWddmMemoryManager(executionEnvironment);
-        executionEnvironment->memoryManager.reset(memoryManager);
+        executionEnvironment.memoryManager.reset(memoryManager);
         // assert we have memory manager
         ASSERT_NE(nullptr, memoryManager);
-        csr.reset(createCommandStream(*executionEnvironment, 0u, 1));
+        csr.reset(createCommandStream(executionEnvironment, 0u, 1));
-        auto &gfxCoreHelper = executionEnvironment->rootDeviceEnvironments[0]->getHelper<GfxCoreHelper>();
+        auto &gfxCoreHelper = executionEnvironment.rootDeviceEnvironments[0]->getHelper<GfxCoreHelper>();
-        osContext = memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment->rootDeviceEnvironments[0])[0],
+        osContext = memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(gfxCoreHelper.getGpgpuEngineInstances(*executionEnvironment.rootDeviceEnvironments[0])[0],
                                                                                                                       preemptionMode));
         osContext->incRefInternal();
     }

shared/source/debugger/debugger_l0.cpp

@@ -45,7 +45,7 @@ void DebuggerL0::initialize() {
         setSingleAddressSpaceSbaTracking(NEO::DebugManager.flags.DebuggerForceSbaTrackingMode.get());
     }
 
-    auto &engines = device->getMemoryManager()->getRegisteredEngines();
+    auto &engines = device->getMemoryManager()->getRegisteredEngines(device->getRootDeviceIndex());
 
     NEO::AllocationProperties properties{device->getRootDeviceIndex(), true, MemoryConstants::pageSize,
                                          NEO::AllocationType::DEBUG_SBA_TRACKING_BUFFER,

shared/source/device/device.cpp

@@ -239,9 +239,10 @@ bool Device::createDeviceImpl() {
         commandStreamReceiver->postInitFlagsSetup();
     }
 
+    auto &registeredEngines = executionEnvironment->memoryManager->getRegisteredEngines(rootDeviceIndex);
     uint32_t defaultEngineIndexWithinMemoryManager = 0;
-    for (auto engineIndex = 0u; engineIndex < executionEnvironment->memoryManager->getRegisteredEnginesCount(); engineIndex++) {
+    for (auto engineIndex = 0u; engineIndex < registeredEngines.size(); engineIndex++) {
-        OsContext *engine = executionEnvironment->memoryManager->getRegisteredEngines()[engineIndex].osContext;
+        OsContext *engine = registeredEngines[engineIndex].osContext;
         if (engine == getDefaultEngine().osContext) {
             defaultEngineIndexWithinMemoryManager = engineIndex;
             break;

shared/source/helpers/common_types.h

@@ -8,11 +8,13 @@
 #pragma once
 #include <cstdint>
 #include <memory>
+#include <unordered_map>
 #include <vector>
 
 namespace NEO {
 struct EngineControl;
 using EngineControlContainer = std::vector<EngineControl>;
+using MultiDeviceEngineControlContainer = std::unordered_map<uint32_t, EngineControlContainer>;
 class Device;
 using DeviceVector = std::vector<std::unique_ptr<Device>>;
 

shared/source/memory_manager/deferrable_allocation_deletion.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2022 Intel Corporation
+ * Copyright (C) 2018-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -19,7 +19,7 @@ DeferrableAllocationDeletion::DeferrableAllocationDeletion(MemoryManager &memory
 bool DeferrableAllocationDeletion::apply() {
     if (graphicsAllocation.isUsed()) {
         bool isStillUsed = false;
-        for (auto &engine : memoryManager.getRegisteredEngines()) {
+        for (auto &engine : memoryManager.getRegisteredEngines(graphicsAllocation.getRootDeviceIndex())) {
             auto contextId = engine.osContext->getContextId();
             if (graphicsAllocation.isUsedByOsContext(contextId)) {
                 if (engine.commandStreamReceiver->testTaskCountReady(engine.commandStreamReceiver->getTagAddress(), graphicsAllocation.getTaskCount(contextId))) {

shared/source/memory_manager/memory_manager.cpp

@@ -81,10 +81,13 @@ MemoryManager::MemoryManager(ExecutionEnvironment &executionEnvironment) : execu
 }
 
 MemoryManager::~MemoryManager() {
-    for (auto &engine : registeredEngines) {
+    for (auto &engineContainer : allRegisteredEngines) {
-        engine.osContext->decRefInternal();
+        for (auto &engine : engineContainer.second) {
+            engine.osContext->decRefInternal();
+        }
+        engineContainer.second.clear();
     }
-    registeredEngines.clear();
+    allRegisteredEngines.clear();
     if (reservedMemory) {
         MemoryManager::alignedFreeWrapper(reservedMemory);
     }
@@ -267,7 +270,7 @@ void MemoryManager::checkGpuUsageAndDestroyGraphicsAllocations(GraphicsAllocatio
             multiContextResourceDestructor->drain(false);
             return;
         }
-        for (auto &engine : getRegisteredEngines()) {
+        for (auto &engine : getRegisteredEngines(gfxAllocation->getRootDeviceIndex())) {
             auto osContextId = engine.osContext->getContextId();
             auto allocationTaskCount = gfxAllocation->getTaskCount(osContextId);
             if (gfxAllocation->isUsedByOsContext(osContextId) &&
@@ -336,7 +339,7 @@ OsContext *MemoryManager::createAndRegisterOsContext(CommandStreamReceiver *comm
 
     UNRECOVERABLE_IF(rootDeviceIndex != osContext->getRootDeviceIndex());
 
-    registeredEngines.emplace_back(commandStreamReceiver, osContext);
+    allRegisteredEngines[rootDeviceIndex].emplace_back(commandStreamReceiver, osContext);
 
     return osContext;
 }
@@ -603,7 +606,7 @@ GraphicsAllocation *MemoryManager::allocateInternalGraphicsMemoryWithHostCopy(ui
 
 bool MemoryManager::mapAuxGpuVA(GraphicsAllocation *graphicsAllocation) {
     bool ret = false;
-    for (auto engine : registeredEngines) {
+    for (auto engine : getRegisteredEngines(graphicsAllocation->getRootDeviceIndex())) {
         if (engine.commandStreamReceiver->pageTableManager.get()) {
             ret = engine.commandStreamReceiver->pageTableManager->updateAuxTable(graphicsAllocation->getGpuAddress(), graphicsAllocation->getDefaultGmm(), true);
             if (!ret) {
@@ -673,10 +676,6 @@ GraphicsAllocation *MemoryManager::allocateGraphicsMemoryForImage(const Allocati
     return allocateGraphicsMemoryForImageImpl(allocationDataWithSize, std::move(gmm));
 }
 
-EngineControlContainer &MemoryManager::getRegisteredEngines() {
-    return registeredEngines;
-}
-
 bool MemoryManager::isExternalAllocation(AllocationType allocationType) {
     if (allocationType == AllocationType::BUFFER ||
         allocationType == AllocationType::BUFFER_HOST_MEMORY ||
@@ -706,9 +705,9 @@ LocalMemoryUsageBankSelector *MemoryManager::getLocalMemoryUsageBankSelector(All
     return internalLocalMemoryUsageBankSelector[rootDeviceIndex].get();
 }
 
-EngineControl *MemoryManager::getRegisteredEngineForCsr(CommandStreamReceiver *commandStreamReceiver) {
+const EngineControl *MemoryManager::getRegisteredEngineForCsr(CommandStreamReceiver *commandStreamReceiver) {
-    EngineControl *engineCtrl = nullptr;
+    const EngineControl *engineCtrl = nullptr;
-    for (auto &engine : registeredEngines) {
+    for (auto &engine : getRegisteredEngines(commandStreamReceiver->getRootDeviceIndex())) {
         if (engine.commandStreamReceiver == commandStreamReceiver) {
             engineCtrl = &engine;
             break;
@@ -718,6 +717,7 @@ EngineControl *MemoryManager::getRegisteredEngineForCsr(CommandStreamReceiver *c
 }
 
 void MemoryManager::unregisterEngineForCsr(CommandStreamReceiver *commandStreamReceiver) {
+    auto &registeredEngines = allRegisteredEngines[commandStreamReceiver->getRootDeviceIndex()];
     auto numRegisteredEngines = registeredEngines.size();
     for (auto i = 0u; i < numRegisteredEngines; i++) {
         if (registeredEngines[i].commandStreamReceiver == commandStreamReceiver) {
@@ -799,41 +799,39 @@ bool MemoryManager::copyMemoryToAllocationBanks(GraphicsAllocation *graphicsAllo
     return true;
 }
 void MemoryManager::waitForEnginesCompletion(GraphicsAllocation &graphicsAllocation) {
-    for (auto &engine : getRegisteredEngines()) {
+    for (auto &engine : getRegisteredEngines(graphicsAllocation.getRootDeviceIndex())) {
-        if (graphicsAllocation.getRootDeviceIndex() == engine.osContext->getRootDeviceIndex()) {
+        auto osContextId = engine.osContext->getContextId();
-            auto osContextId = engine.osContext->getContextId();
+        auto allocationTaskCount = graphicsAllocation.getTaskCount(osContextId);
-            auto allocationTaskCount = graphicsAllocation.getTaskCount(osContextId);
+        if (graphicsAllocation.isUsedByOsContext(osContextId) &&
-            if (graphicsAllocation.isUsedByOsContext(osContextId) &&
+            engine.commandStreamReceiver->getTagAllocation() != nullptr &&
-                engine.commandStreamReceiver->getTagAllocation() != nullptr &&
+            allocationTaskCount > *engine.commandStreamReceiver->getTagAddress()) {
-                allocationTaskCount > *engine.commandStreamReceiver->getTagAddress()) {
+            engine.commandStreamReceiver->waitForCompletionWithTimeout(WaitParams{false, false, TimeoutControls::maxTimeout}, allocationTaskCount);
-                engine.commandStreamReceiver->waitForCompletionWithTimeout(WaitParams{false, false, TimeoutControls::maxTimeout}, allocationTaskCount);
-            }
         }
     }
 }
 
 bool MemoryManager::allocInUse(GraphicsAllocation &graphicsAllocation) {
-    for (auto &engine : getRegisteredEngines()) {
+    for (auto &engine : getRegisteredEngines(graphicsAllocation.getRootDeviceIndex())) {
-        if (graphicsAllocation.getRootDeviceIndex() == engine.osContext->getRootDeviceIndex()) {
+        auto osContextId = engine.osContext->getContextId();
-            auto osContextId = engine.osContext->getContextId();
+        auto allocationTaskCount = graphicsAllocation.getTaskCount(osContextId);
-            auto allocationTaskCount = graphicsAllocation.getTaskCount(osContextId);
+        if (graphicsAllocation.isUsedByOsContext(osContextId) &&
-            if (graphicsAllocation.isUsedByOsContext(osContextId) &&
+            engine.commandStreamReceiver->getTagAllocation() != nullptr &&
-                engine.commandStreamReceiver->getTagAllocation() != nullptr &&
+            allocationTaskCount > *engine.commandStreamReceiver->getTagAddress()) {
-                allocationTaskCount > *engine.commandStreamReceiver->getTagAddress()) {
+            return true;
-                return true;
-            }
         }
     }
     return false;
 }
 
 void MemoryManager::cleanTemporaryAllocationListOnAllEngines(bool waitForCompletion) {
-    for (auto &engine : getRegisteredEngines()) {
+    for (auto &engineContainer : allRegisteredEngines) {
-        auto csr = engine.commandStreamReceiver;
+        for (auto &engine : engineContainer.second) {
-        if (waitForCompletion) {
+            auto csr = engine.commandStreamReceiver;
-            csr->waitForCompletionWithTimeout(WaitParams{false, false, 0}, csr->peekLatestSentTaskCount());
+            if (waitForCompletion) {
+                csr->waitForCompletionWithTimeout(WaitParams{false, false, 0}, csr->peekLatestSentTaskCount());
+            }
+            csr->getInternalAllocationStorage()->cleanAllocationList(*csr->getTagAddress(), AllocationUsage::TEMPORARY_ALLOCATION);
         }
-        csr->getInternalAllocationStorage()->cleanAllocationList(*csr->getTagAddress(), AllocationUsage::TEMPORARY_ALLOCATION);
     }
 }
 
@@ -997,16 +995,15 @@ bool MemoryManager::isKernelBinaryReuseEnabled() {
 
 OsContext *MemoryManager::getDefaultEngineContext(uint32_t rootDeviceIndex, DeviceBitfield subdevicesBitfield) {
     OsContext *defaultContext = nullptr;
-    for (auto engineIndex = 0u; engineIndex < this->getRegisteredEnginesCount(); engineIndex++) {
+    for (auto &engine : getRegisteredEngines(rootDeviceIndex)) {
-        OsContext *engine = this->getRegisteredEngines()[engineIndex].osContext;
+        auto osContext = engine.osContext;
-        if ((engine->getRootDeviceIndex() == rootDeviceIndex) &&
+        if (osContext->isDefaultContext() && osContext->getDeviceBitfield() == subdevicesBitfield) {
-            (engine->isDefaultContext() && engine->getDeviceBitfield() == subdevicesBitfield)) {
+            defaultContext = osContext;
-            defaultContext = engine;
             break;
         }
     }
     if (!defaultContext) {
-        defaultContext = registeredEngines[defaultEngineIndex[rootDeviceIndex]].osContext;
+        defaultContext = getRegisteredEngines(rootDeviceIndex)[defaultEngineIndex[rootDeviceIndex]].osContext;
     }
     return defaultContext;
 }

shared/source/memory_manager/memory_manager.h

@@ -213,10 +213,11 @@ class MemoryManager {
 
     MOCKABLE_VIRTUAL OsContext *createAndRegisterOsContext(CommandStreamReceiver *commandStreamReceiver,
                                                            const EngineDescriptor &engineDescriptor);
-    uint32_t getRegisteredEnginesCount() const { return static_cast<uint32_t>(registeredEngines.size()); }
+    EngineControlContainer &getRegisteredEngines(uint32_t rootDeviceIndex) { return allRegisteredEngines[rootDeviceIndex]; }
-    EngineControlContainer &getRegisteredEngines();
+    const MultiDeviceEngineControlContainer &getRegisteredEngines() const { return allRegisteredEngines; }
-    EngineControl *getRegisteredEngineForCsr(CommandStreamReceiver *commandStreamReceiver);
+    const EngineControl *getRegisteredEngineForCsr(CommandStreamReceiver *commandStreamReceiver);
     void unregisterEngineForCsr(CommandStreamReceiver *commandStreamReceiver);
+
     HostPtrManager *getHostPtrManager() const { return hostPtrManager.get(); }
     void setDefaultEngineIndex(uint32_t rootDeviceIndex, uint32_t engineIndex) { defaultEngineIndex[rootDeviceIndex] = engineIndex; }
     OsContext *getDefaultEngineContext(uint32_t rootDeviceIndex, DeviceBitfield subdevicesBitfield);
@@ -326,7 +327,7 @@ class MemoryManager {
     std::vector<uint32_t> defaultEngineIndex;
     bool supportsMultiStorageResources = true;
     ExecutionEnvironment &executionEnvironment;
-    EngineControlContainer registeredEngines;
+    MultiDeviceEngineControlContainer allRegisteredEngines;
     std::unique_ptr<HostPtrManager> hostPtrManager;
     uint32_t latestContextId = std::numeric_limits<uint32_t>::max();
     std::map<uint32_t, uint32_t> rootDeviceIndexToContextId; // This map will contain initial value of latestContextId for each rootDeviceIndex

shared/source/os_interface/linux/drm_buffer_object.cpp

@@ -77,11 +77,11 @@ bool BufferObjectHandleWrapper::canCloseBoHandle() {
     return controlBlock->refCount == 1;
 }
 
-BufferObject::BufferObject(Drm *drm, uint64_t patIndex, int handle, size_t size, size_t maxOsContextCount)
+BufferObject::BufferObject(uint32_t rootDeviceIndex, Drm *drm, uint64_t patIndex, int handle, size_t size, size_t maxOsContextCount)
-    : BufferObject(drm, patIndex, BufferObjectHandleWrapper{handle}, size, maxOsContextCount) {}
+    : BufferObject(rootDeviceIndex, drm, patIndex, BufferObjectHandleWrapper{handle}, size, maxOsContextCount) {}
 
-BufferObject::BufferObject(Drm *drm, uint64_t patIndex, BufferObjectHandleWrapper &&handle, size_t size, size_t maxOsContextCount)
+BufferObject::BufferObject(uint32_t rootDeviceIndex, Drm *drm, uint64_t patIndex, BufferObjectHandleWrapper &&handle, size_t size, size_t maxOsContextCount)
-    : drm(drm), refCount(1), handle(std::move(handle)), size(size) {
+    : drm(drm), refCount(1), rootDeviceIndex(rootDeviceIndex), handle(std::move(handle)), size(size) {
 
     auto ioctlHelper = drm->getIoctlHelper();
     this->tilingMode = ioctlHelper->getDrmParamValue(DrmParam::TilingNone);

shared/source/os_interface/linux/drm_buffer_object.h

@@ -81,8 +81,8 @@ class BufferObjectHandleWrapper {
 
 class BufferObject {
   public:
-    BufferObject(Drm *drm, uint64_t patIndex, int handle, size_t size, size_t maxOsContextCount);
+    BufferObject(uint32_t rootDeviceIndex, Drm *drm, uint64_t patIndex, int handle, size_t size, size_t maxOsContextCount);
-    BufferObject(Drm *drm, uint64_t patIndex, BufferObjectHandleWrapper &&handle, size_t size, size_t maxOsContextCount);
+    BufferObject(uint32_t rootDeviceIndex, Drm *drm, uint64_t patIndex, BufferObjectHandleWrapper &&handle, size_t size, size_t maxOsContextCount);
 
     MOCKABLE_VIRTUAL ~BufferObject() = default;
 
@@ -168,9 +168,6 @@ class BufferObject {
     void requireExplicitResidency(bool required) {
         requiresExplicitResidency = required;
     }
-    void setRootDeviceIndex(uint32_t index) {
-        rootDeviceIndex = index;
-    }
     uint32_t getRootDeviceIndex() {
         return rootDeviceIndex;
     }

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -265,7 +265,7 @@ bool DrmMemoryManager::setMemAdvise(GraphicsAllocation *gfxAllocation, MemAdvise
 
 bool DrmMemoryManager::setMemPrefetch(GraphicsAllocation *gfxAllocation, SubDeviceIdsVec &subDeviceIds, uint32_t rootDeviceIndex) {
     auto drmAllocation = static_cast<DrmAllocation *>(gfxAllocation);
-    auto osContextLinux = static_cast<OsContextLinux *>(registeredEngines[defaultEngineIndex[rootDeviceIndex]].osContext);
+    auto osContextLinux = getDefaultOsContext(rootDeviceIndex);
 
     for (auto subDeviceId : subDeviceIds) {
         auto vmHandleId = subDeviceId;
@@ -295,7 +295,7 @@ NEO::BufferObject *DrmMemoryManager::allocUserptr(uintptr_t address, size_t size
 
     auto patIndex = drm.getPatIndex(nullptr, AllocationType::EXTERNAL_HOST_PTR, CacheRegion::Default, CachePolicy::WriteBack, false);
 
-    auto res = new (std::nothrow) BufferObject(&drm, patIndex, userptr.handle, size, maxOsContextCount);
+    auto res = new (std::nothrow) BufferObject(rootDeviceIndex, &drm, patIndex, userptr.handle, size, maxOsContextCount);
     if (!res) {
         DEBUG_BREAK_IF(true);
         return nullptr;
@@ -308,7 +308,7 @@ NEO::BufferObject *DrmMemoryManager::allocUserptr(uintptr_t address, size_t size
 void DrmMemoryManager::emitPinningRequest(BufferObject *bo, const AllocationData &allocationData) const {
     auto rootDeviceIndex = allocationData.rootDeviceIndex;
     if (forcePinEnabled && pinBBs.at(rootDeviceIndex) != nullptr && allocationData.flags.forcePin && allocationData.size >= this->pinThreshold) {
-        pinBBs.at(rootDeviceIndex)->pin(&bo, 1, registeredEngines[defaultEngineIndex[rootDeviceIndex]].osContext, 0, getDefaultDrmContextId(rootDeviceIndex));
+        pinBBs.at(rootDeviceIndex)->pin(&bo, 1, getDefaultOsContext(rootDeviceIndex), 0, getDefaultDrmContextId(rootDeviceIndex));
     }
 }
 
@@ -609,7 +609,7 @@ GraphicsAllocation *DrmMemoryManager::allocatePhysicalDeviceMemory(const Allocat
 
     auto patIndex = drm.getPatIndex(gmm.get(), allocationData.type, CacheRegion::Default, CachePolicy::WriteBack, false);
 
-    std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(&drm, patIndex, create.handle, bufferSize, maxOsContextCount));
+    std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(allocationData.rootDeviceIndex, &drm, patIndex, create.handle, bufferSize, maxOsContextCount));
 
     auto allocation = new DrmAllocation(allocationData.rootDeviceIndex, allocationData.type, bo.get(), nullptr, 0u, bufferSize, MemoryPool::SystemCpuInaccessible);
     allocation->setDefaultGmm(gmm.release());
@@ -639,7 +639,7 @@ GraphicsAllocation *DrmMemoryManager::allocateMemoryByKMD(const AllocationData &
 
     auto patIndex = drm.getPatIndex(gmm.get(), allocationData.type, CacheRegion::Default, CachePolicy::WriteBack, false);
 
-    std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(&drm, patIndex, create.handle, bufferSize, maxOsContextCount));
+    std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(allocationData.rootDeviceIndex, &drm, patIndex, create.handle, bufferSize, maxOsContextCount));
     bo->setAddress(gpuRange);
 
     auto allocation = new DrmAllocation(allocationData.rootDeviceIndex, allocationData.type, bo.get(), nullptr, gpuRange, bufferSize, MemoryPool::SystemCpuInaccessible);
@@ -673,7 +673,7 @@ GraphicsAllocation *DrmMemoryManager::allocateGraphicsMemoryForImageImpl(const A
 
     auto patIndex = drm.getPatIndex(gmm.get(), allocationData.type, CacheRegion::Default, CachePolicy::WriteBack, false);
 
-    std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new (std::nothrow) BufferObject(&drm, patIndex, create.handle, allocationData.imgInfo->size, maxOsContextCount));
+    std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new (std::nothrow) BufferObject(allocationData.rootDeviceIndex, &drm, patIndex, create.handle, allocationData.imgInfo->size, maxOsContextCount));
     if (!bo) {
         return nullptr;
     }
@@ -829,8 +829,7 @@ GraphicsAllocation *DrmMemoryManager::createGraphicsAllocationFromMultipleShared
             auto patIndex = drm.getPatIndex(nullptr, properties.allocationType, CacheRegion::Default, CachePolicy::WriteBack, false);
             auto boHandleWrapper = reuseSharedAllocation ? BufferObjectHandleWrapper{boHandle} : tryToGetBoHandleWrapperWithSharedOwnership(boHandle);
 
-            bo = new (std::nothrow) BufferObject(&drm, patIndex, boHandle, size, maxOsContextCount);
+            bo = new (std::nothrow) BufferObject(properties.rootDeviceIndex, &drm, patIndex, boHandle, size, maxOsContextCount);
-            bo->setRootDeviceIndex(properties.rootDeviceIndex);
             i++;
         }
         bos.push_back(bo);
@@ -971,7 +970,7 @@ GraphicsAllocation *DrmMemoryManager::createGraphicsAllocationFromSharedHandle(o
         auto patIndex = drm.getPatIndex(nullptr, properties.allocationType, CacheRegion::Default, CachePolicy::WriteBack, false);
         auto boHandleWrapper = reuseSharedAllocation ? BufferObjectHandleWrapper{boHandle} : tryToGetBoHandleWrapperWithSharedOwnership(boHandle);
 
-        bo = new (std::nothrow) BufferObject(&drm, patIndex, std::move(boHandleWrapper), size, maxOsContextCount);
+        bo = new (std::nothrow) BufferObject(properties.rootDeviceIndex, &drm, patIndex, std::move(boHandleWrapper), size, maxOsContextCount);
 
         if (!bo) {
             return nullptr;
@@ -1000,7 +999,6 @@ GraphicsAllocation *DrmMemoryManager::createGraphicsAllocationFromSharedHandle(o
 
         bo->setAddress(gpuRange);
         bo->setUnmapSize(size);
-        bo->setRootDeviceIndex(properties.rootDeviceIndex);
 
         printDebugString(DebugManager.flags.PrintBOCreateDestroyResult.get(), stdout,
                          "Created BO-%d range: %llx - %llx, size: %lld from PRIME_FD_TO_HANDLE\n",
@@ -1109,9 +1107,9 @@ void DrmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation,
     }
     DrmAllocation *drmAlloc = static_cast<DrmAllocation *>(gfxAllocation);
     this->unregisterAllocation(gfxAllocation);
-
+    auto rootDeviceIndex = gfxAllocation->getRootDeviceIndex();
-    for (auto &engine : this->registeredEngines) {
+    for (auto &engine : getRegisteredEngines(rootDeviceIndex)) {
-        auto memoryOperationsInterface = static_cast<DrmMemoryOperationsHandler *>(executionEnvironment.rootDeviceEnvironments[gfxAllocation->getRootDeviceIndex()]->memoryOperationsInterface.get());
+        auto memoryOperationsInterface = static_cast<DrmMemoryOperationsHandler *>(executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface.get());
         memoryOperationsInterface->evictWithinOsContext(engine.osContext, *gfxAllocation);
     }
 
@@ -1218,7 +1216,7 @@ AllocationStatus DrmMemoryManager::populateOsHandles(OsHandleStorage &handleStor
     }
 
     if (validateHostPtrMemory) {
-        int result = pinBBs.at(rootDeviceIndex)->validateHostPtr(allocatedBos, numberOfBosAllocated, registeredEngines[defaultEngineIndex[rootDeviceIndex]].osContext, 0, getDefaultDrmContextId(rootDeviceIndex));
+        int result = pinBBs.at(rootDeviceIndex)->validateHostPtr(allocatedBos, numberOfBosAllocated, getDefaultOsContext(rootDeviceIndex), 0, getDefaultDrmContextId(rootDeviceIndex));
 
         if (result == EFAULT) {
             for (uint32_t i = 0; i < numberOfBosAllocated; i++) {
@@ -1310,8 +1308,11 @@ int DrmMemoryManager::obtainFdFromHandle(int boHandle, uint32_t rootDeviceIndex)
     return openFd.fileDescriptor;
 }
 
+OsContextLinux *DrmMemoryManager::getDefaultOsContext(uint32_t rootDeviceIndex) const {
+    return static_cast<OsContextLinux *>(allRegisteredEngines.at(rootDeviceIndex)[defaultEngineIndex[rootDeviceIndex]].osContext);
+}
 uint32_t DrmMemoryManager::getDefaultDrmContextId(uint32_t rootDeviceIndex) const {
-    auto osContextLinux = static_cast<OsContextLinux *>(registeredEngines[defaultEngineIndex[rootDeviceIndex]].osContext);
+    auto osContextLinux = getDefaultOsContext(rootDeviceIndex);
     return osContextLinux->getDrmContextIds()[0];
 }
 
@@ -1372,11 +1373,12 @@ std::vector<GraphicsAllocation *> &DrmMemoryManager::getLocalMemAllocs(uint32_t
 bool DrmMemoryManager::makeAllocationResident(GraphicsAllocation *allocation) {
     if (DebugManager.flags.MakeEachAllocationResident.get() == 1) {
         auto drmAllocation = static_cast<DrmAllocation *>(allocation);
-        for (uint32_t i = 0; getDrm(allocation->getRootDeviceIndex()).getVirtualMemoryAddressSpace(i) > 0u; i++) {
+        auto rootDeviceIndex = allocation->getRootDeviceIndex();
-            if (drmAllocation->makeBOsResident(registeredEngines[defaultEngineIndex[allocation->getRootDeviceIndex()]].osContext, i, nullptr, true)) {
+        for (uint32_t i = 0; getDrm(rootDeviceIndex).getVirtualMemoryAddressSpace(i) > 0u; i++) {
+            if (drmAllocation->makeBOsResident(getDefaultOsContext(rootDeviceIndex), i, nullptr, true)) {
                 return false;
             }
-            getDrm(allocation->getRootDeviceIndex()).waitForBind(i);
+            getDrm(rootDeviceIndex).waitForBind(i);
         }
     }
     return true;
@@ -1784,8 +1786,9 @@ GraphicsAllocation *DrmMemoryManager::allocateGraphicsMemoryInDevicePool(const A
     return allocation.release();
 }
 
-BufferObject *DrmMemoryManager::createBufferObjectInMemoryRegion(Drm *drm, Gmm *gmm, AllocationType allocationType, uint64_t gpuAddress,
+BufferObject *DrmMemoryManager::createBufferObjectInMemoryRegion(uint32_t rootDeviceIndex, Gmm *gmm, AllocationType allocationType, uint64_t gpuAddress,
                                                                  size_t size, uint32_t memoryBanks, size_t maxOsContextCount, int32_t pairHandle) {
+    auto drm = &getDrm(rootDeviceIndex);
     auto memoryInfo = drm->getMemoryInfo();
     if (!memoryInfo) {
         return nullptr;
@@ -1807,7 +1810,7 @@ BufferObject *DrmMemoryManager::createBufferObjectInMemoryRegion(Drm *drm, Gmm *
 
     auto patIndex = drm->getPatIndex(gmm, allocationType, CacheRegion::Default, CachePolicy::WriteBack, false);
 
-    auto bo = new (std::nothrow) BufferObject(drm, patIndex, handle, size, maxOsContextCount);
+    auto bo = new (std::nothrow) BufferObject(rootDeviceIndex, drm, patIndex, handle, size, maxOsContextCount);
     if (!bo) {
         return nullptr;
     }
@@ -1854,7 +1857,7 @@ bool DrmMemoryManager::createDrmAllocation(Drm *drm, DrmAllocation *allocation,
         }
         auto gmm = allocation->getGmm(handleId);
         auto boSize = alignUp(gmm->gmmResourceInfo->getSizeAllocation(), MemoryConstants::pageSize64k);
-        bos[handleId] = createBufferObjectInMemoryRegion(drm, gmm, allocation->getAllocationType(), boAddress, boSize, memoryBanks, maxOsContextCount, pairHandle);
+        bos[handleId] = createBufferObjectInMemoryRegion(allocation->getRootDeviceIndex(), gmm, allocation->getAllocationType(), boAddress, boSize, memoryBanks, maxOsContextCount, pairHandle);
         if (nullptr == bos[handleId]) {
             return false;
         }
@@ -1941,7 +1944,7 @@ bool DrmMemoryManager::allocationTypeForCompletionFence(AllocationType allocatio
 void DrmMemoryManager::waitOnCompletionFence(GraphicsAllocation *allocation) {
     auto allocationType = allocation->getAllocationType();
     if (allocationTypeForCompletionFence(allocationType)) {
-        for (auto &engine : getRegisteredEngines()) {
+        for (auto &engine : getRegisteredEngines(allocation->getRootDeviceIndex())) {
             OsContext *osContext = engine.osContext;
             CommandStreamReceiver *csr = engine.commandStreamReceiver;
 
@@ -1990,7 +1993,7 @@ DrmAllocation *DrmMemoryManager::createAllocWithAlignment(const AllocationData &
         cpuPointer = alignUp(cpuPointer, alignment);
 
         auto pointerDiff = ptrDiff(cpuPointer, cpuBasePointer);
-        std::unique_ptr<BufferObject, BufferObject::Deleter> bo(this->createBufferObjectInMemoryRegion(&drm, nullptr, allocationData.type,
+        std::unique_ptr<BufferObject, BufferObject::Deleter> bo(this->createBufferObjectInMemoryRegion(allocationData.rootDeviceIndex, nullptr, allocationData.type,
                                                                                                        reinterpret_cast<uintptr_t>(cpuPointer), alignedSize, 0u, maxOsContextCount, -1));
 
         if (!bo) {
@@ -2176,7 +2179,7 @@ GraphicsAllocation *DrmMemoryManager::createSharedUnifiedMemoryAllocation(const
 
         auto patIndex = drm.getPatIndex(nullptr, allocationData.type, CacheRegion::Default, CachePolicy::WriteBack, false);
 
-        std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(&drm, patIndex, handle, currentSize, maxOsContextCount));
+        std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(allocationData.rootDeviceIndex, &drm, patIndex, handle, currentSize, maxOsContextCount));
 
         if (!ioctlHelper->setVmBoAdvise(bo->peekHandle(), vmAdviseAttribute, nullptr)) {
             this->munmapFunction(cpuBasePointer, totalSizeToAlloc);
@@ -2262,7 +2265,7 @@ DrmAllocation *DrmMemoryManager::createUSMHostAllocationFromSharedHandle(osHandl
     }
 
     if (hasMappedPtr) {
-        auto bo = new BufferObject(&drm, patIndex, openFd.handle, properties.size, maxOsContextCount);
+        auto bo = new BufferObject(properties.rootDeviceIndex, &drm, patIndex, openFd.handle, properties.size, maxOsContextCount);
         bo->setAddress(properties.gpuAddress);
 
         auto gmmHelper = getGmmHelper(properties.rootDeviceIndex);
@@ -2273,7 +2276,7 @@ DrmAllocation *DrmMemoryManager::createUSMHostAllocationFromSharedHandle(osHandl
 
     const bool useBooMmap = drm.getMemoryInfo() && properties.useMmapObject;
     if (!useBooMmap) {
-        auto bo = new BufferObject(&drm, patIndex, openFd.handle, properties.size, maxOsContextCount);
+        auto bo = new BufferObject(properties.rootDeviceIndex, &drm, patIndex, openFd.handle, properties.size, maxOsContextCount);
         bo->setAddress(properties.gpuAddress);
 
         auto gmmHelper = getGmmHelper(properties.rootDeviceIndex);
@@ -2292,7 +2295,7 @@ DrmAllocation *DrmMemoryManager::createUSMHostAllocationFromSharedHandle(osHandl
         void *cpuPointer = nullptr;
         size_t size = lseekFunction(handle, 0, SEEK_END);
 
-        bo = new BufferObject(&drm, patIndex, boHandle, size, maxOsContextCount);
+        bo = new BufferObject(properties.rootDeviceIndex, &drm, patIndex, boHandle, size, maxOsContextCount);
 
         if (properties.allocationType == AllocationType::GPU_TIMESTAMP_DEVICE_BUFFER) {
             cpuPointer = this->mmapFunction(0, size + MemoryConstants::pageSize64k, PROT_NONE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
@@ -2331,7 +2334,6 @@ DrmAllocation *DrmMemoryManager::createUSMHostAllocationFromSharedHandle(osHandl
         emitPinningRequest(bo, allocationData);
 
         bo->setUnmapSize(size);
-        bo->setRootDeviceIndex(properties.rootDeviceIndex);
 
         printDebugString(DebugManager.flags.PrintBOCreateDestroyResult.get(), stdout,
                          "Created BO-%d range: %llx - %llx, size: %lld from PRIME_FD_TO_HANDLE\n",

shared/source/os_interface/linux/drm_memory_manager.h

@@ -19,6 +19,7 @@ class BufferObject;
 class Drm;
 class DrmGemCloseWorker;
 class DrmAllocation;
+class OsContextLinux;
 
 enum class gemCloseWorkerMode;
 
@@ -68,7 +69,7 @@ class DrmMemoryManager : public MemoryManager {
     MOCKABLE_VIRTUAL int obtainFdFromHandle(int boHandle, uint32_t rootDeviceindex);
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override;
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override;
-    MOCKABLE_VIRTUAL BufferObject *createBufferObjectInMemoryRegion(Drm *drm, Gmm *gmm, AllocationType allocationType, uint64_t gpuAddress, size_t size,
+    MOCKABLE_VIRTUAL BufferObject *createBufferObjectInMemoryRegion(uint32_t rootDeviceIndex, Gmm *gmm, AllocationType allocationType, uint64_t gpuAddress, size_t size,
                                                                     uint32_t memoryBanks, size_t maxOsContextCount, int32_t pairHandle);
 
     bool hasPageFaultsEnabled(const Device &neoDevice) override;
@@ -106,6 +107,7 @@ class DrmMemoryManager : public MemoryManager {
     MOCKABLE_VIRTUAL void releaseGpuRange(void *address, size_t size, uint32_t rootDeviceIndex);
     void emitPinningRequest(BufferObject *bo, const AllocationData &allocationData) const;
     uint32_t getDefaultDrmContextId(uint32_t rootDeviceIndex) const;
+    OsContextLinux *getDefaultOsContext(uint32_t rootDeviceIndex) const;
     size_t getUserptrAlignment();
 
     GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const AllocationData &allocationData) override;

shared/source/os_interface/linux/drm_memory_operations_handler_bind.cpp

@@ -160,7 +160,7 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::evictUnusedAllocations(bo
 }
 
 MemoryOperationsStatus DrmMemoryOperationsHandlerBind::evictUnusedAllocationsImpl(std::vector<GraphicsAllocation *> &allocationsForEviction, bool waitForCompletion) {
-    const auto &engines = this->rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines();
+    const auto &engines = this->rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines(this->rootDeviceIndex);
     std::vector<GraphicsAllocation *> evictCandidates;
 
     for (auto subdeviceIndex = 0u; subdeviceIndex < GfxCoreHelper::getSubDevicesCount(rootDeviceEnvironment.getHardwareInfo()); subdeviceIndex++) {
@@ -172,8 +172,7 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::evictUnusedAllocationsImp
             }
 
             for (const auto &engine : engines) {
-                if (this->rootDeviceIndex == engine.commandStreamReceiver->getRootDeviceIndex() &&
+                if (engine.osContext->getDeviceBitfield().test(subdeviceIndex)) {
-                    engine.osContext->getDeviceBitfield().test(subdeviceIndex)) {
                     if (allocation->isAlwaysResident(engine.osContext->getContextId())) {
                         evict = false;
                         break;
@@ -200,8 +199,7 @@ MemoryOperationsStatus DrmMemoryOperationsHandlerBind::evictUnusedAllocationsImp
 
         for (auto &allocationToEvict : evictCandidates) {
             for (const auto &engine : engines) {
-                if (this->rootDeviceIndex == engine.commandStreamReceiver->getRootDeviceIndex() &&
+                if (engine.osContext->getDeviceBitfield().test(subdeviceIndex)) {
-                    engine.osContext->getDeviceBitfield().test(subdeviceIndex)) {
                     DeviceBitfield deviceBitfield;
                     deviceBitfield.set(subdeviceIndex);
                     this->evictImpl(engine.osContext, *allocationToEvict, deviceBitfield);

shared/source/os_interface/linux/drm_neo.cpp

@@ -666,14 +666,11 @@ void Drm::setNewResourceBoundToVM(BufferObject *bo, uint32_t vmHandleId) {
     if (!this->rootDeviceEnvironment.getProductHelper().isTlbFlushRequired() && isAligned(bo->peekAddress(), MemoryConstants::pageSize2Mb)) {
         return;
     }
-    const auto &engines = this->rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines();
+    const auto &engines = this->rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines(bo->getRootDeviceIndex());
     for (const auto &engine : engines) {
         if (engine.osContext->getDeviceBitfield().test(vmHandleId)) {
             auto osContextLinux = static_cast<OsContextLinux *>(engine.osContext);
-
+            osContextLinux->setNewResourceBound();
-            if (&osContextLinux->getDrm() == this) {
-                osContextLinux->setNewResourceBound();
-            }
         }
     }
 }

shared/source/os_interface/windows/wddm/wddm.cpp

@@ -553,9 +553,13 @@ bool Wddm::mapGpuVirtualAddress(Gmm *gmm, D3DKMT_HANDLE handle, D3DGPU_VIRTUAL_A
     bool ret = true;
     auto &productHelper = rootDeviceEnvironment.getHelper<ProductHelper>();
     if (gmm->isCompressionEnabled && productHelper.isPageTableManagerSupported(*rootDeviceEnvironment.getHardwareInfo())) {
-        for (auto engine : rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines()) {
+        for (auto &engineContainer : rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines()) {
-            if (engine.commandStreamReceiver->pageTableManager.get()) {
+            if (rootDeviceEnvironment.executionEnvironment.rootDeviceEnvironments[engineContainer.first].get() == &rootDeviceEnvironment) {
-                ret &= engine.commandStreamReceiver->pageTableManager->updateAuxTable(gpuPtr, gmm, true);
+                for (auto &engine : engineContainer.second) {
+                    if (engine.commandStreamReceiver->pageTableManager.get()) {
+                        ret &= engine.commandStreamReceiver->pageTableManager->updateAuxTable(gpuPtr, gmm, true);
+                    }
+                }
             }
         }
     }

shared/source/os_interface/windows/wddm_memory_manager.cpp

@@ -603,7 +603,8 @@ void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation
     WddmAllocation *input = static_cast<WddmAllocation *>(gfxAllocation);
     DEBUG_BREAK_IF(!validateAllocation(input));
 
-    for (auto &engine : this->registeredEngines) {
+    auto &registeredEngines = getRegisteredEngines(gfxAllocation->getRootDeviceIndex());
+    for (auto &engine : registeredEngines) {
         auto &residencyController = static_cast<OsContextWin *>(engine.osContext)->getResidencyController();
         auto lock = residencyController.acquireLock();
         residencyController.removeFromTrimCandidateListIfUsed(input, true);
@@ -613,7 +614,7 @@ void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation
     auto hwInfo = executionEnvironment.rootDeviceEnvironments[gfxAllocation->getRootDeviceIndex()]->getHardwareInfo();
     auto &productHelper = executionEnvironment.rootDeviceEnvironments[gfxAllocation->getRootDeviceIndex()]->getHelper<ProductHelper>();
     if (gfxAllocation->isCompressionEnabled() && productHelper.isPageTableManagerSupported(*hwInfo)) {
-        for (auto engine : registeredEngines) {
+        for (auto &engine : registeredEngines) {
             if (engine.commandStreamReceiver->pageTableManager.get()) {
                 [[maybe_unused]] auto status = engine.commandStreamReceiver->pageTableManager->updateAuxTable(input->getGpuAddress(), defaultGmm, false);
                 DEBUG_BREAK_IF(!status);
@@ -658,7 +659,7 @@ void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation
 
 void WddmMemoryManager::handleFenceCompletion(GraphicsAllocation *allocation) {
     auto wddmAllocation = static_cast<WddmAllocation *>(allocation);
-    for (auto &engine : this->registeredEngines) {
+    for (auto &engine : getRegisteredEngines(allocation->getRootDeviceIndex())) {
         const auto lastFenceValue = wddmAllocation->getResidencyData().getFenceValueForContextId(engine.osContext->getContextId());
         if (lastFenceValue != 0u) {
             const auto &monitoredFence = static_cast<OsContextWin *>(engine.osContext)->getResidencyController().getMonitoredFence();
@@ -679,9 +680,11 @@ bool WddmMemoryManager::tryDeferDeletions(const D3DKMT_HANDLE *handles, uint32_t
 }
 
 bool WddmMemoryManager::isMemoryBudgetExhausted() const {
-    for (auto &engine : this->registeredEngines) {
+    for (auto &engineContainer : allRegisteredEngines) {
-        if (static_cast<OsContextWin *>(engine.osContext)->getResidencyController().isMemoryBudgetExhausted()) {
+        for (auto &engine : engineContainer.second) {
-            return true;
+            if (static_cast<OsContextWin *>(engine.osContext)->getResidencyController().isMemoryBudgetExhausted()) {
+                return true;
+            }
         }
     }
     return false;

shared/test/common/fixtures/linux/memory_allocator_multi_device_fixture_linux.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2022 Intel Corporation
+ * Copyright (C) 2020-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -13,7 +13,7 @@ using namespace NEO;
 
 void MemoryAllocatorMultiDeviceSystemSpecificFixture::setUp(ExecutionEnvironment &executionEnvironment) {
     auto memoryManager = static_cast<TestedDrmMemoryManager *>(executionEnvironment.memoryManager.get());
-    auto bufferObject = new (std::nothrow) BufferObject(&memoryManager->getDrm(0u), 3, 0, 10, MemoryManager::maxOsContextCount);
+    auto bufferObject = new (std::nothrow) BufferObject(0u, &memoryManager->getDrm(0u), 3, 0, 10, MemoryManager::maxOsContextCount);
     memoryManager->pushSharedBufferObject(bufferObject);
 }
 

shared/test/common/mocks/linux/mock_drm_allocation.h

@@ -38,7 +38,7 @@ class MockBufferObject : public BufferObject {
 
     std::optional<int> execReturnValue;
     std::vector<ExecParams> passedExecParams{};
-    MockBufferObject(Drm *drm) : BufferObject(drm, CommonConstants::unsupportedPatIndex, 0, 0, 1) {
+    MockBufferObject(uint32_t rootDeviceIndex, Drm *drm) : BufferObject(rootDeviceIndex, drm, CommonConstants::unsupportedPatIndex, 0, 0, 1) {
     }
     int exec(uint32_t used, size_t startOffset, unsigned int flags, bool requiresCoherency, OsContext *osContext, uint32_t vmHandleId, uint32_t drmContextId,
              BufferObject *const residency[], size_t residencyCount, ExecObject *execObjectsStorage, uint64_t completionGpuAddress, TaskCountType completionValue) override {
@@ -58,8 +58,8 @@ class MockDrmAllocation : public DrmAllocation {
     using DrmAllocation::memoryPool;
     using DrmAllocation::registeredBoBindHandles;
 
-    MockDrmAllocation(AllocationType allocationType, MemoryPool pool)
+    MockDrmAllocation(uint32_t rootDeviceIndex, AllocationType allocationType, MemoryPool pool)
-        : DrmAllocation(0, allocationType, nullptr, nullptr, 0, static_cast<size_t>(0), pool) {
+        : DrmAllocation(rootDeviceIndex, allocationType, nullptr, nullptr, 0, static_cast<size_t>(0), pool) {
     }
 
     MockDrmAllocation(AllocationType allocationType, MemoryPool pool, BufferObjects &bos)

shared/test/common/mocks/linux/mock_drm_memory_manager.cpp

@@ -62,7 +62,7 @@ BufferObject *TestedDrmMemoryManager::findAndReferenceSharedBufferObject(int boH
     if (failOnfindAndReferenceSharedBufferObject) {
         DrmMockCustom drmMock(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]);
         auto patIndex = drmMock.getPatIndex(nullptr, AllocationType::BUFFER, CacheRegion::Default, CachePolicy::WriteBack, false);
-        return new (std::nothrow) BufferObject(&drmMock, patIndex, boHandle, 4096u, 2u);
+        return new (std::nothrow) BufferObject(rootDeviceIndex, &drmMock, patIndex, boHandle, 4096u, 2u);
     }
     return MemoryManagerCreate<DrmMemoryManager>::findAndReferenceSharedBufferObject(boHandle, rootDeviceIndex);
 }

shared/test/common/mocks/linux/mock_drm_memory_manager.h

@@ -87,8 +87,8 @@ class TestedDrmMemoryManager : public MemoryManagerCreate<DrmMemoryManager> {
     using DrmMemoryManager::unMapPhysicalToVirtualMemory;
     using DrmMemoryManager::waitOnCompletionFence;
     using MemoryManager::allocateGraphicsMemoryInDevicePool;
+    using MemoryManager::allRegisteredEngines;
     using MemoryManager::heapAssigner;
-    using MemoryManager::registeredEngines;
 
     TestedDrmMemoryManager(ExecutionEnvironment &executionEnvironment);
     TestedDrmMemoryManager(bool enableLocalMemory,

shared/test/common/mocks/mock_device.cpp

@@ -89,7 +89,7 @@ void MockDevice::resetCommandStreamReceiver(CommandStreamReceiver *newCsr, uint3
 
     registeredEngine.commandStreamReceiver = newCsr;
     allEngines[engineIndex].commandStreamReceiver = newCsr;
-    memoryManager->getRegisteredEngines().emplace_back(registeredEngine);
+    memoryManager->getRegisteredEngines(rootDeviceIndex).emplace_back(registeredEngine);
     osContext->incRefInternal();
     newCsr->setupContext(*osContext);
     osContext->ensureContextInitialized();

shared/test/common/mocks/mock_memory_manager.cpp

@@ -277,7 +277,7 @@ OsContext *MockMemoryManagerOsAgnosticContext::createAndRegisterOsContext(Comman
                                                                           const EngineDescriptor &engineDescriptor) {
     auto osContext = new OsContext(commandStreamReceiver->getRootDeviceIndex(), 0, engineDescriptor);
     osContext->incRefInternal();
-    registeredEngines.emplace_back(commandStreamReceiver, osContext);
+    allRegisteredEngines[commandStreamReceiver->getRootDeviceIndex()].emplace_back(commandStreamReceiver, osContext);
     return osContext;
 }
 
@@ -286,7 +286,7 @@ OsContext *MockMemoryManagerWithDebuggableOsContext::createAndRegisterOsContext(
     auto osContext = new MockOsContext(0, engineDescriptor);
     osContext->debuggableContext = true;
     osContext->incRefInternal();
-    registeredEngines.emplace_back(commandStreamReceiver, osContext);
+    allRegisteredEngines[commandStreamReceiver->getRootDeviceIndex()].emplace_back(commandStreamReceiver, osContext);
     return osContext;
 }
 

shared/test/common/mocks/mock_memory_manager.h

@@ -32,6 +32,7 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
     using MemoryManager::allocateGraphicsMemoryInPreferredPool;
     using MemoryManager::allocateGraphicsMemoryWithAlignment;
     using MemoryManager::allocateGraphicsMemoryWithProperties;
+    using MemoryManager::allRegisteredEngines;
     using MemoryManager::createGraphicsAllocation;
     using MemoryManager::createStorageInfoFromProperties;
     using MemoryManager::defaultEngineIndex;
@@ -44,7 +45,6 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
     using MemoryManager::overrideAllocationData;
     using MemoryManager::pageFaultManager;
     using MemoryManager::prefetchManager;
-    using MemoryManager::registeredEngines;
     using MemoryManager::supportsMultiStorageResources;
     using MemoryManager::useNonSvmHostPtrAlloc;
     using OsAgnosticMemoryManager::allocateGraphicsMemoryForImageFromHostPtr;

shared/test/common/os_interface/linux/drm_buffer_object_fixture.h

@@ -24,10 +24,10 @@ class TestedBufferObject : public BufferObject {
     using BufferObject::handle;
     using BufferObject::tilingMode;
 
-    TestedBufferObject(Drm *drm) : BufferObject(drm, 3, 1, 0, 1) {
+    TestedBufferObject(uint32_t rootDeviceIndex, Drm *drm) : TestedBufferObject(rootDeviceIndex, drm, 0) {
     }
 
-    TestedBufferObject(Drm *drm, size_t size) : BufferObject(drm, 3, 1, size, 1) {
+    TestedBufferObject(uint32_t rootDeviceIndex, Drm *drm, size_t size) : BufferObject(rootDeviceIndex, drm, 3, 1, size, 1) {
     }
 
     void fillExecObject(ExecObject &execObject, OsContext *osContext, uint32_t vmHandleId, uint32_t drmContextId) override {
@@ -73,6 +73,7 @@ class DrmBufferObjectFixture {
     TestedBufferObject *bo = nullptr;
     ExecObject execObjectsStorage[256]{};
     std::unique_ptr<OsContextLinux> osContext;
+    const uint32_t rootDeviceIndex = 0u;
 
     void setUp() {
         this->mock = std::make_unique<DrmClass>(*executionEnvironment.rootDeviceEnvironments[0]);
@@ -80,7 +81,7 @@ class DrmBufferObjectFixture {
         executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
         osContext.reset(new OsContextLinux(*this->mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor()));
         this->mock->reset();
-        bo = new TestedBufferObject(this->mock.get());
+        bo = new TestedBufferObject(0, this->mock.get());
         ASSERT_NE(nullptr, bo);
     }
 

shared/test/common/os_interface/linux/drm_command_stream_fixture.h

@@ -171,13 +171,13 @@ class DrmCommandStreamEnhancedTemplate : public ::testing::Test {
         friend DrmCommandStreamEnhancedTemplate<DrmType>;
 
       protected:
-        MockBufferObject(Drm *drm, size_t size) : BufferObject(drm, CommonConstants::unsupportedPatIndex, 1, 0, 16u) {
+        MockBufferObject(uint32_t rootDeviceIndex, Drm *drm, size_t size) : BufferObject(rootDeviceIndex, drm, CommonConstants::unsupportedPatIndex, 1, 0, 16u) {
             this->size = alignUp(size, 4096);
         }
     };
 
     MockBufferObject *createBO(size_t size) {
-        return new MockBufferObject(this->mock, size);
+        return new MockBufferObject(0, this->mock, size);
     }
 };
 
@@ -245,20 +245,6 @@ class DrmCommandStreamEnhancedWithFailingExecTemplate : public ::testing::Test {
     const std::vector<BufferObject *> &getResidencyVector() const {
         return static_cast<const TestedDrmCommandStreamReceiver<GfxFamily> *>(csr)->residency;
     }
-
-  protected:
-    class MockBufferObject : public BufferObject {
-        friend DrmCommandStreamEnhancedTemplate<T>;
-
-      protected:
-        MockBufferObject(Drm *drm, size_t size) : BufferObject(drm, CommonConstants::unsupportedPatIndex, 1, 0, 16u) {
-            this->size = alignUp(size, 4096);
-        }
-    };
-
-    MockBufferObject *createBO(size_t size) {
-        return new MockBufferObject(this->mock, size);
-    }
 };
 
 using DrmCommandStreamEnhancedWithFailingExec = DrmCommandStreamEnhancedWithFailingExecTemplate<DrmMockCustom>;

shared/test/common/os_interface/linux/drm_memory_manager_fixture.cpp

@@ -74,7 +74,11 @@ void DrmMemoryManagerFixture::tearDown() {
     mock->testIoctls();
     mock->reset();
 
-    int enginesCount = static_cast<int>(device->getMemoryManager()->getRegisteredEnginesCount());
+    int enginesCount = 0;
+
+    for (auto &engineContainer : memoryManager->getRegisteredEngines()) {
+        enginesCount += engineContainer.second.size();
+    }
 
     mock->ioctlExpected.contextDestroy = enginesCount;
     mock->ioctlExpected.gemClose = enginesCount;

shared/test/unit_test/debugger/linux/test_l0_debugger_linux.cpp

@@ -75,8 +75,8 @@ TEST_F(L0DebuggerSharedLinuxTest, givenNoOSInterfaceThenRegisterElfAndLinkWithAl
     debugData.vIsaSize = 10;
     drmMock->registeredDataSize = 0;
 
-    MockDrmAllocation isaAllocation(AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
+    MockDrmAllocation isaAllocation(neoDevice->getRootDeviceIndex(), AllocationType::KERNEL_ISA, MemoryPool::System4KBPages);
-    MockBufferObject bo(drmMock, 3, 0, 0, 1);
+    MockBufferObject bo(neoDevice->getRootDeviceIndex(), drmMock, 3, 0, 0, 1);
     isaAllocation.bufferObjects[0] = &bo;
 
     neoDevice->getL0Debugger()->registerElfAndLinkWithAllocation(&debugData, &isaAllocation);

shared/test/unit_test/debugger/test_l0_debugger_single_address_space.cpp

@@ -63,7 +63,7 @@ HWTEST_F(SingleAddressSpaceFixture, givenSingleAddressSpaceWhenDebuggerIsCreated
     EXPECT_EQ(0u, debugger->getSbaTrackingGpuVa());
     std::vector<NEO::GraphicsAllocation *> allocations;
 
-    auto &allEngines = pDevice->getMemoryManager()->getRegisteredEngines();
+    auto &allEngines = pDevice->getMemoryManager()->getRegisteredEngines(rootDeviceIndex);
 
     for (auto &engine : allEngines) {
         auto sbaAllocation = debugger->getSbaTrackingBuffer(engine.osContext->getContextId());

shared/test/unit_test/direct_submission/linux/drm_direct_submission_tests.cpp

@@ -432,7 +432,7 @@ HWTEST_F(DrmDirectSubmissionTest, givenNoCompletionFenceSupportWhenSubmittingThe
     auto initialBO = ringBuffer->getBufferObjectToModify(0);
 
     auto drm = executionEnvironment.rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Drm>();
-    MockBufferObject mockBO(drm);
+    MockBufferObject mockBO(0, drm);
     ringBuffer->getBufferObjectToModify(0) = &mockBO;
 
     for (auto i = 0; i < 2; i++) {
@@ -459,7 +459,7 @@ HWTEST_F(DrmDirectSubmissionTest, givenNoCompletionFenceSupportAndExecFailureWhe
     auto initialBO = ringBuffer->getBufferObjectToModify(0);
 
     auto drm = executionEnvironment.rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Drm>();
-    MockBufferObject mockBO(drm);
+    MockBufferObject mockBO(0, drm);
     ringBuffer->getBufferObjectToModify(0) = &mockBO;
 
     mockBO.execReturnValue = ENXIO;
@@ -484,7 +484,7 @@ HWTEST_F(DrmDirectSubmissionTest, givenCompletionFenceSupportAndExecFailureWhenS
     auto ringBuffer = static_cast<DrmAllocation *>(drmDirectSubmission.ringBuffers[drmDirectSubmission.currentRingBuffer].ringBuffer);
     auto initialBO = ringBuffer->getBufferObjectToModify(0);
 
-    MockBufferObject mockBO(drm);
+    MockBufferObject mockBO(0, drm);
     ringBuffer->getBufferObjectToModify(0) = &mockBO;
 
     mockBO.execReturnValue = 1;
@@ -518,7 +518,7 @@ HWTEST_F(DrmDirectSubmissionTest, givenTile0AndCompletionFenceSupportWhenSubmitt
     auto ringBuffer = static_cast<DrmAllocation *>(drmDirectSubmission.ringBuffers[drmDirectSubmission.currentRingBuffer].ringBuffer);
     auto initialBO = ringBuffer->getBufferObjectToModify(0);
 
-    MockBufferObject mockBO(drm);
+    MockBufferObject mockBO(0, drm);
     ringBuffer->getBufferObjectToModify(0) = &mockBO;
 
     for (auto i = 0u; i < 2; i++) {
@@ -557,7 +557,7 @@ HWTEST_F(DrmDirectSubmissionTest, givenTile1AndCompletionFenceSupportWhenSubmitt
     auto ringBuffer = static_cast<DrmAllocation *>(drmDirectSubmission.ringBuffers[drmDirectSubmission.currentRingBuffer].ringBuffer);
     auto initialBO = ringBuffer->getBufferObjectToModify(0);
 
-    MockBufferObject mockBO(drm);
+    MockBufferObject mockBO(0, drm);
     ringBuffer->getBufferObjectToModify(0) = &mockBO;
 
     for (auto i = 0u; i < 2; i++) {
@@ -602,7 +602,7 @@ HWTEST_F(DrmDirectSubmissionTest, givenTwoTilesAndCompletionFenceSupportWhenSubm
     auto ringBuffer = static_cast<DrmAllocation *>(drmDirectSubmission.ringBuffers[drmDirectSubmission.currentRingBuffer].ringBuffer);
     auto initialBO = ringBuffer->getBufferObjectToModify(0);
 
-    MockBufferObject mockBO(drm);
+    MockBufferObject mockBO(0, drm);
     ringBuffer->getBufferObjectToModify(0) = &mockBO;
 
     for (auto i = 0u; i < 2; i++) {

shared/test/unit_test/helpers/timestamp_packet_tests.cpp

@@ -215,14 +215,14 @@ struct DeviceTimestampPacketTests : public ::testing::Test, DeviceFixture {
 };
 
 HWTEST_F(DeviceTimestampPacketTests, givenCommandStreamReceiverHwWhenObtainingPreferredTagPoolSizeThenReturnCorrectValue) {
-    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines()[0].osContext;
+    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines(mockRootDeviceIndex)[0].osContext;
 
     CommandStreamReceiverHw<FamilyType> csr(*executionEnvironment, 0, osContext.getDeviceBitfield());
     EXPECT_EQ(2048u, csr.getPreferredTagPoolSize());
 }
 
 HWTEST_F(DeviceTimestampPacketTests, givenDebugFlagSetWhenCreatingAllocatorThenUseCorrectSize) {
-    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines()[0].osContext;
+    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines(mockRootDeviceIndex)[0].osContext;
 
     {
         CommandStreamReceiverHw<FamilyType> csr(*executionEnvironment, 0, osContext.getDeviceBitfield());
@@ -269,14 +269,14 @@ HWTEST_F(DeviceTimestampPacketTests, givenDebugFlagSetWhenCreatingAllocatorThenU
 }
 
 HWCMDTEST_F(IGFX_XE_HP_CORE, DeviceTimestampPacketTests, givenInvalidDebugFlagSetWhenCreatingCsrThenExceptionIsThrown) {
-    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines()[0].osContext;
+    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines(mockRootDeviceIndex)[0].osContext;
     DebugManager.flags.OverrideTimestampPacketSize.set(12);
 
     EXPECT_ANY_THROW(CommandStreamReceiverHw<FamilyType> csr(*executionEnvironment, 0, osContext.getDeviceBitfield()));
 }
 
 HWTEST_F(DeviceTimestampPacketTests, givenTagAlignmentWhenCreatingAllocatorThenGpuAddressIsAligned) {
-    auto csr = executionEnvironment->memoryManager->getRegisteredEngines()[0].commandStreamReceiver;
+    auto csr = executionEnvironment->memoryManager->getRegisteredEngines(mockRootDeviceIndex)[0].commandStreamReceiver;
 
     auto &gfxCoreHelper = pDevice->getGfxCoreHelper();
 
@@ -292,7 +292,7 @@ HWTEST_F(DeviceTimestampPacketTests, givenTagAlignmentWhenCreatingAllocatorThenG
 HWTEST_F(DeviceTimestampPacketTests, givenDebugFlagSetWhenCreatingTimestampPacketAllocatorThenDisableReusingAndLimitPoolSize) {
     DebugManagerStateRestore restore;
     DebugManager.flags.DisableTimestampPacketOptimizations.set(true);
-    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines()[0].osContext;
+    OsContext &osContext = *executionEnvironment->memoryManager->getRegisteredEngines(mockRootDeviceIndex)[0].osContext;
 
     CommandStreamReceiverHw<FamilyType> csr(*executionEnvironment, 0, osContext.getDeviceBitfield());
     csr.setupContext(osContext);

shared/test/unit_test/memory_manager/deferrable_allocation_deletion_tests.cpp

@@ -216,14 +216,15 @@ TEST_F(DeferrableAllocationDeletionTest, givenNotUsedAllocationWhenDeletionIsApp
 }
 
 TEST_F(DeferrableAllocationDeletionTest, givenAllocationUsedByUnregisteredEngineWhenDeletionIsAppliedThenReturnTrue) {
-    auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});
+    auto rootDeviceIndex = device->getRootDeviceIndex();
+    auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{rootDeviceIndex, MemoryConstants::pageSize});
     allocation->updateTaskCount(2u, defaultOsContextId);
     EXPECT_TRUE(allocation->isUsed());
     DeferrableAllocationDeletion deletion{*memoryManager, *allocation};
 
     device.reset();
     executionEnvironment->rootDeviceEnvironments.clear();
-    EXPECT_EQ(0u, memoryManager->registeredEngines.size());
+    EXPECT_EQ(0u, memoryManager->getRegisteredEngines(rootDeviceIndex).size());
     EXPECT_TRUE(allocation->isUsed());
 
     memoryManager->freeGraphicsMemoryCalled = 0u;

shared/test/unit_test/memory_manager/host_ptr_manager_tests.cpp

@@ -904,7 +904,7 @@ HWTEST_F(HostPtrAllocationTest, givenOverlappingFragmentsWhenCheckIsCalledThenWa
     TaskCountType taskCountReady = 2;
     TaskCountType taskCountNotReady = 1;
 
-    auto &engines = memoryManager->getRegisteredEngines();
+    auto &engines = memoryManager->getRegisteredEngines(mockRootDeviceIndex);
     EXPECT_EQ(1u, engines.size());
 
     auto csr0 = static_cast<MockCommandStreamReceiver *>(engines[0].commandStreamReceiver);

shared/test/unit_test/memory_manager/memory_manager_tests.cpp

@@ -85,7 +85,7 @@ TEST(MemoryManagerTest, givenMemoryManagerWhenGettingDefaultContextThenCorrectCo
     EXPECT_NE(nullptr, osContext2);
     EXPECT_EQ(osContext1, executionEnvironment.memoryManager->getDefaultEngineContext(0, 1));
     EXPECT_EQ(osContext2, executionEnvironment.memoryManager->getDefaultEngineContext(0, 3));
-    EXPECT_EQ(mockMemoryManager->getRegisteredEngines()[mockMemoryManager->defaultEngineIndex[0]].osContext,
+    EXPECT_EQ(mockMemoryManager->getRegisteredEngines(0)[mockMemoryManager->defaultEngineIndex[0]].osContext,
               executionEnvironment.memoryManager->getDefaultEngineContext(0, 2));
 }
 
@@ -124,8 +124,8 @@ TEST(MemoryManagerTest, givenMultipleDevicesMemoryManagerWhenGettingDefaultConte
 
     EXPECT_EQ(osContext3, executionEnvironment.memoryManager->getDefaultEngineContext(1, 1));
     EXPECT_EQ(osContext4, executionEnvironment.memoryManager->getDefaultEngineContext(1, 3));
-    EXPECT_EQ(mockMemoryManager->getRegisteredEngines()[mockMemoryManager->defaultEngineIndex[1]].osContext,
+    EXPECT_EQ(mockMemoryManager->getRegisteredEngines(1)[mockMemoryManager->defaultEngineIndex[1]].osContext,
-              executionEnvironment.memoryManager->getDefaultEngineContext(0, 2));
+              executionEnvironment.memoryManager->getDefaultEngineContext(1, 2));
 }
 
 TEST(MemoryManagerTest, givenFailureOnRegisterSystemMemoryAllocationWhenAllocatingMemoryThenNullptrIsReturned) {

shared/test/unit_test/os_interface/linux/drm_buffer_object_tests.cpp

@@ -115,7 +115,7 @@ TEST_F(DrmBufferObjectTest, whenExecFailsThenPinFails) {
     mock->ioctlRes = -1;
     this->mock->errnoValue = EINVAL;
 
-    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(this->mock.get()));
+    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(rootDeviceIndex, this->mock.get()));
     ASSERT_NE(nullptr, boToPin.get());
 
     bo->setAddress(reinterpret_cast<uint64_t>(buff.get()));
@@ -131,7 +131,7 @@ TEST_F(DrmBufferObjectTest, whenExecFailsThenValidateHostPtrFails) {
     mock->ioctlRes = -1;
     this->mock->errnoValue = EINVAL;
 
-    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(this->mock.get()));
+    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(rootDeviceIndex, this->mock.get()));
     ASSERT_NE(nullptr, boToPin.get());
 
     bo->setAddress(reinterpret_cast<uint64_t>(buff.get()));
@@ -146,7 +146,7 @@ TEST_F(DrmBufferObjectTest, givenResidentBOWhenPrintExecutionBufferIsSetToTrueTh
     DebugManager.flags.PrintExecutionBuffer.set(true);
 
     std::unique_ptr<uint32_t[]> buff(new uint32_t[1024]);
-    std::unique_ptr<BufferObject> bo(new TestedBufferObject(this->mock.get()));
+    std::unique_ptr<BufferObject> bo(new TestedBufferObject(rootDeviceIndex, this->mock.get()));
     ASSERT_NE(nullptr, bo.get());
     bo->setAddress(reinterpret_cast<uint64_t>(buff.get()));
     BufferObject *boArray[1] = {bo.get()};
@@ -234,13 +234,13 @@ TEST(DrmBufferObjectSimpleTest, givenInvalidBoWhenValidateHostptrIsCalledThenErr
     executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     ASSERT_NE(nullptr, mock.get());
-    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(0u, mock.get()));
     ASSERT_NE(nullptr, bo.get());
 
     // fail DRM_IOCTL_I915_GEM_EXECBUFFER2 in pin
     mock->ioctlRes = -1;
 
-    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(mock.get()));
+    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(0u, mock.get()));
     ASSERT_NE(nullptr, boToPin.get());
 
     bo->setAddress(reinterpret_cast<uint64_t>(buff.get()));
@@ -259,13 +259,13 @@ TEST(DrmBufferObjectSimpleTest, givenInvalidBoWhenPinIsCalledThenErrorIsReturned
     executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     ASSERT_NE(nullptr, mock.get());
-    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(0u, mock.get()));
     ASSERT_NE(nullptr, bo.get());
 
     // fail DRM_IOCTL_I915_GEM_EXECBUFFER2 in pin
     mock->ioctlRes = -1;
 
-    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(mock.get()));
+    std::unique_ptr<BufferObject> boToPin(new TestedBufferObject(0u, mock.get()));
     ASSERT_NE(nullptr, boToPin.get());
 
     bo->setAddress(reinterpret_cast<uint64_t>(buff.get()));
@@ -280,25 +280,26 @@ TEST(DrmBufferObjectSimpleTest, givenInvalidBoWhenPinIsCalledThenErrorIsReturned
 TEST(DrmBufferObjectSimpleTest, givenBufferObjectWhenConstructedWithASizeThenTheSizeIsInitialized) {
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> drmMock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
-    std::unique_ptr<BufferObject> bo(new BufferObject(drmMock.get(), 3, 1, 0x1000, 1));
+    std::unique_ptr<BufferObject> bo(new BufferObject(0u, drmMock.get(), 3, 1, 0x1000, 1));
 
     EXPECT_EQ(0x1000u, bo->peekSize());
 }
 
 TEST(DrmBufferObjectSimpleTest, givenArrayOfBosWhenPinnedThenAllBosArePinned) {
+    const auto rootDeviceIndex = 0u;
     std::unique_ptr<uint32_t[]> buff(new uint32_t[256]);
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> mock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
     ASSERT_NE(nullptr, mock.get());
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
 
-    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(rootDeviceIndex, mock.get()));
     ASSERT_NE(nullptr, bo.get());
     mock->ioctlRes = 0;
 
-    std::unique_ptr<TestedBufferObject> boToPin(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> boToPin(new TestedBufferObject(rootDeviceIndex, mock.get()));
-    std::unique_ptr<TestedBufferObject> boToPin2(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> boToPin2(new TestedBufferObject(rootDeviceIndex, mock.get()));
-    std::unique_ptr<TestedBufferObject> boToPin3(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> boToPin3(new TestedBufferObject(rootDeviceIndex, mock.get()));
 
     ASSERT_NE(nullptr, boToPin.get());
     ASSERT_NE(nullptr, boToPin2.get());
@@ -320,19 +321,20 @@ TEST(DrmBufferObjectSimpleTest, givenArrayOfBosWhenPinnedThenAllBosArePinned) {
 }
 
 TEST(DrmBufferObjectSimpleTest, givenArrayOfBosWhenValidatedThenAllBosArePinned) {
+    const auto rootDeviceIndex = 0u;
     std::unique_ptr<uint32_t[]> buff(new uint32_t[256]);
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> mock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
     ASSERT_NE(nullptr, mock.get());
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
 
-    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(rootDeviceIndex, mock.get()));
     ASSERT_NE(nullptr, bo.get());
     mock->ioctlRes = 0;
 
-    std::unique_ptr<TestedBufferObject> boToPin(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> boToPin(new TestedBufferObject(rootDeviceIndex, mock.get()));
-    std::unique_ptr<TestedBufferObject> boToPin2(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> boToPin2(new TestedBufferObject(rootDeviceIndex, mock.get()));
-    std::unique_ptr<TestedBufferObject> boToPin3(new TestedBufferObject(mock.get()));
+    std::unique_ptr<TestedBufferObject> boToPin3(new TestedBufferObject(rootDeviceIndex, mock.get()));
 
     ASSERT_NE(nullptr, boToPin.get());
     ASSERT_NE(nullptr, boToPin2.get());
@@ -358,7 +360,7 @@ TEST_F(DrmBufferObjectTest, givenDeleterWhenBufferObjectIsCreatedAndDeletedThenC
     mock->ioctlExpected.reset();
 
     {
-        std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(mock.get(), 3, 1, 0x1000, 1));
+        std::unique_ptr<BufferObject, BufferObject::Deleter> bo(new BufferObject(0u, mock.get(), 3, 1, 0x1000, 1));
     }
 
     EXPECT_EQ(1, mock->ioctlCnt.gemClose);
@@ -379,8 +381,8 @@ TEST(DrmBufferObject, givenPerContextVmRequiredWhenBoCreatedThenBindInfoIsInitia
     device->getExecutionEnvironment()->calculateMaxOsContextCount();
     DrmMock drm(*(device->getExecutionEnvironment()->rootDeviceEnvironments[0].get()));
     EXPECT_TRUE(drm.isPerContextVMRequired());
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex()).size();
-    MockBufferObject bo(&drm, 3, 0, 0, osContextCount);
+    MockBufferObject bo(device->getRootDeviceIndex(), &drm, 3, 0, 0, osContextCount);
 
     EXPECT_EQ(osContextCount, bo.bindInfo.size());
 
@@ -409,8 +411,9 @@ TEST(DrmBufferObject, givenDrmIoctlReturnsErrorNotSupportedThenBufferObjectRetur
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContextCount = engines.size();
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
 
     std::unique_ptr<OsContextLinux> osContext;
     osContext.reset(new OsContextLinux(*drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor()));
@@ -437,13 +440,14 @@ TEST(DrmBufferObject, givenPerContextVmRequiredWhenBoBoundAndUnboundThenCorrectB
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContextCount = engines.size();
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
 
     EXPECT_EQ(osContextCount, bo.bindInfo.size());
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto contextId = osContextCount / 2;
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContext = engines[contextId].osContext;
     osContext->ensureContextInitialized();
 
     bo.bind(osContext, 0);
@@ -480,13 +484,14 @@ TEST(DrmBufferObject, givenPrintBOBindingResultWhenBOBindAndUnbindSucceedsThenPr
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContextCount = engines.size();
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
 
     EXPECT_EQ(osContextCount, bo.bindInfo.size());
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto contextId = osContextCount / 2;
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContext = engines[contextId].osContext;
     osContext->ensureContextInitialized();
 
     testing::internal::CaptureStdout();
@@ -537,13 +542,14 @@ TEST(DrmBufferObject, givenPrintBOBindingResultWhenBOBindAndUnbindFailsThenPrint
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContextCount = engines.size();
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
 
     EXPECT_EQ(osContextCount, bo.bindInfo.size());
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto contextId = osContextCount / 2;
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContext = engines[contextId].osContext;
     osContext->ensureContextInitialized();
 
     testing::internal::CaptureStderr();
@@ -589,10 +595,11 @@ TEST(DrmBufferObject, givenDrmWhenBindOperationFailsThenFenceValueNotGrow) {
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContextCount = engines.size();
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto contextId = osContextCount / 2;
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContext = engines[contextId].osContext;
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
     drm->bindBufferObject(osContext, 0, &bo);
 
     EXPECT_EQ(drm->fenceVal[0], initFenceValue);
@@ -620,10 +627,11 @@ TEST(DrmBufferObject, givenDrmWhenBindOperationSucceedsThenFenceValueGrow) {
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContextCount = engines.size();
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto contextId = osContextCount / 2;
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContext = engines[contextId].osContext;
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
     drm->bindBufferObject(osContext, 0, &bo);
 
     EXPECT_EQ(drm->fenceVal[0], initFenceValue + 1);
@@ -651,10 +659,11 @@ TEST(DrmBufferObject, givenDrmWhenUnBindOperationFailsThenFenceValueNotGrow) {
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContextCount = engines.size();
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto contextId = osContextCount / 2;
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContext = engines[contextId].osContext;
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
     drm->unbindBufferObject(osContext, 0, &bo);
 
     EXPECT_EQ(drm->fenceVal[0], initFenceValue);
@@ -682,10 +691,11 @@ TEST(DrmBufferObject, givenDrmWhenUnBindOperationSucceedsThenFenceValueGrow) {
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
-    auto contextId = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount() / 2;
+    auto &engines = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines(device->getRootDeviceIndex());
-    auto osContext = device->getExecutionEnvironment()->memoryManager->getRegisteredEngines()[contextId].osContext;
+    auto osContextCount = engines.size();
-    auto osContextCount = device->getExecutionEnvironment()->memoryManager->getRegisteredEnginesCount();
+    auto contextId = osContextCount / 2;
-    MockBufferObject bo(drm, 3, 0, 0, osContextCount);
+    auto osContext = engines[contextId].osContext;
+    MockBufferObject bo(device->getRootDeviceIndex(), drm, 3, 0, 0, osContextCount);
     drm->unbindBufferObject(osContext, 0, &bo);
 
     EXPECT_EQ(drm->fenceVal[0], initFenceValue + 1);
@@ -695,7 +705,7 @@ TEST(DrmBufferObject, whenBindExtHandleAddedThenItIsStored) {
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
     DrmMockResources drm(*executionEnvironment->rootDeviceEnvironments[0]);
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0, &drm, 3, 0, 0, 1);
     bo.addBindExtHandle(4);
 
     EXPECT_EQ(1u, bo.bindExtHandles.size());
@@ -709,7 +719,7 @@ TEST(DrmBufferObject, whenMarkForCapturedCalledThenIsMarkedForCaptureReturnsTrue
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
     DrmMockResources drm(*executionEnvironment->rootDeviceEnvironments[0]);
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0, &drm, 3, 0, 0, 1);
     EXPECT_FALSE(bo.isMarkedForCapture());
 
     bo.markForCapture();
@@ -749,7 +759,7 @@ TEST_F(DrmBufferObjectTest, given47bitAddressWhenSetThenIsAddressNotCanonized) {
 
     uint64_t address = maxNBitValue(47) - maxNBitValue(5);
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0, &drm, 3, 0, 0, 1);
     bo.setAddress(address);
     auto boAddress = bo.peekAddress();
     EXPECT_EQ(boAddress, address);
@@ -764,7 +774,7 @@ TEST_F(DrmBufferObjectTest, given48bitAddressWhenSetThenAddressIsCanonized) {
     uint64_t address = maxNBitValue(48) - maxNBitValue(5);
     uint64_t expectedAddress = std::numeric_limits<uint64_t>::max() - maxNBitValue(5);
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0, &drm, 3, 0, 0, 1);
     bo.setAddress(address);
     auto boAddress = bo.peekAddress();
     EXPECT_EQ(boAddress, expectedAddress);
@@ -776,7 +786,7 @@ TEST_F(DrmBufferObjectTest, givenBoIsCreatedWhenPageFaultIsSupportedThenExplicit
 
     for (auto isPageFaultSupported : {false, true}) {
         drm.pageFaultSupported = isPageFaultSupported;
-        MockBufferObject bo(&drm, 3, 0, 0, 1);
+        MockBufferObject bo(0, &drm, 3, 0, 0, 1);
         EXPECT_EQ(isPageFaultSupported, bo.isExplicitResidencyRequired());
     }
 }
@@ -784,7 +794,7 @@ TEST_F(DrmBufferObjectTest, givenBoIsCreatedWhenPageFaultIsSupportedThenExplicit
 TEST_F(DrmBufferObjectTest, whenBoRequiresExplicitResidencyThenTheCorrespondingQueryReturnsCorrectValue) {
     MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
     DrmMock drm(*(executionEnvironment.rootDeviceEnvironments[0].get()));
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0, &drm, 3, 0, 0, 1);
 
     for (auto required : {false, true}) {
         bo.requireExplicitResidency(required);

shared/test/unit_test/os_interface/linux/drm_command_stream_mm_tests.cpp

@@ -104,8 +104,8 @@ HWTEST_TEMPLATED_F(DrmCommandStreamMemExecTest, GivenDrmSupportsVmBindAndComplet
     mock->isVmBindAvailableCall.callParent = false;
     mock->isVmBindAvailableCall.returnValue = true;
 
-    TestedBufferObject bo(mock, 128);
+    TestedBufferObject bo(rootDeviceIndex, mock, 128);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
     cmdBuffer.bufferObjects[0] = &bo;
     uint8_t buff[128];
 
@@ -138,8 +138,8 @@ HWTEST_TEMPLATED_F(DrmCommandStreamMemExecTest, GivenDrmSupportsVmBindAndNotComp
     mock->isVmBindAvailableCall.callParent = false;
     mock->isVmBindAvailableCall.returnValue = true;
 
-    TestedBufferObject bo(mock, 128);
+    TestedBufferObject bo(rootDeviceIndex, mock, 128);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
     cmdBuffer.bufferObjects[0] = &bo;
     uint8_t buff[128];
 
@@ -173,8 +173,8 @@ HWTEST_TEMPLATED_F(DrmCommandStreamMemExecTest, GivenDrmSupportsCompletionFenceA
     mock->isVmBindAvailableCall.callParent = false;
     mock->isVmBindAvailableCall.returnValue = false;
 
-    TestedBufferObject bo(mock, 128);
+    TestedBufferObject bo(rootDeviceIndex, mock, 128);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
     cmdBuffer.bufferObjects[0] = &bo;
     uint8_t buff[128];
 

shared/test/unit_test/os_interface/linux/drm_command_stream_tests.cpp

@@ -82,8 +82,9 @@ HWTEST_TEMPLATED_F(DrmCommandStreamTest, GivenExecBufferErrorWhenFlushInternalTh
     mock->execBufferResult = -1;
     mock->baseErrno = false;
     mock->errnoRetVal = EWOULDBLOCK;
-    TestedBufferObject bo(mock, 128);
+    auto rootDeviceIndex = csr->getRootDeviceIndex();
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    TestedBufferObject bo(rootDeviceIndex, mock, 128);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
     cmdBuffer.bufferObjects[0] = &bo;
     uint8_t buff[128]{};
 

shared/test/unit_test/os_interface/linux/drm_command_stream_tests_1.cpp

@@ -1066,9 +1066,9 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, givenPrintBOsForSubmitWhenPrint
     DebugManagerStateRestore restorer;
     DebugManager.flags.PrintBOsForSubmit.set(true);
 
-    MockDrmAllocation allocation1(AllocationType::BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation allocation1(rootDeviceIndex, AllocationType::BUFFER, MemoryPool::System4KBPages);
-    MockDrmAllocation allocation2(AllocationType::BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation allocation2(rootDeviceIndex, AllocationType::BUFFER, MemoryPool::System4KBPages);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
 
     csr->makeResident(allocation1);
     csr->makeResident(allocation2);
@@ -1102,9 +1102,9 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, givenPrintBOsForSubmitAndFailur
     DebugManagerStateRestore restorer;
     DebugManager.flags.PrintBOsForSubmit.set(true);
 
-    MockDrmAllocation allocation1(AllocationType::BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation allocation1(rootDeviceIndex, AllocationType::BUFFER, MemoryPool::System4KBPages);
-    MockDrmAllocation allocation2(AllocationType::BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation allocation2(rootDeviceIndex, AllocationType::BUFFER, MemoryPool::System4KBPages);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
 
     cmdBuffer.makeBOsResidentResult = ENOSPC;
 
@@ -1127,9 +1127,9 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, givenPrintBOsForSubmitAndFailur
     DebugManagerStateRestore restorer;
     DebugManager.flags.PrintBOsForSubmit.set(true);
 
-    MockDrmAllocation allocation1(AllocationType::BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation allocation1(rootDeviceIndex, AllocationType::BUFFER, MemoryPool::System4KBPages);
-    MockDrmAllocation allocation2(AllocationType::BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation allocation2(rootDeviceIndex, AllocationType::BUFFER, MemoryPool::System4KBPages);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(rootDeviceIndex, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
 
     allocation1.makeBOsResidentResult = ENOSPC;
 

shared/test/unit_test/os_interface/linux/drm_command_stream_xehp_and_later_prelim_tests.cpp

@@ -442,15 +442,15 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DrmImplicitScalingCommandStreamTest, givenTwoTilesW
     auto csr = createCsr<FamilyType>();
 
     auto size = 1024u;
-    auto multiStorageBo0 = new BufferObject(drm, 3, 30, 0, 1);
+    auto multiStorageBo0 = new BufferObject(0u, drm, 3, 30, 0, 1);
-    auto multiStorageBo1 = new BufferObject(drm, 3, 31, 0, 1);
+    auto multiStorageBo1 = new BufferObject(0u, drm, 3, 31, 0, 1);
     BufferObjects multiStorageBos{multiStorageBo0, multiStorageBo1};
     auto multiStorageAllocation = new DrmAllocation(0, AllocationType::UNKNOWN, multiStorageBos, nullptr, 0u, size, MemoryPool::LocalMemory);
     multiStorageAllocation->storageInfo.memoryBanks = 0b11;
     csr->CommandStreamReceiver::makeResident(*multiStorageAllocation);
 
-    auto tileInstancedBo0 = new BufferObject(drm, 3, 40, 0, 1);
+    auto tileInstancedBo0 = new BufferObject(0u, drm, 3, 40, 0, 1);
-    auto tileInstancedBo1 = new BufferObject(drm, 3, 41, 0, 1);
+    auto tileInstancedBo1 = new BufferObject(0u, drm, 3, 41, 0, 1);
     BufferObjects tileInstancedBos{tileInstancedBo0, tileInstancedBo1};
     auto tileInstancedAllocation = new DrmAllocation(0, AllocationType::UNKNOWN, tileInstancedBos, nullptr, 0u, size, MemoryPool::LocalMemory);
     tileInstancedAllocation->storageInfo.memoryBanks = 0b11;
@@ -514,10 +514,10 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DrmImplicitScalingCommandStreamTest, whenForceExecu
                                          gemCloseWorkerMode::gemCloseWorkerActive);
     csr->setupContext(*osContext);
 
-    auto tileInstancedBo0 = new BufferObject(drm, 3, 40, 0, 1);
+    auto tileInstancedBo0 = new BufferObject(0u, drm, 3, 40, 0, 1);
-    auto tileInstancedBo1 = new BufferObject(drm, 3, 41, 0, 1);
+    auto tileInstancedBo1 = new BufferObject(0u, drm, 3, 41, 0, 1);
-    auto tileInstancedBo2 = new BufferObject(drm, 3, 42, 0, 1);
+    auto tileInstancedBo2 = new BufferObject(0u, drm, 3, 42, 0, 1);
-    auto tileInstancedBo3 = new BufferObject(drm, 3, 43, 0, 1);
+    auto tileInstancedBo3 = new BufferObject(0u, drm, 3, 43, 0, 1);
     BufferObjects tileInstancedBos{tileInstancedBo0, tileInstancedBo1, tileInstancedBo2, tileInstancedBo3};
     auto tileInstancedAllocation = new DrmAllocation(0, AllocationType::UNKNOWN, tileInstancedBos, nullptr, 0u, 1024u, MemoryPool::LocalMemory);
     tileInstancedAllocation->storageInfo.memoryBanks = 0b11;
@@ -557,10 +557,10 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DrmImplicitScalingCommandStreamTest, whenForceExecu
                                          gemCloseWorkerMode::gemCloseWorkerActive);
     csr->setupContext(*osContext);
 
-    auto tileInstancedBo0 = new BufferObject(drm, 3, 40, 0, 1);
+    auto tileInstancedBo0 = new BufferObject(0u, drm, 3, 40, 0, 1);
-    auto tileInstancedBo1 = new BufferObject(drm, 3, 41, 0, 1);
+    auto tileInstancedBo1 = new BufferObject(0u, drm, 3, 41, 0, 1);
-    auto tileInstancedBo2 = new BufferObject(drm, 3, 42, 0, 1);
+    auto tileInstancedBo2 = new BufferObject(0u, drm, 3, 42, 0, 1);
-    auto tileInstancedBo3 = new BufferObject(drm, 3, 43, 0, 1);
+    auto tileInstancedBo3 = new BufferObject(0u, drm, 3, 43, 0, 1);
     BufferObjects tileInstancedBos{tileInstancedBo0, tileInstancedBo1, tileInstancedBo2, tileInstancedBo3};
     auto tileInstancedAllocation = new DrmAllocation(0, AllocationType::UNKNOWN, tileInstancedBos, nullptr, 0u, 1024u, MemoryPool::LocalMemory);
     tileInstancedAllocation->storageInfo.memoryBanks = 0b11;
@@ -606,7 +606,7 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DrmImplicitScalingCommandStreamTest, givenDisabledI
     csr->setupContext(*osContext);
 
     const auto size = 1024u;
-    BufferObject *bufferObject = new BufferObject(drm, 3, 30, 0, 1);
+    BufferObject *bufferObject = new BufferObject(0u, drm, 3, 30, 0, 1);
     BufferObjects bufferObjects{bufferObject};
     auto allocation = new DrmAllocation(0, AllocationType::UNKNOWN, bufferObjects, nullptr, 0u, size, MemoryPool::LocalMemory);
     csr->CommandStreamReceiver::makeResident(*allocation);
@@ -640,7 +640,7 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DrmImplicitScalingCommandStreamTest, givenMultiTile
     csr->setupContext(*osContext);
 
     const auto size = 1024u;
-    BufferObject *bufferObject = new BufferObject(drm, 3, 30, 0, 1);
+    BufferObject *bufferObject = new BufferObject(0u, drm, 3, 30, 0, 1);
     BufferObjects bufferObjects{bufferObject};
     auto allocation = new DrmAllocation(0, AllocationType::UNKNOWN, bufferObjects, nullptr, 0u, size, MemoryPool::LocalMemory);
     csr->CommandStreamReceiver::makeResident(*allocation);

shared/test/unit_test/os_interface/linux/drm_command_stream_xehp_and_later_tests.cpp

@@ -82,8 +82,8 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DrmCommandStreamMultiTileMemExecTest, GivenDrmSuppo
     mock->isVmBindAvailableCall.callParent = false;
     mock->isVmBindAvailableCall.returnValue = true;
 
-    TestedBufferObject bo(mock, 128);
+    TestedBufferObject bo(0, mock, 128);
-    MockDrmAllocation cmdBuffer(AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
+    MockDrmAllocation cmdBuffer(0u, AllocationType::COMMAND_BUFFER, MemoryPool::System4KBPages);
     cmdBuffer.bufferObjects[0] = &bo;
     uint8_t buff[128];
 

shared/test/unit_test/os_interface/linux/drm_debug_prelim_tests.cpp

@@ -27,6 +27,7 @@ struct DrmDebugPrelimTest : public ::testing::Test {
         executionEnvironment->rootDeviceEnvironments[0]->initGmm();
         executionEnvironment->initializeMemoryManager();
     }
+    const uint32_t rootDeviceIndex = 0u;
 
   protected:
     std::unique_ptr<ExecutionEnvironment> executionEnvironment;
@@ -242,7 +243,7 @@ TEST(DrmPrelimTest, givenContextDebugNotAvailableWhenCheckedForSupportThenTrueIs
 TEST_F(DrmDebugPrelimTest, givenAddedBindExtHandlesInBoWhenBindingWithinDefaultEngineContextThenExtensionsArePassedToVmBindIoctl) {
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(rootDeviceIndex, &drm, 3, 0, 0, 1);
     bo.addBindExtHandle(4);
     bo.addBindExtHandle(5);
 
@@ -262,7 +263,7 @@ TEST_F(DrmDebugPrelimTest, givenAddedBindExtHandlesInBoWhenBindingWithinDefaultE
 TEST_F(DrmDebugPrelimTest, givenAddedBindExtHandlesInBoWhenBindingWithinInternalContextThenExtensionsAreNotPassedToVmBindIoctl) {
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(rootDeviceIndex, &drm, 3, 0, 0, 1);
     bo.addBindExtHandle(4);
     bo.addBindExtHandle(5);
 
@@ -276,7 +277,7 @@ TEST_F(DrmDebugPrelimTest, givenAddedBindExtHandlesInBoWhenBindingWithinInternal
 TEST_F(DrmDebugPrelimTest, givenAddedBindExtHandlesInBoWhenBindingWithinCopyEngineContextThenExtensionsAreNotPassedToVmBindIoctl) {
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(rootDeviceIndex, &drm, 3, 0, 0, 1);
     bo.addBindExtHandle(4);
     bo.addBindExtHandle(5);
 
@@ -293,7 +294,7 @@ HWTEST_F(DrmDebugPrelimTest, givenAddedBindExtHandlesInBoWhenUnbindingThenExtens
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
     drm.queryAndSetVmBindPatIndexProgrammingSupport();
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0u, &drm, 3, 0, 0, 1);
     bo.addBindExtHandle(4);
     bo.addBindExtHandle(5);
 

shared/test/unit_test/os_interface/linux/drm_gem_close_worker_tests.cpp

@@ -57,6 +57,7 @@ class DrmGemCloseWorkerFixture {
     DrmMemoryManager *mm;
     DrmMockForWorker *drmMock;
     uint32_t deadCnt = deadCntInit;
+    const uint32_t rootDeviceIndex = 0u;
 
     void setUp() {
         this->drmMock = new DrmMockForWorker(*executionEnvironment.rootDeviceEnvironments[0]);
@@ -64,9 +65,9 @@ class DrmGemCloseWorkerFixture {
         auto hwInfo = executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo();
         drmMock->setupIoctlHelper(hwInfo->platform.eProductFamily);
 
-        executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
+        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
-        executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drmMock));
+        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drmMock));
-        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drmMock, 0u);
+        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drmMock, 0u);
 
         this->mm = new DrmMemoryManager(gemCloseWorkerMode::gemCloseWorkerInactive,
                                         false,
@@ -101,7 +102,7 @@ TEST_F(DrmGemCloseWorkerTests, WhenClosingGemThenSucceeds) {
     this->drmMock->gemCloseExpected = 1;
 
     auto worker = new DrmGemCloseWorker(*mm);
-    auto bo = new BufferObject(this->drmMock, 3, 1, 0, 1);
+    auto bo = new BufferObject(rootDeviceIndex, this->drmMock, 3, 1, 0, 1);
 
     worker->push(bo);
 
@@ -112,7 +113,7 @@ TEST_F(DrmGemCloseWorkerTests, GivenMultipleThreadsWhenClosingGemThenSucceeds) {
     this->drmMock->gemCloseExpected = -1;
 
     auto worker = new DrmGemCloseWorker(*mm);
-    auto bo = new BufferObject(this->drmMock, 3, 1, 0, 1);
+    auto bo = new BufferObject(rootDeviceIndex, this->drmMock, 3, 1, 0, 1);
 
     worker->push(bo);
 
@@ -132,7 +133,7 @@ TEST_F(DrmGemCloseWorkerTests, GivenMultipleThreadsAndCloseFalseWhenClosingGemTh
     this->drmMock->gemCloseExpected = -1;
 
     auto worker = new DrmGemCloseWorker(*mm);
-    auto bo = new BufferObject(this->drmMock, 3, 1, 0, 1);
+    auto bo = new BufferObject(rootDeviceIndex, this->drmMock, 3, 1, 0, 1);
 
     worker->push(bo);
     worker->close(false);
@@ -150,7 +151,7 @@ TEST_F(DrmGemCloseWorkerTests, givenAllocationWhenAskedForUnreferenceWithForceFl
     this->drmMock->gemCloseExpected = 1;
 
     auto worker = new DrmGemCloseWorker(*mm);
-    auto bo = new BufferObject(this->drmMock, 3, 1, 0, 1);
+    auto bo = new BufferObject(rootDeviceIndex, this->drmMock, 3, 1, 0, 1);
 
     bo->reference();
     worker->push(bo);

shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_prelim_tests.cpp

@@ -29,7 +29,7 @@
 #include "gtest/gtest.h"
 
 TEST_F(DrmMemoryManagerLocalMemoryWithCustomPrelimMockTest, givenDrmMemoryManagerWithLocalMemoryWhenLockResourceIsCalledOnBufferObjectThenReturnPtr) {
-    BufferObject bo(mock, 3, 1, 1024, 1);
+    BufferObject bo(0, mock, 3, 1, 1024, 1);
 
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, &bo, nullptr, 0u, 0u, MemoryPool::LocalMemory);
     EXPECT_EQ(&bo, drmAllocation.getBO());
@@ -54,7 +54,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenDrmMemoryManagerWithPrelimSup
     auto gpuAddress = 0x1234u;
     auto size = MemoryConstants::pageSize64k;
 
-    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(&memoryManager->getDrm(0),
+    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(0u,
                                                                                             nullptr,
                                                                                             AllocationType::BUFFER,
                                                                                             gpuAddress,
@@ -1585,7 +1585,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenPrintBOCreateDestroyResultFla
     mock->queryEngineInfo();
 
     testing::internal::CaptureStdout();
-    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(&memoryManager->getDrm(0),
+    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(0u,
                                                                                             nullptr,
                                                                                             AllocationType::BUFFER,
                                                                                             gpuAddress,
@@ -2170,7 +2170,7 @@ TEST_F(DrmMemoryManagerTestPrelim, givenDrmMemoryManagerWhenLockUnlockIsCalledOn
     this->ioctlResExt = {mock->ioctlCnt.total, -1};
     mock->ioctlResExt = &ioctlResExt;
 
-    BufferObject bo(mock, 3, 1, 0, 1);
+    BufferObject bo(0, mock, 3, 1, 0, 1);
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, &bo, nullptr, 0u, 0u, MemoryPool::LocalMemory);
     EXPECT_NE(nullptr, drmAllocation.getBO());
 
@@ -2185,7 +2185,7 @@ TEST_F(DrmMemoryManagerTestPrelim, givenDrmMemoryManagerWhenLockUnlockIsCalledOn
     mock->ioctlExpected.gemMmapOffset = 2;
     mock->failOnMmapOffset = true;
 
-    BufferObject bo(mock, 3, 1, 0, 1);
+    BufferObject bo(0, mock, 3, 1, 0, 1);
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, &bo, nullptr, 0u, 0u, MemoryPool::LocalMemory);
     EXPECT_NE(nullptr, drmAllocation.getBO());
 
@@ -2564,8 +2564,8 @@ TEST(AllocationInfoLogging, givenDrmGraphicsAllocationWithMultipleBOsWhenGetting
     executionEnvironment->rootDeviceEnvironments[0]->initGmm();
 
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
-    BufferObject bo0(&drm, 3, 0, 0, 1), bo1(&drm, 3, 1, 0, 1),
+    BufferObject bo0(0, &drm, 3, 0, 0, 1), bo1(0, &drm, 3, 1, 0, 1),
-        bo2(&drm, 3, 2, 0, 1), bo3(&drm, 3, 3, 0, 1);
+        bo2(0, &drm, 3, 2, 0, 1), bo3(0, &drm, 3, 3, 0, 1);
     BufferObjects bos{&bo0, &bo1, &bo2, &bo3};
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, bos, nullptr, 0u, 0u, MemoryPool::LocalMemory);
 

shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_upstream_tests.cpp

@@ -119,7 +119,7 @@ HWTEST2_F(DrmMemoryManagerLocalMemoryTest, givenDrmMemoryManagerWhenCreateBuffer
     auto gpuAddress = 0x1234u;
     auto size = MemoryConstants::pageSize64k;
 
-    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(&memoryManager->getDrm(0),
+    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(rootDeviceIndex,
                                                                                             nullptr,
                                                                                             AllocationType::BUFFER,
                                                                                             gpuAddress,
@@ -440,7 +440,7 @@ class DrmMemoryManagerLocalMemoryMemoryBankMock : public TestedDrmMemoryManager
                                               ExecutionEnvironment &executionEnvironment) : TestedDrmMemoryManager(enableLocalMemory, allowForcePin, validateHostPtrMemory, executionEnvironment) {
     }
 
-    BufferObject *createBufferObjectInMemoryRegion(Drm *drm,
+    BufferObject *createBufferObjectInMemoryRegion(uint32_t rootDeviceIndex,
                                                    Gmm *gmm,
                                                    AllocationType allocationType,
                                                    uint64_t gpuAddress,
@@ -640,7 +640,7 @@ TEST_F(DrmMemoryManagerLocalMemoryTest, givenDrmMemoryManagerWithLocalMemoryWhen
 }
 
 TEST_F(DrmMemoryManagerLocalMemoryWithCustomMockTest, givenDrmMemoryManagerWithLocalMemoryWhenLockResourceIsCalledOnBufferObjectThenReturnPtr) {
-    BufferObject bo(mock, 3, 1, 1024, 0);
+    BufferObject bo(0, mock, 3, 1, 1024, 0);
 
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, &bo, nullptr, 0u, 0u, MemoryPool::LocalMemory);
     EXPECT_EQ(&bo, drmAllocation.getBO());
@@ -844,7 +844,7 @@ TEST_F(DrmMemoryManagerTestImpl, givenDrmMemoryManagerWhenLockUnlockIsCalledOnAl
     this->ioctlResExt = {mockExp->ioctlCnt.total, -1};
     mockExp->ioctlResExt = &ioctlResExt;
 
-    BufferObject bo(mockExp, 3, 1, 0, 0);
+    BufferObject bo(0, mockExp, 3, 1, 0, 0);
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, &bo, nullptr, 0u, 0u, MemoryPool::LocalMemory);
     EXPECT_NE(nullptr, drmAllocation.getBO());
 
@@ -860,7 +860,7 @@ TEST_F(DrmMemoryManagerTestImpl, givenDrmMemoryManagerWhenLockUnlockIsCalledOnAl
     mockExp->returnIoctlExtraErrorValue = true;
     mockExp->failOnMmapOffset = true;
 
-    BufferObject bo(mockExp, 3, 1, 0, 0);
+    BufferObject bo(0, mockExp, 3, 1, 0, 0);
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, &bo, nullptr, 0u, 0u, MemoryPool::LocalMemory);
     EXPECT_NE(nullptr, drmAllocation.getBO());
 

shared/test/unit_test/os_interface/linux/drm_query_prelim_tests.cpp

@@ -375,7 +375,7 @@ TEST(DrmBufferObjectTestPrelim, givenBufferObjectSetToColourWithBindWhenBindingT
     executionEnvironment->rootDeviceEnvironments[0]->initGmm();
     executionEnvironment->initializeMemoryManager();
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
-    BufferObjectMock bo(&drm, 3, 1, 0, 1);
+    BufferObjectMock bo(0u, &drm, 3, 1, 0, 1);
     bo.setColourWithBind();
     bo.setColourChunk(MemoryConstants::pageSize64k);
     bo.addColouringAddress(0xffeeffee);
@@ -432,7 +432,7 @@ TEST(DrmBufferObjectTestPrelim, givenContextWhenQueryingVmIdThenIoctlIsCalled) {
     executionEnvironment->rootDeviceEnvironments[0]->initGmm();
     executionEnvironment->initializeMemoryManager();
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
-    BufferObjectMock bo(&drm, 3, 1, 0, 1);
+    BufferObjectMock bo(0u, &drm, 3, 1, 0, 1);
     OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     osContext.ensureContextInitialized();
 
@@ -449,7 +449,7 @@ TEST(DrmBufferObjectTestPrelim, givenBufferObjectMarkedForCaptureWhenBindingThen
     executionEnvironment->rootDeviceEnvironments[0]->initGmm();
     executionEnvironment->initializeMemoryManager();
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
-    BufferObjectMock bo(&drm, 3, 1, 0, 1);
+    BufferObjectMock bo(0, &drm, 3, 1, 0, 1);
     OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     osContext.ensureContextInitialized();
     bo.markForCapture();
@@ -467,7 +467,7 @@ TEST(DrmBufferObjectTestPrelim, givenNoActiveDirectSubmissionAndForceUseImmediat
     executionEnvironment->rootDeviceEnvironments[0]->initGmm();
     executionEnvironment->initializeMemoryManager();
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
-    BufferObjectMock bo(&drm, 3, 1, 0, 1);
+    BufferObjectMock bo(0u, &drm, 3, 1, 0, 1);
     OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     osContext.ensureContextInitialized();
 
@@ -483,7 +483,7 @@ TEST(DrmBufferObjectTestPrelim, whenBindingThenImmediateFlagIsSetAndExtensionLis
     executionEnvironment->initializeMemoryManager();
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
     drm.setDirectSubmissionActive(true);
-    BufferObjectMock bo(&drm, 3, 1, 0, 1);
+    BufferObjectMock bo(0u, &drm, 3, 1, 0, 1);
     OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     osContext.ensureContextInitialized();
     osContext.setDirectSubmissionActive();

shared/test/unit_test/os_interface/linux/drm_query_tests.cpp

@@ -72,6 +72,6 @@ TEST(DrmQueryTest, givenDrmAllocationWhenShouldAllocationFaultIsCalledThenReturn
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
     DrmMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::MemoryNull);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::MemoryNull);
     EXPECT_FALSE(allocation.shouldAllocationPageFault(&drm));
 }

shared/test/unit_test/os_interface/linux/drm_query_topology_prelim_tests.cpp

@@ -364,11 +364,11 @@ TEST(DrmQueryTest, givenUseKmdMigrationWhenShouldAllocationFaultIsCalledOnFaulta
         AllocationType::UNIFIED_SHARED_MEMORY};
 
     for (auto allocationType : allocationTypesThatShouldFault) {
-        MockDrmAllocation allocation(allocationType, MemoryPool::MemoryNull);
+        MockDrmAllocation allocation(0u, allocationType, MemoryPool::MemoryNull);
         EXPECT_TRUE(allocation.shouldAllocationPageFault(&drm));
     }
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::MemoryNull);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::MemoryNull);
     EXPECT_FALSE(allocation.shouldAllocationPageFault(&drm));
 }
 
@@ -388,7 +388,7 @@ TEST(DrmQueryTest, givenDrmAllocationWhenShouldAllocationFaultIsCalledOnNonFault
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
     drm.pageFaultSupported = false;
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::MemoryNull);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::MemoryNull);
     EXPECT_FALSE(allocation.shouldAllocationPageFault(&drm));
 }
 
@@ -400,6 +400,6 @@ TEST(DrmQueryTest, givenEnableImplicitMigrationOnFaultableHardwareWhenShouldAllo
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
     drm.pageFaultSupported = true;
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::MemoryNull);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::MemoryNull);
     EXPECT_TRUE(allocation.shouldAllocationPageFault(&drm));
 }

shared/test/unit_test/os_interface/linux/drm_residency_handler_prelim_tests.cpp

@@ -93,7 +93,7 @@ struct DrmMemoryOperationsHandlerBindFixture : public ::testing::Test {
         operationHandler = static_cast<MockDrmMemoryOperationsHandlerBind *>(executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface.get());
         memoryManagerBackup = executionEnvironment->memoryManager.release();
         executionEnvironment->memoryManager.reset(memoryManager.get());
-        memoryManager->registeredEngines = memoryManagerBackup->getRegisteredEngines();
+        memoryManager->allRegisteredEngines = memoryManagerBackup->getRegisteredEngines();
     }
     void SetUp() override {
         setUp(false);
@@ -102,7 +102,9 @@ struct DrmMemoryOperationsHandlerBindFixture : public ::testing::Test {
     void TearDown() override {
         executionEnvironment->memoryManager.release();
         executionEnvironment->memoryManager.reset(memoryManagerBackup);
-        memoryManager->getRegisteredEngines().clear();
+        for (auto &engineContainer : memoryManager->allRegisteredEngines) {
+            engineContainer.second.clear();
+        }
     }
 
   protected:
@@ -123,7 +125,7 @@ TEST_F(DrmMemoryOperationsHandlerBindMultiRootDeviceTest, whenSetNewResourceBoun
         using OsContextLinux::lastFlushedTlbFlushCounter;
     };
 
-    BufferObject mockBo(mock, 3, 1, 0, 1);
+    BufferObject mockBo(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
     mock->setNewResourceBoundToVM(&mockBo, 1u);
 
     for (const auto &engine : device->getAllEngines()) {
@@ -142,7 +144,8 @@ TEST_F(DrmMemoryOperationsHandlerBindMultiRootDeviceTest, whenSetNewResourceBoun
     }
 
     auto mock2 = executionEnvironment->rootDeviceEnvironments[1u]->osInterface->getDriverModel()->as<DrmQueryMock>();
-    mock2->setNewResourceBoundToVM(&mockBo, 0u);
+    BufferObject mockBo2(devices[1]->getRootDeviceIndex(), mock, 3, 1, 0, 1);
+    mock2->setNewResourceBoundToVM(&mockBo2, 0u);
 
     for (const auto &engine : devices[1]->getAllEngines()) {
         auto osContexLinux = static_cast<MockOsContextLinux *>(engine.osContext);
@@ -218,7 +221,7 @@ struct DrmMemoryOperationsHandlerBindFixture2 : public ::testing::Test {
         operationHandler = static_cast<MockDrmMemoryOperationsHandlerBind *>(executionEnvironment->rootDeviceEnvironments[1]->memoryOperationsInterface.get());
         memoryManagerBackup = executionEnvironment->memoryManager.release();
         executionEnvironment->memoryManager.reset(memoryManager.get());
-        memoryManager->registeredEngines = memoryManagerBackup->getRegisteredEngines();
+        memoryManager->allRegisteredEngines = memoryManagerBackup->getRegisteredEngines();
     }
     void SetUp() override {
         setUp(false);
@@ -227,7 +230,9 @@ struct DrmMemoryOperationsHandlerBindFixture2 : public ::testing::Test {
     void TearDown() override {
         executionEnvironment->memoryManager.release();
         executionEnvironment->memoryManager.reset(memoryManagerBackup);
-        memoryManager->getRegisteredEngines().clear();
+        for (auto &engineContainer : memoryManager->allRegisteredEngines) {
+            engineContainer.second.clear();
+        }
     }
 
   protected:
@@ -537,7 +542,7 @@ TEST_F(DrmMemoryOperationsHandlerBindTest, givenMakeBOsResidentFailsThenMakeResi
 
     auto size = 1024u;
     BufferObjects bos;
-    BufferObject mockBo(mock, 3, 1, 0, 1);
+    BufferObject mockBo(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
     bos.push_back(&mockBo);
 
     auto allocation = new MockDrmAllocationBOsResident(0, AllocationType::UNKNOWN, bos, nullptr, 0u, size, MemoryPool::LocalMemory);
@@ -682,8 +687,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportAndMultiSubdevice
     DebugManager.flags.UseVmBind.set(1);
     mock->bindAvailable = true;
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    BufferObject boToPin(mock, 3, 2, 0, 1);
+    BufferObject boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
 
     pinBB.pin(&boToPinPtr, 1u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -696,8 +701,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportAndMultiSubdevice
     DebugManager.flags.UseVmBind.set(1);
     mock->bindAvailable = true;
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    BufferObject boToPin(mock, 3, 2, 0, 1);
+    BufferObject boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
 
     pinBB.validateHostPtr(&boToPinPtr, 1u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -710,8 +715,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportAndMultiSubdevice
     DebugManager.flags.UseVmBind.set(1);
     mock->bindAvailable = true;
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    BufferObject boToPin(mock, 3, 2, 0, 1);
+    BufferObject boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
     uint32_t vmHandleId = 1u;
 
@@ -727,9 +732,9 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportAndMultiSubdevice
     mock->bindAvailable = true;
     mock->context.vmBindReturn = -1;
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    BufferObject boToPin(mock, 3, 2, 0, 1);
+    BufferObject boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
-    BufferObject boToPin2(mock, 3, 3, 0, 1);
+    BufferObject boToPin2(device->getRootDeviceIndex(), mock, 3, 3, 0, 1);
     BufferObject *boToPinPtr[] = {&boToPin, &boToPin2};
 
     auto ret = pinBB.validateHostPtr(boToPinPtr, 2u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -763,10 +768,11 @@ HWTEST_F(DrmMemoryOperationsHandlerBindWithPerContextVms, givenVmBindMultipleSub
     uint32_t vmIdForContext0 = 0;
     uint32_t vmIdForContext1 = 0;
 
-    memoryManager->registeredEngines = EngineControlContainer{this->device->allEngines};
+    auto &engines = memoryManager->allRegisteredEngines[this->device->getRootDeviceIndex()];
-    memoryManager->registeredEngines.insert(memoryManager->registeredEngines.end(), this->device->getSubDevice(0)->getAllEngines().begin(), this->device->getSubDevice(0)->getAllEngines().end());
+    engines = EngineControlContainer{this->device->allEngines};
-    memoryManager->registeredEngines.insert(memoryManager->registeredEngines.end(), this->device->getSubDevice(1)->getAllEngines().begin(), this->device->getSubDevice(1)->getAllEngines().end());
+    engines.insert(engines.end(), this->device->getSubDevice(0)->getAllEngines().begin(), this->device->getSubDevice(0)->getAllEngines().end());
-    for (auto engine : memoryManager->registeredEngines) {
+    engines.insert(engines.end(), this->device->getSubDevice(1)->getAllEngines().begin(), this->device->getSubDevice(1)->getAllEngines().end());
+    for (auto &engine : engines) {
         engine.osContext->incRefInternal();
         if (engine.osContext->isDefaultContext()) {
 
@@ -844,17 +850,17 @@ HWTEST_F(DrmMemoryOperationsHandlerBindWithPerContextVms, givenVmBindMultipleRoo
     uint32_t vmIdForDevice0 = 0;
     uint32_t vmIdForDevice0Subdevice0 = 0;
     uint32_t vmIdForDevice1 = 0;
+    uint32_t vmIdForDevice1Subdevice0 = 0;
 
-    memoryManager->registeredEngines = EngineControlContainer{this->device->allEngines};
+    {
-    memoryManager->registeredEngines.insert(memoryManager->registeredEngines.end(), this->device->getSubDevice(0)->getAllEngines().begin(), this->device->getSubDevice(0)->getAllEngines().end());
-    memoryManager->registeredEngines.insert(memoryManager->registeredEngines.end(), device1->getAllEngines().begin(), device1->getAllEngines().end());
-
-    for (auto engine : memoryManager->registeredEngines) {
-        engine.osContext->incRefInternal();
-        if (engine.osContext->isDefaultContext()) {
-
-            if (engine.osContext->getRootDeviceIndex() == 0) {
 
+        auto &engines = memoryManager->allRegisteredEngines[this->device->getRootDeviceIndex()];
+        engines = EngineControlContainer{this->device->allEngines};
+        engines.insert(engines.end(), this->device->getSubDevice(0)->getAllEngines().begin(), this->device->getSubDevice(0)->getAllEngines().end());
+        engines.insert(engines.end(), this->device->getSubDevice(1)->getAllEngines().begin(), this->device->getSubDevice(1)->getAllEngines().end());
+        for (auto &engine : engines) {
+            engine.osContext->incRefInternal();
+            if (engine.osContext->isDefaultContext()) {
                 if (engine.osContext->getDeviceBitfield().to_ulong() == 3) {
                     auto osContexLinux = static_cast<OsContextLinux *>(engine.osContext);
                     vmIdForDevice0 = osContexLinux->getDrmVmIds()[0];
@@ -862,18 +868,34 @@ HWTEST_F(DrmMemoryOperationsHandlerBindWithPerContextVms, givenVmBindMultipleRoo
                     auto osContexLinux = static_cast<OsContextLinux *>(engine.osContext);
                     vmIdForDevice0Subdevice0 = osContexLinux->getDrmVmIds()[0];
                 }
-            } else {
-                if (engine.osContext->getDeviceBitfield().to_ulong() == 3) {
-                    auto osContexLinux = static_cast<OsContextLinux *>(engine.osContext);
-                    vmIdForDevice1 = osContexLinux->getDrmVmIds()[0];
-                }
             }
         }
     }
 
+    {
+        auto &engines = memoryManager->allRegisteredEngines[device1->getRootDeviceIndex()];
+        engines = EngineControlContainer{this->device->allEngines};
+        engines.insert(engines.end(), device1->getSubDevice(0)->getAllEngines().begin(), device1->getSubDevice(0)->getAllEngines().end());
+        engines.insert(engines.end(), device1->getSubDevice(1)->getAllEngines().begin(), device1->getSubDevice(1)->getAllEngines().end());
+        for (auto &engine : engines) {
+            engine.osContext->incRefInternal();
+            if (engine.osContext->isDefaultContext()) {
+
+                if (engine.osContext->getDeviceBitfield().to_ulong() == 3) {
+                    auto osContexLinux = static_cast<OsContextLinux *>(engine.osContext);
+                    vmIdForDevice1 = osContexLinux->getDrmVmIds()[0];
+
+                } else if (engine.osContext->getDeviceBitfield().to_ulong() == 1) {
+                    auto osContexLinux = static_cast<OsContextLinux *>(engine.osContext);
+                    vmIdForDevice1Subdevice0 = osContexLinux->getDrmVmIds()[0];
+                }
+            }
+        }
+    }
     EXPECT_NE(0u, vmIdForDevice0);
     EXPECT_NE(0u, vmIdForDevice0Subdevice0);
     EXPECT_NE(0u, vmIdForDevice1);
+    EXPECT_NE(0u, vmIdForDevice1Subdevice0);
 
     AllocationData allocationData;
     allocationData.size = 13u;
@@ -905,7 +927,7 @@ HWTEST_F(DrmMemoryOperationsHandlerBindWithPerContextVms, givenVmBindMultipleRoo
     EXPECT_EQ(vmBindCalledBefore + 1, mock1->context.vmBindCalled);
 
     EXPECT_FALSE(mock->context.receivedVmBind.has_value());
-    EXPECT_EQ(vmIdForDevice1, mock1->context.receivedVmBind->vmId);
+    EXPECT_EQ(vmIdForDevice1Subdevice0, mock1->context.receivedVmBind->vmId);
 }
 
 HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenDirectSubmissionWhenPinBOThenVmBindIsCalledInsteadOfExec) {
@@ -913,8 +935,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenDirectSubmissionWhenPinBOThenV
     mock->bindAvailable = true;
     device->getDefaultEngine().osContext->setDirectSubmissionActive();
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    BufferObject boToPin(mock, 3, 2, 0, 1);
+    BufferObject boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
 
     pinBB.pin(&boToPinPtr, 1u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -928,8 +950,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenDirectSubmissionAndValidateHos
     mock->bindAvailable = true;
     device->getDefaultEngine().osContext->setDirectSubmissionActive();
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    BufferObject boToPin(mock, 3, 2, 0, 1);
+    BufferObject boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
 
     pinBB.validateHostPtr(&boToPinPtr, 1u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -946,8 +968,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportWhenPinBOThenAllo
     DebugManager.flags.UseVmBind.set(1);
     mock->bindAvailable = true;
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    MockBO boToPin(mock, 3, 2, 0, 1);
+    MockBO boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
 
     auto ret = pinBB.pin(&boToPinPtr, 1u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -965,8 +987,8 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportWhenPinBOAndVmBin
     mock->bindAvailable = true;
     mock->context.vmBindReturn = -1;
 
-    BufferObject pinBB(mock, 3, 1, 0, 1);
+    BufferObject pinBB(device->getRootDeviceIndex(), mock, 3, 1, 0, 1);
-    MockBO boToPin(mock, 3, 2, 0, 1);
+    MockBO boToPin(device->getRootDeviceIndex(), mock, 3, 2, 0, 1);
     BufferObject *boToPinPtr = &boToPin;
 
     auto ret = pinBB.pin(&boToPinPtr, 1u, device->getDefaultEngine().osContext, 0u, 0u);
@@ -1011,7 +1033,7 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenPatIndexProgrammingEnabledWhen
 
     uint64_t gpuAddress = 0x123000;
     size_t size = 1;
-    BufferObject bo(mock, static_cast<uint64_t>(MockGmmClientContextBase::MockPatIndex::cached), 0, 1, 1);
+    BufferObject bo(0, mock, static_cast<uint64_t>(MockGmmClientContextBase::MockPatIndex::cached), 0, 1, 1);
     DrmAllocation allocation(0, 1, AllocationType::BUFFER, &bo, nullptr, gpuAddress, size, MemoryPool::System4KBPages);
 
     auto allocationPtr = static_cast<GraphicsAllocation *>(&allocation);
@@ -1074,7 +1096,7 @@ HWTEST_F(DrmMemoryOperationsHandlerBindTest, givenPatIndexErrorAndUncachedDebugF
 
     uint64_t gpuAddress = 0x123000;
     size_t size = 1;
-    BufferObject bo(mock, static_cast<uint64_t>(MockGmmClientContextBase::MockPatIndex::cached), 0, 1, 1);
+    BufferObject bo(0, mock, static_cast<uint64_t>(MockGmmClientContextBase::MockPatIndex::cached), 0, 1, 1);
     DrmAllocation allocation(0, 1, AllocationType::BUFFER, &bo, nullptr, gpuAddress, size, MemoryPool::System4KBPages);
 
     EXPECT_ANY_THROW(mock->getPatIndex(allocation.getDefaultGmm(), allocation.getAllocationType(), CacheRegion::Default, CachePolicy::WriteBack, false));

shared/test/unit_test/os_interface/linux/drm_vm_bind_prelim_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2022 Intel Corporation
+ * Copyright (C) 2018-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -22,7 +22,7 @@ TEST(DrmVmBindTest, givenBoRequiringImmediateBindWhenBindingThenImmediateFlagIsP
     executionEnvironment->initializeMemoryManager();
     DrmQueryMock drm{*executionEnvironment->rootDeviceEnvironments[0]};
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0, &drm, 3, 0, 0, 1);
     bo.requireImmediateBinding(true);
 
     OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
@@ -40,7 +40,7 @@ TEST(DrmVmBindTest, givenBoRequiringExplicitResidencyWhenBindingThenMakeResident
     drm.pageFaultSupported = true;
 
     for (auto requireResidency : {false, true}) {
-        MockBufferObject bo(&drm, 3, 0, 0, 1);
+        MockBufferObject bo(0, &drm, 3, 0, 0, 1);
         bo.requireExplicitResidency(requireResidency);
 
         OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
@@ -78,7 +78,7 @@ TEST(DrmVmBindTest, givenBoNotRequiringExplicitResidencyWhenCallingWaitForBindTh
     drm.pageFaultSupported = true;
 
     for (auto requireResidency : {false, true}) {
-        MockBufferObject bo(&drm, 3, 0, 0, 1);
+        MockBufferObject bo(0, &drm, 3, 0, 0, 1);
         bo.requireExplicitResidency(requireResidency);
 
         OsContextLinux osContext(drm, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
@@ -111,8 +111,8 @@ TEST(DrmVmBindTest, givenUseKmdMigrationWhenCallingBindBoOnUnifiedSharedMemoryTh
     osContext.ensureContextInitialized();
     uint32_t vmHandleId = 0;
 
-    MockBufferObject bo(&drm, 3, 0, 0, 1);
+    MockBufferObject bo(0u, &drm, 3, 0, 0, 1);
-    MockDrmAllocation allocation(AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory);
+    MockDrmAllocation allocation(0u, AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory);
     allocation.bufferObjects[0] = &bo;
 
     allocation.bindBO(&bo, &osContext, vmHandleId, nullptr, true);

shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp

@@ -238,8 +238,8 @@ TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseFailsThenDont
     drm->ioctlCallsCount = 0;
     drm->ioctlRetVal = -1;
 
-    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
+    MockBufferObject bo(0u, drm.get(), 3, 0, 0, 1);
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::LocalMemory);
     allocation.bufferObjects[0] = &bo;
 
     MemAdviseFlags memAdviseFlags{};
@@ -252,9 +252,9 @@ TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseFailsThenDont
 }
 
 TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseWithNonAtomicIsCalledThenUpdateTheCorrespondingVmAdviceForBufferObject) {
-    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
+    MockBufferObject bo(0u, drm.get(), 3, 0, 0, 1);
     drm->ioctlCallsCount = 0;
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::LocalMemory);
     allocation.bufferObjects[0] = &bo;
 
     MemAdviseFlags memAdviseFlags{};
@@ -270,8 +270,8 @@ TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseWithNonAtomic
 
 TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseWithDevicePreferredLocationIsCalledThenUpdateTheCorrespondingVmAdviceForBufferObject) {
     drm->ioctlCallsCount = 0;
-    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
+    MockBufferObject bo(0u, drm.get(), 3, 0, 0, 1);
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::LocalMemory);
     allocation.bufferObjects[0] = &bo;
     allocation.storageInfo.memoryBanks = 0x1;
     allocation.setNumHandles(1);
@@ -289,8 +289,8 @@ TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseWithDevicePre
 
 TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemPrefetchSucceedsThenReturnTrue) {
     SubDeviceIdsVec subDeviceIds{0};
-    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
+    MockBufferObject bo(0u, drm.get(), 3, 0, 0, 1);
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::LocalMemory);
     allocation.bufferObjects[0] = &bo;
 
     drm->ioctlRetVal = 0;
@@ -299,8 +299,8 @@ TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemPrefetchSucceedsThe
 
 TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemPrefetchFailsThenReturnFalse) {
     SubDeviceIdsVec subDeviceIds{0};
-    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
+    MockBufferObject bo(0u, drm.get(), 3, 0, 0, 1);
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::LocalMemory);
     allocation.bufferObjects[0] = &bo;
 
     drm->ioctlRetVal = EINVAL;

shared/test/unit_test/os_interface/linux/file_logger_linux_tests.cpp

@@ -27,13 +27,13 @@ TEST(FileLogger, GivenLogAllocationMemoryPoolFlagThenLogsCorrectInfo) {
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
     DrmMock drm(*executionEnvironment->rootDeviceEnvironments[0]);
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::System64KBPages);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::System64KBPages);
     auto gmmHelper = executionEnvironment->rootDeviceEnvironments[0]->getGmmHelper();
     auto canonizedGpuAddress = gmmHelper->canonize(0x12345);
 
     allocation.setCpuPtrAndGpuAddress(&allocation, canonizedGpuAddress);
 
-    MockBufferObject bo(&drm);
+    MockBufferObject bo(0u, &drm);
     bo.handle.setBoHandle(4);
 
     allocation.bufferObjects[0] = &bo;
@@ -78,13 +78,13 @@ TEST(FileLogger, givenLogAllocationStdoutWhenLogAllocationThenLogToStdoutInstead
     auto executionEnvironment = std::make_unique<MockExecutionEnvironment>();
     DrmMock drm(*executionEnvironment->rootDeviceEnvironments[0]);
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::System64KBPages);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::System64KBPages);
     auto gmmHelper = executionEnvironment->rootDeviceEnvironments[0]->getGmmHelper();
     auto canonizedGpuAddress = gmmHelper->canonize(0x12345);
 
     allocation.setCpuPtrAndGpuAddress(&allocation, canonizedGpuAddress);
 
-    MockBufferObject bo(&drm);
+    MockBufferObject bo(0u, &drm);
     bo.handle.setBoHandle(4);
 
     allocation.bufferObjects[0] = &bo;
@@ -128,7 +128,7 @@ TEST(FileLogger, GivenDrmAllocationWithoutBOThenNoHandleLogged) {
     // Log file not created
     bool logFileCreated = fileExists(fileLogger.getLogFileName());
     EXPECT_FALSE(logFileCreated);
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::System64KBPages);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::System64KBPages);
 
     fileLogger.logAllocation(&allocation);
 
@@ -159,7 +159,7 @@ TEST(FileLogger, GivenLogAllocationMemoryPoolFlagSetFalseThenAllocationIsNotLogg
     bool logFileCreated = fileExists(fileLogger.getLogFileName());
     EXPECT_FALSE(logFileCreated);
 
-    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::System64KBPages);
+    MockDrmAllocation allocation(0u, AllocationType::BUFFER, MemoryPool::System64KBPages);
 
     fileLogger.logAllocation(&allocation);
 

shared/test/unit_test/os_interface/windows/device_command_stream_tests.cpp

@@ -1035,7 +1035,7 @@ HWTEST_F(WddmCsrCompressionTests, givenDisabledCompressionWhenInitializedThenDon
     setCompressionEnabled(false, false);
     myMockWddm = static_cast<WddmMock *>(executionEnvironment->rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Wddm>());
     MockWddmCsr<FamilyType> mockWddmCsr(*executionEnvironment, 1, device->getDeviceBitfield());
-    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines()) {
+    for (auto engine : executionEnvironment->memoryManager->getRegisteredEngines(device->getRootDeviceIndex())) {
         EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
     }
 }
@@ -1059,14 +1059,14 @@ HWTEST_F(WddmCsrCompressionTests, givenDisabledCompressionWhenFlushingThenDontIn
     device->resetCommandStreamReceiver(mockWddmCsr);
 
     auto memoryManager = executionEnvironment->memoryManager.get();
-    for (auto engine : memoryManager->getRegisteredEngines()) {
+    for (auto engine : memoryManager->getRegisteredEngines(device->getRootDeviceIndex())) {
         EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
     }
 
     auto graphicsAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{mockWddmCsr->getRootDeviceIndex(), MemoryConstants::pageSize});
     IndirectHeap cs(graphicsAllocation);
 
-    for (auto engine : memoryManager->getRegisteredEngines()) {
+    for (auto engine : memoryManager->getRegisteredEngines(device->getRootDeviceIndex())) {
         EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
     }
 
@@ -1074,7 +1074,7 @@ HWTEST_F(WddmCsrCompressionTests, givenDisabledCompressionWhenFlushingThenDontIn
 
     mockWddmCsr->flushTask(cs, 0u, &cs, &cs, &cs, 0u, dispatchFlags, *device);
 
-    for (auto engine : memoryManager->getRegisteredEngines()) {
+    for (auto engine : memoryManager->getRegisteredEngines(device->getRootDeviceIndex())) {
         EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
     }