diff --git a/unit_tests/command_stream/command_stream_receiver_tests.cpp b/unit_tests/command_stream/command_stream_receiver_tests.cpp
index 8e607d7c1d..3bf7a9e33c 100644
--- a/unit_tests/command_stream/command_stream_receiver_tests.cpp
+++ b/unit_tests/command_stream/command_stream_receiver_tests.cpp
@@ -12,6 +12,7 @@
 #include "core/unit_tests/helpers/debug_manager_state_restore.h"
 #include "runtime/aub_mem_dump/aub_services.h"
 #include "runtime/command_stream/command_stream_receiver.h"
+#include "runtime/command_stream/scratch_space_controller.h"
 #include "runtime/helpers/timestamp_packet.h"
 #include "runtime/mem_obj/buffer.h"
 #include "runtime/memory_manager/internal_allocation_storage.h"
@@ -21,6 +22,7 @@
 #include "runtime/utilities/tag_allocator.h"
 #include "test.h"
 #include "unit_tests/fixtures/device_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/gen_common/matchers.h"
 #include "unit_tests/helpers/unit_test_helper.h"
 #include "unit_tests/mocks/mock_buffer.h"
@@ -640,22 +642,16 @@ TEST(CommandStreamReceiverDeviceIndexTest, givenOsContextWithNoDeviceBitfieldWhe
     EXPECT_EQ(0u, csr.getDeviceIndex());
 }
 
