diff --git a/opencl/test/unit_test/command_stream/aub_command_stream_receiver_2_tests.cpp b/opencl/test/unit_test/command_stream/aub_command_stream_receiver_2_tests.cpp
index 657ba30c66..3fea1ed297 100644
--- a/opencl/test/unit_test/command_stream/aub_command_stream_receiver_2_tests.cpp
+++ b/opencl/test/unit_test/command_stream/aub_command_stream_receiver_2_tests.cpp
@@ -473,7 +473,7 @@ HWTEST_F(AubCommandStreamReceiverNoHostPtrTests, givenAubCommandStreamReceiverWh
 
     auto imgInfo = MockGmm::initImgInfo(imgDesc, 0, nullptr);
 
-    AllocationProperties allocProperties = MemObjHelper::getAllocationPropertiesWithImageInfo(aubCsr->getRootDeviceIndex(), imgInfo, true, {}, *hwInfo, {});
+    AllocationProperties allocProperties = MemObjHelper::getAllocationPropertiesWithImageInfo(aubCsr->getRootDeviceIndex(), imgInfo, true, {}, *hwInfo, osContext.getDeviceBitfield());
 
     auto imageAllocation = memoryManager->allocateGraphicsMemoryInPreferredPool(allocProperties, nullptr);
     ASSERT_NE(nullptr, imageAllocation);
diff --git a/opencl/test/unit_test/memory_manager/host_ptr_manager_tests.cpp b/opencl/test/unit_test/memory_manager/host_ptr_manager_tests.cpp
index 9c25834428..ea1df178a1 100644
--- a/opencl/test/unit_test/memory_manager/host_ptr_manager_tests.cpp
+++ b/opencl/test/unit_test/memory_manager/host_ptr_manager_tests.cpp
@@ -9,6 +9,7 @@
 #include "shared/source/helpers/aligned_memory.h"
 #include "shared/source/helpers/constants.h"
 #include "shared/source/helpers/ptr_math.h"
+#include "shared/source/os_interface/os_context.h"
 
 #include "opencl/test/unit_test/fixtures/memory_manager_fixture.h"
 #include "opencl/test/unit_test/mocks/mock_allocation_properties.h"
@@ -788,10 +789,10 @@ TEST_F(HostPtrAllocationTest, givenTwoAllocationsThatSharesOneFragmentWhenOneIsD
     void *cpuPtr2 = ptrOffset(cpuPtr1, MemoryConstants::pageSize);
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
-    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, 2 * MemoryConstants::pageSize - 1}, cpuPtr1);
+    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, 2 * MemoryConstants::pageSize - 1, csr->getOsContext().getDeviceBitfield()}, cpuPtr1);
     EXPECT_EQ(2u, hostPtrManager->getFragmentCount());
 
-    auto graphicsAllocation2 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr2);
+    auto graphicsAllocation2 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr2);
     EXPECT_EQ(3u, hostPtrManager->getFragmentCount());
     memoryManager->freeGraphicsMemory(graphicsAllocation1);
     EXPECT_EQ(2u, hostPtrManager->getFragmentCount());
@@ -820,7 +821,7 @@ TEST_F(HostPtrAllocationTest, whenOverlappedFragmentIsBiggerThenStoredAndStoredF
     void *cpuPtr1 = (void *)0x100004;
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
-    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr1);
+    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr1);
     EXPECT_EQ(2u, hostPtrManager->getFragmentCount());
 
     EXPECT_NE(nullptr, graphicsAllocation1);
@@ -874,8 +875,8 @@ HWTEST_F(HostPtrAllocationTest, givenOverlappingFragmentsWhenCheckIsCalledThenWa
     void *cpuPtr = reinterpret_cast<void *>(0x100004);
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
-    auto graphicsAllocation0 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr);
-    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr);
+    auto graphicsAllocation0 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr);
+    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr);
 
     auto storage0 = new MockInternalAllocationStorage(*csr0);
     auto storage1 = new MockInternalAllocationStorage(*csr1);
@@ -913,7 +914,7 @@ TEST_F(HostPtrAllocationTest, whenOverlappedFragmentIsBiggerThenStoredAndStoredF
     void *cpuPtr1 = (void *)0x100004;
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
-    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr1);
+    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr1);
     EXPECT_EQ(2u, hostPtrManager->getFragmentCount());
 
     EXPECT_NE(nullptr, graphicsAllocation1);
