Extend key to host ptr fragments container

use separated entries per root device

Related-To: NEO-3691

Change-Id: I26b85b8852b23e6a4d290da689174c59343536b3
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>

opencl/source/mem_obj/buffer.cpp

@@ -187,7 +187,7 @@ Buffer *Buffer::create(Context *context,
         memoryManager->isLocalMemorySupported(rootDeviceIndex),
         HwHelper::get(context->getDevice(0)->getHardwareInfo().platform.eRenderCoreFamily).obtainRenderBufferCompressionPreference(context->getDevice(0)->getHardwareInfo(), size));
 
-    checkMemory(memoryProperties, size, hostPtr, errcodeRet, alignementSatisfied, copyMemoryFromHostPtr, memoryManager);
+    checkMemory(memoryProperties, size, hostPtr, errcodeRet, alignementSatisfied, copyMemoryFromHostPtr, memoryManager, rootDeviceIndex);
 
     if (errcodeRet != CL_SUCCESS) {
         return nullptr;
@@ -384,7 +384,8 @@ void Buffer::checkMemory(MemoryProperties memoryProperties,
                          cl_int &errcodeRet,
                          bool &alignementSatisfied,
                          bool &copyMemoryFromHostPtr,
-                         MemoryManager *memoryManager) {
+                         MemoryManager *memoryManager,
+                         uint32_t rootDeviceIndex) {
     errcodeRet = CL_SUCCESS;
     alignementSatisfied = true;
     copyMemoryFromHostPtr = false;
@@ -403,7 +404,7 @@ void Buffer::checkMemory(MemoryProperties memoryProperties,
 
     if (memoryProperties.flags.useHostPtr) {
         if (hostPtr) {
-            auto fragment = memoryManager->getHostPtrManager()->getFragment(hostPtr);
+            auto fragment = memoryManager->getHostPtrManager()->getFragment({hostPtr, rootDeviceIndex});
             if (fragment && fragment->driverAllocation) {
                 errcodeRet = CL_INVALID_HOST_PTR;
                 return;

opencl/source/mem_obj/buffer.h

@@ -178,7 +178,8 @@ class Buffer : public MemObj {
                             cl_int &errcodeRet,
                             bool &isZeroCopy,
                             bool &copyMemoryFromHostPtr,
-                            MemoryManager *memMngr);
+                            MemoryManager *memMngr,
+                            uint32_t rootDeviceIndex);
     static GraphicsAllocation::AllocationType getGraphicsAllocationType(const MemoryProperties &properties, Context &context,
                                                                         bool renderCompressedBuffers, bool localMemoryEnabled,
                                                                         bool preferCompression);

opencl/source/memory_manager/os_agnostic_memory_manager.cpp

@@ -180,16 +180,17 @@ void OsAgnosticMemoryManager::addAllocationToHostPtrManager(GraphicsAllocation *
     fragment.fragmentSize = alignUp(gfxAllocation->getUnderlyingBufferSize(), MemoryConstants::pageSize);
     fragment.osInternalStorage = new OsHandle();
     fragment.residency = new ResidencyData();
-    hostPtrManager->storeFragment(fragment);
+    hostPtrManager->storeFragment(gfxAllocation->getRootDeviceIndex(), fragment);
 }
 
 void OsAgnosticMemoryManager::removeAllocationFromHostPtrManager(GraphicsAllocation *gfxAllocation) {
     auto buffer = gfxAllocation->getUnderlyingBuffer();
-    auto fragment = hostPtrManager->getFragment(buffer);
+    auto rootDeviceIndex = gfxAllocation->getRootDeviceIndex();
+    auto fragment = hostPtrManager->getFragment({buffer, rootDeviceIndex});
     if (fragment && fragment->driverAllocation) {
         OsHandle *osStorageToRelease = fragment->osInternalStorage;
         ResidencyData *residencyDataToRelease = fragment->residency;
-        if (hostPtrManager->releaseHostPtr(buffer)) {
+        if (hostPtrManager->releaseHostPtr(rootDeviceIndex, buffer)) {
             delete osStorageToRelease;
             delete residencyDataToRelease;
         }
@@ -265,7 +266,7 @@ MemoryManager::AllocationStatus OsAgnosticMemoryManager::populateOsHandles(OsHan
             newFragment.fragmentSize = handleStorage.fragmentStorageData[i].fragmentSize;
             newFragment.osInternalStorage = handleStorage.fragmentStorageData[i].osHandleStorage;
             newFragment.residency = handleStorage.fragmentStorageData[i].residency;
-            hostPtrManager->storeFragment(newFragment);
+            hostPtrManager->storeFragment(rootDeviceIndex, newFragment);
         }
     }
     return AllocationStatus::Success;

opencl/test/unit_test/memory_manager/host_ptr_manager_tests.cpp

@@ -21,11 +21,15 @@
 
 using namespace NEO;
 
-TEST(HostPtrManager, GivenAlignedPointerAndAlignedSizeWhenGettingAllocationRequirementsThenOneFragmentIsReturned) {
+struct HostPtrManagerTest : ::testing::Test {
+    const uint32_t rootDeviceIndex = 1u;
+};
+
+TEST_F(HostPtrManagerTest, GivenAlignedPointerAndAlignedSizeWhenGettingAllocationRequirementsThenOneFragmentIsReturned) {
     auto size = MemoryConstants::pageSize * 10;
     void *ptr = (void *)0x1000;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
 
     EXPECT_EQ(1u, reqs.requiredFragmentsCount);
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::MIDDLE);
@@ -39,13 +43,15 @@ TEST(HostPtrManager, GivenAlignedPointerAndAlignedSizeWhenGettingAllocationRequi
 
     EXPECT_EQ(nullptr, reqs.allocationFragments[1].allocationPtr);
     EXPECT_EQ(nullptr, reqs.allocationFragments[2].allocationPtr);
+
+    EXPECT_EQ(rootDeviceIndex, reqs.rootDeviceIndex);
 }
 
-TEST(HostPtrManager, GivenAlignedPointerAndNotAlignedSizeWhenGettingAllocationRequirementsThenTwoFragmentsAreReturned) {
+TEST_F(HostPtrManagerTest, GivenAlignedPointerAndNotAlignedSizeWhenGettingAllocationRequirementsThenTwoFragmentsAreReturned) {
     auto size = MemoryConstants::pageSize * 10 - 1;
     void *ptr = (void *)0x1000;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(2u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::MIDDLE);
@@ -64,11 +70,11 @@ TEST(HostPtrManager, GivenAlignedPointerAndNotAlignedSizeWhenGettingAllocationRe
     EXPECT_EQ(0u, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenNotAlignedPointerAndNotAlignedSizeWhenGettingAllocationRequirementsThenThreeFragmentsAreReturned) {
+TEST_F(HostPtrManagerTest, GivenNotAlignedPointerAndNotAlignedSizeWhenGettingAllocationRequirementsThenThreeFragmentsAreReturned) {
     auto size = MemoryConstants::pageSize * 10 - 1;
     void *ptr = (void *)0x1045;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(3u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::LEADING);
@@ -91,11 +97,11 @@ TEST(HostPtrManager, GivenNotAlignedPointerAndNotAlignedSizeWhenGettingAllocatio
     EXPECT_EQ(MemoryConstants::pageSize, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenNotAlignedPointerAndNotAlignedSizeWithinOnePageWhenGettingAllocationRequirementsThenOneFragmentIsReturned) {
+TEST_F(HostPtrManagerTest, GivenNotAlignedPointerAndNotAlignedSizeWithinOnePageWhenGettingAllocationRequirementsThenOneFragmentIsReturned) {
     auto size = 200;
     void *ptr = (void *)0x1045;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(1u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::LEADING);
@@ -116,11 +122,11 @@ TEST(HostPtrManager, GivenNotAlignedPointerAndNotAlignedSizeWithinOnePageWhenGet
     EXPECT_EQ(0u, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenNotAlignedPointerAndNotAlignedSizeWithinTwoPagesWhenGettingAllocationRequirementsThenTwoFragmentsAreReturned) {
+TEST_F(HostPtrManagerTest, GivenNotAlignedPointerAndNotAlignedSizeWithinTwoPagesWhenGettingAllocationRequirementsThenTwoFragmentsAreReturned) {
     auto size = MemoryConstants::pageSize;
     void *ptr = (void *)0x1045;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(2u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::LEADING);
@@ -142,11 +148,11 @@ TEST(HostPtrManager, GivenNotAlignedPointerAndNotAlignedSizeWithinTwoPagesWhenGe
     EXPECT_EQ(0u, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenAlignedPointerAndAlignedSizeOfOnePageWhenGettingAllocationRequirementsThenOnlyMiddleFragmentIsReturned) {
+TEST_F(HostPtrManagerTest, GivenAlignedPointerAndAlignedSizeOfOnePageWhenGettingAllocationRequirementsThenOnlyMiddleFragmentIsReturned) {
     auto size = MemoryConstants::pageSize * 10;
     void *ptr = (void *)0x1000;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(1u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::MIDDLE);
@@ -167,11 +173,11 @@ TEST(HostPtrManager, GivenAlignedPointerAndAlignedSizeOfOnePageWhenGettingAlloca
     EXPECT_EQ(0u, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenNotAlignedPointerAndSizeThatFitsToPageWhenGettingAllocationRequirementsThenLeadingAndMiddleFragmentsAreReturned) {
+TEST_F(HostPtrManagerTest, GivenNotAlignedPointerAndSizeThatFitsToPageWhenGettingAllocationRequirementsThenLeadingAndMiddleFragmentsAreReturned) {
     auto size = MemoryConstants::pageSize * 10 - 1;
     void *ptr = (void *)0x1001;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(2u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::LEADING);
@@ -193,11 +199,11 @@ TEST(HostPtrManager, GivenNotAlignedPointerAndSizeThatFitsToPageWhenGettingAlloc
     EXPECT_EQ(0u, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenAlignedPointerAndPageSizeWhenGettingAllocationRequirementsThenOnlyMiddleFragmentIsReturned) {
+TEST_F(HostPtrManagerTest, GivenAlignedPointerAndPageSizeWhenGettingAllocationRequirementsThenOnlyMiddleFragmentIsReturned) {
     auto size = MemoryConstants::pageSize;
     void *ptr = (void *)0x1000;
 
-    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    AllocationRequirements reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     EXPECT_EQ(1u, reqs.requiredFragmentsCount);
 
     EXPECT_EQ(reqs.allocationFragments[0].fragmentPosition, FragmentPosition::MIDDLE);
@@ -218,10 +224,10 @@ TEST(HostPtrManager, GivenAlignedPointerAndPageSizeWhenGettingAllocationRequirem
     EXPECT_EQ(0u, reqs.allocationFragments[2].allocationSize);
 }
 
-TEST(HostPtrManager, GivenAllocationRequirementsForMiddleAllocationThatIsNotStoredInManagerWhenGettingAllocationRequirementsThenNullptrIsReturned) {
+TEST_F(HostPtrManagerTest, GivenAllocationRequirementsForMiddleAllocationThatIsNotStoredInManagerWhenGettingAllocationRequirementsThenNullptrIsReturned) {
     auto size = MemoryConstants::pageSize;
     void *ptr = (void *)0x1000;
-    auto reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     MockHostPtrManager hostPtrManager;
 
     auto gpuAllocationFragments = hostPtrManager.populateAlreadyAllocatedFragments(reqs);
@@ -233,7 +239,7 @@ TEST(HostPtrManager, GivenAllocationRequirementsForMiddleAllocationThatIsNotStor
     EXPECT_EQ(nullptr, gpuAllocationFragments.fragmentStorageData[2].cpuPtr);
 }
 
-TEST(HostPtrManager, GivenAllocationRequirementsForMiddleAllocationThatIsStoredInManagerWhenGettingAllocationRequirementsThenProperAllocationIsReturnedAndRefCountIncreased) {
+TEST_F(HostPtrManagerTest, GivenAllocationRequirementsForMiddleAllocationThatIsStoredInManagerWhenGettingAllocationRequirementsThenProperAllocationIsReturnedAndRefCountIncreased) {
 
     MockHostPtrManager hostPtrManager;
     FragmentStorage allocationFragment;
@@ -244,9 +250,9 @@ TEST(HostPtrManager, GivenAllocationRequirementsForMiddleAllocationThatIsStoredI
     allocationFragment.fragmentSize = ptrSize;
     allocationFragment.osInternalStorage = osInternalStorage;
 
-    hostPtrManager.storeFragment(allocationFragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, allocationFragment);
 
-    auto reqs = MockHostPtrManager::getAllocationRequirements(cpuPtr, ptrSize);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, cpuPtr, ptrSize);
 
     auto gpuAllocationFragments = hostPtrManager.populateAlreadyAllocatedFragments(reqs);
 
@@ -257,11 +263,11 @@ TEST(HostPtrManager, GivenAllocationRequirementsForMiddleAllocationThatIsStoredI
     EXPECT_EQ(nullptr, gpuAllocationFragments.fragmentStorageData[2].osHandleStorage);
     EXPECT_EQ(nullptr, gpuAllocationFragments.fragmentStorageData[2].cpuPtr);
 
-    auto fragment = hostPtrManager.getFragment(cpuPtr);
+    auto fragment = hostPtrManager.getFragment({cpuPtr, rootDeviceIndex});
     EXPECT_EQ(2, fragment->refCount);
 }
 
-TEST(HostPtrManager, GivenAllocationRequirementsForAllocationWithinSizeOfStoredAllocationInManagerWhenGettingAllocationRequirementsThenProperAllocationIsReturned) {
+TEST_F(HostPtrManagerTest, GivenAllocationRequirementsForAllocationWithinSizeOfStoredAllocationInManagerWhenGettingAllocationRequirementsThenProperAllocationIsReturned) {
 
     MockHostPtrManager hostPtrManager;
     FragmentStorage allocationFragment;
@@ -272,9 +278,9 @@ TEST(HostPtrManager, GivenAllocationRequirementsForAllocationWithinSizeOfStoredA
     allocationFragment.fragmentSize = ptrSize;
     allocationFragment.osInternalStorage = osInternalStorage;
 
-    hostPtrManager.storeFragment(allocationFragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, allocationFragment);
 
-    auto reqs = MockHostPtrManager::getAllocationRequirements(cpuPtr, MemoryConstants::pageSize);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, cpuPtr, MemoryConstants::pageSize);
 
     auto gpuAllocationFragments = hostPtrManager.populateAlreadyAllocatedFragments(reqs);
 
@@ -285,11 +291,11 @@ TEST(HostPtrManager, GivenAllocationRequirementsForAllocationWithinSizeOfStoredA
     EXPECT_EQ(nullptr, gpuAllocationFragments.fragmentStorageData[2].osHandleStorage);
     EXPECT_EQ(nullptr, gpuAllocationFragments.fragmentStorageData[2].cpuPtr);
 
-    auto fragment = hostPtrManager.getFragment(cpuPtr);
+    auto fragment = hostPtrManager.getFragment({cpuPtr, rootDeviceIndex});
     EXPECT_EQ(2, fragment->refCount);
 }
 
-TEST(HostPtrManager, WhenStoringFragmentThenContainerCountIsIncremented) {
+TEST_F(HostPtrManagerTest, WhenStoringFragmentThenContainerCountIsIncremented) {
     MockHostPtrManager hostPtrManager;
 
     FragmentStorage allocationFragment;
@@ -297,30 +303,30 @@ TEST(HostPtrManager, WhenStoringFragmentThenContainerCountIsIncremented) {
     EXPECT_EQ(allocationFragment.fragmentSize, 0u);
     EXPECT_EQ(allocationFragment.refCount, 0);
 
-    hostPtrManager.storeFragment(allocationFragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, allocationFragment);
 
     EXPECT_EQ(1u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, WhenStoringFragmentTwiceThenContainerCountIsIncrementedOnce) {
+TEST_F(HostPtrManagerTest, WhenStoringFragmentTwiceThenContainerCountIsIncrementedOnce) {
     MockHostPtrManager hostPtrManager;
 
     FragmentStorage allocationFragment;
 
-    hostPtrManager.storeFragment(allocationFragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, allocationFragment);
-    hostPtrManager.storeFragment(allocationFragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, allocationFragment);
 
     EXPECT_EQ(1u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenEmptyHostPtrManagerWhenAskingForFragmentThenNullptrIsReturned) {
+TEST_F(HostPtrManagerTest, GivenEmptyHostPtrManagerWhenAskingForFragmentThenNullptrIsReturned) {
     MockHostPtrManager hostPtrManager;
-    auto fragment = hostPtrManager.getFragment((void *)0x10121);
+    auto fragment = hostPtrManager.getFragment({(void *)0x10121, rootDeviceIndex});
     EXPECT_EQ(nullptr, fragment);
     EXPECT_EQ(0u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenNonEmptyHostPtrManagerWhenAskingForFragmentThenProperFragmentIsReturnedWithRefCountOne) {
+TEST_F(HostPtrManagerTest, GivenNonEmptyHostPtrManagerWhenAskingForFragmentThenProperFragmentIsReturnedWithRefCountOne) {
     MockHostPtrManager hostPtrManager;
     FragmentStorage fragment;
     void *cpuPtr = (void *)0x10121;
@@ -329,8 +335,8 @@ TEST(HostPtrManager, GivenNonEmptyHostPtrManagerWhenAskingForFragmentThenProperF
     fragment.fragmentSize = fragmentSize;
     fragment.refCount = 0;
 
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
-    auto retFragment = hostPtrManager.getFragment(cpuPtr);
+    auto retFragment = hostPtrManager.getFragment({cpuPtr, rootDeviceIndex});
 
     EXPECT_NE(retFragment, &fragment);
     EXPECT_EQ(1, retFragment->refCount);
@@ -339,7 +345,7 @@ TEST(HostPtrManager, GivenNonEmptyHostPtrManagerWhenAskingForFragmentThenProperF
     EXPECT_EQ(1u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenHostPtrManagerFilledTwiceWithTheSamePointerWhenAskingForFragmentThenProperFragmentIsReturnedWithRefCountTwo) {
+TEST_F(HostPtrManagerTest, GivenHostPtrManagerFilledTwiceWithTheSamePointerWhenAskingForFragmentThenProperFragmentIsReturnedWithRefCountTwo) {
     MockHostPtrManager hostPtrManager;
     FragmentStorage fragment;
     void *cpuPtr = (void *)0x10121;
@@ -348,9 +354,9 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledTwiceWithTheSamePointerWhenAskingF
     fragment.fragmentSize = fragmentSize;
     fragment.refCount = 0;
 
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
-    auto retFragment = hostPtrManager.getFragment(cpuPtr);
+    auto retFragment = hostPtrManager.getFragment({cpuPtr, rootDeviceIndex});
 
     EXPECT_NE(retFragment, &fragment);
     EXPECT_EQ(2, retFragment->refCount);
@@ -359,7 +365,7 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledTwiceWithTheSamePointerWhenAskingF
     EXPECT_EQ(1u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenFragmentIsBeingReleasedThenManagerMaintainsProperRefferenceCount) {
+TEST_F(HostPtrManagerTest, GivenHostPtrManagerFilledWithFragmentsWhenFragmentIsBeingReleasedThenManagerMaintainsProperRefferenceCount) {
     MockHostPtrManager hostPtrManager;
     FragmentStorage fragment;
     void *cpuPtr = (void *)0x1000;
@@ -368,28 +374,28 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenFragmentIsBeingRe
     fragment.fragmentCpuPointer = cpuPtr;
     fragment.fragmentSize = fragmentSize;
 
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     ASSERT_EQ(1u, hostPtrManager.getFragmentCount());
 
-    auto fragmentReadyForRelease = hostPtrManager.releaseHostPtr(cpuPtr);
+    auto fragmentReadyForRelease = hostPtrManager.releaseHostPtr(rootDeviceIndex, cpuPtr);
     EXPECT_FALSE(fragmentReadyForRelease);
 
-    auto retFragment = hostPtrManager.getFragment(cpuPtr);
+    auto retFragment = hostPtrManager.getFragment({cpuPtr, rootDeviceIndex});
 
     EXPECT_EQ(1, retFragment->refCount);
 
-    fragmentReadyForRelease = hostPtrManager.releaseHostPtr(cpuPtr);
+    fragmentReadyForRelease = hostPtrManager.releaseHostPtr(rootDeviceIndex, cpuPtr);
 
     EXPECT_TRUE(fragmentReadyForRelease);
 
-    retFragment = hostPtrManager.getFragment(cpuPtr);
+    retFragment = hostPtrManager.getFragment({cpuPtr, rootDeviceIndex});
 
     EXPECT_EQ(nullptr, retFragment);
     EXPECT_EQ(0u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenOsHandleStorageWhenAskedToStoreTheFragmentThenFragmentIsStoredProperly) {
+TEST_F(HostPtrManagerTest, GivenOsHandleStorageWhenAskedToStoreTheFragmentThenFragmentIsStoredProperly) {
     OsHandleStorage storage;
     void *cpu1 = (void *)0x1000;
     void *cpu2 = (void *)0x2000;
@@ -407,23 +413,23 @@ TEST(HostPtrManager, GivenOsHandleStorageWhenAskedToStoreTheFragmentThenFragment
 
     EXPECT_EQ(0u, hostPtrManager.getFragmentCount());
 
-    hostPtrManager.storeFragment(storage.fragmentStorageData[0]);
+    hostPtrManager.storeFragment(rootDeviceIndex, storage.fragmentStorageData[0]);
-    hostPtrManager.storeFragment(storage.fragmentStorageData[1]);
+    hostPtrManager.storeFragment(rootDeviceIndex, storage.fragmentStorageData[1]);
 
     EXPECT_EQ(2u, hostPtrManager.getFragmentCount());
 
-    hostPtrManager.releaseHandleStorage(storage);
+    hostPtrManager.releaseHandleStorage(rootDeviceIndex, storage);
 
     EXPECT_EQ(0u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenHostPtrFilledWith3TripleFragmentsWhenAskedForPopulationThenAllFragmentsAreResued) {
+TEST_F(HostPtrManagerTest, GivenHostPtrFilledWith3TripleFragmentsWhenAskedForPopulationThenAllFragmentsAreResued) {
     void *cpuPtr = (void *)0x1001;
     auto fragmentSize = MemoryConstants::pageSize * 10;
 
     MockHostPtrManager hostPtrManager;
 
-    auto reqs = hostPtrManager.getAllocationRequirements(cpuPtr, fragmentSize);
+    auto reqs = hostPtrManager.getAllocationRequirements(rootDeviceIndex, cpuPtr, fragmentSize);
     ASSERT_EQ(3u, reqs.requiredFragmentsCount);
 
     FragmentStorage fragments[maxFragmentsCount];
@@ -431,7 +437,7 @@ TEST(HostPtrManager, GivenHostPtrFilledWith3TripleFragmentsWhenAskedForPopulatio
     for (int i = 0; i < maxFragmentsCount; i++) {
         fragments[i].fragmentCpuPointer = const_cast<void *>(reqs.allocationFragments[i].allocationPtr);
         fragments[i].fragmentSize = reqs.allocationFragments[i].allocationSize;
-        hostPtrManager.storeFragment(fragments[i]);
+        hostPtrManager.storeFragment(rootDeviceIndex, fragments[i]);
     }
 
     EXPECT_EQ(3u, hostPtrManager.getFragmentCount());
@@ -442,108 +448,132 @@ TEST(HostPtrManager, GivenHostPtrFilledWith3TripleFragmentsWhenAskedForPopulatio
     for (int i = 0; i < maxFragmentsCount; i++) {
         EXPECT_EQ(OsHandles.fragmentStorageData[i].cpuPtr, reqs.allocationFragments[i].allocationPtr);
         EXPECT_EQ(OsHandles.fragmentStorageData[i].fragmentSize, reqs.allocationFragments[i].allocationSize);
-        auto fragment = hostPtrManager.getFragment(const_cast<void *>(reqs.allocationFragments[i].allocationPtr));
+        auto fragment = hostPtrManager.getFragment({const_cast<void *>(reqs.allocationFragments[i].allocationPtr),
+                                                    rootDeviceIndex});
         ASSERT_NE(nullptr, fragment);
         EXPECT_EQ(2, fragment->refCount);
         EXPECT_EQ(OsHandles.fragmentStorageData[i].cpuPtr, fragment->fragmentCpuPointer);
     }
 
     for (int i = 0; i < maxFragmentsCount; i++) {
-        hostPtrManager.releaseHostPtr(fragments[i].fragmentCpuPointer);
+        hostPtrManager.releaseHostPtr(rootDeviceIndex, fragments[i].fragmentCpuPointer);
     }
     EXPECT_EQ(3u, hostPtrManager.getFragmentCount());
     for (int i = 0; i < maxFragmentsCount; i++) {
-        auto fragment = hostPtrManager.getFragment(const_cast<void *>(reqs.allocationFragments[i].allocationPtr));
+        auto fragment = hostPtrManager.getFragment({const_cast<void *>(reqs.allocationFragments[i].allocationPtr),
+                                                    rootDeviceIndex});
         ASSERT_NE(nullptr, fragment);
         EXPECT_EQ(1, fragment->refCount);
     }
     for (int i = 0; i < maxFragmentsCount; i++) {
-        hostPtrManager.releaseHostPtr(fragments[i].fragmentCpuPointer);
+        hostPtrManager.releaseHostPtr(rootDeviceIndex, fragments[i].fragmentCpuPointer);
     }
     EXPECT_EQ(0u, hostPtrManager.getFragmentCount());
 }
 
-TEST(HostPtrManager, GivenFragmentSizeZeroWhenGettingFragmentThenNullptrIsReturned) {
+TEST_F(HostPtrManagerTest, GivenFragmentSizeZeroWhenGettingFragmentThenNullptrIsReturned) {
     HostPtrManager hostPtrManager;
 
     auto ptr1 = (void *)0x010000;
     FragmentStorage fragment1;
     fragment1.fragmentCpuPointer = ptr1;
     fragment1.fragmentSize = 0;
-    hostPtrManager.storeFragment(fragment1);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment1);
 
     auto ptr2 = (void *)0x040000;
     FragmentStorage fragment2;
     fragment2.fragmentCpuPointer = ptr2;
     fragment2.fragmentSize = 0;
-    hostPtrManager.storeFragment(fragment2);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment2);
 
     auto cptr1 = (void *)0x00F000;
-    auto frag1 = hostPtrManager.getFragment(cptr1);
+    auto frag1 = hostPtrManager.getFragment({cptr1, rootDeviceIndex});
     EXPECT_EQ(frag1, nullptr);
 
     auto cptr2 = (void *)0x010000;
-    auto frag2 = hostPtrManager.getFragment(cptr2);
+    auto frag2 = hostPtrManager.getFragment({cptr2, rootDeviceIndex});
     EXPECT_NE(frag2, nullptr);
 
     auto cptr3 = (void *)0x010001;
-    auto frag3 = hostPtrManager.getFragment(cptr3);
+    auto frag3 = hostPtrManager.getFragment({cptr3, rootDeviceIndex});
     EXPECT_EQ(frag3, nullptr);
 
     auto cptr4 = (void *)0x020000;
-    auto frag4 = hostPtrManager.getFragment(cptr4);
+    auto frag4 = hostPtrManager.getFragment({cptr4, rootDeviceIndex});
     EXPECT_EQ(frag4, nullptr);
 
     auto cptr5 = (void *)0x040000;
-    auto frag5 = hostPtrManager.getFragment(cptr5);
+    auto frag5 = hostPtrManager.getFragment({cptr5, rootDeviceIndex});
     EXPECT_NE(frag5, nullptr);
 
     auto cptr6 = (void *)0x040001;
-    auto frag6 = hostPtrManager.getFragment(cptr6);
+    auto frag6 = hostPtrManager.getFragment({cptr6, rootDeviceIndex});
     EXPECT_EQ(frag6, nullptr);
 
     auto cptr7 = (void *)0x060000;
-    auto frag7 = hostPtrManager.getFragment(cptr7);
+    auto frag7 = hostPtrManager.getFragment({cptr7, rootDeviceIndex});
     EXPECT_EQ(frag7, nullptr);
 }
 
-TEST(HostPtrManager, GivenFragmentSizeNonZeroWhenGettingFragmentThenCorrectAllocationIsReturned) {
+TEST_F(HostPtrManagerTest, GivenFragmentSizeNonZeroWhenGettingFragmentThenCorrectAllocationIsReturned) {
     MockHostPtrManager hostPtrManager;
+    uint32_t rootDeviceIndex2 = 2u;
 
     auto size1 = MemoryConstants::pageSize;
 
-    auto ptr1 = (void *)0x010000;
+    auto ptr11 = (void *)0x010000;
-    FragmentStorage fragment1;
+    FragmentStorage fragment11;
-    fragment1.fragmentCpuPointer = ptr1;
+    fragment11.fragmentCpuPointer = ptr11;
-    fragment1.fragmentSize = size1;
+    fragment11.fragmentSize = size1;
-    hostPtrManager.storeFragment(fragment1);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment11);
 
-    auto ptr2 = (void *)0x040000;
+    auto ptr12 = (void *)0x020000;
-    FragmentStorage fragment2;
+    FragmentStorage fragment12;
-    fragment2.fragmentCpuPointer = ptr2;
+    fragment12.fragmentCpuPointer = ptr12;
-    fragment2.fragmentSize = size1;
+    fragment12.fragmentSize = size1;
-    hostPtrManager.storeFragment(fragment2);
+    hostPtrManager.storeFragment(rootDeviceIndex2, fragment12);
+
+    auto ptr21 = (void *)0x040000;
+    FragmentStorage fragment21;
+    fragment21.fragmentCpuPointer = ptr21;
+    fragment21.fragmentSize = size1;
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment21);
+
+    auto ptr22 = (void *)0x060000;
+    FragmentStorage fragment22;
+    fragment22.fragmentCpuPointer = ptr22;
+    fragment22.fragmentSize = size1;
+    hostPtrManager.storeFragment(rootDeviceIndex2, fragment22);
 
     auto cptr1 = (void *)0x010060;
-    auto frag1 = hostPtrManager.getFragment(cptr1);
+    auto frag11 = hostPtrManager.getFragment({cptr1, rootDeviceIndex});
-    EXPECT_NE(frag1, nullptr);
+    EXPECT_NE(frag11, nullptr);
+    auto frag12 = hostPtrManager.getFragment({cptr1, rootDeviceIndex2});
+    EXPECT_EQ(frag12, nullptr);
 
     auto cptr2 = (void *)0x020000;
-    auto frag2 = hostPtrManager.getFragment(cptr2);
+    auto frag21 = hostPtrManager.getFragment({cptr2, rootDeviceIndex});
-    EXPECT_EQ(frag2, nullptr);
+    EXPECT_EQ(frag21, nullptr);
+    auto frag22 = hostPtrManager.getFragment({cptr2, rootDeviceIndex2});
+    EXPECT_NE(frag22, nullptr);
 
     auto cptr3 = (void *)0x040060;
-    auto frag3 = hostPtrManager.getFragment(cptr3);
+    auto frag31 = hostPtrManager.getFragment({cptr3, rootDeviceIndex});
-    EXPECT_NE(frag3, nullptr);
+    EXPECT_NE(frag31, nullptr);
+    auto frag32 = hostPtrManager.getFragment({cptr3, rootDeviceIndex2});
+    EXPECT_EQ(frag32, nullptr);
 
     auto cptr4 = (void *)0x060000;
-    auto frag4 = hostPtrManager.getFragment(cptr4);
+    auto frag41 = hostPtrManager.getFragment({cptr4, rootDeviceIndex});
-    EXPECT_EQ(frag4, nullptr);
+    EXPECT_EQ(frag41, nullptr);
+    auto frag42 = hostPtrManager.getFragment({cptr4, rootDeviceIndex2});
+    EXPECT_NE(frag42, nullptr);
 
     AllocationRequirements requiredAllocations;
+    requiredAllocations.rootDeviceIndex = rootDeviceIndex;
     auto ptr3 = (void *)0x040000;
     auto size3 = MemoryConstants::pageSize * 2;
-    requiredAllocations = hostPtrManager.getAllocationRequirements(ptr3, size3);
+    requiredAllocations = hostPtrManager.getAllocationRequirements(rootDeviceIndex, ptr3, size3);
     auto catchme = false;
     try {
         OsHandleStorage st = hostPtrManager.populateAlreadyAllocatedFragments(requiredAllocations);
@@ -554,109 +584,126 @@ TEST(HostPtrManager, GivenFragmentSizeNonZeroWhenGettingFragmentThenCorrectAlloc
     EXPECT_TRUE(catchme);
 }
 
-TEST(HostPtrManager, WhenCheckingForOverlapsThenCorrectStatusIsReturned) {
+TEST_F(HostPtrManagerTest, WhenCheckingForOverlapsThenCorrectStatusIsReturned) {
     MockHostPtrManager hostPtrManager;
+    uint32_t rootDeviceIndex2 = 2u;
 
     auto size1 = MemoryConstants::pageSize;
 
-    auto ptr1 = (void *)0x010000;
+    auto ptr11 = (void *)0x010000;
-    FragmentStorage fragment1;
+    FragmentStorage fragment11;
-    fragment1.fragmentCpuPointer = ptr1;
+    fragment11.fragmentCpuPointer = ptr11;
-    fragment1.fragmentSize = size1;
+    fragment11.fragmentSize = size1;
-    hostPtrManager.storeFragment(fragment1);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment11);
 
-    auto ptr2 = (void *)0x040000;
+    auto ptr12 = (void *)0x020000;
-    FragmentStorage fragment2;
+    FragmentStorage fragment12;
-    fragment2.fragmentCpuPointer = ptr2;
+    fragment12.fragmentCpuPointer = ptr12;
-    fragment2.fragmentSize = size1;
+    fragment12.fragmentSize = size1;
-    hostPtrManager.storeFragment(fragment2);
+    hostPtrManager.storeFragment(rootDeviceIndex2, fragment12);
+
+    auto ptr21 = (void *)0x040000;
+    FragmentStorage fragment21;
+    fragment21.fragmentCpuPointer = ptr21;
+    fragment21.fragmentSize = size1;
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment21);
+
+    auto ptr22 = (void *)0x060000;
+    FragmentStorage fragment22;
+    fragment22.fragmentCpuPointer = ptr22;
+    fragment22.fragmentSize = size1;
+    hostPtrManager.storeFragment(rootDeviceIndex2, fragment22);
 
     OverlapStatus overlappingStatus;
     auto cptr1 = (void *)0x010060;
-    auto frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr1, 1u, overlappingStatus);
+    auto frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr1, 1u, overlappingStatus);
     EXPECT_NE(frag1, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT);
 
-    frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(ptr1, size1, overlappingStatus);
+    frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptr11, size1, overlappingStatus);
     EXPECT_NE(frag1, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_WITH_EXACT_SIZE_AS_STORED_FRAGMENT);
 
-    frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(ptr1, size1 - 1, overlappingStatus);
+    frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptr11, size1 - 1, overlappingStatus);
     EXPECT_NE(frag1, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT);
 
+    frag1 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex2, ptr11, size1, overlappingStatus);
+    EXPECT_EQ(frag1, nullptr);
+    EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER);
+
     auto cptr2 = (void *)0x020000;
-    auto frag2 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr2, 1u, overlappingStatus);
+    auto frag2 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr2, 1u, overlappingStatus);
     EXPECT_EQ(frag2, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER);
 
     auto cptr3 = (void *)0x040060;
-    auto frag3 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr3, 1u, overlappingStatus);
+    auto frag3 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr3, 1u, overlappingStatus);
     EXPECT_NE(frag3, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT);
 
     auto cptr4 = (void *)0x060000;
-    auto frag4 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr4, 1u, overlappingStatus);
+    auto frag4 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr4, 1u, overlappingStatus);
     EXPECT_EQ(frag4, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER);
 
     auto cptr5 = (void *)0x040000;
-    auto frag5 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr5, size1 - 1, overlappingStatus);
+    auto frag5 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr5, size1 - 1, overlappingStatus);
     EXPECT_NE(frag5, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT);
 
     auto cptr6 = (void *)0x040000;
-    auto frag6 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr6, size1 + 1, overlappingStatus);
+    auto frag6 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr6, size1 + 1, overlappingStatus);
     EXPECT_EQ(frag6, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT);
 
     auto cptr7 = (void *)0x03FFF0;
-    auto frag7 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr7, 2 * size1, overlappingStatus);
+    auto frag7 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr7, 2 * size1, overlappingStatus);
     EXPECT_EQ(frag7, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT);
 
     auto cptr8 = (void *)0x040000;
-    auto frag8 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr8, size1, overlappingStatus);
+    auto frag8 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr8, size1, overlappingStatus);
     EXPECT_NE(frag8, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_WITH_EXACT_SIZE_AS_STORED_FRAGMENT);
 
     auto cptr9 = (void *)0x010060;
-    auto frag9 = hostPtrManager.getFragmentAndCheckForOverlaps(cptr9, 2 * size1, overlappingStatus);
+    auto frag9 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, cptr9, 2 * size1, overlappingStatus);
     EXPECT_EQ(frag9, nullptr);
     EXPECT_EQ(overlappingStatus, OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT);
 }
 
-TEST(HostPtrManager, GivenHostPtrManagerFilledWithBigFragmentWhenAskedForFragmnetInTheMiddleOfBigFragmentThenBigFragmentIsReturned) {
+TEST_F(HostPtrManagerTest, GivenHostPtrManagerFilledWithBigFragmentWhenAskedForFragmnetInTheMiddleOfBigFragmentThenBigFragmentIsReturned) {
     auto bigSize = 10 * MemoryConstants::pageSize;
     auto bigPtr = (void *)0x01000;
     FragmentStorage fragment;
     fragment.fragmentCpuPointer = bigPtr;
     fragment.fragmentSize = bigSize;
     MockHostPtrManager hostPtrManager;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     EXPECT_EQ(1u, hostPtrManager.getFragmentCount());
 
     auto ptrInTheMiddle = (void *)0x2000;
     auto smallSize = MemoryConstants::pageSize;
 
-    auto storedBigFragment = hostPtrManager.getFragment(bigPtr);
+    auto storedBigFragment = hostPtrManager.getFragment({bigPtr, rootDeviceIndex});
 
-    auto fragment2 = hostPtrManager.getFragment(ptrInTheMiddle);
+    auto fragment2 = hostPtrManager.getFragment({ptrInTheMiddle, rootDeviceIndex});
     EXPECT_EQ(storedBigFragment, fragment2);
 
     OverlapStatus overlapStatus;
-    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrInTheMiddle, smallSize, overlapStatus);
+    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrInTheMiddle, smallSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT, overlapStatus);
     EXPECT_EQ(fragment3, storedBigFragment);
 
     auto ptrOutside = (void *)0x1000000;
     auto outsideSize = 1;
 
-    auto perfectMatchFragment = hostPtrManager.getFragmentAndCheckForOverlaps(bigPtr, bigSize, overlapStatus);
+    auto perfectMatchFragment = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, bigPtr, bigSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_WITH_EXACT_SIZE_AS_STORED_FRAGMENT, overlapStatus);
     EXPECT_EQ(perfectMatchFragment, storedBigFragment);
 
-    auto oustideFragment = hostPtrManager.getFragmentAndCheckForOverlaps(ptrOutside, outsideSize, overlapStatus);
+    auto oustideFragment = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrOutside, outsideSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, oustideFragment);
 
@@ -664,26 +711,26 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledWithBigFragmentWhenAskedForFragmne
     auto ptrPartial = (void *)(((uintptr_t)bigPtr + bigSize) - 100);
     auto partialBigSize = MemoryConstants::pageSize * 100;
 
-    auto partialFragment = hostPtrManager.getFragmentAndCheckForOverlaps(ptrPartial, partialBigSize, overlapStatus);
+    auto partialFragment = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrPartial, partialBigSize, overlapStatus);
     EXPECT_EQ(nullptr, partialFragment);
     EXPECT_EQ(OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT, overlapStatus);
 }
 
-TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenCheckedForOverlappingThenProperOverlappingStatusIsReturned) {
+TEST_F(HostPtrManagerTest, GivenHostPtrManagerFilledWithFragmentsWhenCheckedForOverlappingThenProperOverlappingStatusIsReturned) {
     auto bigPtr = (void *)0x04000;
     auto bigSize = 10 * MemoryConstants::pageSize;
     FragmentStorage fragment;
     fragment.fragmentCpuPointer = bigPtr;
     fragment.fragmentSize = bigSize;
     MockHostPtrManager hostPtrManager;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     EXPECT_EQ(1u, hostPtrManager.getFragmentCount());
 
     auto ptrNonOverlapingPriorToBigPtr = (void *)0x2000;
     auto smallSize = MemoryConstants::pageSize;
 
     OverlapStatus overlapStatus;
-    auto fragment2 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingPriorToBigPtr, smallSize, overlapStatus);
+    auto fragment2 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingPriorToBigPtr, smallSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment2);
 
@@ -691,27 +738,27 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenCheckedForOverlap
     auto checkMatch = (uintptr_t)ptrNonOverlapingPriorToBigPtrByPage + smallSize;
     EXPECT_EQ(checkMatch, (uintptr_t)bigPtr);
 
-    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingPriorToBigPtrByPage, smallSize, overlapStatus);
+    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingPriorToBigPtrByPage, smallSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment3);
 
-    auto fragment4 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingPriorToBigPtrByPage, smallSize + 1, overlapStatus);
+    auto fragment4 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingPriorToBigPtrByPage, smallSize + 1, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT, overlapStatus);
     EXPECT_EQ(nullptr, fragment4);
 }
 
-TEST(HostPtrManager, GivenEmptyHostPtrManagerWhenAskedForOverlapingThenNoOverlappingIsReturned) {
+TEST_F(HostPtrManagerTest, GivenEmptyHostPtrManagerWhenAskedForOverlapingThenNoOverlappingIsReturned) {
     MockHostPtrManager hostPtrManager;
     auto bigPtr = (void *)0x04000;
     auto bigSize = 10 * MemoryConstants::pageSize;
 
     OverlapStatus overlapStatus;
-    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(bigPtr, bigSize, overlapStatus);
+    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, bigPtr, bigSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment3);
 }
 
-TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenAskedForOverlpaingThenProperStatusIsReturned) {
+TEST_F(HostPtrManagerTest, GivenHostPtrManagerFilledWithFragmentsWhenAskedForOverlpaingThenProperStatusIsReturned) {
     auto bigPtr1 = (void *)0x01000;
     auto bigPtr2 = (void *)0x03000;
     auto bigSize = MemoryConstants::pageSize;
@@ -719,33 +766,33 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenAskedForOverlpain
     fragment.fragmentCpuPointer = bigPtr1;
     fragment.fragmentSize = bigSize;
     MockHostPtrManager hostPtrManager;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     fragment.fragmentCpuPointer = bigPtr2;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     EXPECT_EQ(2u, hostPtrManager.getFragmentCount());
 
     auto ptrNonOverlapingInTheMiddleOfBigPtrs = (void *)0x2000;
     auto ptrNonOverlapingAfterBigPtr = (void *)0x4000;
     auto ptrNonOverlapingBeforeBigPtr = (void *)0;
     OverlapStatus overlapStatus;
-    auto fragment1 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingInTheMiddleOfBigPtrs, bigSize, overlapStatus);
+    auto fragment1 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingInTheMiddleOfBigPtrs, bigSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment1);
 
-    auto fragment2 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingInTheMiddleOfBigPtrs, bigSize * 5, overlapStatus);
+    auto fragment2 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingInTheMiddleOfBigPtrs, bigSize * 5, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT, overlapStatus);
     EXPECT_EQ(nullptr, fragment2);
 
-    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingAfterBigPtr, bigSize * 5, overlapStatus);
+    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingAfterBigPtr, bigSize * 5, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment3);
 
-    auto fragment4 = hostPtrManager.getFragmentAndCheckForOverlaps(ptrNonOverlapingBeforeBigPtr, bigSize, overlapStatus);
+    auto fragment4 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, ptrNonOverlapingBeforeBigPtr, bigSize, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment4);
 }
 
-TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenAskedForOverlapingThenProperOverlapingStatusIsReturned) {
+TEST_F(HostPtrManagerTest, GivenHostPtrManagerFilledWithFragmentsWhenAskedForOverlapingThenProperOverlapingStatusIsReturned) {
     auto bigPtr1 = (void *)0x10000;
     auto bigPtr2 = (void *)0x03000;
     auto bigPtr3 = (void *)0x11000;
@@ -758,29 +805,29 @@ TEST(HostPtrManager, GivenHostPtrManagerFilledWithFragmentsWhenAskedForOverlapin
     fragment.fragmentCpuPointer = bigPtr1;
     fragment.fragmentSize = bigSize1;
     MockHostPtrManager hostPtrManager;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     fragment.fragmentCpuPointer = bigPtr2;
     fragment.fragmentSize = bigSize2;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
     fragment.fragmentCpuPointer = bigPtr3;
     fragment.fragmentSize = bigSize3;
-    hostPtrManager.storeFragment(fragment);
+    hostPtrManager.storeFragment(rootDeviceIndex, fragment);
 
     EXPECT_EQ(3u, hostPtrManager.getFragmentCount());
 
     OverlapStatus overlapStatus;
-    auto fragment1 = hostPtrManager.getFragmentAndCheckForOverlaps(bigPtr1, bigSize1 + 1, overlapStatus);
+    auto fragment1 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, bigPtr1, bigSize1 + 1, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT, overlapStatus);
     EXPECT_EQ(nullptr, fragment1);
 
     auto priorToBig1 = (void *)0x9999;
 
-    auto fragment2 = hostPtrManager.getFragmentAndCheckForOverlaps(priorToBig1, 1, overlapStatus);
+    auto fragment2 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, priorToBig1, 1, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER, overlapStatus);
     EXPECT_EQ(nullptr, fragment2);
 
     auto middleOfBig3 = (void *)0x11111;
-    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(middleOfBig3, 1, overlapStatus);
+    auto fragment3 = hostPtrManager.getFragmentAndCheckForOverlaps(rootDeviceIndex, middleOfBig3, 1, overlapStatus);
     EXPECT_EQ(OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT, overlapStatus);
     EXPECT_NE(nullptr, fragment3);
 }
@@ -809,12 +856,12 @@ TEST_F(HostPtrAllocationTest, whenPrepareOsHandlesForAllocationThenPopulateAsMan
     void *cpuPtr = reinterpret_cast<void *>(0x100001);
     size_t allocationSize = MemoryConstants::pageSize / 2;
     for (uint32_t expectedFragmentCount = 1; expectedFragmentCount <= 3; expectedFragmentCount++, allocationSize += MemoryConstants::pageSize) {
-        auto requirements = hostPtrManager->getAllocationRequirements(cpuPtr, allocationSize);
+        auto requirements = hostPtrManager->getAllocationRequirements(csr->getRootDeviceIndex(), cpuPtr, allocationSize);
         EXPECT_EQ(expectedFragmentCount, requirements.requiredFragmentsCount);
         auto osStorage = hostPtrManager->prepareOsStorageForAllocation(*memoryManager, allocationSize, cpuPtr, 0);
         EXPECT_EQ(expectedFragmentCount, osStorage.fragmentCount);
         EXPECT_EQ(expectedFragmentCount, hostPtrManager->getFragmentCount());
-        hostPtrManager->releaseHandleStorage(osStorage);
+        hostPtrManager->releaseHandleStorage(csr->getRootDeviceIndex(), osStorage);
         memoryManager->cleanOsHandles(osStorage, 0);
         EXPECT_EQ(0u, hostPtrManager->getFragmentCount());
     }
@@ -830,9 +877,9 @@ TEST_F(HostPtrAllocationTest, whenOverlappedFragmentIsBiggerThenStoredAndStoredF
 
     EXPECT_NE(nullptr, graphicsAllocation1);
 
-    auto fragment1 = hostPtrManager->getFragment(alignDown(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment1 = hostPtrManager->getFragment({alignDown(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment1);
-    auto fragment2 = hostPtrManager->getFragment(alignUp(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment2 = hostPtrManager->getFragment({alignUp(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment2);
 
     uint32_t taskCountReady = 2;
@@ -902,6 +949,7 @@ HWTEST_F(HostPtrAllocationTest, givenOverlappingFragmentsWhenCheckIsCalledThenWa
     requirements.allocationFragments[0].allocationPtr = alignDown(cpuPtr, MemoryConstants::pageSize);
     requirements.allocationFragments[0].allocationSize = MemoryConstants::pageSize * 10;
     requirements.allocationFragments[0].fragmentPosition = FragmentPosition::NONE;
+    requirements.rootDeviceIndex = csr0->getRootDeviceIndex();
 
     hostPtrManager->checkAllocationsForOverlapping(*memoryManager, &requirements);
 
@@ -923,9 +971,9 @@ TEST_F(HostPtrAllocationTest, whenOverlappedFragmentIsBiggerThenStoredAndStoredF
 
     EXPECT_NE(nullptr, graphicsAllocation1);
 
-    auto fragment1 = hostPtrManager->getFragment(alignDown(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment1 = hostPtrManager->getFragment({alignDown(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment1);
-    auto fragment2 = hostPtrManager->getFragment(alignUp(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment2 = hostPtrManager->getFragment({alignUp(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment2);
 
     uint32_t taskCountReady = 2;
@@ -944,6 +992,7 @@ TEST_F(HostPtrAllocationTest, whenOverlappedFragmentIsBiggerThenStoredAndStoredF
     requirements.allocationFragments[0].allocationPtr = alignDown(cpuPtr1, MemoryConstants::pageSize);
     requirements.allocationFragments[0].allocationSize = MemoryConstants::pageSize * 10;
     requirements.allocationFragments[0].fragmentPosition = FragmentPosition::NONE;
+    requirements.rootDeviceIndex = csr->getRootDeviceIndex();
 
     RequirementsStatus status = hostPtrManager->checkAllocationsForOverlapping(*memoryManager, &requirements);
 
@@ -965,13 +1014,13 @@ TEST_F(HostPtrAllocationTest, GivenAllocationsWithoutBiggerOverlapWhenChckingFor
     EXPECT_NE(nullptr, graphicsAllocation1);
     EXPECT_NE(nullptr, graphicsAllocation2);
 
-    auto fragment1 = hostPtrManager->getFragment(alignDown(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment1 = hostPtrManager->getFragment({alignDown(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment1);
-    auto fragment2 = hostPtrManager->getFragment(alignUp(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment2 = hostPtrManager->getFragment({alignUp(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment2);
-    auto fragment3 = hostPtrManager->getFragment(alignDown(cpuPtr2, MemoryConstants::pageSize));
+    auto fragment3 = hostPtrManager->getFragment({alignDown(cpuPtr2, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment3);
-    auto fragment4 = hostPtrManager->getFragment(alignUp(cpuPtr2, MemoryConstants::pageSize));
+    auto fragment4 = hostPtrManager->getFragment({alignUp(cpuPtr2, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment4);
 
     AllocationRequirements requirements;
@@ -1005,9 +1054,9 @@ TEST_F(HostPtrAllocationTest, GivenAllocationsWithBiggerOverlapWhenChckingForOve
 
     EXPECT_NE(nullptr, graphicsAllocation1);
 
-    auto fragment1 = hostPtrManager->getFragment(alignDown(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment1 = hostPtrManager->getFragment({alignDown(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment1);
-    auto fragment2 = hostPtrManager->getFragment(alignUp(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment2 = hostPtrManager->getFragment({alignUp(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment2);
 
     uint32_t taskCountReady = 1;
@@ -1031,3 +1080,44 @@ TEST_F(HostPtrAllocationTest, GivenAllocationsWithBiggerOverlapWhenChckingForOve
 
     EXPECT_EQ(RequirementsStatus::SUCCESS, status);
 }
+
+TEST(HostPtrEntryKeyTest, givenTwoHostPtrEntryKeysWhenComparingThemThenKeyWithLowerRootDeviceIndexIsLower) {
+
+    auto hostPtr0 = reinterpret_cast<void *>(0x100);
+    auto hostPtr1 = reinterpret_cast<void *>(0x200);
+    auto hostPtr2 = reinterpret_cast<void *>(0x300);
+
+    HostPtrEntryKey key0{hostPtr1, 0u};
+    HostPtrEntryKey key1{hostPtr1, 1u};
+
+    EXPECT_TRUE(key0 < key1);
+    EXPECT_FALSE(key1 < key0);
+
+    key0.ptr = hostPtr0;
+    EXPECT_TRUE(key0 < key1);
+    EXPECT_FALSE(key1 < key0);
+
+    key0.ptr = hostPtr2;
+    EXPECT_TRUE(key0 < key1);
+    EXPECT_FALSE(key1 < key0);
+}
+
+TEST(HostPtrEntryKeyTest, givenTwoHostPtrEntryKeysWithSameRootDeviceIndexWhenComparingThemThenKeyWithLowerPtrIsLower) {
+    auto hostPtr0 = reinterpret_cast<void *>(0x100);
+    auto hostPtr1 = reinterpret_cast<void *>(0x200);
+
+    HostPtrEntryKey key0{hostPtr0, 1u};
+    HostPtrEntryKey key1{hostPtr1, 1u};
+
+    EXPECT_TRUE(key0 < key1);
+    EXPECT_FALSE(key1 < key0);
+}
+TEST(HostPtrEntryKeyTest, givenTwoSameHostPtrEntryKeysWithSameRootDeviceIndexWhenComparingThemThenTheyAreEqual) {
+    auto hostPtr = reinterpret_cast<void *>(0x100);
+
+    HostPtrEntryKey key0{hostPtr, 1u};
+    HostPtrEntryKey key1{hostPtr, 1u};
+
+    EXPECT_FALSE(key0 < key1);
+    EXPECT_FALSE(key1 < key0);
+}

opencl/test/unit_test/memory_manager/memory_manager_tests.cpp

@@ -182,7 +182,7 @@ TEST_F(MemoryAllocatorTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationToH
 
     MockGraphicsAllocation gfxAllocation(cpuPtr, size);
     memoryManager->addAllocationToHostPtrManager(&gfxAllocation);
-    auto fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    auto fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), csr->getRootDeviceIndex()});
     EXPECT_NE(fragment, nullptr);
     EXPECT_TRUE(fragment->driverAllocation);
     EXPECT_EQ(fragment->refCount, 1);
@@ -193,22 +193,22 @@ TEST_F(MemoryAllocatorTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationToH
 
     FragmentStorage fragmentStorage = {};
     fragmentStorage.fragmentCpuPointer = cpuPtr;
-    memoryManager->getHostPtrManager()->storeFragment(fragmentStorage);
+    memoryManager->getHostPtrManager()->storeFragment(csr->getRootDeviceIndex(), fragmentStorage);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), csr->getRootDeviceIndex()});
     EXPECT_EQ(fragment->refCount, 2);
 
     fragment->driverAllocation = false;
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), csr->getRootDeviceIndex()});
     EXPECT_EQ(fragment->refCount, 2);
     fragment->driverAllocation = true;
 
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), csr->getRootDeviceIndex()});
     EXPECT_EQ(fragment->refCount, 1);
 
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), csr->getRootDeviceIndex()});
     EXPECT_EQ(fragment, nullptr);
 }
 
@@ -262,7 +262,7 @@ TEST_F(MemoryAllocatorTest, GivenAlignedHostPtrWithAlignedSizeWhenAllocatingGrap
     EXPECT_NE(nullptr, graphicsAllocation);
 
     EXPECT_EQ(1u, hostPtrManager->getFragmentCount());
-    auto fragmentData = hostPtrManager->getFragment(ptr);
+    auto fragmentData = hostPtrManager->getFragment({ptr, device->getRootDeviceIndex()});
 
     ASSERT_NE(nullptr, fragmentData);
 
@@ -322,7 +322,7 @@ TEST_F(MemoryAllocatorTest, WhenPopulatingOsHandleThenOneFragmentIsReturned) {
     EXPECT_NE(nullptr, storage.fragmentStorageData[0].osHandleStorage);
     EXPECT_EQ(nullptr, storage.fragmentStorageData[1].osHandleStorage);
     EXPECT_EQ(nullptr, storage.fragmentStorageData[2].osHandleStorage);
-    memoryManager->getHostPtrManager()->releaseHandleStorage(storage);
+    memoryManager->getHostPtrManager()->releaseHandleStorage(csr->getRootDeviceIndex(), storage);
     memoryManager->cleanOsHandles(storage, 0);
 }
 
@@ -336,7 +336,7 @@ TEST_F(MemoryAllocatorTest, givenOsHandleStorageWhenOsHandlesAreCleanedAndAubMan
     OsHandleStorage storage;
     storage.fragmentStorageData[0].cpuPtr = (void *)0x1000;
     mockMemoryManager.populateOsHandles(storage, 0);
-    mockMemoryManager.getHostPtrManager()->releaseHandleStorage(storage);
+    mockMemoryManager.getHostPtrManager()->releaseHandleStorage(csr->getRootDeviceIndex(), storage);
     mockMemoryManager.cleanOsHandles(storage, 0);
 
     EXPECT_EQ(nullptr, mockAubCenter->aubManager);
@@ -360,8 +360,8 @@ TEST_F(MemoryAllocatorTest, givenOsHandleStorageAndFreeMemoryEnabledWhenOsHandle
 
     OsHandleStorage storage;
     storage.fragmentStorageData[0].cpuPtr = reinterpret_cast<void *>(0x1000);
-    mockMemoryManager.populateOsHandles(storage, 0);
+    mockMemoryManager.populateOsHandles(storage, rootDeviceIndex);
-    mockMemoryManager.getHostPtrManager()->releaseHandleStorage(storage);
+    mockMemoryManager.getHostPtrManager()->releaseHandleStorage(rootDeviceIndex, storage);
     mockMemoryManager.cleanOsHandles(storage, rootDeviceIndex);
 
     EXPECT_FALSE(mockManager0->freeMemoryCalled);
@@ -374,7 +374,7 @@ TEST_F(MemoryAllocatorTest, GivenEmptyMemoryManagerAndMisalingedHostPtrWithHugeS
 
     MockMemoryManager mockMemoryManager(*executionEnvironment);
     auto hostPtrManager = static_cast<MockHostPtrManager *>(mockMemoryManager.getHostPtrManager());
-    auto reqs = MockHostPtrManager::getAllocationRequirements(cpuPtr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(device->getRootDeviceIndex(), cpuPtr, size);
 
     ASSERT_EQ(3u, reqs.requiredFragmentsCount);
 
@@ -1419,13 +1419,13 @@ TEST_F(MemoryManagerWithCsrTest, GivenAllocationsInHostPtrManagerReadyForCleanin
     EXPECT_NE(nullptr, graphicsAllocation1);
     EXPECT_NE(nullptr, graphicsAllocation2);
 
-    auto fragment1 = hostPtrManager->getFragment(alignDown(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment1 = hostPtrManager->getFragment({alignDown(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment1);
-    auto fragment2 = hostPtrManager->getFragment(alignUp(cpuPtr1, MemoryConstants::pageSize));
+    auto fragment2 = hostPtrManager->getFragment({alignUp(cpuPtr1, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment2);
-    auto fragment3 = hostPtrManager->getFragment(alignDown(cpuPtr2, MemoryConstants::pageSize));
+    auto fragment3 = hostPtrManager->getFragment({alignDown(cpuPtr2, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment3);
-    auto fragment4 = hostPtrManager->getFragment(alignUp(cpuPtr2, MemoryConstants::pageSize));
+    auto fragment4 = hostPtrManager->getFragment({alignUp(cpuPtr2, MemoryConstants::pageSize), csr->getRootDeviceIndex()});
     EXPECT_NE(nullptr, fragment4);
 
     uint32_t taskCountReady = 1;

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

@@ -1019,7 +1019,7 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, makeResidentTwice) {
 HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, makeResidentTwiceWhenFragmentStorage) {
     auto ptr = (void *)0x1001;
     auto size = MemoryConstants::pageSize * 10;
-    auto reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(csr->getRootDeviceIndex(), ptr, size);
     auto allocation = mm->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, size}, ptr);
 
     ASSERT_EQ(3u, allocation->fragmentsStorage.fragmentCount);
@@ -1115,7 +1115,7 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, GivenAllocationCreatedFromThree
     auto ptr = (void *)0x1001;
     auto size = MemoryConstants::pageSize * 10;
 
-    auto reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(csr->getRootDeviceIndex(), ptr, size);
 
     auto allocation = mm->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, size}, ptr);
 
@@ -1145,7 +1145,7 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, GivenAllocationsContainingDiffe
     auto size = MemoryConstants::pageSize;
     auto size2 = 100u;
 
-    auto reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(csr->getRootDeviceIndex(), ptr, size);
 
     auto allocation = mm->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, size}, ptr);
 
@@ -1172,7 +1172,7 @@ HWTEST_TEMPLATED_F(DrmCommandStreamEnhancedTest, GivenAllocationsContainingDiffe
     csr->getResidencyAllocations().clear();
 
     auto allocation2 = mm->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, size2}, ptr);
-    reqs = MockHostPtrManager::getAllocationRequirements(ptr, size2);
+    reqs = MockHostPtrManager::getAllocationRequirements(csr->getRootDeviceIndex(), ptr, size2);
 
     ASSERT_EQ(1u, allocation2->fragmentsStorage.fragmentCount);
     ASSERT_EQ(1u, reqs.requiredFragmentsCount);

opencl/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp

@@ -80,9 +80,10 @@ TEST_F(DrmMemoryManagerTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationTo
     void *cpuPtr = (void *)0x30000;
     size_t size = 0x1000;
 
-    DrmAllocation gfxAllocation(0, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, cpuPtr, size, (osHandle)1u, MemoryPool::MemoryNull);
+    const uint32_t rootDeviceIndex = 0u;
+    DrmAllocation gfxAllocation(rootDeviceIndex, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, cpuPtr, size, (osHandle)1u, MemoryPool::MemoryNull);
     memoryManager->addAllocationToHostPtrManager(&gfxAllocation);
-    auto fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    auto fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), rootDeviceIndex});
     EXPECT_NE(fragment, nullptr);
     EXPECT_TRUE(fragment->driverAllocation);
     EXPECT_EQ(fragment->refCount, 1);
@@ -95,22 +96,22 @@ TEST_F(DrmMemoryManagerTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationTo
 
     FragmentStorage fragmentStorage = {};
     fragmentStorage.fragmentCpuPointer = cpuPtr;
-    memoryManager->getHostPtrManager()->storeFragment(fragmentStorage);
+    memoryManager->getHostPtrManager()->storeFragment(rootDeviceIndex, fragmentStorage);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), rootDeviceIndex});
     EXPECT_EQ(fragment->refCount, 2);
 
     fragment->driverAllocation = false;
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), rootDeviceIndex});
     EXPECT_EQ(fragment->refCount, 2);
     fragment->driverAllocation = true;
 
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), rootDeviceIndex});
     EXPECT_EQ(fragment->refCount, 1);
 
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment({gfxAllocation.getUnderlyingBuffer(), rootDeviceIndex});
     EXPECT_EQ(fragment, nullptr);
 }
 
@@ -1013,7 +1014,7 @@ TEST_F(DrmMemoryManagerTest, GivenMisalignedHostPtrAndMultiplePagesSizeWhenAsked
 
     ASSERT_EQ(3u, hostPtrManager->getFragmentCount());
 
-    auto reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
 
     for (int i = 0; i < maxFragmentsCount; i++) {
         ASSERT_NE(nullptr, graphicsAllocation->fragmentsStorage.fragmentStorageData[i].osHandleStorage->bo);
@@ -3372,8 +3373,8 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMem
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
 
     EXPECT_EQ(0u, hostPtrManager->getFragmentCount());
-    EXPECT_EQ(nullptr, hostPtrManager->getFragment(handleStorage.fragmentStorageData[1].cpuPtr));
+    EXPECT_EQ(nullptr, hostPtrManager->getFragment({handleStorage.fragmentStorageData[1].cpuPtr, rootDeviceIndex}));
-    EXPECT_EQ(nullptr, hostPtrManager->getFragment(handleStorage.fragmentStorageData[2].cpuPtr));
+    EXPECT_EQ(nullptr, hostPtrManager->getFragment({handleStorage.fragmentStorageData[2].cpuPtr, rootDeviceIndex}));
 
     handleStorage.fragmentStorageData[0].freeTheFragment = false;
     handleStorage.fragmentStorageData[1].freeTheFragment = true;
@@ -3407,7 +3408,7 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenEnabledValidateHostMem
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
     EXPECT_EQ(1u, hostPtrManager->getFragmentCount());
-    EXPECT_NE(nullptr, hostPtrManager->getFragment(handleStorage.fragmentStorageData[0].cpuPtr));
+    EXPECT_NE(nullptr, hostPtrManager->getFragment({handleStorage.fragmentStorageData[0].cpuPtr, device->getRootDeviceIndex()}));
 
     handleStorage.fragmentStorageData[0].freeTheFragment = true;
     memoryManager->cleanOsHandles(handleStorage, rootDeviceIndex);

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

@@ -743,12 +743,12 @@ TEST_F(WddmCommandStreamTest, givenTwoTemporaryAllocationsWhenCleanTemporaryAllo
     auto alignedPtr = alignDown(host_ptr, MemoryConstants::pageSize);
     auto alignedPtr2 = alignDown(host_ptr2, MemoryConstants::pageSize);
 
-    auto fragment = hostPtrManager->getFragment(alignedPtr2);
+    auto fragment = hostPtrManager->getFragment({alignedPtr2, csr->getRootDeviceIndex()});
     ASSERT_NE(nullptr, fragment);
 
     EXPECT_EQ(alignedPtr2, fragment->fragmentCpuPointer);
 
-    auto fragment2 = hostPtrManager->getFragment(alignedPtr);
+    auto fragment2 = hostPtrManager->getFragment({alignedPtr, csr->getRootDeviceIndex()});
     EXPECT_EQ(nullptr, fragment2);
     // destroy remaining allocation
     csr->waitForTaskCountAndCleanAllocationList(100, TEMPORARY_ALLOCATION);

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

@@ -434,11 +434,12 @@ TEST_F(WddmMemoryManagerTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationT
     uint64_t gpuPtr = 0x123;
 
     MockWddmAllocation gfxAllocation;
+    HostPtrEntryKey key{cpuPtr, gfxAllocation.getRootDeviceIndex()};
     gfxAllocation.cpuPtr = cpuPtr;
     gfxAllocation.size = size;
     gfxAllocation.gpuPtr = gpuPtr;
     memoryManager->addAllocationToHostPtrManager(&gfxAllocation);
-    auto fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    auto fragment = memoryManager->getHostPtrManager()->getFragment(key);
     EXPECT_NE(fragment, nullptr);
     EXPECT_TRUE(fragment->driverAllocation);
     EXPECT_EQ(fragment->refCount, 1);
@@ -452,22 +453,22 @@ TEST_F(WddmMemoryManagerTest, GivenGraphicsAllocationWhenAddAndRemoveAllocationT
 
     FragmentStorage fragmentStorage = {};
     fragmentStorage.fragmentCpuPointer = cpuPtr;
-    memoryManager->getHostPtrManager()->storeFragment(fragmentStorage);
+    memoryManager->getHostPtrManager()->storeFragment(gfxAllocation.getRootDeviceIndex(), fragmentStorage);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment(key);
     EXPECT_EQ(fragment->refCount, 2);
 
     fragment->driverAllocation = false;
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment(key);
     EXPECT_EQ(fragment->refCount, 2);
     fragment->driverAllocation = true;
 
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment(key);
     EXPECT_EQ(fragment->refCount, 1);
 
     memoryManager->removeAllocationFromHostPtrManager(&gfxAllocation);
-    fragment = memoryManager->getHostPtrManager()->getFragment(gfxAllocation.getUnderlyingBuffer());
+    fragment = memoryManager->getHostPtrManager()->getFragment(key);
     EXPECT_EQ(fragment, nullptr);
 }
 
@@ -842,13 +843,13 @@ TEST_F(WddmMemoryManagerTest, GivenOffsetsWhenAllocatingGpuMemHostThenAllocatedO
 
     auto hostPtrManager = memoryManager->getHostPtrManager();
 
-    auto fragment = hostPtrManager->getFragment(ptr);
+    auto fragment = hostPtrManager->getFragment({ptr, rootDeviceIndex});
     ASSERT_NE(nullptr, fragment);
     EXPECT_TRUE(fragment->refCount == 1);
     EXPECT_NE(fragment->osInternalStorage, nullptr);
 
     // offset by 3 pages, not in boundary
-    auto fragment2 = hostPtrManager->getFragment((char *)ptr + 3 * 4096);
+    auto fragment2 = hostPtrManager->getFragment({reinterpret_cast<char *>(ptr) + 3 * 4096, rootDeviceIndex});
 
     EXPECT_EQ(nullptr, fragment2);
 
@@ -858,7 +859,7 @@ TEST_F(WddmMemoryManagerTest, GivenOffsetsWhenAllocatingGpuMemHostThenAllocatedO
     // Should be same cpu ptr and gpu ptr
     EXPECT_EQ(offsetPtr, gpuAllocation2->getUnderlyingBuffer());
 
-    auto fragment3 = hostPtrManager->getFragment(offsetPtr);
+    auto fragment3 = hostPtrManager->getFragment({offsetPtr, rootDeviceIndex});
     ASSERT_NE(nullptr, fragment3);
 
     EXPECT_TRUE(fragment3->refCount == 2);
@@ -868,14 +869,14 @@ TEST_F(WddmMemoryManagerTest, GivenOffsetsWhenAllocatingGpuMemHostThenAllocatedO
 
     memoryManager->freeGraphicsMemory(gpuAllocation2);
 
-    auto fragment4 = hostPtrManager->getFragment(ptr);
+    auto fragment4 = hostPtrManager->getFragment({ptr, rootDeviceIndex});
     ASSERT_NE(nullptr, fragment4);
 
     EXPECT_TRUE(fragment4->refCount == 1);
 
     memoryManager->freeGraphicsMemory(gpuAllocation);
 
-    fragment4 = hostPtrManager->getFragment(ptr);
+    fragment4 = hostPtrManager->getFragment({ptr, rootDeviceIndex});
     EXPECT_EQ(nullptr, fragment4);
 
     alignedFree(ptr);
@@ -890,7 +891,7 @@ TEST_F(WddmMemoryManagerTest, WhenAllocatingGpuMemThenOsInternalStorageIsPopulat
     ASSERT_NE(nullptr, gpuAllocation);
     EXPECT_EQ(ptr, gpuAllocation->getUnderlyingBuffer());
 
-    auto fragment = memoryManager->getHostPtrManager()->getFragment(ptr);
+    auto fragment = memoryManager->getHostPtrManager()->getFragment({ptr, rootDeviceIndex});
     ASSERT_NE(nullptr, fragment);
     EXPECT_TRUE(fragment->refCount == 1);
     EXPECT_NE(fragment->osInternalStorage->handle, 0);
@@ -1189,7 +1190,7 @@ TEST_F(BufferWithWddmMemory, GivenMisalignedHostPtrAndMultiplePagesSizeWhenAsked
 
     ASSERT_EQ(3u, hostPtrManager->getFragmentCount());
 
-    auto reqs = MockHostPtrManager::getAllocationRequirements(ptr, size);
+    auto reqs = MockHostPtrManager::getAllocationRequirements(context.getDevice(0)->getRootDeviceIndex(), ptr, size);
 
     for (int i = 0; i < maxFragmentsCount; i++) {
         EXPECT_NE((D3DKMT_HANDLE) nullptr, graphicsAllocation->fragmentsStorage.fragmentStorageData[i].osHandleStorage->handle);
@@ -1261,7 +1262,7 @@ TEST_F(BufferWithWddmMemory, givenFragmentsThatAreNotInOrderWhenGraphicsAllocati
     fragment.fragmentSize = size;
     fragment.osInternalStorage = handleStorage.fragmentStorageData[0].osHandleStorage;
     fragment.osInternalStorage->gpuPtr = gpuAdress;
-    memoryManager->getHostPtrManager()->storeFragment(fragment);
+    memoryManager->getHostPtrManager()->storeFragment(rootDeviceIndex, fragment);
 
     AllocationData allocationData;
     allocationData.size = size;
@@ -1296,7 +1297,7 @@ TEST_F(BufferWithWddmMemory, givenFragmentsThatAreNotInOrderWhenGraphicsAllocati
     fragment.fragmentSize = size;
     fragment.osInternalStorage = handleStorage.fragmentStorageData[0].osHandleStorage;
     fragment.osInternalStorage->gpuPtr = gpuAdress;
-    memoryManager->getHostPtrManager()->storeFragment(fragment);
+    memoryManager->getHostPtrManager()->storeFragment(rootDeviceIndex, fragment);
 
     auto offset = 80;
     auto allocationPtr = ptrOffset(ptr, offset);
@@ -1811,16 +1812,16 @@ TEST_F(WddmMemoryManagerTest2, givenReadOnlyMemoryPassedToPopulateOsHandlesWhenC
 
     EXPECT_CALL(*wddm, createAllocationsAndMapGpuVa(::testing::_)).WillOnce(::testing::Return(STATUS_GRAPHICS_NO_VIDEO_MEMORY));
 
-    auto result = memoryManager->populateOsHandles(handleStorage, 0);
+    auto result = memoryManager->populateOsHandles(handleStorage, mockRootDeviceIndex);
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
 
     EXPECT_EQ(MemoryManager::AllocationStatus::InvalidHostPointer, result);
     auto numberOfStoredFragments = hostPtrManager->getFragmentCount();
     EXPECT_EQ(0u, numberOfStoredFragments);
-    EXPECT_EQ(nullptr, hostPtrManager->getFragment(handleStorage.fragmentStorageData[1].cpuPtr));
+    EXPECT_EQ(nullptr, hostPtrManager->getFragment({handleStorage.fragmentStorageData[1].cpuPtr, mockRootDeviceIndex}));
 
     handleStorage.fragmentStorageData[1].freeTheFragment = true;
-    memoryManager->cleanOsHandles(handleStorage, 0);
+    memoryManager->cleanOsHandles(handleStorage, mockRootDeviceIndex);
 }
 
 TEST(WddmMemoryManagerCleanupTest, givenUsedTagAllocationInWddmMemoryManagerWhenCleanupMemoryManagerThenDontAccessCsr) {

shared/source/memory_manager/host_ptr_defines.h

@@ -9,6 +9,7 @@
 
 #include <cinttypes>
 #include <cstdlib>
+#include <limits>
 
 namespace NEO {
 
@@ -49,6 +50,7 @@ struct AllocationRequirements {
     PartialAllocation allocationFragments[maxFragmentsCount];
     uint64_t totalRequiredSize = 0u;
     uint32_t requiredFragmentsCount = 0u;
+    uint32_t rootDeviceIndex = std::numeric_limits<uint32_t>::max();
 };
 
 struct FragmentStorage {

shared/source/memory_manager/host_ptr_manager.cpp

@@ -11,40 +11,36 @@
 
 using namespace NEO;
 
-HostPtrFragmentsContainer::iterator HostPtrManager::findElement(const void *ptr) {
+HostPtrFragmentsContainer::iterator HostPtrManager::findElement(HostPtrEntryKey key) {
-    auto nextElement = partialAllocations.lower_bound(ptr);
+    auto nextElement = partialAllocations.lower_bound(key);
     auto element = nextElement;
     if (element != partialAllocations.end()) {
+
         auto &storedFragment = element->second;
-        if (storedFragment.fragmentCpuPointer <= ptr) {
+        if (element->first.rootDeviceIndex == key.rootDeviceIndex && storedFragment.fragmentCpuPointer == key.ptr) {
             return element;
-        } else if (element != partialAllocations.begin()) {
-            element--;
-            auto &storedFragment = element->second;
-            auto storedEndAddress = (uintptr_t)storedFragment.fragmentCpuPointer + storedFragment.fragmentSize;
-            if (storedFragment.fragmentSize == 0) {
-                storedEndAddress++;
-            }
-            if ((uintptr_t)ptr < (uintptr_t)storedEndAddress) {
-                return element;
-            }
         }
-    } else if (element != partialAllocations.begin()) {
+    }
+    if (element != partialAllocations.begin()) {
         element--;
+        if (element->first.rootDeviceIndex != key.rootDeviceIndex) {
+            return partialAllocations.end();
+        }
         auto &storedFragment = element->second;
-        auto storedEndAddress = (uintptr_t)storedFragment.fragmentCpuPointer + storedFragment.fragmentSize;
+        auto storedEndAddress = reinterpret_cast<uintptr_t>(storedFragment.fragmentCpuPointer) + storedFragment.fragmentSize;
         if (storedFragment.fragmentSize == 0) {
             storedEndAddress++;
         }
-        if ((uintptr_t)ptr < (uintptr_t)storedEndAddress) {
+        if (reinterpret_cast<uintptr_t>(key.ptr) < storedEndAddress) {
             return element;
         }
     }
     return partialAllocations.end();
 }
 
-AllocationRequirements HostPtrManager::getAllocationRequirements(const void *inputPtr, size_t size) {
+AllocationRequirements HostPtrManager::getAllocationRequirements(uint32_t rootDeviceIndex, const void *inputPtr, size_t size) {
     AllocationRequirements requiredAllocations;
+    requiredAllocations.rootDeviceIndex = rootDeviceIndex;
 
     auto allocationCount = 0;
     auto wholeAllocationSize = alignSizeWholePage(inputPtr, size);
@@ -93,7 +89,8 @@ OsHandleStorage HostPtrManager::populateAlreadyAllocatedFragments(AllocationRequ
     OsHandleStorage handleStorage;
     for (unsigned int i = 0; i < requirements.requiredFragmentsCount; i++) {
         OverlapStatus overlapStatus = OverlapStatus::FRAGMENT_NOT_CHECKED;
-        FragmentStorage *fragmentStorage = getFragmentAndCheckForOverlaps(const_cast<void *>(requirements.allocationFragments[i].allocationPtr), requirements.allocationFragments[i].allocationSize, overlapStatus);
+        FragmentStorage *fragmentStorage = getFragmentAndCheckForOverlaps(requirements.rootDeviceIndex, const_cast<void *>(requirements.allocationFragments[i].allocationPtr),
+                                                                          requirements.allocationFragments[i].allocationSize, overlapStatus);
         if (overlapStatus == OverlapStatus::FRAGMENT_WITHIN_STORED_FRAGMENT) {
             UNRECOVERABLE_IF(fragmentStorage == nullptr);
             fragmentStorage->refCount++;
@@ -122,43 +119,44 @@ OsHandleStorage HostPtrManager::populateAlreadyAllocatedFragments(AllocationRequ
     return handleStorage;
 }
 
-void HostPtrManager::storeFragment(FragmentStorage &fragment) {
+void HostPtrManager::storeFragment(uint32_t rootDeviceIndex, FragmentStorage &fragment) {
     std::lock_guard<decltype(allocationsMutex)> lock(allocationsMutex);
-    auto element = findElement(fragment.fragmentCpuPointer);
+    HostPtrEntryKey key{fragment.fragmentCpuPointer, rootDeviceIndex};
+    auto element = findElement(key);
     if (element != partialAllocations.end()) {
         element->second.refCount++;
     } else {
         fragment.refCount++;
-        partialAllocations.insert(std::pair<const void *, FragmentStorage>(fragment.fragmentCpuPointer, fragment));
+        partialAllocations.insert(std::pair<HostPtrEntryKey, FragmentStorage>(key, fragment));
     }
 }
 
-void HostPtrManager::storeFragment(AllocationStorageData &storageData) {
+void HostPtrManager::storeFragment(uint32_t rootDeviceIndex, AllocationStorageData &storageData) {
     FragmentStorage fragment;
     fragment.fragmentCpuPointer = const_cast<void *>(storageData.cpuPtr);
     fragment.fragmentSize = storageData.fragmentSize;
     fragment.osInternalStorage = storageData.osHandleStorage;
     fragment.residency = storageData.residency;
-    storeFragment(fragment);
+    storeFragment(rootDeviceIndex, fragment);
 }
 
 std::unique_lock<std::recursive_mutex> HostPtrManager::obtainOwnership() {
     return std::unique_lock<std::recursive_mutex>(allocationsMutex);
 }
 
-void HostPtrManager::releaseHandleStorage(OsHandleStorage &fragments) {
+void HostPtrManager::releaseHandleStorage(uint32_t rootDeviceIndex, OsHandleStorage &fragments) {
     for (int i = 0; i < maxFragmentsCount; i++) {
         if (fragments.fragmentStorageData[i].fragmentSize || fragments.fragmentStorageData[i].cpuPtr) {
-            fragments.fragmentStorageData[i].freeTheFragment = releaseHostPtr(fragments.fragmentStorageData[i].cpuPtr);
+            fragments.fragmentStorageData[i].freeTheFragment = releaseHostPtr(rootDeviceIndex, fragments.fragmentStorageData[i].cpuPtr);
         }
     }
 }
 
-bool HostPtrManager::releaseHostPtr(const void *ptr) {
+bool HostPtrManager::releaseHostPtr(uint32_t rootDeviceIndex, const void *ptr) {
     std::lock_guard<decltype(allocationsMutex)> lock(allocationsMutex);
     bool fragmentReadyToBeReleased = false;
 
-    auto element = findElement(ptr);
+    auto element = findElement({ptr, rootDeviceIndex});
 
     DEBUG_BREAK_IF(element == partialAllocations.end());
 
@@ -171,9 +169,9 @@ bool HostPtrManager::releaseHostPtr(const void *ptr) {
     return fragmentReadyToBeReleased;
 }
 
-FragmentStorage *HostPtrManager::getFragment(const void *inputPtr) {
+FragmentStorage *HostPtrManager::getFragment(HostPtrEntryKey key) {
     std::lock_guard<decltype(allocationsMutex)> lock(allocationsMutex);
-    auto element = findElement(inputPtr);
+    auto element = findElement(key);
     if (element != partialAllocations.end()) {
         return &element->second;
     }
@@ -181,10 +179,10 @@ FragmentStorage *HostPtrManager::getFragment(const void *inputPtr) {
 }
 
 //for given inputs see if any allocation overlaps
-FragmentStorage *HostPtrManager::getFragmentAndCheckForOverlaps(const void *inPtr, size_t size, OverlapStatus &overlappingStatus) {
+FragmentStorage *HostPtrManager::getFragmentAndCheckForOverlaps(uint32_t rootDeviceIndex, const void *inPtr, size_t size, OverlapStatus &overlappingStatus) {
     std::lock_guard<decltype(allocationsMutex)> lock(allocationsMutex);
     void *inputPtr = const_cast<void *>(inPtr);
-    auto nextElement = partialAllocations.lower_bound(inputPtr);
+    auto nextElement = partialAllocations.lower_bound({inputPtr, rootDeviceIndex});
     auto element = nextElement;
     overlappingStatus = OverlapStatus::FRAGMENT_NOT_OVERLAPING_WITH_ANY_OTHER;
 
@@ -193,6 +191,9 @@ FragmentStorage *HostPtrManager::getFragmentAndCheckForOverlaps(const void *inPt
     }
 
     if (element != partialAllocations.end()) {
+        if (element->first.rootDeviceIndex != rootDeviceIndex) {
+            return nullptr;
+        }
         auto &storedFragment = element->second;
         if (storedFragment.fragmentCpuPointer == inputPtr && storedFragment.fragmentSize == size) {
             overlappingStatus = OverlapStatus::FRAGMENT_WITH_EXACT_SIZE_AS_STORED_FRAGMENT;
@@ -213,6 +214,9 @@ FragmentStorage *HostPtrManager::getFragmentAndCheckForOverlaps(const void *inPt
         }
         //next fragment doesn't have to be after the inputPtr
         if (nextElement != partialAllocations.end()) {
+            if (nextElement->first.rootDeviceIndex != rootDeviceIndex) {
+                return nullptr;
+            }
             auto &storedNextElement = nextElement->second;
             auto storedNextEndAddress = (uintptr_t)storedNextElement.fragmentCpuPointer + storedNextElement.fragmentSize;
             auto storedNextStartAddress = (uintptr_t)storedNextElement.fragmentCpuPointer;
@@ -242,7 +246,7 @@ FragmentStorage *HostPtrManager::getFragmentAndCheckForOverlaps(const void *inPt
 
 OsHandleStorage HostPtrManager::prepareOsStorageForAllocation(MemoryManager &memoryManager, size_t size, const void *ptr, uint32_t rootDeviceIndex) {
     std::lock_guard<decltype(allocationsMutex)> lock(allocationsMutex);
-    auto requirements = HostPtrManager::getAllocationRequirements(ptr, size);
+    auto requirements = HostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
     UNRECOVERABLE_IF(checkAllocationsForOverlapping(memoryManager, &requirements) == RequirementsStatus::FATAL);
     auto osStorage = populateAlreadyAllocatedFragments(requirements);
     if (osStorage.fragmentCount > 0) {
@@ -262,20 +266,23 @@ RequirementsStatus HostPtrManager::checkAllocationsForOverlapping(MemoryManager
     for (unsigned int i = 0; i < requirements->requiredFragmentsCount; i++) {
         OverlapStatus overlapStatus = OverlapStatus::FRAGMENT_NOT_CHECKED;
 
-        getFragmentAndCheckForOverlaps(requirements->allocationFragments[i].allocationPtr, requirements->allocationFragments[i].allocationSize, overlapStatus);
+        getFragmentAndCheckForOverlaps(requirements->rootDeviceIndex, requirements->allocationFragments[i].allocationPtr,
+                                       requirements->allocationFragments[i].allocationSize, overlapStatus);
         if (overlapStatus == OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT) {
             // clean temporary allocations
             memoryManager.cleanTemporaryAllocationListOnAllEngines(false);
 
             // check overlapping again
-            getFragmentAndCheckForOverlaps(requirements->allocationFragments[i].allocationPtr, requirements->allocationFragments[i].allocationSize, overlapStatus);
+            getFragmentAndCheckForOverlaps(requirements->rootDeviceIndex, requirements->allocationFragments[i].allocationPtr,
+                                           requirements->allocationFragments[i].allocationSize, overlapStatus);
             if (overlapStatus == OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT) {
 
                 // Wait for completion
                 memoryManager.cleanTemporaryAllocationListOnAllEngines(true);
 
                 // check overlapping last time
-                getFragmentAndCheckForOverlaps(requirements->allocationFragments[i].allocationPtr, requirements->allocationFragments[i].allocationSize, overlapStatus);
+                getFragmentAndCheckForOverlaps(requirements->rootDeviceIndex, requirements->allocationFragments[i].allocationPtr,
+                                               requirements->allocationFragments[i].allocationSize, overlapStatus);
                 if (overlapStatus == OverlapStatus::FRAGMENT_OVERLAPING_AND_BIGGER_THEN_STORED_FRAGMENT) {
                     status = RequirementsStatus::FATAL;
                     break;

shared/source/memory_manager/host_ptr_manager.h

@@ -13,25 +13,35 @@
 
 namespace NEO {
 
-using HostPtrFragmentsContainer = std::map<const void *, FragmentStorage>;
+struct HostPtrEntryKey {
+    const void *ptr = nullptr;
+    uint32_t rootDeviceIndex = std::numeric_limits<uint32_t>::max();
+
+    bool operator<(const HostPtrEntryKey &key) const {
+        return rootDeviceIndex < key.rootDeviceIndex ||
+               ((ptr < key.ptr) && (rootDeviceIndex == key.rootDeviceIndex));
+    }
+};
+
+using HostPtrFragmentsContainer = std::map<HostPtrEntryKey, FragmentStorage>;
 class MemoryManager;
 class HostPtrManager {
   public:
-    FragmentStorage *getFragment(const void *inputPtr);
+    FragmentStorage *getFragment(HostPtrEntryKey key);
     OsHandleStorage prepareOsStorageForAllocation(MemoryManager &memoryManager, size_t size, const void *ptr, uint32_t rootDeviceIndex);
-    void releaseHandleStorage(OsHandleStorage &fragments);
+    void releaseHandleStorage(uint32_t rootDeviceIndex, OsHandleStorage &fragments);
-    bool releaseHostPtr(const void *ptr);
+    bool releaseHostPtr(uint32_t rootDeviceIndex, const void *ptr);
-    void storeFragment(AllocationStorageData &storageData);
+    void storeFragment(uint32_t rootDeviceIndex, AllocationStorageData &storageData);
-    void storeFragment(FragmentStorage &fragment);
+    void storeFragment(uint32_t rootDeviceIndex, FragmentStorage &fragment);
     std::unique_lock<std::recursive_mutex> obtainOwnership();
 
   protected:
-    static AllocationRequirements getAllocationRequirements(const void *inputPtr, size_t size);
+    static AllocationRequirements getAllocationRequirements(uint32_t rootDeviceIndex, const void *inputPtr, size_t size);
     OsHandleStorage populateAlreadyAllocatedFragments(AllocationRequirements &requirements);
-    FragmentStorage *getFragmentAndCheckForOverlaps(const void *inputPtr, size_t size, OverlapStatus &overlappingStatus);
+    FragmentStorage *getFragmentAndCheckForOverlaps(uint32_t rootDeviceIndex, const void *inputPtr, size_t size, OverlapStatus &overlappingStatus);
     RequirementsStatus checkAllocationsForOverlapping(MemoryManager &memoryManager, AllocationRequirements *requirements);
 
-    HostPtrFragmentsContainer::iterator findElement(const void *ptr);
+    HostPtrFragmentsContainer::iterator findElement(HostPtrEntryKey key);
     HostPtrFragmentsContainer partialAllocations;
     std::recursive_mutex allocationsMutex;
 };

shared/source/memory_manager/memory_manager.cpp

@@ -123,7 +123,7 @@ GraphicsAllocation *MemoryManager::allocateGraphicsMemoryForImageFromHostPtr(con
 }
 
 void MemoryManager::cleanGraphicsMemoryCreatedFromHostPtr(GraphicsAllocation *graphicsAllocation) {
-    hostPtrManager->releaseHandleStorage(graphicsAllocation->fragmentsStorage);
+    hostPtrManager->releaseHandleStorage(graphicsAllocation->getRootDeviceIndex(), graphicsAllocation->fragmentsStorage);
     cleanOsHandles(graphicsAllocation->fragmentsStorage, graphicsAllocation->getRootDeviceIndex());
 }
 

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -572,16 +572,16 @@ void DrmMemoryManager::addAllocationToHostPtrManager(GraphicsAllocation *gfxAllo
     fragment.osInternalStorage = new OsHandle();
     fragment.residency = new ResidencyData();
     fragment.osInternalStorage->bo = drmMemory->getBO();
-    hostPtrManager->storeFragment(fragment);
+    hostPtrManager->storeFragment(gfxAllocation->getRootDeviceIndex(), fragment);
 }
 
 void DrmMemoryManager::removeAllocationFromHostPtrManager(GraphicsAllocation *gfxAllocation) {
     auto buffer = gfxAllocation->getUnderlyingBuffer();
-    auto fragment = hostPtrManager->getFragment(buffer);
+    auto fragment = hostPtrManager->getFragment({buffer, gfxAllocation->getRootDeviceIndex()});
     if (fragment && fragment->driverAllocation) {
         OsHandle *osStorageToRelease = fragment->osInternalStorage;
         ResidencyData *residencyDataToRelease = fragment->residency;
-        if (hostPtrManager->releaseHostPtr(buffer)) {
+        if (hostPtrManager->releaseHostPtr(gfxAllocation->getRootDeviceIndex(), buffer)) {
             delete osStorageToRelease;
             delete residencyDataToRelease;
         }
@@ -669,7 +669,7 @@ MemoryManager::AllocationStatus DrmMemoryManager::populateOsHandles(OsHandleStor
     }
 
     for (uint32_t i = 0; i < numberOfBosAllocated; i++) {
-        hostPtrManager->storeFragment(handleStorage.fragmentStorageData[indexesOfAllocatedBos[i]]);
+        hostPtrManager->storeFragment(rootDeviceIndex, handleStorage.fragmentStorageData[indexesOfAllocatedBos[i]]);
     }
     return AllocationStatus::Success;
 }

shared/source/os_interface/windows/wddm_memory_manager.cpp

@@ -347,15 +347,15 @@ void WddmMemoryManager::addAllocationToHostPtrManager(GraphicsAllocation *gfxAll
     fragment.osInternalStorage->handle = wddmMemory->getDefaultHandle();
     fragment.osInternalStorage->gmm = gfxAllocation->getDefaultGmm();
     fragment.residency = &wddmMemory->getResidencyData();
-    hostPtrManager->storeFragment(fragment);
+    hostPtrManager->storeFragment(gfxAllocation->getRootDeviceIndex(), fragment);
 }
 
 void WddmMemoryManager::removeAllocationFromHostPtrManager(GraphicsAllocation *gfxAllocation) {
     auto buffer = gfxAllocation->getUnderlyingBuffer();
-    auto fragment = hostPtrManager->getFragment(buffer);
+    auto fragment = hostPtrManager->getFragment({buffer, gfxAllocation->getRootDeviceIndex()});
     if (fragment && fragment->driverAllocation) {
         OsHandle *osStorageToRelease = fragment->osInternalStorage;
-        if (hostPtrManager->releaseHostPtr(buffer)) {
+        if (hostPtrManager->releaseHostPtr(gfxAllocation->getRootDeviceIndex(), buffer)) {
             delete osStorageToRelease;
         }
     }
@@ -497,7 +497,7 @@ MemoryManager::AllocationStatus WddmMemoryManager::populateOsHandles(OsHandleSto
     UNRECOVERABLE_IF(result != STATUS_SUCCESS);
 
     for (uint32_t i = 0; i < allocatedFragmentsCounter; i++) {
-        hostPtrManager->storeFragment(handleStorage.fragmentStorageData[allocatedFragmentIndexes[i]]);
+        hostPtrManager->storeFragment(rootDeviceIndex, handleStorage.fragmentStorageData[allocatedFragmentIndexes[i]]);
     }
 
     return AllocationStatus::Success;
@@ -532,7 +532,7 @@ void WddmMemoryManager::cleanOsHandles(OsHandleStorage &handleStorage, uint32_t
 void WddmMemoryManager::obtainGpuAddressFromFragments(WddmAllocation *allocation, OsHandleStorage &handleStorage) {
     if (this->force32bitAllocations && (handleStorage.fragmentCount > 0)) {
         auto hostPtr = allocation->getUnderlyingBuffer();
-        auto fragment = hostPtrManager->getFragment(hostPtr);
+        auto fragment = hostPtrManager->getFragment({hostPtr, allocation->getRootDeviceIndex()});
         if (fragment && fragment->driverAllocation) {
             auto gpuPtr = handleStorage.fragmentStorageData[0].osHandleStorage->gpuPtr;
             for (uint32_t i = 1; i < handleStorage.fragmentCount; i++) {