Move PageTableManager from rootDeviceEnvironment to CommandStreamReceiver

Also add init of pageTable registers to blitter path Signed-off-by: Katarzyna Cencelewska <katarzyna.cencelewska@intel.com>
2026-01-08 14:02:58 +08:00 · 2021-09-29 15:59:41 +00:00
parent 897420236a
commit 30b3f5cdb0
20 changed files with 510 additions and 76 deletions
--- a/opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_3_tests.cpp
+++ b/opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_3_tests.cpp
@@ -1844,7 +1844,9 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCommandStreamReceiverWhenFlus
    pDevice->resetCommandStreamReceiver(csr);

    MockGmmPageTableMngr *pageTableManager = new MockGmmPageTableMngr();
-    pDevice->getExecutionEnvironment()->rootDeviceEnvironments[pDevice->getRootDeviceIndex()]->pageTableManager.reset(pageTableManager);
+    csr->pageTableManager.reset(pageTableManager);
+    MockGmmPageTableMngr *pageTableManager2 = new MockGmmPageTableMngr();
+    csr2->pageTableManager.reset(pageTableManager2);

    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(csr, ::testing::_)).Times(1);
    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(csr2, ::testing::_)).Times(0);
@@ -1865,7 +1867,7 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCommandStreamReceiverWhenFlus

    csr->flushTask(cs, 0u, cs, cs, cs, 0u, dispatchFlags, *pDevice);

-    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(csr2, ::testing::_)).Times(1);
+    EXPECT_CALL(*pageTableManager2, initContextAuxTableRegister(csr2, ::testing::_)).Times(1);
    pDevice->resetCommandStreamReceiver(csr2);
    csr2->flushTask(cs, 0u, cs, cs, cs, 0u, dispatchFlags, *pDevice);
    EXPECT_TRUE(csr2->pageTableManagerInitialized);
@@ -1873,12 +1875,141 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCommandStreamReceiverWhenFlus
    memoryManager->freeGraphicsMemory(graphicsAllocation);
 }