@@ -951,10 +952,10 @@ TEST_F(HostPtrAllocationTest, checkAllocationsForOverlappingWithoutBiggerOverlap
     void *cpuPtr2 = (void *)0x101008;
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
-    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr1);
+    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr1);
     EXPECT_EQ(2u, hostPtrManager->getFragmentCount());
 
-    auto graphicsAllocation2 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize * 3}, cpuPtr2);
+    auto graphicsAllocation2 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize * 3, csr->getOsContext().getDeviceBitfield()}, cpuPtr2);
     EXPECT_EQ(4u, hostPtrManager->getFragmentCount());
 
     EXPECT_NE(nullptr, graphicsAllocation1);
@@ -994,7 +995,7 @@ TEST_F(HostPtrAllocationTest, checkAllocationsForOverlappingWithBiggerOverlapUnt
 
     void *cpuPtr1 = (void *)0x100004;
 
-    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize}, cpuPtr1);
+    auto graphicsAllocation1 = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{csr->getRootDeviceIndex(), false, MemoryConstants::pageSize, csr->getOsContext().getDeviceBitfield()}, cpuPtr1);
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
     EXPECT_EQ(2u, hostPtrManager->getFragmentCount());
 
diff --git a/opencl/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.inl b/opencl/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.inl
index 5e4acaef73..17a894f361 100644
--- a/opencl/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.inl
+++ b/opencl/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.inl
@@ -13,6 +13,7 @@
 #include "opencl/test/unit_test/helpers/unit_test_helper.h"
 #include "opencl/test/unit_test/mocks/mock_allocation_properties.h"
 #include "opencl/test/unit_test/mocks/mock_execution_environment.h"
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
 #include "opencl/test/unit_test/mocks/mock_memory_manager.h"
 #include "test.h"
 
@@ -29,7 +30,7 @@ using MemoryManagerGetAlloctionDataTests = ::testing::Test;
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenHostMemoryAllocationTypeAndAllocateMemoryFlagAndNullptrWhenAllocationDataIsQueriedThenCorrectFlagsAndSizeAreSet) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::BUFFER_HOST_MEMORY, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::BUFFER_HOST_MEMORY, false, mockDeviceBitfield);
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
 