-TEST(CommandStreamReceiverRootDeviceIndexTest, commandStreamGraphicsAllocationsHaveCorrectRootDeviceIndex) {
-    const uint32_t expectedRootDeviceIndex = 101;
+using CommandStreamReceiverMultiRootDeviceTest = MultiRootDeviceFixture;
 
-    // Setup
-    auto executionEnvironment = platformImpl->peekExecutionEnvironment();
-    executionEnvironment->prepareRootDeviceEnvironments(2 * expectedRootDeviceIndex);
-    auto memoryManager = new MockMemoryManager(false, false, *executionEnvironment);
-    executionEnvironment->memoryManager.reset(memoryManager);
-    std::unique_ptr<MockDevice> device(Device::create<MockDevice>(executionEnvironment, expectedRootDeviceIndex));
+TEST_F(CommandStreamReceiverMultiRootDeviceTest, commandStreamGraphicsAllocationsHaveCorrectRootDeviceIndex) {
     auto commandStreamReceiver = &device->getGpgpuCommandStreamReceiver();
 
     ASSERT_NE(nullptr, commandStreamReceiver);
     EXPECT_EQ(expectedRootDeviceIndex, commandStreamReceiver->getRootDeviceIndex());
 
     // Linear stream / Command buffer
-    GraphicsAllocation *allocation = memoryManager->allocateGraphicsMemoryWithProperties({expectedRootDeviceIndex, 128u, GraphicsAllocation::AllocationType::COMMAND_BUFFER});
+    GraphicsAllocation *allocation = mockMemoryManager->allocateGraphicsMemoryWithProperties({expectedRootDeviceIndex, 128u, GraphicsAllocation::AllocationType::COMMAND_BUFFER});
     LinearStream commandStream{allocation};
 
     commandStreamReceiver->ensureCommandBufferAllocation(commandStream, 100u, 0u);
@@ -667,7 +663,7 @@ TEST(CommandStreamReceiverRootDeviceIndexTest, commandStreamGraphicsAllocationsH
     EXPECT_NE(allocation, commandStream.getGraphicsAllocation());
     EXPECT_EQ(0u, commandStream.getMaxAvailableSpace() % MemoryConstants::pageSize64k);
     EXPECT_EQ(expectedRootDeviceIndex, commandStream.getGraphicsAllocation()->getRootDeviceIndex());
-    memoryManager->freeGraphicsMemory(commandStream.getGraphicsAllocation());
+    mockMemoryManager->freeGraphicsMemory(commandStream.getGraphicsAllocation());
 
     // Debug surface
     auto debugSurface = commandStreamReceiver->allocateDebugSurface(MemoryConstants::pageSize);
@@ -682,7 +678,7 @@ TEST(CommandStreamReceiverRootDeviceIndexTest, commandStreamGraphicsAllocationsH
         ASSERT_NE(nullptr, heap);
         ASSERT_NE(nullptr, heap->getGraphicsAllocation());
         EXPECT_EQ(expectedRootDeviceIndex, heap->getGraphicsAllocation()->getRootDeviceIndex());
-        memoryManager->freeGraphicsMemory(heap->getGraphicsAllocation());
+        mockMemoryManager->freeGraphicsMemory(heap->getGraphicsAllocation());
         delete heap;
     }
 
@@ -702,4 +698,14 @@ TEST(CommandStreamReceiverRootDeviceIndexTest, commandStreamGraphicsAllocationsH
     EXPECT_TRUE(commandStreamReceiver->createAllocationForHostSurface(surface, false));
     ASSERT_NE(nullptr, surface.getAllocation());
     EXPECT_EQ(expectedRootDeviceIndex, surface.getAllocation()->getRootDeviceIndex());
+
+    // Scratch allocation
+    auto scratchController = commandStreamReceiver->getScratchSpaceController();
+    bool cfeStateDirty = false;
+    bool stateBaseAddressDirty = false;
+    std::vector<unsigned char> surfaceHeap(0x1000);
+    scratchController->setRequiredScratchSpace(surfaceHeap.data(), 0x1000u, 0u, 0u, *device->getDefaultEngine().osContext, stateBaseAddressDirty, cfeStateDirty);
+    auto scratchAllocation = scratchController->getScratchSpaceAllocation();
+    ASSERT_NE(nullptr, scratchAllocation);
+    EXPECT_EQ(expectedRootDeviceIndex, scratchAllocation->getRootDeviceIndex());
 }
diff --git a/unit_tests/command_stream/experimental_command_buffer_tests.cpp b/unit_tests/command_stream/experimental_command_buffer_tests.cpp
index 40a7145443..dde8a46c99 100644
--- a/unit_tests/command_stream/experimental_command_buffer_tests.cpp
+++ b/unit_tests/command_stream/experimental_command_buffer_tests.cpp
@@ -11,6 +11,7 @@
 #include "runtime/memory_manager/internal_allocation_storage.h"
 #include "runtime/memory_manager/memory_manager.h"
 #include "test.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/fixtures/ult_command_stream_receiver_fixture.h"
 #include "unit_tests/mocks/mock_experimental_command_buffer.h"
 #include "unit_tests/mocks/mock_memory_manager.h"
@@ -354,15 +355,9 @@ HWTEST_F(ExperimentalCommandBufferTest, givenEnabledExperimentalCmdBufferWhenCom
     EXPECT_STREQ(output.c_str(), "");
 }
 
-TEST(ExperimentalCommandBufferRootDeviceIndexTest, experimentalCommandBufferGraphicsAllocationsHaveCorrectRootDeviceIndex) {
-    const uint32_t expectedRootDeviceIndex = 101;
+using ExperimentalCommandBufferRootDeviceIndexTest = MultiRootDeviceFixture;
 
-    // Setup
-    auto executionEnvironment = platformImpl->peekExecutionEnvironment();
-    executionEnvironment->prepareRootDeviceEnvironments(2 * expectedRootDeviceIndex);
-    auto memoryManager = new MockMemoryManager(false, false, *executionEnvironment);
-    executionEnvironment->memoryManager.reset(memoryManager);
-    std::unique_ptr<MockDevice> device(Device::create<MockDevice>(executionEnvironment, expectedRootDeviceIndex));
+TEST_F(ExperimentalCommandBufferRootDeviceIndexTest, experimentalCommandBufferGraphicsAllocationsHaveCorrectRootDeviceIndex) {
     auto experimentalCommandBuffer = std::make_unique<MockExperimentalCommandBuffer>(&device->getGpgpuCommandStreamReceiver());
 
     ASSERT_NE(nullptr, experimentalCommandBuffer);
diff --git a/unit_tests/d3d_sharing/d3d9_tests.cpp b/unit_tests/d3d_sharing/d3d9_tests.cpp
index a916c4434d..2fb73af31d 100644
--- a/unit_tests/d3d_sharing/d3d9_tests.cpp
+++ b/unit_tests/d3d_sharing/d3d9_tests.cpp
@@ -14,6 +14,7 @@
 #include "runtime/sharings/d3d/cl_d3d_api.h"
 #include "runtime/sharings/d3d/d3d_surface.h"
 #include "test.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/fixtures/platform_fixture.h"
 #include "unit_tests/mocks/mock_command_queue.h"
 #include "unit_tests/mocks/mock_context.h"
@@ -28,6 +29,42 @@ DXGI_ADAPTER_DESC MockD3DSharingFunctions<D3DTypesHelper::D3D9>::mockDxgiDesc =
 template <>
 IDXGIAdapter *MockD3DSharingFunctions<D3DTypesHelper::D3D9>::getDxgiDescAdapterRequested = nullptr;
 
+class MockMM : public OsAgnosticMemoryManager {
+  public:
+    MockMM(const ExecutionEnvironment &executionEnvironment) : OsAgnosticMemoryManager(const_cast<ExecutionEnvironment &>(executionEnvironment)){};
+    GraphicsAllocation *createGraphicsAllocationFromSharedHandle(osHandle handle, const AllocationProperties &properties, bool requireSpecificBitness) override {
+        auto alloc = OsAgnosticMemoryManager::createGraphicsAllocationFromSharedHandle(handle, properties, requireSpecificBitness);
+        alloc->setDefaultGmm(forceGmm);
+        gmmOwnershipPassed = true;
+        return alloc;
+    }
+    GraphicsAllocation *allocateGraphicsMemoryForImage(const AllocationData &allocationData) override {
+        auto gmm = std::make_unique<Gmm>(*allocationData.imgInfo, StorageInfo{});
+        AllocationProperties properties(allocationData.rootDeviceIndex, nullptr, false, GraphicsAllocation::AllocationType::SHARED_IMAGE, false);
+        auto alloc = OsAgnosticMemoryManager::createGraphicsAllocationFromSharedHandle(1, properties, false);
+        alloc->setDefaultGmm(forceGmm);
+        gmmOwnershipPassed = true;
+        return alloc;
+    }
+
+    void *lockResourceImpl(GraphicsAllocation &allocation) override {
+        lockResourceCalled++;
+        EXPECT_EQ(expectedLockingAllocation, &allocation);
+        return lockResourceReturnValue;
+    }
+    void unlockResourceImpl(GraphicsAllocation &allocation) override {
+        unlockResourceCalled++;
+        EXPECT_EQ(expectedLockingAllocation, &allocation);
+    }
+
+    int32_t lockResourceCalled = 0;
+    int32_t unlockResourceCalled = 0;
+    GraphicsAllocation *expectedLockingAllocation = nullptr;
+    void *lockResourceReturnValue = nullptr;
+    Gmm *forceGmm = nullptr;
+    bool gmmOwnershipPassed = false;
+};
+
 class D3D9Tests : public PlatformFixture, public ::testing::Test {
   public:
     typedef typename D3DTypesHelper::D3D9 D3D9;
@@ -38,42 +75,6 @@ class D3D9Tests : public PlatformFixture, public ::testing::Test {
     typedef typename D3D9::D3DTexture2dDesc D3DTexture2dDesc;
     typedef typename D3D9::D3DTexture2d D3DTexture2d;
 
-    class MockMM : public OsAgnosticMemoryManager {
-      public:
-        MockMM(const ExecutionEnvironment &executionEnvironment) : OsAgnosticMemoryManager(const_cast<ExecutionEnvironment &>(executionEnvironment)){};
-        GraphicsAllocation *createGraphicsAllocationFromSharedHandle(osHandle handle, const AllocationProperties &properties, bool requireSpecificBitness) override {
-            auto alloc = OsAgnosticMemoryManager::createGraphicsAllocationFromSharedHandle(handle, properties, requireSpecificBitness);
-            alloc->setDefaultGmm(forceGmm);
-            gmmOwnershipPassed = true;
-            return alloc;
-        }
-        GraphicsAllocation *allocateGraphicsMemoryForImage(const AllocationData &allocationData) override {
-            auto gmm = std::make_unique<Gmm>(*allocationData.imgInfo, StorageInfo{});
-            AllocationProperties properties(allocationData.rootDeviceIndex, nullptr, false, GraphicsAllocation::AllocationType::SHARED_IMAGE, false);
-            auto alloc = OsAgnosticMemoryManager::createGraphicsAllocationFromSharedHandle(1, properties, false);
-            alloc->setDefaultGmm(forceGmm);
-            gmmOwnershipPassed = true;
-            return alloc;
-        }
-
-        void *lockResourceImpl(GraphicsAllocation &allocation) override {
-            lockResourceCalled++;
-            EXPECT_EQ(expectedLockingAllocation, &allocation);
-            return lockResourceReturnValue;
-        }
-        void unlockResourceImpl(GraphicsAllocation &allocation) override {
-            unlockResourceCalled++;
-            EXPECT_EQ(expectedLockingAllocation, &allocation);
-        }
-
-        int32_t lockResourceCalled = 0;
-        int32_t unlockResourceCalled = 0;
-        GraphicsAllocation *expectedLockingAllocation = nullptr;
-        void *lockResourceReturnValue = nullptr;
-        Gmm *forceGmm = nullptr;
-        bool gmmOwnershipPassed = false;
-    };
-
     void setupMockGmm() {
         cl_image_desc imgDesc = {};
         imgDesc.image_height = 10;
@@ -1069,4 +1070,70 @@ TEST_P(D3D9ImageFormatTests, validFormat) {
     EXPECT_EQ(imgFormat.image_channel_order, expectedClChannelOrder);
     EXPECT_TRUE(oclPlane == expectedOclPlane);
 }
-} // namespace NEO
+
+using D3D9MultiRootDeviceTest = MultiRootDeviceFixture;
+
+TEST_F(D3D9MultiRootDeviceTest, givenD3DHandleIsNullWhenCreatingSharedSurfaceAndRootDeviceIndexIsSpecifiedThenAllocationHasCorrectRootDeviceIndex) {
+    cl_image_desc imgDesc = {};
+    imgDesc.image_height = 10;
+    imgDesc.image_width = 10;
+    imgDesc.image_depth = 1;
+    imgDesc.image_type = CL_MEM_OBJECT_IMAGE2D;
+    auto imgInfo = MockGmm::initImgInfo(imgDesc, 0, nullptr);
+    auto gmm = MockGmm::queryImgParams(imgInfo).release();
+
+    auto memoryManager = std::make_unique<MockMM>(*device->executionEnvironment);
+    memoryManager->forceGmm = gmm;
+
+    auto mockSharingFcns = new NiceMock<MockD3DSharingFunctions<D3DTypesHelper::D3D9>>();
+    mockSharingFcns->mockTexture2dDesc.Format = D3DFMT_R32F;
+    mockSharingFcns->mockTexture2dDesc.Height = 10;
+    mockSharingFcns->mockTexture2dDesc.Width = 10;
+
+    cl_dx9_surface_info_khr surfaceInfo = {};
+    surfaceInfo.shared_handle = reinterpret_cast<HANDLE>(0);
+
+    ON_CALL(*mockSharingFcns, getTexture2dDesc(_, _)).WillByDefault(SetArgPointee<0>(mockSharingFcns->mockTexture2dDesc));
+
+    MockContext ctx(device.get());
+    ctx.setSharingFunctions(mockSharingFcns);
+    ctx.memoryManager = memoryManager.get();
+
+    auto sharedImg = std::unique_ptr<Image>(D3DSurface::create(&ctx, &surfaceInfo, CL_MEM_READ_WRITE, 0, 0, nullptr));
+    ASSERT_NE(nullptr, sharedImg.get());
+    ASSERT_NE(nullptr, sharedImg->getGraphicsAllocation());
+    EXPECT_EQ(expectedRootDeviceIndex, sharedImg->getGraphicsAllocation()->getRootDeviceIndex());
+}
+
+TEST_F(D3D9MultiRootDeviceTest, givenD3DHandleIsNotNullWhenCreatingSharedSurfaceAndRootDeviceIndexIsSpecifiedThenAllocationHasCorrectRootDeviceIndex) {
+    cl_image_desc imgDesc = {};
+    imgDesc.image_height = 10;
+    imgDesc.image_width = 10;
+    imgDesc.image_depth = 1;
+    imgDesc.image_type = CL_MEM_OBJECT_IMAGE2D;
+    auto imgInfo = MockGmm::initImgInfo(imgDesc, 0, nullptr);
+    auto gmm = MockGmm::queryImgParams(imgInfo).release();
+
+    auto memoryManager = std::make_unique<MockMM>(*device->executionEnvironment);
+    memoryManager->forceGmm = gmm;
+
+    auto mockSharingFcns = new NiceMock<MockD3DSharingFunctions<D3DTypesHelper::D3D9>>();
+    mockSharingFcns->mockTexture2dDesc.Format = D3DFMT_R32F;
+    mockSharingFcns->mockTexture2dDesc.Height = 10;
+    mockSharingFcns->mockTexture2dDesc.Width = 10;
+
+    cl_dx9_surface_info_khr surfaceInfo = {};
+    surfaceInfo.shared_handle = reinterpret_cast<HANDLE>(1);
+
+    ON_CALL(*mockSharingFcns, getTexture2dDesc(_, _)).WillByDefault(SetArgPointee<0>(mockSharingFcns->mockTexture2dDesc));
+
+    MockContext ctx(device.get());
+    ctx.setSharingFunctions(mockSharingFcns);
+    ctx.memoryManager = memoryManager.get();
+
+    auto sharedImg = std::unique_ptr<Image>(D3DSurface::create(&ctx, &surfaceInfo, CL_MEM_READ_WRITE, 0, 0, nullptr));
+    ASSERT_NE(nullptr, sharedImg.get());
+    ASSERT_NE(nullptr, sharedImg->getGraphicsAllocation());
+    EXPECT_EQ(expectedRootDeviceIndex, sharedImg->getGraphicsAllocation()->getRootDeviceIndex());
+}
+} // namespace NEO
\ No newline at end of file
diff --git a/unit_tests/fixtures/CMakeLists.txt b/unit_tests/fixtures/CMakeLists.txt
index 9e3a204028..b8fb8f1b2c 100644
--- a/unit_tests/fixtures/CMakeLists.txt
+++ b/unit_tests/fixtures/CMakeLists.txt
@@ -34,6 +34,7 @@ set(IGDRCL_SRCS_tests_fixtures
   ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_fixture.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_fixture.h
   ${CMAKE_CURRENT_SOURCE_DIR}/mock_aub_center_fixture.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/multi_root_device_fixture.h
   ${CMAKE_CURRENT_SOURCE_DIR}/platform_fixture.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/platform_fixture.h
   ${CMAKE_CURRENT_SOURCE_DIR}/run_kernel_fixture.h
diff --git a/unit_tests/fixtures/multi_root_device_fixture.h b/unit_tests/fixtures/multi_root_device_fixture.h
new file mode 100644
index 0000000000..faedb216a4
--- /dev/null
+++ b/unit_tests/fixtures/multi_root_device_fixture.h
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2019 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+#include "core/unit_tests/helpers/debug_manager_state_restore.h"
+#include "unit_tests/mocks/mock_context.h"
+#include "unit_tests/mocks/mock_device.h"
+#include "unit_tests/mocks/mock_memory_manager.h"
+
+namespace NEO {
+class MultiRootDeviceFixture : public ::testing::Test {
+  public:
+    void SetUp() override {
+        DebugManager.flags.CreateMultipleRootDevices.set(2 * expectedRootDeviceIndex);
+        device.reset(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr, expectedRootDeviceIndex));
+        context.reset(new MockContext(device.get()));
+        mockMemoryManager = reinterpret_cast<MockMemoryManager *>(device->getMemoryManager());
+    }
+
+    const uint32_t expectedRootDeviceIndex = 1;
+    DebugManagerStateRestore restorer;
+    std::unique_ptr<MockDevice> device;
+    std::unique_ptr<MockContext> context;
+    MockMemoryManager *mockMemoryManager;
+};
+}; // namespace NEO
diff --git a/unit_tests/kernel/kernel_reflection_surface_tests.cpp b/unit_tests/kernel/kernel_reflection_surface_tests.cpp
index 2099438adf..6bce1c5cf6 100644
--- a/unit_tests/kernel/kernel_reflection_surface_tests.cpp
+++ b/unit_tests/kernel/kernel_reflection_surface_tests.cpp
@@ -13,6 +13,7 @@
 #include "unit_tests/fixtures/execution_model_fixture.h"
 #include "unit_tests/fixtures/execution_model_kernel_fixture.h"
 #include "unit_tests/fixtures/image_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/gen_common/matchers.h"
 #include "unit_tests/helpers/gtest_helpers.h"
 #include "unit_tests/mocks/mock_context.h"
@@ -2108,3 +2109,118 @@ TEST_F(ReflectionSurfaceConstantValuesPatchingTest, GivenBlockWithConstantMemory
 
     delete parentKernel;
 }
+
+using KernelReflectionMultiDeviceTest = MultiRootDeviceFixture;
+
+TEST_F(KernelReflectionMultiDeviceTest, ObtainKernelReflectionSurfaceWithoutKernelArgs) {
+    MockProgram program(*device->getExecutionEnvironment());
+    KernelInfo *blockInfo = new KernelInfo;
+    KernelInfo &info = *blockInfo;
+    cl_queue_properties properties[1] = {0};
+    DeviceQueue devQueue(context.get(), device.get(), properties[0]);
+
+    SPatchExecutionEnvironment environment = {};
+    environment.HasDeviceEnqueue = 1;
+    info.patchInfo.executionEnvironment = &environment;
+
+    SKernelBinaryHeaderCommon kernelHeader;
+    info.heapInfo.pKernelHeader = &kernelHeader;
+
+    SPatchDataParameterStream dataParameterStream;
+    dataParameterStream.Size = 0;
+    dataParameterStream.DataParameterStreamSize = 0;
+    info.patchInfo.dataParameterStream = &dataParameterStream;
+
+    SPatchBindingTableState bindingTableState;
+    bindingTableState.Count = 0;
+    bindingTableState.Offset = 0;
+    bindingTableState.Size = 0;
+    bindingTableState.SurfaceStateOffset = 0;
+    info.patchInfo.bindingTableState = &bindingTableState;
+
+    MockKernel kernel(&program, info, *device.get());
+
+    EXPECT_TRUE(kernel.isParentKernel);
+
+    program.blockKernelManager->addBlockKernelInfo(blockInfo);
+
+    kernel.createReflectionSurface();
+    auto reflectionSurface = kernel.getKernelReflectionSurface();
+    ASSERT_NE(nullptr, reflectionSurface);
+    EXPECT_EQ(expectedRootDeviceIndex, reflectionSurface->getRootDeviceIndex());
+
+    kernel.patchReflectionSurface<true>(&devQueue, nullptr);
+
+    uint64_t undefinedOffset = MockKernel::ReflectionSurfaceHelperPublic::undefinedOffset;
+
+    EXPECT_EQ(undefinedOffset, MockKernel::ReflectionSurfaceHelperPublic::defaultQueue.offset);
+    EXPECT_EQ(undefinedOffset, MockKernel::ReflectionSurfaceHelperPublic::devQueue.offset);
+    EXPECT_EQ(undefinedOffset, MockKernel::ReflectionSurfaceHelperPublic::eventPool.offset);
+
+    EXPECT_EQ(0u, MockKernel::ReflectionSurfaceHelperPublic::defaultQueue.size);
+    EXPECT_EQ(0u, MockKernel::ReflectionSurfaceHelperPublic::devQueue.size);
+    EXPECT_EQ(0u, MockKernel::ReflectionSurfaceHelperPublic::eventPool.size);
+}
+
+TEST_F(KernelReflectionMultiDeviceTest, ObtainKernelReflectionSurfaceWithDeviceQueueKernelArg) {
+    MockProgram program(*device->getExecutionEnvironment());
+
+    KernelInfo *blockInfo = new KernelInfo;
+    KernelInfo &info = *blockInfo;
+    cl_queue_properties properties[1] = {0};
+    DeviceQueue devQueue(context.get(), device.get(), properties[0]);
+
+    uint32_t devQueueCurbeOffset = 16;
+    uint32_t devQueueCurbeSize = 4;
+
+    SPatchExecutionEnvironment environment = {};
+    environment.HasDeviceEnqueue = 1;
+    info.patchInfo.executionEnvironment = &environment;
+
+    SKernelBinaryHeaderCommon kernelHeader;
+    info.heapInfo.pKernelHeader = &kernelHeader;
+
+    SPatchDataParameterStream dataParameterStream;
+    dataParameterStream.Size = 0;
+    dataParameterStream.DataParameterStreamSize = 0;
+    info.patchInfo.dataParameterStream = &dataParameterStream;
+
+    SPatchBindingTableState bindingTableState;
+    bindingTableState.Count = 0;
+    bindingTableState.Offset = 0;
+    bindingTableState.Size = 0;
+    bindingTableState.SurfaceStateOffset = 0;
+    info.patchInfo.bindingTableState = &bindingTableState;
+
+    KernelArgInfo argInfo;
+    argInfo.isDeviceQueue = true;
+
+    info.kernelArgInfo.resize(1);
+    info.kernelArgInfo[0] = argInfo;
+
+    info.kernelArgInfo[0].kernelArgPatchInfoVector.resize(1);
+    info.kernelArgInfo[0].kernelArgPatchInfoVector[0].crossthreadOffset = devQueueCurbeOffset;
+    info.kernelArgInfo[0].kernelArgPatchInfoVector[0].size = devQueueCurbeSize;
+
+    MockKernel kernel(&program, info, *device.get());
+
+    EXPECT_TRUE(kernel.isParentKernel);
+
+    program.blockKernelManager->addBlockKernelInfo(blockInfo);
+
+    kernel.createReflectionSurface();
+    auto reflectionSurface = kernel.getKernelReflectionSurface();
+    ASSERT_NE(nullptr, reflectionSurface);
+    EXPECT_EQ(expectedRootDeviceIndex, reflectionSurface->getRootDeviceIndex());
+
+    kernel.patchReflectionSurface<true>(&devQueue, nullptr);
+    uint64_t undefinedOffset = MockKernel::ReflectionSurfaceHelperPublic::undefinedOffset;
+
+    EXPECT_EQ(undefinedOffset, MockKernel::ReflectionSurfaceHelperPublic::defaultQueue.offset);
+    EXPECT_EQ(devQueueCurbeOffset, MockKernel::ReflectionSurfaceHelperPublic::devQueue.offset);
+    EXPECT_EQ(undefinedOffset, MockKernel::ReflectionSurfaceHelperPublic::eventPool.offset);
+
+    EXPECT_EQ(0u, MockKernel::ReflectionSurfaceHelperPublic::defaultQueue.size);
+    EXPECT_EQ(4u, MockKernel::ReflectionSurfaceHelperPublic::devQueue.size);
+    EXPECT_EQ(0u, MockKernel::ReflectionSurfaceHelperPublic::eventPool.size);
+}
diff --git a/unit_tests/kernel/kernel_tests.cpp b/unit_tests/kernel/kernel_tests.cpp
index 77e5054939..bc25866579 100644
--- a/unit_tests/kernel/kernel_tests.cpp
+++ b/unit_tests/kernel/kernel_tests.cpp
@@ -28,6 +28,7 @@
 #include "unit_tests/fixtures/device_host_queue_fixture.h"
 #include "unit_tests/fixtures/execution_model_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/helpers/gtest_helpers.h"
 #include "unit_tests/libult/ult_command_stream_receiver.h"
 #include "unit_tests/mocks/mock_command_queue.h"
@@ -2948,3 +2949,25 @@ HWCMDTEST_F(IGFX_GEN8_CORE, DeviceQueueHwTest, whenSlbEndOffsetGreaterThanZeroTh
 
     free(slbCopy);
 }
+
+using KernelMultiRootDeviceTest = MultiRootDeviceFixture;
+
+TEST_F(KernelMultiRootDeviceTest, privateSurfaceHasCorrectRootDeviceIndex) {
+    auto kernelInfo = std::make_unique<KernelInfo>();
+
+    // setup private memory
+    SPatchAllocateStatelessPrivateSurface tokenSPS;
+    tokenSPS.SurfaceStateHeapOffset = 64;
+    tokenSPS.DataParamOffset = 40;
+    tokenSPS.DataParamSize = 8;
+    tokenSPS.PerThreadPrivateMemorySize = 112;
+    kernelInfo->patchInfo.pAllocateStatelessPrivateSurface = &tokenSPS;
+
+    MockProgram program(*device->getExecutionEnvironment(), context.get(), false);
+    std::unique_ptr<MockKernel> kernel(new MockKernel(&program, *kernelInfo, *device.get()));
+    kernel->initialize();
+
+    auto privateSurface = kernel->getPrivateSurface();
+    ASSERT_NE(nullptr, privateSurface);
+    EXPECT_EQ(expectedRootDeviceIndex, privateSurface->getRootDeviceIndex());
+}
diff --git a/unit_tests/mem_obj/buffer_tests.cpp b/unit_tests/mem_obj/buffer_tests.cpp
index 4fa3c86117..a9ee352605 100644
--- a/unit_tests/mem_obj/buffer_tests.cpp
+++ b/unit_tests/mem_obj/buffer_tests.cpp
@@ -26,6 +26,7 @@
 #include "test.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/gen_common/matchers.h"
 #include "unit_tests/helpers/hw_parse.h"
 #include "unit_tests/helpers/unit_test_helper.h"
@@ -2419,3 +2420,16 @@ TEST_F(BufferTransferTests, givenBufferWhenTransferFromHostPtrCalledThenCopyRequ
 
     EXPECT_TRUE(memcmp(expectedBufferMemory, buffer->getCpuAddress(), copySize[0]) == 0);
 }
+
+using MultiRootDeviceBufferTest = MultiRootDeviceFixture;
+
+TEST_F(MultiRootDeviceBufferTest, bufferGraphicsAllocationHasCorrectRootDeviceIndex) {
+    cl_int retVal = 0;
+    cl_mem_flags flags = CL_MEM_READ_WRITE;
+
+    std::unique_ptr<Buffer> buffer(Buffer::create(context.get(), flags, MemoryConstants::pageSize, nullptr, retVal));
+
+    auto graphicsAllocation = buffer->getGraphicsAllocation();
+    ASSERT_NE(nullptr, graphicsAllocation);
+    EXPECT_EQ(expectedRootDeviceIndex, graphicsAllocation->getRootDeviceIndex());
+}
diff --git a/unit_tests/mem_obj/image_tests.cpp b/unit_tests/mem_obj/image_tests.cpp
index 170ff5d75e..bb0c93938b 100644
--- a/unit_tests/mem_obj/image_tests.cpp
+++ b/unit_tests/mem_obj/image_tests.cpp
@@ -18,6 +18,7 @@
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/image_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/helpers/kernel_binary_helper.h"
 #include "unit_tests/helpers/unit_test_helper.h"
 #include "unit_tests/mem_obj/image_compression_fixture.h"
@@ -1565,3 +1566,13 @@ HWTEST_F(HwImageTest, givenImageHwWithUnifiedSurfaceAndMcsWhenSettingParamsForMu
 
     EXPECT_TRUE(mockImage->setAuxParamsForMCSCCSCalled);
 }
+
+using ImageMultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(ImageMultiRootDeviceTests, imageAllocationHasCorrectRootDeviceIndex) {
+    std::unique_ptr<Image> image(ImageHelper<Image3dDefaults>::create(context.get()));
+
+    auto graphicsAllocation = image->getGraphicsAllocation();
+    ASSERT_NE(nullptr, graphicsAllocation);
+    EXPECT_EQ(expectedRootDeviceIndex, graphicsAllocation->getRootDeviceIndex());
+}
\ No newline at end of file
diff --git a/unit_tests/mem_obj/mem_obj_tests.cpp b/unit_tests/mem_obj/mem_obj_tests.cpp
index 13e0969522..3545e40174 100644
--- a/unit_tests/mem_obj/mem_obj_tests.cpp
+++ b/unit_tests/mem_obj/mem_obj_tests.cpp
@@ -14,6 +14,7 @@
 #include "runtime/memory_manager/allocations_list.h"
 #include "runtime/os_interface/os_context.h"
 #include "runtime/platform/platform.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/mocks/mock_context.h"
 #include "unit_tests/mocks/mock_deferred_deleter.h"
 #include "unit_tests/mocks/mock_device.h"
@@ -490,3 +491,18 @@ TEST(MemObj, givenMemObjNotUsingHostPtrWhenGettingBasePtrTwiceReturnSameMapPtr)
     EXPECT_EQ(mapPtr, mapAllocation->getUnderlyingBuffer());
     EXPECT_EQ(mapPtr, memObj.getAllocatedMapPtr());
 }
+
+using MemObjMultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(MemObjMultiRootDeviceTests, memObjMapAllocationHasCorrectRootDeviceIndex) {
+    MemoryPropertiesFlags memoryProperties = MemoryPropertiesFlagsParser::createMemoryPropertiesFlags(CL_MEM_READ_WRITE, 0, 0);
+    MemObj memObj(context.get(), CL_MEM_OBJECT_BUFFER, memoryProperties, CL_MEM_READ_WRITE, 0,
+                  1, nullptr, nullptr, nullptr, true, false, false);
+
+    void *mapPtr = memObj.getBasePtrForMap(device->getRootDeviceIndex());
+    EXPECT_NE(nullptr, mapPtr);
+
+    auto mapAllocation = memObj.getMapAllocation();
+    ASSERT_NE(nullptr, mapAllocation);
+    EXPECT_EQ(expectedRootDeviceIndex, mapAllocation->getRootDeviceIndex());
+}
\ No newline at end of file
diff --git a/unit_tests/mem_obj/pipe_tests.cpp b/unit_tests/mem_obj/pipe_tests.cpp
index 4d92cd7d12..5ffba353f2 100644
--- a/unit_tests/mem_obj/pipe_tests.cpp
+++ b/unit_tests/mem_obj/pipe_tests.cpp
@@ -10,6 +10,7 @@
 #include "test.h"
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/memory_management_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/mocks/mock_context.h"
 
 using namespace NEO;
@@ -112,4 +113,17 @@ TEST_F(PipeTest, givenPipeWithDifferentCpuAndGpuAddressesWhenSetArgPipeThenUseGp
     EXPECT_EQ(valueToPatch, pipeAllocation->getGpuAddressToPatch());
 
     delete pipe;
+}
+
+using MultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(MultiRootDeviceTests, pipeGraphicsAllocationHasCorrectRootDeviceIndex) {
+    int errCode = CL_SUCCESS;
+
+    std::unique_ptr<Pipe> pipe(Pipe::create(context.get(), CL_MEM_READ_ONLY, 1, 20, nullptr, errCode));
+    EXPECT_EQ(CL_SUCCESS, errCode);
+    ASSERT_NE(nullptr, pipe.get());
+    auto graphicsAllocation = pipe->getGraphicsAllocation();
+    ASSERT_NE(nullptr, graphicsAllocation);
+    EXPECT_EQ(expectedRootDeviceIndex, graphicsAllocation->getRootDeviceIndex());
 }
\ No newline at end of file
diff --git a/unit_tests/memory_manager/memory_manager_tests.cpp b/unit_tests/memory_manager/memory_manager_tests.cpp
index dc3d6d473a..a9ffead15c 100644
--- a/unit_tests/memory_manager/memory_manager_tests.cpp
+++ b/unit_tests/memory_manager/memory_manager_tests.cpp
@@ -24,6 +24,7 @@
 #include "unit_tests/fixtures/device_fixture.h"
 #include "unit_tests/fixtures/memory_allocator_fixture.h"
 #include "unit_tests/fixtures/memory_manager_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/helpers/execution_environment_helper.h"
 #include "unit_tests/helpers/variable_backup.h"
 #include "unit_tests/mocks/mock_context.h"
@@ -1942,3 +1943,24 @@ HWTEST_F(MemoryAllocatorTest, givenMemoryManagerWhenEnableHostPtrTrackingFlagIsS
         EXPECT_EQ(executionEnvironment->getHardwareInfo()->capabilityTable.hostPtrTrackingEnabled, memoryManager->isHostPointerTrackingEnabled());
     }
 }
+
+using MemoryManagerMultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(MemoryManagerMultiRootDeviceTests, globalsSurfaceHasCorrectRootDeviceIndex) {
+    if (device->getMemoryManager()->isLimitedRange(expectedRootDeviceIndex)) {
+        delete context->svmAllocsManager;
+        context->svmAllocsManager = nullptr;
+    }
+
+    std::vector<unsigned char> initData(1024, 0x5B);
+    GraphicsAllocation *allocation = allocateGlobalsSurface(context.get(), device.get(), initData.size(), false, true, initData.data());
+
+    ASSERT_NE(nullptr, allocation);
+    EXPECT_EQ(expectedRootDeviceIndex, allocation->getRootDeviceIndex());
+
+    if (device->getMemoryManager()->isLimitedRange(expectedRootDeviceIndex)) {
+        device->getMemoryManager()->freeGraphicsMemory(allocation);
+    } else {
+        context->getSVMAllocsManager()->freeSVMAlloc(allocation->getUnderlyingBuffer());
+    }
+}
diff --git a/unit_tests/mocks/mock_device.h b/unit_tests/mocks/mock_device.h
index ab5a64a7a1..5946a93981 100644
--- a/unit_tests/mocks/mock_device.h
+++ b/unit_tests/mocks/mock_device.h
@@ -98,13 +98,13 @@ class MockDevice : public RootDevice {
     }
 
     template <typename T>
-    static T *createWithNewExecutionEnvironment(const HardwareInfo *pHwInfo) {
+    static T *createWithNewExecutionEnvironment(const HardwareInfo *pHwInfo, uint32_t rootDeviceIndex = 0) {
         ExecutionEnvironment *executionEnvironment = new ExecutionEnvironment();
         auto numRootDevices = DebugManager.flags.CreateMultipleRootDevices.get() ? DebugManager.flags.CreateMultipleRootDevices.get() : 1u;
         executionEnvironment->prepareRootDeviceEnvironments(numRootDevices);
         pHwInfo = pHwInfo ? pHwInfo : platformDevices[0];
         executionEnvironment->setHwInfo(pHwInfo);
-        return createWithExecutionEnvironment<T>(pHwInfo, executionEnvironment, 0u);
+        return createWithExecutionEnvironment<T>(pHwInfo, executionEnvironment, rootDeviceIndex);
     }
     bool initializeRootCommandStreamReceiver() override {
         if (callBaseInitializeRootCommandStreamReceiver) {
@@ -129,7 +129,7 @@ class MockDevice : public RootDevice {
 };
 
 template <>
-inline Device *MockDevice::createWithNewExecutionEnvironment<Device>(const HardwareInfo *pHwInfo) {
+inline Device *MockDevice::createWithNewExecutionEnvironment<Device>(const HardwareInfo *pHwInfo, uint32_t rootDeviceIndex) {
     auto executionEnvironment = new ExecutionEnvironment();
     executionEnvironment->prepareRootDeviceEnvironments(1);
     MockAubCenterFixture::setMockAubCenter(*executionEnvironment->rootDeviceEnvironments[0]);
diff --git a/unit_tests/program/kernel_info_tests.cpp b/unit_tests/program/kernel_info_tests.cpp
index a50e3712c6..614cf2478c 100644
--- a/unit_tests/program/kernel_info_tests.cpp
+++ b/unit_tests/program/kernel_info_tests.cpp
@@ -8,6 +8,7 @@
 #include "runtime/execution_environment/execution_environment.h"
 #include "runtime/memory_manager/os_agnostic_memory_manager.h"
 #include "runtime/program/kernel_info.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/mocks/mock_execution_environment.h"
 
 #include "gtest/gtest.h"
@@ -247,3 +248,22 @@ TEST(KernelInfo, givenKernelInfoWhenStoreNonTransformableArgThenArgInfoIsNotTran
     const auto &argInfo = kernelInfo->kernelArgInfo[argumentNumber];
     EXPECT_FALSE(argInfo.isTransformable);
 }
+
+using KernelInfoMultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(KernelInfoMultiRootDeviceTests, kernelAllocationHasCorrectRootDeviceIndex) {
+    KernelInfo kernelInfo;
+    SKernelBinaryHeaderCommon kernelHeader;
+    const size_t heapSize = 0x40;
+    char heap[heapSize];
+    kernelHeader.KernelHeapSize = heapSize;
+    kernelInfo.heapInfo.pKernelHeader = &kernelHeader;
+    kernelInfo.heapInfo.pKernelHeap = &heap;
+
+    auto retVal = kernelInfo.createKernelAllocation(expectedRootDeviceIndex, mockMemoryManager);
+    EXPECT_TRUE(retVal);
+    auto allocation = kernelInfo.kernelAllocation;
+    ASSERT_NE(nullptr, allocation);
+    EXPECT_EQ(expectedRootDeviceIndex, allocation->getRootDeviceIndex());
+    mockMemoryManager->checkGpuUsageAndDestroyGraphicsAllocations(allocation);
+}
diff --git a/unit_tests/program/printf_handler_tests.cpp b/unit_tests/program/printf_handler_tests.cpp
index 353204cdb2..783ff269f5 100644
--- a/unit_tests/program/printf_handler_tests.cpp
+++ b/unit_tests/program/printf_handler_tests.cpp
@@ -7,6 +7,7 @@
 
 #include "runtime/program/printf_handler.h"
 #include "unit_tests/fixtures/device_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/mocks/mock_context.h"
 #include "unit_tests/mocks/mock_device.h"
 #include "unit_tests/mocks/mock_graphics_allocation.h"
@@ -72,6 +73,7 @@ TEST(PrintfHandlerTest, givenPreparedPrintfHandlerWhenGetSurfaceIsCalledThenResu
     delete pProgram;
     delete device;
 }
+
 TEST(PrintfHandlerTest, givenParentKernelWihoutPrintfAndBlockKernelWithPrintfWhenPrintfHandlerCreateCalledThenResaultIsAnObject) {
 
     std::unique_ptr<MockDevice> device(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
@@ -155,3 +157,28 @@ TEST(PrintfHandlerTest, GivenEmptyMultiDispatchInfoWhenCreatingPrintfHandlerThen
     auto printfHandler = PrintfHandler::create(multiDispatchInfo, device);
     EXPECT_EQ(nullptr, printfHandler);
 }
+
+using PrintfHandlerMultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(PrintfHandlerMultiRootDeviceTests, printfSurfaceHasCorrectRootDeviceIndex) {
+    auto printfSurface = std::make_unique<SPatchAllocateStatelessPrintfSurface>();
+    printfSurface->DataParamOffset = 0;
+    printfSurface->DataParamSize = 8;
+
+    auto kernelInfo = std::make_unique<KernelInfo>();
+    kernelInfo->patchInfo.pAllocateStatelessPrintfSurface = printfSurface.get();
+
+    auto program = std::make_unique<MockProgram>(*device->getExecutionEnvironment(), context.get(), false);
+
+    uint64_t crossThread[10];
+    auto kernel = std::make_unique<MockKernel>(program.get(), *kernelInfo, *device);
+    kernel->setCrossThreadData(&crossThread, sizeof(uint64_t) * 8);
+
+    MockMultiDispatchInfo multiDispatchInfo(kernel.get());
+    std::unique_ptr<PrintfHandler> printfHandler(PrintfHandler::create(multiDispatchInfo, *device));
+    printfHandler->prepareDispatch(multiDispatchInfo);
+    auto surface = printfHandler->getSurface();
+
+    ASSERT_NE(nullptr, surface);
+    EXPECT_EQ(expectedRootDeviceIndex, surface->getRootDeviceIndex());
+}
diff --git a/unit_tests/program/program_tests.cpp b/unit_tests/program/program_tests.cpp
index e06f12b3af..3b6587f5ba 100644
--- a/unit_tests/program/program_tests.cpp
+++ b/unit_tests/program/program_tests.cpp
@@ -27,6 +27,7 @@
 #include "runtime/program/create.inl"
 #include "test.h"
 #include "unit_tests/fixtures/device_fixture.h"
+#include "unit_tests/fixtures/multi_root_device_fixture.h"
 #include "unit_tests/global_environment.h"
 #include "unit_tests/helpers/kernel_binary_helper.h"
 #include "unit_tests/libult/ult_command_stream_receiver.h"
@@ -78,9 +79,9 @@ std::vector<const char *> KernelNames{
     "CopyBuffer",
 };
 
-class MockExecutionEnvironment : public ExecutionEnvironment {
+class MockCompIfaceExecutionEnvironment : public ExecutionEnvironment {
   public:
-    MockExecutionEnvironment(CompilerInterface *compilerInterface) : compilerInterface(compilerInterface) {}
+    MockCompIfaceExecutionEnvironment(CompilerInterface *compilerInterface) : compilerInterface(compilerInterface) {}
 
     CompilerInterface *getCompilerInterface() override {
         return compilerInterface;
@@ -694,7 +695,7 @@ TEST_P(ProgramFromSourceTest, CreateWithSource_Build) {
     pMockProgram->SetBuildStatus(CL_BUILD_NONE);
 
     // fail build - CompilerInterface cannot be obtained
-    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockExecutionEnvironment>(nullptr);
+    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockCompIfaceExecutionEnvironment>(nullptr);
     auto p2 = std::make_unique<MockProgram>(*noCompilerInterfaceExecutionEnvironment);
     retVal = p2->build(0, nullptr, nullptr, nullptr, nullptr, false);
     EXPECT_EQ(CL_OUT_OF_HOST_MEMORY, retVal);
@@ -959,7 +960,7 @@ TEST_P(ProgramFromSourceTest, CreateWithSource_Compile) {
     delete p3;
 
     // fail compilation - CompilerInterface cannot be obtained
-    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockExecutionEnvironment>(nullptr);
+    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockCompIfaceExecutionEnvironment>(nullptr);
     auto p2 = std::make_unique<MockProgram>(*noCompilerInterfaceExecutionEnvironment);
     retVal = p2->compile(0, nullptr, nullptr, 0, nullptr, nullptr, nullptr, nullptr);
     EXPECT_EQ(CL_OUT_OF_HOST_MEMORY, retVal);
@@ -1004,7 +1005,7 @@ struct MockCompilerInterfaceCaptureBuildOptions : CompilerInterface {
 
 TEST_P(ProgramFromSourceTest, CompileProgramWithInternalFlags) {
     auto cip = std::make_unique<MockCompilerInterfaceCaptureBuildOptions>();
-    MockExecutionEnvironment executionEnvironment(cip.get());
+    MockCompIfaceExecutionEnvironment executionEnvironment(cip.get());
     auto program = std::make_unique<SucceedingGenBinaryProgram>(executionEnvironment);
     cl_device_id deviceId = pContext->getDevice(0);
     Device *pDevice = castToObject<Device>(deviceId);
@@ -1186,7 +1187,7 @@ TEST_P(ProgramFromSourceTest, CreateWithSource_Link) {
 }
 
 TEST_P(ProgramFromSourceTest, CreateWithSource_CreateLibrary) {
-    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockExecutionEnvironment>(nullptr);
+    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockCompIfaceExecutionEnvironment>(nullptr);
     auto p = std::make_unique<MockProgram>(*noCompilerInterfaceExecutionEnvironment);
     cl_program program = pProgram;
 
@@ -2181,7 +2182,7 @@ TEST_F(ProgramTests, ValidBinaryWithIGCVersionEqual0) {
 }
 
 TEST_F(ProgramTests, RebuildBinaryButNoCompilerInterface) {
-    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockExecutionEnvironment>(nullptr);
+    auto noCompilerInterfaceExecutionEnvironment = std::make_unique<MockCompIfaceExecutionEnvironment>(nullptr);
     auto program = std::make_unique<MockProgram>(*noCompilerInterfaceExecutionEnvironment);
     EXPECT_NE(nullptr, program);
     cl_device_id deviceId = pContext->getDevice(0);
@@ -2216,7 +2217,7 @@ TEST_F(ProgramTests, RebuildBinaryWithRebuildError) {
     };
 
     auto cip = std::make_unique<MyCompilerInterface>();
-    MockExecutionEnvironment executionEnvironment(cip.get());
+    MockCompIfaceExecutionEnvironment executionEnvironment(cip.get());
     auto program = std::make_unique<MockProgram>(executionEnvironment);
     cl_device_id deviceId = pContext->getDevice(0);
     Device *pDevice = castToObject<Device>(deviceId);
@@ -2240,7 +2241,7 @@ TEST_F(ProgramTests, RebuildBinaryWithRebuildError) {
 
 TEST_F(ProgramTests, BuildProgramWithReraFlag) {
     auto cip = std::make_unique<MockCompilerInterfaceCaptureBuildOptions>();
-    MockExecutionEnvironment executionEnvironment(cip.get());
+    MockCompIfaceExecutionEnvironment executionEnvironment(cip.get());
     auto program = std::make_unique<SucceedingGenBinaryProgram>(executionEnvironment);
     cl_device_id deviceId = pContext->getDevice(0);
     Device *pDevice = castToObject<Device>(deviceId);
@@ -3229,3 +3230,27 @@ TEST_F(ProgramBinTest, GivenDebugDataAvailableWhenLinkingProgramThenDebugDataIsS
 
     EXPECT_NE(nullptr, pProgram->getDebugData());
 }
+
+using ProgramMultiRootDeviceTests = MultiRootDeviceFixture;
+
+TEST_F(ProgramMultiRootDeviceTests, privateSurfaceHasCorrectRootDeviceIndex) {
+    auto program = std::make_unique<MockProgram>(*device->getExecutionEnvironment(), context.get(), false);
+
+    auto privateSurfaceBlock = std::make_unique<SPatchAllocateStatelessPrivateSurface>();
+    privateSurfaceBlock->DataParamOffset = 0;
+    privateSurfaceBlock->DataParamSize = 8;
+    privateSurfaceBlock->Size = 8;
+    privateSurfaceBlock->SurfaceStateHeapOffset = 0;
+    privateSurfaceBlock->Token = 0;
+    privateSurfaceBlock->PerThreadPrivateMemorySize = 1000;
+
+    auto infoBlock = std::make_unique<KernelInfo>();
+    infoBlock->patchInfo.pAllocateStatelessPrivateSurface = privateSurfaceBlock.get();
+
+    program->blockKernelManager->addBlockKernelInfo(infoBlock.release());
+    program->allocateBlockPrivateSurfaces(device->getRootDeviceIndex());
+
+    auto privateSurface = program->getBlockKernelManager()->getPrivateSurface(0);
+    EXPECT_NE(nullptr, privateSurface);
+    EXPECT_EQ(expectedRootDeviceIndex, privateSurface->getRootDeviceIndex());
+}
\ No newline at end of file