+HWTEST_F(CommandStreamReceiverFlushTaskTests, givenPageTableManagerPointerWhenCallBlitBufferThenPageTableManagerInitializedForProperCsr) {
+    auto bcsCsr = new MockCsrHw2<FamilyType>(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
+    auto bcsCsr2 = new MockCsrHw2<FamilyType>(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
+    pDevice->resetCommandStreamReceiver(bcsCsr);
+
+    MockGmmPageTableMngr *pageTableManager = new MockGmmPageTableMngr();
+    bcsCsr->pageTableManager.reset(pageTableManager);
+    MockGmmPageTableMngr *pageTableManager2 = new MockGmmPageTableMngr();
+    bcsCsr2->pageTableManager.reset(pageTableManager2);
+
+    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(bcsCsr, ::testing::_)).Times(1);
+    EXPECT_CALL(*pageTableManager2, initContextAuxTableRegister(bcsCsr2, ::testing::_)).Times(0);
+
+    auto memoryManager = pDevice->getMemoryManager();
+    auto graphicsAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{pDevice->getRootDeviceIndex(), MemoryConstants::pageSize});
+
+    EXPECT_FALSE(bcsCsr->pageTableManagerInitialized);
+    EXPECT_FALSE(bcsCsr2->pageTableManagerInitialized);
+
+    auto blitProperties = BlitProperties::constructPropertiesForCopy(graphicsAllocation,               //dstAllocation
+                                                                     graphicsAllocation,               //srcAllocation
+                                                                     0,                                //dstOffset
+                                                                     0,                                //srcOffset
+                                                                     0,                                //copySize
+                                                                     0,                                //srcRowPitch
+                                                                     0,                                //srcSlicePitch
+                                                                     0,                                //dstRowPitch
+                                                                     0,                                //dstSlicePitch
+                                                                     bcsCsr->getClearColorAllocation() //clearColorAllocation
+    );
+    BlitPropertiesContainer container;
+    container.push_back(blitProperties);
+
+    bcsCsr->blitBuffer(container, true, false, *pDevice);
+
+    EXPECT_TRUE(bcsCsr->pageTableManagerInitialized);
+    EXPECT_FALSE(bcsCsr2->pageTableManagerInitialized);
+
+    EXPECT_CALL(*pageTableManager2, initContextAuxTableRegister(bcsCsr2, ::testing::_)).Times(1);
+    pDevice->resetCommandStreamReceiver(bcsCsr2);
+    bcsCsr2->blitBuffer(container, true, false, *pDevice);
+
+    EXPECT_TRUE(bcsCsr2->pageTableManagerInitialized);
+
+    memoryManager->freeGraphicsMemory(graphicsAllocation);
+}
+
+HWTEST_F(CommandStreamReceiverFlushTaskTests, givenPageTableManagerPointerWhenCallBlitBufferAndPageTableManagerInitializedThenNotInitializeAgain) {
+    auto bcsCsr = new MockCsrHw2<FamilyType>(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
+    pDevice->resetCommandStreamReceiver(bcsCsr);
+
+    MockGmmPageTableMngr *pageTableManager = new MockGmmPageTableMngr();
+    bcsCsr->pageTableManager.reset(pageTableManager);
+
+    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(bcsCsr, ::testing::_)).Times(1);
+
+    auto memoryManager = pDevice->getMemoryManager();
+    auto graphicsAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{pDevice->getRootDeviceIndex(), MemoryConstants::pageSize});
+
+    EXPECT_FALSE(bcsCsr->pageTableManagerInitialized);
+
+    auto blitProperties = BlitProperties::constructPropertiesForCopy(graphicsAllocation,               //dstAllocation
+                                                                     graphicsAllocation,               //srcAllocation
+                                                                     0,                                //dstOffset
+                                                                     0,                                //srcOffset
+                                                                     0,                                //copySize
+                                                                     0,                                //srcRowPitch
+                                                                     0,                                //srcSlicePitch
+                                                                     0,                                //dstRowPitch
+                                                                     0,                                //dstSlicePitch
+                                                                     bcsCsr->getClearColorAllocation() //clearColorAllocation
+    );
+    BlitPropertiesContainer container;
+    container.push_back(blitProperties);
+
+    bcsCsr->blitBuffer(container, true, false, *pDevice);
+
+    EXPECT_TRUE(bcsCsr->pageTableManagerInitialized);
+
+    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(bcsCsr, ::testing::_)).Times(0);
+    bcsCsr->blitBuffer(container, true, false, *pDevice);
+
+    memoryManager->freeGraphicsMemory(graphicsAllocation);
+}
+
+HWTEST_F(CommandStreamReceiverFlushTaskTests, givenNullPageTableManagerWhenCallBlitBufferThenPageTableManagerIsNotInitialized) {
+    auto bcsCsr = new MockCsrHw2<FamilyType>(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
+    auto bcsCsr2 = new MockCsrHw2<FamilyType>(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
+    pDevice->resetCommandStreamReceiver(bcsCsr);
+
+    bcsCsr->pageTableManager.reset(nullptr);
+    bcsCsr2->pageTableManager.reset(nullptr);
+
+    auto memoryManager = pDevice->getMemoryManager();
+    auto graphicsAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{pDevice->getRootDeviceIndex(), MemoryConstants::pageSize});
+
+    EXPECT_FALSE(bcsCsr->pageTableManagerInitialized);
+    EXPECT_FALSE(bcsCsr2->pageTableManagerInitialized);
+
+    auto blitProperties = BlitProperties::constructPropertiesForCopy(graphicsAllocation,               //dstAllocation
+                                                                     graphicsAllocation,               //srcAllocation
+                                                                     0,                                //dstOffset
+                                                                     0,                                //srcOffset
+                                                                     0,                                //copySize
+                                                                     0,                                //srcRowPitch
+                                                                     0,                                //srcSlicePitch
+                                                                     0,                                //dstRowPitch
+                                                                     0,                                //dstSlicePitch
+                                                                     bcsCsr->getClearColorAllocation() //clearColorAllocation
+    );
+    BlitPropertiesContainer container;
+    container.push_back(blitProperties);
+
+    bcsCsr->blitBuffer(container, true, false, *pDevice);
+
+    EXPECT_FALSE(bcsCsr->pageTableManagerInitialized);
+    EXPECT_FALSE(bcsCsr2->pageTableManagerInitialized);
+
+    pDevice->resetCommandStreamReceiver(bcsCsr2);
+    bcsCsr2->blitBuffer(container, true, false, *pDevice);
+
+    EXPECT_FALSE(bcsCsr2->pageTableManagerInitialized);
+
+    bcsCsr2->pageTableManagerInitialized = true;
+    EXPECT_NO_THROW(bcsCsr2->blitBuffer(container, true, false, *pDevice));
+
+    memoryManager->freeGraphicsMemory(graphicsAllocation);
+}
+
 HWTEST_F(CommandStreamReceiverFlushTaskTests, givenCommandStreamReceiverWhenInitializingPageTableManagerRegisterFailsThenPageTableManagerIsNotInitialized) {
    auto csr = new MockCsrHw2<FamilyType>(*pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
    pDevice->resetCommandStreamReceiver(csr);

    MockGmmPageTableMngr *pageTableManager = new MockGmmPageTableMngr();
-    pDevice->getExecutionEnvironment()->rootDeviceEnvironments[pDevice->getRootDeviceIndex()]->pageTableManager.reset(pageTableManager);
+    csr->pageTableManager.reset(pageTableManager);

    EXPECT_CALL(*pageTableManager, initContextAuxTableRegister(csr, ::testing::_)).Times(2).WillRepeatedly(::testing::Return(GMM_ERROR));

--- a/opencl/test/unit_test/command_stream/command_stream_receiver_tests.cpp
+++ b/opencl/test/unit_test/command_stream/command_stream_receiver_tests.cpp
@@ -1506,43 +1506,39 @@ TEST_F(CommandStreamReceiverMultiRootDeviceTest, WhenCreatingCommandStreamGraphi

 using CommandStreamReceiverPageTableManagerTest = ::testing::Test;

-TEST_F(CommandStreamReceiverPageTableManagerTest, givenNonDefaultEngineTypeWhenNeedsPageTableManagerIsCalledThenFalseIsReturned) {
+TEST_F(CommandStreamReceiverPageTableManagerTest, givenNonRegularEngineWhenNeedsPageTableManagerIsCalledThenFalseIsReturned) {
    MockExecutionEnvironment executionEnvironment;
    executionEnvironment.initializeMemoryManager();
    DeviceBitfield deviceBitfield(1);
    MockCommandStreamReceiver commandStreamReceiver(executionEnvironment, 0u, deviceBitfield);
-    auto hwInfo = executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo();
-    auto defaultEngineType = getChosenEngineType(*hwInfo);
-    auto engineType = aub_stream::EngineType::ENGINE_BCS;
-    EXPECT_NE(defaultEngineType, engineType);
-    EXPECT_EQ(nullptr, executionEnvironment.rootDeviceEnvironments[0]->pageTableManager.get());
-    EXPECT_FALSE(commandStreamReceiver.needsPageTableManager(engineType));
+    auto engineUsage = EngineUsage::Internal;
+    EXPECT_EQ(nullptr, commandStreamReceiver.pageTableManager.get());
+    EXPECT_FALSE(commandStreamReceiver.needsPageTableManager(engineUsage));
 }

-TEST_F(CommandStreamReceiverPageTableManagerTest, givenDefaultEngineTypeAndExistingPageTableManagerWhenNeedsPageTableManagerIsCalledThenFalseIsReturned) {
+TEST_F(CommandStreamReceiverPageTableManagerTest, givenRegularEngineTypeAndExistingPageTableManagerWhenNeedsPageTableManagerIsCalledThenFalseIsReturned) {
    MockExecutionEnvironment executionEnvironment;
    executionEnvironment.initializeMemoryManager();
    DeviceBitfield deviceBitfield(1);
    MockCommandStreamReceiver commandStreamReceiver(executionEnvironment, 0u, deviceBitfield);
-    auto hwInfo = executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo();
-    auto defaultEngineType = getChosenEngineType(*hwInfo);
+    auto engineUsage = EngineUsage::Regular;

    GmmPageTableMngr *dummyPageTableManager = reinterpret_cast<GmmPageTableMngr *>(0x1234);

-    executionEnvironment.rootDeviceEnvironments[0]->pageTableManager.reset(dummyPageTableManager);
-    EXPECT_FALSE(commandStreamReceiver.needsPageTableManager(defaultEngineType));
-    executionEnvironment.rootDeviceEnvironments[0]->pageTableManager.release();
+    commandStreamReceiver.pageTableManager.reset(dummyPageTableManager);
+    EXPECT_FALSE(commandStreamReceiver.needsPageTableManager(engineUsage));
+    commandStreamReceiver.pageTableManager.release();
 }

-TEST_F(CommandStreamReceiverPageTableManagerTest, givenDefaultEngineTypeAndNonExisitingPageTableManagerWhenNeedsPageTableManagerIsCalledThenSupportOfPageTableManagerIsReturned) {
+TEST_F(CommandStreamReceiverPageTableManagerTest, givenRegularEngineAndNonExisitingPageTableManagerWhenNeedsPageTableManagerIsCalledThenSupportOfPageTableManagerIsReturned) {
    MockExecutionEnvironment executionEnvironment;
    executionEnvironment.initializeMemoryManager();
    DeviceBitfield deviceBitfield(1);
    MockCommandStreamReceiver commandStreamReceiver(executionEnvironment, 0u, deviceBitfield);
    auto hwInfo = executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo();
-    auto defaultEngineType = getChosenEngineType(*hwInfo);
+    auto engineUsage = EngineUsage::Regular;
    bool supportsPageTableManager = HwInfoConfig::get(hwInfo->platform.eProductFamily)->isPageTableManagerSupported(*hwInfo);
-    EXPECT_EQ(nullptr, executionEnvironment.rootDeviceEnvironments[0]->pageTableManager.get());
+    EXPECT_EQ(nullptr, commandStreamReceiver.pageTableManager.get());

-    EXPECT_EQ(supportsPageTableManager, commandStreamReceiver.needsPageTableManager(defaultEngineType));
+    EXPECT_EQ(supportsPageTableManager, commandStreamReceiver.needsPageTableManager(engineUsage));
 }
--- a/opencl/test/unit_test/device/device_tests.cpp
+++ b/opencl/test/unit_test/device/device_tests.cpp
@@ -551,7 +551,7 @@ HWTEST_F(DeviceHwTest, givenHwHelperInputWhenInitializingCsrThenCreatePageTableM
        executionEnvironment.rootDeviceEnvironments[i]->setHwInfo(&localHwInfo);
    }
    executionEnvironment.initializeMemoryManager();
-    auto defaultEngineType = getChosenEngineType(localHwInfo);
+    auto regularEngineUsage = EngineUsage::Regular;
    std::unique_ptr<MockDevice> device;
    device.reset(MockDevice::createWithExecutionEnvironment<MockDevice>(&localHwInfo, &executionEnvironment, 0));
    auto &csr0 = device->getUltCommandStreamReceiver<FamilyType>();
@@ -562,14 +562,14 @@ HWTEST_F(DeviceHwTest, givenHwHelperInputWhenInitializingCsrThenCreatePageTableM
    hwInfo->capabilityTable.ftrRenderCompressedImages = false;
    device.reset(MockDevice::createWithExecutionEnvironment<MockDevice>(&localHwInfo, &executionEnvironment, 1));
    auto &csr1 = device->getUltCommandStreamReceiver<FamilyType>();
-    EXPECT_EQ(csr1.needsPageTableManager(defaultEngineType), csr1.createPageTableManagerCalled);
+    EXPECT_EQ(csr1.needsPageTableManager(regularEngineUsage), csr1.createPageTableManagerCalled);

    hwInfo = executionEnvironment.rootDeviceEnvironments[2]->getMutableHardwareInfo();
    hwInfo->capabilityTable.ftrRenderCompressedBuffers = false;
    hwInfo->capabilityTable.ftrRenderCompressedImages = true;
    device.reset(MockDevice::createWithExecutionEnvironment<MockDevice>(&localHwInfo, &executionEnvironment, 2));
    auto &csr2 = device->getUltCommandStreamReceiver<FamilyType>();
-    EXPECT_EQ(csr2.needsPageTableManager(defaultEngineType), csr2.createPageTableManagerCalled);
+    EXPECT_EQ(csr2.needsPageTableManager(regularEngineUsage), csr2.createPageTableManagerCalled);
 }

 HWTEST_F(DeviceHwTest, givenDeviceCreationWhenCsrFailsToCreateGlobalSyncAllocationThenReturnNull) {
--- a/opencl/test/unit_test/memory_manager/memory_manager_tests.cpp
+++ b/opencl/test/unit_test/memory_manager/memory_manager_tests.cpp
@@ -45,6 +45,7 @@
 #include "opencl/test/unit_test/mocks/mock_context.h"
 #include "opencl/test/unit_test/mocks/mock_csr.h"
 #include "opencl/test/unit_test/mocks/mock_gmm.h"
+#include "opencl/test/unit_test/mocks/mock_gmm_page_table_mngr.h"
 #include "opencl/test/unit_test/mocks/mock_kernel.h"
 #include "opencl/test/unit_test/mocks/mock_mdi.h"
 #include "opencl/test/unit_test/mocks/mock_memory_manager.h"
@@ -2755,6 +2756,147 @@ HWTEST_F(MemoryAllocatorTest, givenMemoryManagerWhenHostPtrTrackingEnabledThenNo
 }

 using PageTableManagerTest = ::testing::Test;
+using namespace ::testing;
+HWTEST_F(PageTableManagerTest, givenPageTableManagerWhenMapAuxGpuVaThenForAllEnginesWithPageTableUpdateAuxTableAreCalled) {
+    ExecutionEnvironment *executionEnvironment = platform()->peekExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(2);
+    for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
+        executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
+        executionEnvironment->rootDeviceEnvironments[i]->initGmm();
+    }
+
+    auto memoryManager = new MockMemoryManager(false, false, *executionEnvironment);
+
+    executionEnvironment->memoryManager.reset(memoryManager);
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto csr2 = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular},
+        {aub_stream::ENGINE_BCS, EngineUsage::Regular},
+    };
+
+    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+    memoryManager->createAndRegisterOsContext(csr2.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[1],
+                                                                                                       PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
+    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
+    auto mockMngr2 = new NiceMock<MockGmmPageTableMngr>();
+
+    memoryManager->getRegisteredEngines()[0].commandStreamReceiver->pageTableManager.reset(mockMngr);
+    memoryManager->getRegisteredEngines()[1].commandStreamReceiver->pageTableManager.reset(mockMngr2);
+
+    MockGraphicsAllocation allocation(1u, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, 0, 0, 0, MemoryPool::MemoryNull);
+    MockGmm gmm(executionEnvironment->rootDeviceEnvironments[allocation.getRootDeviceIndex()]->getGmmClientContext());
+    gmm.isCompressionEnabled = true;
+    allocation.setDefaultGmm(&gmm);
+    GMM_DDI_UPDATEAUXTABLE givenDdiUpdateAuxTable = {};
+    GMM_DDI_UPDATEAUXTABLE givenDdiUpdateAuxTable2 = {};
+    GMM_DDI_UPDATEAUXTABLE expectedDdiUpdateAuxTable = {};
+    expectedDdiUpdateAuxTable.BaseGpuVA = allocation.getGpuAddress();
+    expectedDdiUpdateAuxTable.BaseResInfo = allocation.getDefaultGmm()->gmmResourceInfo->peekGmmResourceInfo();
+    expectedDdiUpdateAuxTable.DoNotWait = true;
+    expectedDdiUpdateAuxTable.Map = true;
+
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(1).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
+    EXPECT_CALL(*mockMngr2, updateAuxTable(_)).Times(1).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg2) {givenDdiUpdateAuxTable2 = *arg2; return GMM_SUCCESS; }));
+
+    bool result = memoryManager->mapAuxGpuVA(&allocation);
+    EXPECT_TRUE(result);
+
+    EXPECT_TRUE(memcmp(&expectedDdiUpdateAuxTable, &givenDdiUpdateAuxTable, sizeof(GMM_DDI_UPDATEAUXTABLE)) == 0);
+    EXPECT_TRUE(memcmp(&expectedDdiUpdateAuxTable, &givenDdiUpdateAuxTable2, sizeof(GMM_DDI_UPDATEAUXTABLE)) == 0);
+}
+HWTEST_F(PageTableManagerTest, givenPageTableManagerWhenUpdateAuxTableGmmErrorThenMapAuxGpuVaReturnFalse) {
+    ExecutionEnvironment *executionEnvironment = platform()->peekExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(2);
+    for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
+        executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
+        executionEnvironment->rootDeviceEnvironments[i]->initGmm();
+    }
+
+    auto memoryManager = new MockMemoryManager(false, false, *executionEnvironment);
+
+    executionEnvironment->memoryManager.reset(memoryManager);
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
+    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
+
+    memoryManager->getRegisteredEngines()[0].commandStreamReceiver->pageTableManager.reset(mockMngr);
+
+    MockGraphicsAllocation allocation(1u, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, 0, 0, 0, MemoryPool::MemoryNull);
+    MockGmm gmm(executionEnvironment->rootDeviceEnvironments[allocation.getRootDeviceIndex()]->getGmmClientContext());
+    gmm.isCompressionEnabled = true;
+    allocation.setDefaultGmm(&gmm);
+
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(1).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) { return GMM_ERROR; }));
+
+    bool result = memoryManager->mapAuxGpuVA(&allocation);
+    EXPECT_FALSE(result);
+}
+
+HWTEST_F(PageTableManagerTest, givenNullPageTableManagerWhenMapAuxGpuVaThenNoThrow) {
+    ExecutionEnvironment *executionEnvironment = platform()->peekExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(2);
+    for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
+        executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
+        executionEnvironment->rootDeviceEnvironments[i]->initGmm();
+    }
+
+    auto memoryManager = new MockMemoryManager(false, false, *executionEnvironment);
+
+    executionEnvironment->memoryManager.reset(memoryManager);
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular},
+    };
+
+    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
+    memoryManager->getRegisteredEngines()[0].commandStreamReceiver->pageTableManager.reset(nullptr);
+
+    MockGraphicsAllocation allocation(1u, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, 0, 0, 0, MemoryPool::MemoryNull);
+    MockGmm gmm(executionEnvironment->rootDeviceEnvironments[allocation.getRootDeviceIndex()]->getGmmClientContext());
+    gmm.isCompressionEnabled = true;
+    allocation.setDefaultGmm(&gmm);
+
+    EXPECT_NO_THROW(memoryManager->mapAuxGpuVA(&allocation));
+}
+
+HWTEST_F(PageTableManagerTest, givenNullPageTableManagerWhenMapAuxGpuVaThenReturnFalse) {
+    ExecutionEnvironment *executionEnvironment = platform()->peekExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(2);
+    for (auto i = 0u; i < executionEnvironment->rootDeviceEnvironments.size(); i++) {
+        executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
+        executionEnvironment->rootDeviceEnvironments[i]->initGmm();
+    }
+
+    auto memoryManager = new MockMemoryManager(false, false, *executionEnvironment);
+
+    executionEnvironment->memoryManager.reset(memoryManager);
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(HwHelper::get(hwInfo.platform.eRenderCoreFamily).getGpgpuEngineInstances(hwInfo)[0],
+                                                                                                      PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
+    for (auto engine : memoryManager->getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(nullptr);
+    }
+
+    MockGraphicsAllocation allocation(1u, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, 0, 0, 0, MemoryPool::MemoryNull);
+
+    bool result = memoryManager->mapAuxGpuVA(&allocation);
+    EXPECT_FALSE(result);
+}

 HWTEST_F(PageTableManagerTest, givenMemoryManagerThatSupportsPageTableManagerWhenMapAuxGpuVAIsCalledThenItReturnsTrue) {
    ExecutionEnvironment *executionEnvironment = platform()->peekExecutionEnvironment();
--- a/opencl/test/unit_test/os_interface/linux/drm_command_stream_tests_1.cpp
+++ b/opencl/test/unit_test/os_interface/linux/drm_command_stream_tests_1.cpp
@@ -862,7 +862,7 @@ HWTEST_TEMPLATED_F(DrmCommandStreamTest, givenDrmCommandStreamReceiverWhenCreate
    executionEnvironment.rootDeviceEnvironments[1]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(new DrmMockCustom()));
    auto csr = std::make_unique<MockDrmCsr<FamilyType>>(executionEnvironment, 1, 1, gemCloseWorkerMode::gemCloseWorkerActive);
    auto pageTableManager = csr->createPageTableManager();
-    EXPECT_EQ(executionEnvironment.rootDeviceEnvironments[1]->pageTableManager.get(), pageTableManager);
+    EXPECT_EQ(csr->pageTableManager.get(), pageTableManager);
 }

 HWTEST_TEMPLATED_F(DrmCommandStreamTest, givenLocalMemoryEnabledWhenCreatingDrmCsrThenEnableBatching) {
@@ -893,12 +893,12 @@ HWTEST_TEMPLATED_F(DrmCommandStreamTest, givenLocalMemoryEnabledWhenCreatingDrmC

 HWTEST_TEMPLATED_F(DrmCommandStreamTest, givenPageTableManagerAndMapTrueWhenUpdateAuxTableIsCalledThenItReturnsTrue) {
    auto mockMngr = new MockGmmPageTableMngr();
-    executionEnvironment.rootDeviceEnvironments[0]->pageTableManager.reset(mockMngr);
+    csr->pageTableManager.reset(mockMngr);
    executionEnvironment.rootDeviceEnvironments[0]->initGmm();
    auto gmm = std::make_unique<MockGmm>(executionEnvironment.rootDeviceEnvironments[0]->getGmmClientContext());
    GMM_DDI_UPDATEAUXTABLE ddiUpdateAuxTable = {};
    EXPECT_CALL(*mockMngr, updateAuxTable(::testing::_)).Times(1).WillOnce(::testing::Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {ddiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
-    auto result = executionEnvironment.rootDeviceEnvironments[0]->pageTableManager->updateAuxTable(0, gmm.get(), true);
+    auto result = csr->pageTableManager->updateAuxTable(0, gmm.get(), true);
    EXPECT_EQ(ddiUpdateAuxTable.BaseGpuVA, 0ull);
    EXPECT_EQ(ddiUpdateAuxTable.BaseResInfo, gmm->gmmResourceInfo->peekHandle());
    EXPECT_EQ(ddiUpdateAuxTable.DoNotWait, true);
@@ -909,12 +909,12 @@ HWTEST_TEMPLATED_F(DrmCommandStreamTest, givenPageTableManagerAndMapTrueWhenUpda

 HWTEST_TEMPLATED_F(DrmCommandStreamTest, givenPageTableManagerAndMapFalseWhenUpdateAuxTableIsCalledThenItReturnsTrue) {
    auto mockMngr = new MockGmmPageTableMngr();
-    executionEnvironment.rootDeviceEnvironments[0]->pageTableManager.reset(mockMngr);
+    csr->pageTableManager.reset(mockMngr);
    executionEnvironment.rootDeviceEnvironments[0]->initGmm();
    auto gmm = std::make_unique<MockGmm>(executionEnvironment.rootDeviceEnvironments[0]->getGmmClientContext());
    GMM_DDI_UPDATEAUXTABLE ddiUpdateAuxTable = {};
    EXPECT_CALL(*mockMngr, updateAuxTable(::testing::_)).Times(1).WillOnce(::testing::Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {ddiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
-    auto result = executionEnvironment.rootDeviceEnvironments[0]->pageTableManager->updateAuxTable(0, gmm.get(), false);
+    auto result = csr->pageTableManager->updateAuxTable(0, gmm.get(), false);
    EXPECT_EQ(ddiUpdateAuxTable.BaseGpuVA, 0ull);
    EXPECT_EQ(ddiUpdateAuxTable.BaseResInfo, gmm->gmmResourceInfo->peekHandle());
    EXPECT_EQ(ddiUpdateAuxTable.DoNotWait, true);
--- a/opencl/test/unit_test/os_interface/windows/device_command_stream_tests.cpp
+++ b/opencl/test/unit_test/os_interface/windows/device_command_stream_tests.cpp
@@ -970,13 +970,13 @@ HWTEST_P(WddmCsrCompressionParameterizedTest, givenEnabledCompressionWhenInitial
    std::unique_ptr<MockDevice> device(Device::create<MockDevice>(executionEnvironment, 1u));
    setCompressionEnabled(compressionEnabled, !compressionEnabled);
    myMockWddm = static_cast<WddmMock *>(executionEnvironment->rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Wddm>());
-    EXPECT_EQ(nullptr, executionEnvironment->rootDeviceEnvironments[index]->pageTableManager.get());

    MockWddmCsr<FamilyType> mockWddmCsr(*executionEnvironment, index, 1);
    mockWddmCsr.createPageTableManager();
-    ASSERT_NE(nullptr, executionEnvironment->rootDeviceEnvironments[index]->pageTableManager.get());

-    auto mockMngr = reinterpret_cast<MockGmmPageTableMngr *>(executionEnvironment->rootDeviceEnvironments[index]->pageTableManager.get());
+    ASSERT_NE(nullptr, mockWddmCsr.pageTableManager.get());
+
+    auto mockMngr = reinterpret_cast<MockGmmPageTableMngr *>(mockWddmCsr.pageTableManager.get());
    EXPECT_EQ(1u, mockMngr->setCsrHanleCalled);
    EXPECT_EQ(&mockWddmCsr, mockMngr->passedCsrHandle);

@@ -995,7 +995,9 @@ HWTEST_F(WddmCsrCompressionTests, givenDisabledCompressionWhenInitializedThenDon
    setCompressionEnabled(false, false);
    myMockWddm = static_cast<WddmMock *>(executionEnvironment->rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Wddm>());
    MockWddmCsr<FamilyType> mockWddmCsr(*executionEnvironment, 1, device->getDeviceBitfield());
-    EXPECT_EQ(nullptr, executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.get());
+    for (auto engine : executionEnvironment->memoryManager.get()->getRegisteredEngines()) {
+        EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
+    }
 }

 INSTANTIATE_TEST_CASE_P(
@@ -1017,19 +1019,24 @@ HWTEST_F(WddmCsrCompressionTests, givenDisabledCompressionWhenFlushingThenDontIn
    device->resetCommandStreamReceiver(mockWddmCsr);

    auto memoryManager = executionEnvironment->memoryManager.get();
-
-    EXPECT_EQ(nullptr, executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.get());
+    for (auto engine : memoryManager->getRegisteredEngines()) {
+        EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
+    }

    auto graphicsAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{mockWddmCsr->getRootDeviceIndex(), MemoryConstants::pageSize});
    IndirectHeap cs(graphicsAllocation);

-    EXPECT_EQ(nullptr, executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.get());
+    for (auto engine : memoryManager->getRegisteredEngines()) {
+        EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
+    }

    DispatchFlags dispatchFlags = DispatchFlagsHelper::createDefaultDispatchFlags();

    mockWddmCsr->flushTask(cs, 0u, cs, cs, cs, 0u, dispatchFlags, *device);

-    EXPECT_EQ(nullptr, executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.get());
+    for (auto engine : memoryManager->getRegisteredEngines()) {
+        EXPECT_EQ(nullptr, engine.commandStreamReceiver->pageTableManager.get());
+    }

    mockWddmCsr->flushBatchedSubmissions();
    memoryManager->freeGraphicsMemory(graphicsAllocation);
--- a/opencl/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp
+++ b/opencl/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp
@@ -12,6 +12,7 @@
 #include "shared/source/helpers/aligned_memory.h"
 #include "shared/source/helpers/array_count.h"
 #include "shared/source/os_interface/device_factory.h"
+#include "shared/source/os_interface/hw_info_config.h"
 #include "shared/source/os_interface/os_library.h"
 #include "shared/source/os_interface/windows/os_context_win.h"
 #include "shared/source/os_interface/windows/wddm_residency_controller.h"
@@ -49,7 +50,20 @@ void WddmMemoryManagerFixture::SetUp() {
    wddm = static_cast<WddmMock *>(Wddm::createWddm(nullptr, *rootDeviceEnvironment));
    if (defaultHwInfo->capabilityTable.ftrRenderCompressedBuffers || defaultHwInfo->capabilityTable.ftrRenderCompressedImages) {
        GMM_TRANSLATIONTABLE_CALLBACKS dummyTTCallbacks = {};
-        rootDeviceEnvironment->pageTableManager.reset(GmmPageTableMngr::create(nullptr, 0, &dummyTTCallbacks));
+
+        auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+        auto hwInfo = *defaultHwInfo;
+        EngineInstancesContainer regularEngines = {
+            {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+        memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                          PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
+        for (auto engine : memoryManager->getRegisteredEngines()) {
+            if (engine.getEngineUsage() == EngineUsage::Regular) {
+                engine.commandStreamReceiver->pageTableManager.reset(GmmPageTableMngr::create(nullptr, 0, &dummyTTCallbacks));
+            }
+        }
    }
    wddm->init();
    constexpr uint64_t heap32Base = (is32bit) ? 0x1000 : 0x800000000000;
@@ -742,6 +756,7 @@ HWTEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerWhenAllocateGraphicsMemory
    if (memoryManager->isLimitedGPU(0)) {
        GTEST_SKIP();
    }
+    rootDeviceEnvironment->executionEnvironment.initializeMemoryManager();
    memoryManager->allocateGraphicsMemoryInNonDevicePool = true;
    auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{mockRootDeviceIndex, true, size, GraphicsAllocation::AllocationType::BUFFER_COMPRESSED, mockDeviceBitfield}, ptr);

@@ -1784,11 +1799,23 @@ TEST_F(MockWddmMemoryManagerTest, givenDisabledAsyncDeleterFlagWhenMemoryManager
 }

 TEST_F(MockWddmMemoryManagerTest, givenPageTableManagerWhenMapAuxGpuVaCalledThenUseWddmToMap) {
+    if (!HwInfoConfig::get(defaultHwInfo->platform.eProductFamily)->isPageTableManagerSupported(*defaultHwInfo)) {
+        GTEST_SKIP();
+    }
    wddm->init();
    WddmMemoryManager memoryManager(*executionEnvironment);
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    memoryManager.createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                     PreemptionHelper::getDefaultPreemptionMode(hwInfo)));

    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
-    executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.reset(mockMngr);
+    for (auto engine : memoryManager.getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    auto allocation = memoryManager.allocateGraphicsMemoryWithProperties(AllocationProperties(1, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, mockDeviceBitfield));

@@ -1799,7 +1826,8 @@ TEST_F(MockWddmMemoryManagerTest, givenPageTableManagerWhenMapAuxGpuVaCalledThen
    expectedDdiUpdateAuxTable.DoNotWait = true;
    expectedDdiUpdateAuxTable.Map = true;

-    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(1).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
+    auto expectedCallCount = static_cast<int>(regularEngines.size());
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(expectedCallCount).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));

    auto result = memoryManager.mapAuxGpuVA(allocation);
    EXPECT_TRUE(result);
@@ -1807,16 +1835,61 @@ TEST_F(MockWddmMemoryManagerTest, givenPageTableManagerWhenMapAuxGpuVaCalledThen
    memoryManager.freeGraphicsMemory(allocation);
 }

-TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenMappedGpuVaThenMapAuxVa) {
+TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenMappedGpuVaAndPageTableNotSupportedThenMapAuxVa) {
+    if (HwInfoConfig::get(defaultHwInfo->platform.eProductFamily)->isPageTableManagerSupported(*defaultHwInfo)) {
+        GTEST_SKIP();
+    }
    auto rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[1].get();
    std::unique_ptr<Gmm> gmm(new Gmm(rootDeviceEnvironment->getGmmClientContext(), reinterpret_cast<void *>(123), 4096u, 0, false));
    gmm->isCompressionEnabled = true;
    D3DGPU_VIRTUAL_ADDRESS gpuVa = 0;
    WddmMock wddm(*executionEnvironment->rootDeviceEnvironments[1].get());
    wddm.init();
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                                            PreemptionHelper::getDefaultPreemptionMode(hwInfo)));

    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
-    rootDeviceEnvironment->pageTableManager.reset(mockMngr);
+    for (auto engine : executionEnvironment->memoryManager.get()->getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }
+
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(0);
+
+    auto hwInfoMock = hardwareInfoTable[wddm.getGfxPlatform()->eProductFamily];
+    ASSERT_NE(nullptr, hwInfoMock);
+    auto result = wddm.mapGpuVirtualAddress(gmm.get(), ALLOCATION_HANDLE, wddm.getGfxPartition().Standard.Base, wddm.getGfxPartition().Standard.Limit, 0u, gpuVa);
+    ASSERT_TRUE(result);
+
+    EXPECT_EQ(GmmHelper::canonize(wddm.getGfxPartition().Standard.Base), gpuVa);
+}
+
+TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenMappedGpuVaAndPageTableSupportedThenMapAuxVa) {
+    if (!HwInfoConfig::get(defaultHwInfo->platform.eProductFamily)->isPageTableManagerSupported(*defaultHwInfo)) {
+        GTEST_SKIP();
+    }
+    auto rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[1].get();
+    std::unique_ptr<Gmm> gmm(new Gmm(rootDeviceEnvironment->getGmmClientContext(), reinterpret_cast<void *>(123), 4096u, 0, false));
+    gmm->isCompressionEnabled = true;
+    D3DGPU_VIRTUAL_ADDRESS gpuVa = 0;
+    WddmMock wddm(*executionEnvironment->rootDeviceEnvironments[1].get());
+    wddm.init();
+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                                            PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
+    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
+    for (auto engine : executionEnvironment->memoryManager.get()->getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    GMM_DDI_UPDATEAUXTABLE givenDdiUpdateAuxTable = {};
    GMM_DDI_UPDATEAUXTABLE expectedDdiUpdateAuxTable = {};
@@ -1825,7 +1898,9 @@ TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenMappedGpuVa
    expectedDdiUpdateAuxTable.DoNotWait = true;
    expectedDdiUpdateAuxTable.Map = true;

-    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(1).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
+    auto expectedCallCount = executionEnvironment->memoryManager.get()->getRegisteredEnginesCount();
+
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(expectedCallCount).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));

    auto hwInfoMock = hardwareInfoTable[wddm.getGfxPlatform()->eProductFamily];
    ASSERT_NE(nullptr, hwInfoMock);
@@ -1836,13 +1911,26 @@ TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenMappedGpuVa
    EXPECT_TRUE(memcmp(&expectedDdiUpdateAuxTable, &givenDdiUpdateAuxTable, sizeof(GMM_DDI_UPDATEAUXTABLE)) == 0);
 }

-TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenReleaseingThenUnmapAuxVa) {
+TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationAndPageTableSupportedWhenReleaseingThenUnmapAuxVa) {
+    if (!HwInfoConfig::get(defaultHwInfo->platform.eProductFamily)->isPageTableManagerSupported(*defaultHwInfo)) {
+        GTEST_SKIP();
+    }
    wddm->init();
    WddmMemoryManager memoryManager(*executionEnvironment);
    D3DGPU_VIRTUAL_ADDRESS gpuVa = 123;

+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    memoryManager.createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                     PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
-    executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.reset(mockMngr);
+    for (auto engine : memoryManager.getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    auto wddmAlloc = static_cast<WddmAllocation *>(memoryManager.allocateGraphicsMemoryWithProperties(AllocationProperties(1, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::INTERNAL_HOST_MEMORY, mockDeviceBitfield)));
    wddmAlloc->setGpuAddress(gpuVa);
@@ -1855,7 +1943,9 @@ TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedAllocationWhenReleaseingT
    expectedDdiUpdateAuxTable.DoNotWait = true;
    expectedDdiUpdateAuxTable.Map = false;

-    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(1).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
+    auto expectedCallCount = memoryManager.getRegisteredEnginesCount();
+
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(expectedCallCount).WillOnce(Invoke([&](const GMM_DDI_UPDATEAUXTABLE *arg) {givenDdiUpdateAuxTable = *arg; return GMM_SUCCESS; }));
    memoryManager.freeGraphicsMemory(wddmAlloc);

    EXPECT_TRUE(memcmp(&expectedDdiUpdateAuxTable, &givenDdiUpdateAuxTable, sizeof(GMM_DDI_UPDATEAUXTABLE)) == 0);
@@ -1865,8 +1955,18 @@ TEST_F(MockWddmMemoryManagerTest, givenNonRenderCompressedAllocationWhenReleasei
    wddm->init();
    WddmMemoryManager memoryManager(*executionEnvironment);

+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    memoryManager.createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                     PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
-    executionEnvironment->rootDeviceEnvironments[1]->pageTableManager.reset(mockMngr);
+    for (auto engine : memoryManager.getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    auto wddmAlloc = static_cast<WddmAllocation *>(memoryManager.allocateGraphicsMemoryWithProperties(MockAllocationProperties{rootDeviceIndex, MemoryConstants::pageSize}));
    wddmAlloc->getDefaultGmm()->isCompressionEnabled = false;
@@ -1884,8 +1984,18 @@ TEST_F(MockWddmMemoryManagerTest, givenNonRenderCompressedAllocationWhenMappedGp
    WddmMock wddm(*rootDeviceEnvironment);
    wddm.init();

+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                                            PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
-    rootDeviceEnvironment->pageTableManager.reset(mockMngr);
+    for (auto engine : executionEnvironment->memoryManager.get()->getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(0);

@@ -1910,9 +2020,19 @@ TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedFlagSetWhenInternalIsUnse
    wddm->init();
    WddmMemoryManager memoryManager(*executionEnvironment);

+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    memoryManager.createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                     PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
    auto rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[1].get();
-    rootDeviceEnvironment->pageTableManager.reset(mockMngr);
+    for (auto engine : memoryManager.getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    auto myGmm = new Gmm(rootDeviceEnvironment->getGmmClientContext(), reinterpret_cast<void *>(123), 4096u, 0, false);
    myGmm->isCompressionEnabled = false;
@@ -1930,13 +2050,27 @@ TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedFlagSetWhenInternalIsUnse
 }

 TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedFlagSetWhenInternalIsSetThenUpdateAuxTable) {
+    if (!HwInfoConfig::get(defaultHwInfo->platform.eProductFamily)->isPageTableManagerSupported(*defaultHwInfo)) {
+        GTEST_SKIP();
+    }
    D3DGPU_VIRTUAL_ADDRESS gpuVa = 0;
    wddm->init();
    WddmMemoryManager memoryManager(*executionEnvironment);

+    auto csr = std::unique_ptr<CommandStreamReceiver>(createCommandStream(*executionEnvironment, 1u, 1));
+    auto hwInfo = *defaultHwInfo;
+    EngineInstancesContainer regularEngines = {
+        {aub_stream::ENGINE_CCS, EngineUsage::Regular}};
+
+    memoryManager.createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor(regularEngines[0],
+                                                                                                     PreemptionHelper::getDefaultPreemptionMode(hwInfo)));
+
    auto mockMngr = new NiceMock<MockGmmPageTableMngr>();
    auto rootDeviceEnvironment = executionEnvironment->rootDeviceEnvironments[1].get();
-    rootDeviceEnvironment->pageTableManager.reset(mockMngr);
+    rootDeviceEnvironment->executionEnvironment.initializeMemoryManager();
+    for (auto engine : memoryManager.getRegisteredEngines()) {
+        engine.commandStreamReceiver->pageTableManager.reset(mockMngr);
+    }

    auto myGmm = new Gmm(rootDeviceEnvironment->getGmmClientContext(), reinterpret_cast<void *>(123), 4096u, 0, false);
    myGmm->isCompressionEnabled = true;
@@ -1946,7 +2080,9 @@ TEST_F(MockWddmMemoryManagerTest, givenRenderCompressedFlagSetWhenInternalIsSetT
    delete wddmAlloc->getDefaultGmm();
    wddmAlloc->setDefaultGmm(myGmm);

-    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(1);
+    auto expectedCallCount = memoryManager.getRegisteredEnginesCount();
+
+    EXPECT_CALL(*mockMngr, updateAuxTable(_)).Times(expectedCallCount);

    auto result = wddm->mapGpuVirtualAddress(myGmm, ALLOCATION_HANDLE, wddm->getGfxPartition().Standard.Base, wddm->getGfxPartition().Standard.Limit, 0u, gpuVa);
    EXPECT_TRUE(result);
--- a/shared/source/command_stream/command_stream_receiver.cpp
+++ b/shared/source/command_stream/command_stream_receiver.cpp
@@ -14,6 +14,7 @@
 #include "shared/source/command_stream/scratch_space_controller.h"
 #include "shared/source/device/device.h"
 #include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/gmm_helper/page_table_mngr.h"
 #include "shared/source/helpers/array_count.h"
 #include "shared/source/helpers/cache_policy.h"
 #include "shared/source/helpers/flush_stamp.h"
@@ -667,14 +668,13 @@ bool CommandStreamReceiver::expectMemory(const void *gfxAddress, const void *src
    return (isMemoryEqual == isEqualMemoryExpected);
 }

-bool CommandStreamReceiver::needsPageTableManager(aub_stream::EngineType engineType) const {
+bool CommandStreamReceiver::needsPageTableManager(EngineUsage engineTypeUsage) const {
    auto hwInfo = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo();
-    auto defaultEngineType = getChosenEngineType(*hwInfo);
-    if (engineType != defaultEngineType) {
+
+    if (engineTypeUsage != EngineUsage::Regular) {
        return false;
    }
-    auto rootDeviceEnvironment = executionEnvironment.rootDeviceEnvironments[rootDeviceIndex].get();
-    if (rootDeviceEnvironment->pageTableManager.get() != nullptr) {
+    if (pageTableManager.get() != nullptr) {
        return false;
    }
    return HwInfoConfig::get(hwInfo->platform.eProductFamily)->isPageTableManagerSupported(*hwInfo);
--- a/shared/source/command_stream/command_stream_receiver.h
+++ b/shared/source/command_stream/command_stream_receiver.h
@@ -104,7 +104,7 @@ class CommandStreamReceiver {
    ResidencyContainer &getEvictionAllocations();

    virtual GmmPageTableMngr *createPageTableManager() { return nullptr; }
-    bool needsPageTableManager(aub_stream::EngineType engineType) const;
+    bool needsPageTableManager(EngineUsage engineTypeUsage) const;

    void waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationUsage);
    MOCKABLE_VIRTUAL void waitForTaskCountAndCleanTemporaryAllocationList(uint32_t requiredTaskCount);
@@ -274,6 +274,8 @@ class CommandStreamReceiver {
        return activePartitions;
    }

+    std::unique_ptr<GmmPageTableMngr> pageTableManager;
+
  protected:
    void cleanupResources();
    void printDeviceIndex();
--- a/shared/source/command_stream/command_stream_receiver_hw_base.inl
+++ b/shared/source/command_stream/command_stream_receiver_hw_base.inl
@@ -340,8 +340,8 @@ CompletionStamp CommandStreamReceiverHw<GfxFamily>::flushTask(
    }

    programEngineModeCommands(commandStreamCSR, dispatchFlags);
-    if (executionEnvironment.rootDeviceEnvironments[device.getRootDeviceIndex()]->pageTableManager.get() && !pageTableManagerInitialized) {
-        pageTableManagerInitialized = executionEnvironment.rootDeviceEnvironments[device.getRootDeviceIndex()]->pageTableManager->initPageTableManagerRegisters(this);
+    if (pageTableManager.get() && !pageTableManagerInitialized) {
+        pageTableManagerInitialized = pageTableManager->initPageTableManagerRegisters(this);
    }

    programHardwareContext(commandStreamCSR);
@@ -1041,6 +1041,10 @@ uint32_t CommandStreamReceiverHw<GfxFamily>::blitBuffer(const BlitPropertiesCont

    programEnginePrologue(commandStream);

+    if (pageTableManager.get() && !pageTableManagerInitialized) {
+        pageTableManagerInitialized = pageTableManager->initPageTableManagerRegisters(this);
+    }
+
    for (auto &blitProperties : blitPropertiesContainer) {
        TimestampPacketHelper::programCsrDependenciesForTimestampPacketContainer<GfxFamily>(commandStream, blitProperties.csrDependencies);
        TimestampPacketHelper::programCsrDependenciesForForTaskCountContainer<GfxFamily>(commandStream, blitProperties.csrDependencies);
--- a/shared/source/device/device.cpp
+++ b/shared/source/device/device.cpp
@@ -317,7 +317,7 @@ bool Device::createEngine(uint32_t deviceCsrIndex, EngineTypeUsage engineTypeUsa
        commandStreamReceiver->initializeDefaultsForInternalEngine();
    }

-    if (commandStreamReceiver->needsPageTableManager(engineType)) {
+    if (commandStreamReceiver->needsPageTableManager(engineUsage)) {
        commandStreamReceiver->createPageTableManager();
    }

--- a/shared/source/execution_environment/root_device_environment.h
+++ b/shared/source/execution_environment/root_device_environment.h
@@ -65,7 +65,6 @@ struct RootDeviceEnvironment {
    std::unique_ptr<SipKernel> sipKernels[static_cast<uint32_t>(SipKernelType::COUNT)];
    std::unique_ptr<GmmHelper> gmmHelper;
    std::unique_ptr<OSInterface> osInterface;
-    std::unique_ptr<GmmPageTableMngr> pageTableManager;
    std::unique_ptr<MemoryOperationsHandler> memoryOperationsInterface;
    std::unique_ptr<AubCenter> aubCenter;
    std::unique_ptr<BindlessHeapsHelper> bindlessHeapsHelper;
--- a/shared/source/gmm_helper/page_table_mngr_impl.cpp
+++ b/shared/source/gmm_helper/page_table_mngr_impl.cpp
@@ -23,7 +23,7 @@ bool GmmPageTableMngr::updateAuxTable(uint64_t gpuVa, Gmm *gmm, bool map) {
    ddiUpdateAuxTable.BaseGpuVA = gpuVa;
    ddiUpdateAuxTable.BaseResInfo = gmm->gmmResourceInfo->peekGmmResourceInfo();
    ddiUpdateAuxTable.DoNotWait = true;
-    ddiUpdateAuxTable.Map = map ? 1u : 0u;
+    ddiUpdateAuxTable.Map = map;

    return updateAuxTable(&ddiUpdateAuxTable) == GMM_STATUS::GMM_SUCCESS;
 }
--- a/shared/source/memory_manager/memory_manager.cpp
+++ b/shared/source/memory_manager/memory_manager.cpp
@@ -495,11 +495,16 @@ GraphicsAllocation *MemoryManager::allocateInternalGraphicsMemoryWithHostCopy(ui
 }

 bool MemoryManager::mapAuxGpuVA(GraphicsAllocation *graphicsAllocation) {
-    auto index = graphicsAllocation->getRootDeviceIndex();
-    if (executionEnvironment.rootDeviceEnvironments[index]->pageTableManager.get()) {
-        return executionEnvironment.rootDeviceEnvironments[index]->pageTableManager->updateAuxTable(graphicsAllocation->getGpuAddress(), graphicsAllocation->getDefaultGmm(), true);
+    bool ret = false;
+    for (auto engine : registeredEngines) {
+        if (engine.commandStreamReceiver->pageTableManager.get()) {
+            ret = engine.commandStreamReceiver->pageTableManager->updateAuxTable(graphicsAllocation->getGpuAddress(), graphicsAllocation->getDefaultGmm(), true);
+            if (!ret) {
+                break;
+            }
+        }
    }
-    return false;
+    return ret;
 }

 GraphicsAllocation *MemoryManager::allocateGraphicsMemory(const AllocationData &allocationData) {
--- a/shared/source/os_interface/linux/drm_command_stream.h
+++ b/shared/source/os_interface/linux/drm_command_stream.h
@@ -63,6 +63,8 @@ class DrmCommandStreamReceiver : public DeviceCommandStreamReceiver<GfxFamily> {

    void printBOsForSubmit(ResidencyContainer &allocationsForResidency, GraphicsAllocation &cmdBufferAllocation);

+    using CommandStreamReceiver::pageTableManager;
+
  protected:
    MOCKABLE_VIRTUAL void flushInternal(const BatchBuffer &batchBuffer, const ResidencyContainer &allocationsForResidency);
    MOCKABLE_VIRTUAL void exec(const BatchBuffer &batchBuffer, uint32_t vmHandleId, uint32_t drmContextId);
--- a/shared/source/os_interface/linux/drm_command_stream.inl
+++ b/shared/source/os_interface/linux/drm_command_stream.inl
@@ -222,7 +222,7 @@ GmmPageTableMngr *DrmCommandStreamReceiver<GfxFamily>::createPageTableManager()
    auto gmmPageTableMngr = GmmPageTableMngr::create(gmmClientContext, TT_TYPE::AUXTT, nullptr);
    gmmPageTableMngr->setCsrHandle(this);

-    rootDeviceEnvironment->pageTableManager.reset(gmmPageTableMngr);
+    this->pageTableManager.reset(gmmPageTableMngr);

    return gmmPageTableMngr;
 }
--- a/shared/source/os_interface/windows/wddm/wddm.cpp
+++ b/shared/source/os_interface/windows/wddm/wddm.cpp
@@ -7,6 +7,7 @@

 #include "shared/source/os_interface/windows/wddm/wddm.h"

+#include "shared/source/command_stream/command_stream_receiver.h"
 #include "shared/source/command_stream/preemption.h"
 #include "shared/source/execution_environment/execution_environment.h"
 #include "shared/source/execution_environment/root_device_environment.h"
@@ -452,12 +453,16 @@ bool Wddm::mapGpuVirtualAddress(Gmm *gmm, D3DKMT_HANDLE handle, D3DGPU_VIRTUAL_A
    }

    kmDafListener->notifyMapGpuVA(featureTable->ftrKmdDaf, getAdapter(), device, handle, MapGPUVA.VirtualAddress, getGdi()->escape);
-
-    if (gmm->isCompressionEnabled && rootDeviceEnvironment.pageTableManager.get()) {
-        return rootDeviceEnvironment.pageTableManager->updateAuxTable(gpuPtr, gmm, true);
+    bool ret = true;
+    if (gmm->isCompressionEnabled && HwInfoConfig::get(gfxPlatform->eProductFamily)->isPageTableManagerSupported(*rootDeviceEnvironment.getHardwareInfo())) {
+        for (auto engine : rootDeviceEnvironment.executionEnvironment.memoryManager.get()->getRegisteredEngines()) {
+            if (engine.commandStreamReceiver->pageTableManager.get()) {
+                ret &= engine.commandStreamReceiver->pageTableManager->updateAuxTable(gpuPtr, gmm, true);
+            }
+        }
    }

-    return true;
+    return ret;
 }

 D3DGPU_VIRTUAL_ADDRESS Wddm::reserveGpuVirtualAddress(D3DGPU_VIRTUAL_ADDRESS minimumAddress,
--- a/shared/source/os_interface/windows/wddm_device_command_stream.h
+++ b/shared/source/os_interface/windows/wddm_device_command_stream.h
@@ -35,6 +35,8 @@ class WddmCommandStreamReceiver : public DeviceCommandStreamReceiver<GfxFamily>
    }
    GmmPageTableMngr *createPageTableManager() override;

+    using CommandStreamReceiver::pageTableManager;
+
  protected:
    void kmDafLockAllocations(ResidencyContainer &allocationsForResidency);

--- a/shared/source/os_interface/windows/wddm_device_command_stream.inl
+++ b/shared/source/os_interface/windows/wddm_device_command_stream.inl
@@ -144,7 +144,7 @@ GmmPageTableMngr *WddmCommandStreamReceiver<GfxFamily>::createPageTableManager()

    GmmPageTableMngr *gmmPageTableMngr = GmmPageTableMngr::create(rootDeviceEnvironment->getGmmClientContext(), TT_TYPE::AUXTT, &ttCallbacks);
    gmmPageTableMngr->setCsrHandle(this);
-    rootDeviceEnvironment->pageTableManager.reset(gmmPageTableMngr);
+    this->pageTableManager.reset(gmmPageTableMngr);
    return gmmPageTableMngr;
 }

--- a/shared/source/os_interface/windows/wddm_memory_manager.cpp
+++ b/shared/source/os_interface/windows/wddm_memory_manager.cpp
@@ -23,6 +23,7 @@
 #include "shared/source/memory_manager/deferred_deleter.h"
 #include "shared/source/memory_manager/host_ptr_manager.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
+#include "shared/source/os_interface/hw_info_config.h"
 #include "shared/source/os_interface/os_interface.h"
 #include "shared/source/os_interface/windows/os_context_win.h"
 #include "shared/source/os_interface/windows/wddm/wddm.h"
@@ -486,11 +487,13 @@ void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation
    }

    auto defaultGmm = gfxAllocation->getDefaultGmm();
-    if (defaultGmm) {
-        auto index = gfxAllocation->getRootDeviceIndex();
-        if (defaultGmm->isCompressionEnabled && executionEnvironment.rootDeviceEnvironments[index]->pageTableManager.get()) {
-            [[maybe_unused]] auto status = executionEnvironment.rootDeviceEnvironments[index]->pageTableManager->updateAuxTable(input->getGpuAddress(), defaultGmm, false);
-            DEBUG_BREAK_IF(!status);
+    auto hwInfo = executionEnvironment.rootDeviceEnvironments[gfxAllocation->getRootDeviceIndex()]->getHardwareInfo();
+    if (defaultGmm && defaultGmm->isCompressionEnabled && HwInfoConfig::get(hwInfo->platform.eProductFamily)->isPageTableManagerSupported(*hwInfo)) {
+        for (auto engine : registeredEngines) {
+            if (engine.commandStreamReceiver->pageTableManager.get()) {
+                [[maybe_unused]] auto status = engine.commandStreamReceiver->pageTableManager->updateAuxTable(input->getGpuAddress(), defaultGmm, false);
+                DEBUG_BREAK_IF(!status);
+            }
        }
    }
    for (auto handleId = 0u; handleId < gfxAllocation->getNumGmms(); handleId++) {