@@ -40,7 +41,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenHostMemoryAllocationTypeAndAlloc
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenNonHostMemoryAllocatoinTypeWhenAllocationDataIsQueriedThenUseSystemMemoryFlagsIsNotSet) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::BUFFER, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::BUFFER, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -52,7 +53,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenNonHostMemoryAllocatoinTypeWhenA
 
 HWTEST_F(MemoryManagerGetAlloctionDataTests, givenCommandBufferAllocationTypeWhenGetAllocationDataIsCalledThenSystemMemoryIsRequested) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::COMMAND_BUFFER, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::COMMAND_BUFFER, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -63,7 +64,7 @@ HWTEST_F(MemoryManagerGetAlloctionDataTests, givenCommandBufferAllocationTypeWhe
 TEST_F(MemoryManagerGetAlloctionDataTests, givenAllocateMemoryFlagTrueWhenHostPtrIsNotNullThenAllocationDataHasHostPtrNulled) {
     AllocationData allocData;
     char memory = 0;
-    AllocationProperties properties(0, true, sizeof(memory), GraphicsAllocation::AllocationType::BUFFER, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, sizeof(memory), GraphicsAllocation::AllocationType::BUFFER, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, &memory, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -74,7 +75,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenAllocateMemoryFlagTrueWhenHostPt
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenBufferTypeWhenAllocationDataIsQueriedThenForcePinFlagIsSet) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::BUFFER, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::BUFFER, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -84,7 +85,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenBufferTypeWhenAllocationDataIsQu
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenBufferHostMemoryTypeWhenAllocationDataIsQueriedThenForcePinFlagIsSet) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::BUFFER_HOST_MEMORY, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::BUFFER_HOST_MEMORY, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -94,7 +95,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenBufferHostMemoryTypeWhenAllocati
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenBufferCompressedTypeWhenAllocationDataIsQueriedThenForcePinFlagIsSet) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::BUFFER_COMPRESSED, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::BUFFER_COMPRESSED, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -104,7 +105,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenBufferCompressedTypeWhenAllocati
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenWriteCombinedTypeWhenAllocationDataIsQueriedThenForcePinFlagIsSet) {
     AllocationData allocData;
-    AllocationProperties properties(0, true, 10, GraphicsAllocation::AllocationType::WRITE_COMBINED, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, true, 10, GraphicsAllocation::AllocationType::WRITE_COMBINED, false, mockDeviceBitfield);
 
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -114,7 +115,7 @@ TEST_F(MemoryManagerGetAlloctionDataTests, givenWriteCombinedTypeWhenAllocationD
 
 TEST_F(MemoryManagerGetAlloctionDataTests, givenDefaultAllocationFlagsWhenAllocationDataIsQueriedThenAllocateMemoryIsFalse) {
     AllocationData allocData;
-    AllocationProperties properties(0, false, 0, GraphicsAllocation::AllocationType::BUFFER_COMPRESSED, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, false, 0, GraphicsAllocation::AllocationType::BUFFER_COMPRESSED, false, mockDeviceBitfield);
     char memory;
     MockMemoryManager mockMemoryManager;
     mockMemoryManager.getAllocationData(allocData, properties, &memory, mockMemoryManager.createStorageInfoFromProperties(properties));
@@ -406,7 +407,7 @@ TEST(MemoryManagerTest, givenEnabled64kbPagesWhenGraphicsMemoryIsAllocatedWithHo
     MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
     MockMemoryManager memoryManager(true, false, executionEnvironment);
     AllocationData allocData;
-    AllocationProperties properties(0, false, 1, GraphicsAllocation::AllocationType::BUFFER_HOST_MEMORY, false, {});
+    AllocationProperties properties(mockRootDeviceIndex, false, 1, GraphicsAllocation::AllocationType::BUFFER_HOST_MEMORY, false, mockDeviceBitfield);
 
     char memory[1];
     memoryManager.getAllocationData(allocData, properties, &memory, memoryManager.createStorageInfoFromProperties(properties));
@@ -665,7 +666,7 @@ TEST(MemoryManagerTest, givenExternalHostMemoryWhenGetAllocationDataIsCalledThen
     AllocationData allocData;
     auto hostPtr = reinterpret_cast<void *>(0x1234);
     MockMemoryManager mockMemoryManager;
-    AllocationProperties properties{0, false, 1, GraphicsAllocation::AllocationType::EXTERNAL_HOST_PTR, false, {}};
+    AllocationProperties properties{mockRootDeviceIndex, false, 1, GraphicsAllocation::AllocationType::EXTERNAL_HOST_PTR, false, mockDeviceBitfield};
     mockMemoryManager.getAllocationData(allocData, properties, hostPtr, mockMemoryManager.createStorageInfoFromProperties(properties));
     EXPECT_TRUE(allocData.flags.useSystemMemory);
     EXPECT_FALSE(allocData.flags.allocateMemory);
@@ -693,7 +694,7 @@ TEST(MemoryManagerTest, givenMapAllocationWhenGetAllocationDataIsCalledThenItHas
     AllocationData allocData;
     auto hostPtr = reinterpret_cast<void *>(0x1234);
     MockMemoryManager mockMemoryManager;
-    AllocationProperties properties{0, false, 1, GraphicsAllocation::AllocationType::MAP_ALLOCATION, false, {}};
+    AllocationProperties properties{mockRootDeviceIndex, false, 1, GraphicsAllocation::AllocationType::MAP_ALLOCATION, false, mockDeviceBitfield};
     mockMemoryManager.getAllocationData(allocData, properties, hostPtr, mockMemoryManager.createStorageInfoFromProperties(properties));
     EXPECT_TRUE(allocData.flags.useSystemMemory);
     EXPECT_FALSE(allocData.flags.allocateMemory);
diff --git a/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp b/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp
index cd758839ee..e38e8244ab 100644
--- a/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp
+++ b/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp
@@ -10,6 +10,7 @@
 #include "shared/source/memory_manager/memory_manager.h"
 
 #include "opencl/test/unit_test/fixtures/memory_allocator_multi_device_fixture.h"
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
 #include "test.h"
 
 using namespace NEO;
@@ -21,7 +22,7 @@ TEST_P(MemoryManagerMultiDeviceTest, givenRootDeviceIndexSpecifiedWhenAllocateGr
                                                                     GraphicsAllocation::AllocationType::KERNEL_ISA};
     for (auto allocationType : allocationTypes) {
         for (uint32_t rootDeviceIndex = 0; rootDeviceIndex < getNumRootDevices(); ++rootDeviceIndex) {
-            AllocationProperties properties{rootDeviceIndex, true, MemoryConstants::pageSize, allocationType, false, false, 0};
+            AllocationProperties properties{rootDeviceIndex, true, MemoryConstants::pageSize, allocationType, false, false, mockDeviceBitfield};
 
             auto gfxAllocation = memoryManager->allocateGraphicsMemoryWithProperties(properties);
             ASSERT_NE(gfxAllocation, nullptr);
diff --git a/opencl/test/unit_test/mocks/mock_allocation_properties.h b/opencl/test/unit_test/mocks/mock_allocation_properties.h
index 6c9795c2aa..355c0d79c6 100644
--- a/opencl/test/unit_test/mocks/mock_allocation_properties.h
+++ b/opencl/test/unit_test/mocks/mock_allocation_properties.h
@@ -8,10 +8,15 @@
 #pragma once
 #include "shared/source/memory_manager/memory_manager.h"
 
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
+
 namespace NEO {
 struct MockAllocationProperties : public AllocationProperties {
-    MockAllocationProperties(uint32_t rootDeviceIndex, size_t size) : AllocationProperties(rootDeviceIndex, size, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY) {}
-    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size) : AllocationProperties(rootDeviceIndex, allocateMemory, size, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, false, {}) {}
-    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, GraphicsAllocation::AllocationType allocationType) : AllocationProperties(rootDeviceIndex, allocateMemory, size, allocationType, false, {}) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, size_t size) : AllocationProperties(rootDeviceIndex, true, size, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, false, mockDeviceBitfield) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, size_t size, DeviceBitfield deviceBitfield) : AllocationProperties(rootDeviceIndex, true, size, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, false, deviceBitfield) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size) : AllocationProperties(rootDeviceIndex, allocateMemory, size, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, false, mockDeviceBitfield) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, DeviceBitfield deviceBitfield) : AllocationProperties(rootDeviceIndex, allocateMemory, size, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, false, deviceBitfield) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, GraphicsAllocation::AllocationType allocationType) : AllocationProperties(rootDeviceIndex, allocateMemory, size, allocationType, false, mockDeviceBitfield) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, GraphicsAllocation::AllocationType allocationType, DeviceBitfield deviceBitfield) : AllocationProperties(rootDeviceIndex, allocateMemory, size, allocationType, false, deviceBitfield) {}
 };
 } // namespace NEO
diff --git a/opencl/test/unit_test/mocks/mock_gmm.h b/opencl/test/unit_test/mocks/mock_gmm.h
index 20d998a0f7..82d269ddc0 100644
--- a/opencl/test/unit_test/mocks/mock_gmm.h
+++ b/opencl/test/unit_test/mocks/mock_gmm.h
@@ -13,6 +13,7 @@
 #include "opencl/source/helpers/surface_formats.h"
 #include "opencl/source/mem_obj/image.h"
 #include "opencl/test/unit_test/mocks/mock_gmm_resource_info.h"
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
 #include "opencl/test/unit_test/mocks/mock_platform.h"
 
 namespace NEO {
@@ -51,7 +52,7 @@ class MockGmm : public Gmm {
         imgDesc.image_width = 5;
         imgDesc.image_height = 5;
         auto imgInfo = MockGmm::initImgInfo(imgDesc, 0, nullptr);
-        return memoryManager.allocateGraphicsMemoryWithProperties({0, true, imgInfo, GraphicsAllocation::AllocationType::IMAGE, {}});
+        return memoryManager.allocateGraphicsMemoryWithProperties({mockRootDeviceIndex, true, imgInfo, GraphicsAllocation::AllocationType::IMAGE, mockDeviceBitfield});
     }
 };
 } // namespace NEO
diff --git a/opencl/test/unit_test/mocks/mock_graphics_allocation.h b/opencl/test/unit_test/mocks/mock_graphics_allocation.h
index 11bc60ca5d..f8fd01bfff 100644
--- a/opencl/test/unit_test/mocks/mock_graphics_allocation.h
+++ b/opencl/test/unit_test/mocks/mock_graphics_allocation.h
@@ -11,6 +11,10 @@
 #include "opencl/source/memory_manager/os_agnostic_memory_manager.h"
 
 namespace NEO {
+
+constexpr uint32_t mockRootDeviceIndex = 0u;
+constexpr DeviceBitfield mockDeviceBitfield(0b1);
+
 class MockGraphicsAllocation : public MemoryAllocation {
   public:
     using MemoryAllocation::aubInfo;
diff --git a/opencl/test/unit_test/mocks/mock_os_context.h b/opencl/test/unit_test/mocks/mock_os_context.h
index 77425e6c17..9b996f5867 100644
--- a/opencl/test/unit_test/mocks/mock_os_context.h
+++ b/opencl/test/unit_test/mocks/mock_os_context.h
@@ -11,6 +11,8 @@
 namespace NEO {
 class MockOsContext : public OsContext {
   public:
+    using OsContext::getDeviceBitfield;
+
     MockOsContext(uint32_t contextId, DeviceBitfield deviceBitfield,
                   aub_stream::EngineType engineType, PreemptionMode preemptionMode,
                   bool lowPriority, bool internalEngine, bool rootDevice)
diff --git a/opencl/test/unit_test/mocks/mock_timestamp_container.h b/opencl/test/unit_test/mocks/mock_timestamp_container.h
index a6b7cf806d..bdff4f785c 100644
--- a/opencl/test/unit_test/mocks/mock_timestamp_container.h
+++ b/opencl/test/unit_test/mocks/mock_timestamp_container.h
@@ -8,6 +8,8 @@
 #pragma once
 #include "shared/source/helpers/timestamp_packet.h"
 
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
+
 namespace NEO {
 
 template <typename TagType = TimestampPacketStorage>
@@ -19,7 +21,7 @@ class MockTagAllocator : public TagAllocator<TagType> {
     using NodeType = typename BaseClass::NodeType;
 
     MockTagAllocator(uint32_t rootDeviceIndex, MemoryManager *memoryManager, size_t tagCount = 10)
-        : BaseClass(rootDeviceIndex, memoryManager, tagCount, MemoryConstants::cacheLineSize, sizeof(TagType), false, {}) {}
+        : BaseClass(rootDeviceIndex, memoryManager, tagCount, MemoryConstants::cacheLineSize, sizeof(TagType), false, mockDeviceBitfield) {}
 
     void returnTag(NodeType *node) override {
         releaseReferenceNodes.push_back(node);
diff --git a/opencl/test/unit_test/profiling/profiling_tests.cpp b/opencl/test/unit_test/profiling/profiling_tests.cpp
index d1384c0ca7..632c4218dc 100644
--- a/opencl/test/unit_test/profiling/profiling_tests.cpp
+++ b/opencl/test/unit_test/profiling/profiling_tests.cpp
@@ -852,8 +852,8 @@ struct FixedGpuAddressTagAllocator : TagAllocator<TagType> {
     };
 
     FixedGpuAddressTagAllocator(CommandStreamReceiver &csr, uint64_t gpuAddress)
-        : TagAllocator<TagType>(0, csr.getMemoryManager(), csr.getPreferredTagPoolSize(), MemoryConstants::cacheLineSize,
-                                sizeof(TagType), false, {}) {
+        : TagAllocator<TagType>(csr.getRootDeviceIndex(), csr.getMemoryManager(), csr.getPreferredTagPoolSize(), MemoryConstants::cacheLineSize,
+                                sizeof(TagType), false, csr.getOsContext().getDeviceBitfield()) {
         auto tag = reinterpret_cast<MockTagNode *>(this->freeTags.peekHead());
         tag->setGpuAddress(gpuAddress);
     }
diff --git a/opencl/test/unit_test/utilities/tag_allocator_tests.cpp b/opencl/test/unit_test/utilities/tag_allocator_tests.cpp
index c8efc8261b..d20513486d 100644
--- a/opencl/test/unit_test/utilities/tag_allocator_tests.cpp
+++ b/opencl/test/unit_test/utilities/tag_allocator_tests.cpp
@@ -7,8 +7,11 @@
 
 #include "shared/source/helpers/timestamp_packet.h"
 #include "shared/source/utilities/tag_allocator.h"
+#include "shared/test/unit_test/helpers/debug_manager_state_restore.h"
+#include "shared/test/unit_test/mocks/ult_device_factory.h"
 
 #include "opencl/test/unit_test/fixtures/memory_allocator_fixture.h"
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
 #include "test.h"
 
 #include "gtest/gtest.h"
@@ -18,7 +21,13 @@
 using namespace NEO;
 
 struct TagAllocatorTest : public Test<MemoryAllocatorFixture> {
-    const DeviceBitfield deviceBitfield{0xef};
+    const DeviceBitfield deviceBitfield{0xf};
+    DebugManagerStateRestore restorer;
+
+    void SetUp() override {
+        DebugManager.flags.CreateMultipleSubDevices.set(4);
+        MemoryAllocatorFixture::SetUp();
+    }
 };
 
 struct TimeStamps {
@@ -417,9 +426,9 @@ TEST_F(TagAllocatorTest, givenTagsOnDeferredListWhenReleasingItThenMoveReadyTags
 }
 
 TEST_F(TagAllocatorTest, givenTagAllocatorWhenGraphicsAllocationIsCreatedThenSetValidllocationType) {
-    TagAllocator<TimestampPacketStorage> timestampPacketAllocator(0, memoryManager, 1, 1, sizeof(TimestampPacketStorage), false, {});
-    TagAllocator<HwTimeStamps> hwTimeStampsAllocator(0, memoryManager, 1, 1, sizeof(HwTimeStamps), false, {});
-    TagAllocator<HwPerfCounter> hwPerfCounterAllocator(0, memoryManager, 1, 1, sizeof(HwPerfCounter), false, {});
+    TagAllocator<TimestampPacketStorage> timestampPacketAllocator(mockRootDeviceIndex, memoryManager, 1, 1, sizeof(TimestampPacketStorage), false, mockDeviceBitfield);
+    TagAllocator<HwTimeStamps> hwTimeStampsAllocator(mockRootDeviceIndex, memoryManager, 1, 1, sizeof(HwTimeStamps), false, mockDeviceBitfield);
+    TagAllocator<HwPerfCounter> hwPerfCounterAllocator(mockRootDeviceIndex, memoryManager, 1, 1, sizeof(HwPerfCounter), false, mockDeviceBitfield);
 
     auto timestampPacketTag = timestampPacketAllocator.getTag();
     auto hwTimeStampsTag = hwTimeStampsAllocator.getTag();
diff --git a/shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp b/shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp
index b09d82df97..b6986f2d61 100644
--- a/shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp
+++ b/shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp
@@ -11,6 +11,7 @@
 #include "shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests_fixture.h"
 #include "shared/test/unit_test/test_macros/test_checks_shared.h"
 
+#include "opencl/test/unit_test/mocks/mock_graphics_allocation.h"
 #include "opencl/test/unit_test/mocks/mock_memory_manager.h"
 
 using namespace NEO;
@@ -209,7 +210,9 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenSetAubWritableIsCalledTh
     void *cmdQ = reinterpret_cast<void *>(0xFFFF);
     auto memoryManager = std::make_unique<MockMemoryManager>(executionEnvironment);
     auto unifiedMemoryManager = std::make_unique<SVMAllocsManager>(memoryManager.get());
-    void *alloc1 = unifiedMemoryManager->createSharedUnifiedMemoryAllocation(0, 10, SVMAllocsManager::UnifiedMemoryProperties(InternalMemoryType::SHARED_UNIFIED_MEMORY), cmdQ);
+    auto properties = SVMAllocsManager::UnifiedMemoryProperties(InternalMemoryType::SHARED_UNIFIED_MEMORY);
+    properties.subdeviceBitfield = mockDeviceBitfield;
+    void *alloc1 = unifiedMemoryManager->createSharedUnifiedMemoryAllocation(mockRootDeviceIndex, 10, properties, cmdQ);
 
     pageFaultManager->baseAubWritable(false, alloc1, unifiedMemoryManager.get());