feature: extend TBX page fault manager from CPU implementation

In TBX mode, the host could not write to host buffers after access from device code due to the lack of a migration mechanism post-initial TBX upload. Migration is unnecessary with real hardware, but required for TBX. This patch introduces a new page fault manager type that extends the original CPU fault manager, enabling automatic migration of host buffers in TBX mode. Refactoring was necessary to avoid diamond inheritance, achieved by using a template parameter as the base class for OS-specific fault managers. Related-To: NEO-12268 Signed-off-by: Jack Myers <jack.myers@intel.com>
2025-12-21 01:04:57 +08:00 · 2024-12-03 00:28:34 +00:00
parent 924ad580bd
commit 51c0e80299
36 changed files with 591 additions and 136 deletions
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.h
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.h
@@ -18,6 +18,7 @@ namespace NEO {

 class AubSubCaptureManager;
 class TbxStream;
+class CpuPageFaultManager;

 template <typename GfxFamily>
 class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFamily> {
@@ -36,6 +37,9 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
        return 2000; // 2s
    }

+    void allowCPUMemoryAccessIfHostBuffer(AllocationType allocType, void *cpuAddress, size_t size);
+    void protectCPUMemoryAccessIfHostBuffer(AllocationType allocType, void *cpuAddress, size_t size);
+
  public:
    using CommandStreamReceiverSimulatedCommonHw<GfxFamily>::initAdditionalMMIO;
    using CommandStreamReceiverSimulatedCommonHw<GfxFamily>::aubManager;
@@ -81,10 +85,12 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa

    void initializeEngine() override;

-    MemoryManager *getMemoryManager() {
+    MOCKABLE_VIRTUAL MemoryManager *getMemoryManager() {
        return CommandStreamReceiver::getMemoryManager();
    }

+    MOCKABLE_VIRTUAL CpuPageFaultManager *getTbxPageFaultManager();
+
    TbxStream tbxStream;
    std::unique_ptr<AubSubCaptureManager> subCaptureManager;
    uint32_t aubDeviceId;
--- a/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
+++ b/shared/source/command_stream/tbx_command_stream_receiver_hw.inl
@@ -26,14 +26,22 @@
 #include "shared/source/helpers/gfx_core_helper.h"
 #include "shared/source/helpers/hw_info.h"
 #include "shared/source/helpers/ptr_math.h"
+#include "shared/source/memory_manager/allocation_type.h"
 #include "shared/source/memory_manager/graphics_allocation.h"
+#include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/os_interface/aub_memory_operations_handler.h"
 #include "shared/source/os_interface/product_helper.h"
+#include "shared/source/page_fault_manager/tbx_page_fault_manager.h"

 #include <cstring>

 namespace NEO {

+template <typename GfxFamily>
+CpuPageFaultManager *TbxCommandStreamReceiverHw<GfxFamily>::getTbxPageFaultManager() {
+    return this->getMemoryManager()->getPageFaultManager();
+}
+
 template <typename GfxFamily>
 TbxCommandStreamReceiverHw<GfxFamily>::TbxCommandStreamReceiverHw(ExecutionEnvironment &executionEnvironment,
                                                                  uint32_t rootDeviceIndex,
@@ -74,6 +82,26 @@ TbxCommandStreamReceiverHw<GfxFamily>::~TbxCommandStreamReceiverHw() {
    this->freeEngineInfo(gttRemap);
 }

+template <typename GfxFamily>
+void TbxCommandStreamReceiverHw<GfxFamily>::allowCPUMemoryAccessIfHostBuffer(AllocationType allocType, void *cpuAddress, size_t size) {
+    if (allocType == AllocationType::bufferHostMemory) {
+        auto faultManager = getTbxPageFaultManager();
+        if (faultManager != nullptr) {
+            faultManager->allowCPUMemoryAccess(cpuAddress, size);
+        }
+    }
+}
+
+template <typename GfxFamily>
+void TbxCommandStreamReceiverHw<GfxFamily>::protectCPUMemoryAccessIfHostBuffer(AllocationType allocType, void *cpuAddress, size_t size) {
+    if (allocType == AllocationType::bufferHostMemory) {
+        auto faultManager = getTbxPageFaultManager();
+        if (faultManager != nullptr) {
+            faultManager->protectCPUMemoryAccess(cpuAddress, size);
+        }
+    }
+}
+
 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::initializeEngine() {
    isEngineInitialized = true;
@@ -431,21 +459,32 @@ void TbxCommandStreamReceiverHw<GfxFamily>::writeMemory(uint64_t gpuAddress, voi

 template <typename GfxFamily>
 bool TbxCommandStreamReceiverHw<GfxFamily>::writeMemory(GraphicsAllocation &gfxAllocation, bool isChunkCopy, uint64_t gpuVaChunkOffset, size_t chunkSize) {
+    uint64_t gpuAddress;
+    void *cpuAddress;
+    size_t size;
+
+    if (!this->getParametersForMemory(gfxAllocation, gpuAddress, cpuAddress, size)) {
+        return false;
+    }
+
+    auto allocType = gfxAllocation.getAllocationType();
+    if (allocType == AllocationType::bufferHostMemory) {
+        auto faultManager = getTbxPageFaultManager();
+        if (faultManager != nullptr) {
+            faultManager->insertAllocation(this, &gfxAllocation, cpuAddress, size);
+        }
+    }
+
    if (!this->isTbxWritable(gfxAllocation)) {
        return false;
    }

+    this->allowCPUMemoryAccessIfHostBuffer(allocType, cpuAddress, size);
+
    if (!isEngineInitialized) {
        initializeEngine();
    }

-    uint64_t gpuAddress;
-    void *cpuAddress;
-    size_t size;
-    if (!this->getParametersForMemory(gfxAllocation, gpuAddress, cpuAddress, size)) {
-        return false;
-    }
-
    if (aubManager) {
        this->writeMemoryWithAubManager(gfxAllocation, isChunkCopy, gpuVaChunkOffset, chunkSize);
    } else {
@@ -460,6 +499,7 @@ bool TbxCommandStreamReceiverHw<GfxFamily>::writeMemory(GraphicsAllocation &gfxA
    if (AubHelper::isOneTimeAubWritableAllocationType(gfxAllocation.getAllocationType())) {
        this->setTbxWritable(false, gfxAllocation);
    }
+    this->protectCPUMemoryAccessIfHostBuffer(allocType, cpuAddress, size);

    return true;
 }
@@ -533,6 +573,7 @@ void TbxCommandStreamReceiverHw<GfxFamily>::processEviction() {

 template <typename GfxFamily>
 SubmissionStatus TbxCommandStreamReceiverHw<GfxFamily>::processResidency(ResidencyContainer &allocationsForResidency, uint32_t handleId) {
+
    for (auto &gfxAllocation : allocationsForResidency) {
        if (dumpTbxNonWritable) {
            this->setTbxWritable(true, *gfxAllocation);
@@ -554,12 +595,16 @@ SubmissionStatus TbxCommandStreamReceiverHw<GfxFamily>::processResidency(Residen

 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::downloadAllocationTbx(GraphicsAllocation &gfxAllocation) {
+
    uint64_t gpuAddress = 0;
    void *cpuAddress = nullptr;
    size_t size = 0;

    this->getParametersForMemory(gfxAllocation, gpuAddress, cpuAddress, size);

+    auto allocType = gfxAllocation.getAllocationType();
+    this->allowCPUMemoryAccessIfHostBuffer(allocType, cpuAddress, size);
+
    if (hardwareContextController) {
        hardwareContextController->readMemory(gpuAddress, cpuAddress, size,
                                              this->getMemoryBank(&gfxAllocation), gfxAllocation.getUsedPageSize());
@@ -671,6 +716,12 @@ void TbxCommandStreamReceiverHw<GfxFamily>::dumpAllocation(GraphicsAllocation &g
 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::removeDownloadAllocation(GraphicsAllocation *alloc) {
    auto lockCSR = this->obtainUniqueOwnership();
+
    this->allocationsForDownload.erase(alloc);
+
+    auto faultManager = getTbxPageFaultManager();
+    if (faultManager != nullptr) {
+        faultManager->removeAllocation(alloc);
+    }
 }
 } // namespace NEO
--- a/shared/source/debug_settings/debug_settings_manager.h
+++ b/shared/source/debug_settings/debug_settings_manager.h
@@ -6,6 +6,7 @@
 */

 #pragma once
+#include "shared/source/helpers/options.h"
 #include "shared/source/helpers/string.h"
 #include "shared/source/utilities/io_functions.h"

@@ -161,6 +162,11 @@ class DebugSettingsManager {
        }
    }

+    inline bool isTbxMode() {
+        auto setCsr = flags.SetCommandStreamReceiver.get();
+        return (setCsr == static_cast<int32_t>(CommandStreamReceiverType::tbx)) || (setCsr == static_cast<int32_t>(CommandStreamReceiverType::tbxWithAub));
+    }
+
  protected:
    std::unique_ptr<SettingsReader> readerImpl;
    bool isLoopAtDriverInitEnabled() const {
--- a/shared/source/helpers/debug_helpers.cpp
+++ b/shared/source/helpers/debug_helpers.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -9,6 +9,7 @@

 #include "shared/source/debug_settings/debug_settings_manager.h"
 #include "shared/source/helpers/abort.h"
+#include "shared/source/helpers/options.h"

 #include <cassert>
 #include <cstdio>
--- a/shared/source/memory_manager/memory_manager.cpp
+++ b/shared/source/memory_manager/memory_manager.cpp
@@ -83,9 +83,11 @@ MemoryManager::MemoryManager(ExecutionEnvironment &executionEnvironment) : execu
        localMemAllocsSize[rootDeviceIndex].store(0u);
    }

-    if (anyLocalMemorySupported) {
-        pageFaultManager = PageFaultManager::create();
-        prefetchManager = PrefetchManager::create();
+    if (anyLocalMemorySupported || debugManager.isTbxMode()) {
+        pageFaultManager = CpuPageFaultManager::create();
+        if (anyLocalMemorySupported) {
+            prefetchManager = PrefetchManager::create();
+        }
    }

    if (debugManager.flags.EnableMultiStorageResources.get() != -1) {
--- a/shared/source/memory_manager/memory_manager.h
+++ b/shared/source/memory_manager/memory_manager.h
@@ -29,7 +29,7 @@ namespace NEO {
 using SubDeviceIdsVec = StackVec<uint32_t, 4>;

 class MultiGraphicsAllocation;
-class PageFaultManager;
+class CpuPageFaultManager;
 class GfxPartition;
 struct ImageInfo;
 struct AllocationData;
@@ -192,7 +192,7 @@ class MemoryManager {
        return deferredDeleter.get();
    }

-    PageFaultManager *getPageFaultManager() const {
+    MOCKABLE_VIRTUAL CpuPageFaultManager *getPageFaultManager() const {
        return pageFaultManager.get();
    }

@@ -410,7 +410,7 @@ class MemoryManager {
    std::vector<std::unique_ptr<LocalMemoryUsageBankSelector>> internalLocalMemoryUsageBankSelector;
    std::vector<std::unique_ptr<LocalMemoryUsageBankSelector>> externalLocalMemoryUsageBankSelector;
    void *reservedMemory = nullptr;
-    std::unique_ptr<PageFaultManager> pageFaultManager;
+    std::unique_ptr<CpuPageFaultManager> pageFaultManager;
    std::unique_ptr<PrefetchManager> prefetchManager;
    OSMemory::ReservedCpuAddressRange reservedCpuAddressRange;
    std::vector<std::unique_ptr<HeapAssigner>> heapAssigners;
--- a/shared/source/page_fault_manager/CMakeLists.txt
+++ b/shared/source/page_fault_manager/CMakeLists.txt
@@ -1,5 +1,5 @@
 #
-# Copyright (C) 2019-2020 Intel Corporation
+# Copyright (C) 2019-2024 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
 #
@@ -8,6 +8,8 @@ set(NEO_CORE_PAGE_FAULT_MANAGER
    ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
    ${CMAKE_CURRENT_SOURCE_DIR}/cpu_page_fault_manager.cpp
    ${CMAKE_CURRENT_SOURCE_DIR}/cpu_page_fault_manager.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/tbx_page_fault_manager.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/tbx_page_fault_manager.h
 )

 set_property(GLOBAL PROPERTY NEO_CORE_PAGE_FAULT_MANAGER ${NEO_CORE_PAGE_FAULT_MANAGER})
--- a/shared/source/page_fault_manager/cpu_page_fault_manager.cpp
+++ b/shared/source/page_fault_manager/cpu_page_fault_manager.cpp
@@ -16,19 +16,26 @@
 #include <algorithm>

 namespace NEO {
-void PageFaultManager::insertAllocation(void *ptr, size_t size, SVMAllocsManager *unifiedMemoryManager, void *cmdQ, const MemoryProperties &memoryProperties) {
+
+void CpuPageFaultManager::insertAllocation(void *ptr, size_t size, SVMAllocsManager *unifiedMemoryManager, void *cmdQ, const MemoryProperties &memoryProperties) {
    auto initialPlacement = MemoryPropertiesHelper::getUSMInitialPlacement(memoryProperties);
    const auto domain = (initialPlacement == GraphicsAllocation::UsmInitialPlacement::CPU) ? AllocationDomain::cpu : AllocationDomain::none;

    std::unique_lock<SpinLock> lock{mtx};
-    this->memoryData.insert(std::make_pair(ptr, PageFaultData{size, unifiedMemoryManager, cmdQ, domain}));
+    PageFaultData faultData{};
+    faultData.faultType = FaultMode::cpu;
+    faultData.size = size;
+    faultData.unifiedMemoryManager = unifiedMemoryManager;
+    faultData.cmdQ = cmdQ;
+    faultData.domain = domain;
+    this->memoryData.insert(std::make_pair(ptr, faultData));
    if (initialPlacement != GraphicsAllocation::UsmInitialPlacement::CPU) {
        this->protectCPUMemoryAccess(ptr, size);
    }
    unifiedMemoryManager->nonGpuDomainAllocs.push_back(ptr);
 }

-void PageFaultManager::removeAllocation(void *ptr) {
+void CpuPageFaultManager::removeAllocation(void *ptr) {
    std::unique_lock<SpinLock> lock{mtx};
    auto alloc = memoryData.find(ptr);
    if (alloc != memoryData.end()) {
@@ -45,7 +52,7 @@ void PageFaultManager::removeAllocation(void *ptr) {
    }
 }

-void PageFaultManager::moveAllocationToGpuDomain(void *ptr) {
+void CpuPageFaultManager::moveAllocationToGpuDomain(void *ptr) {
    std::unique_lock<SpinLock> lock{mtx};
    auto alloc = memoryData.find(ptr);
    if (alloc != memoryData.end()) {
@@ -61,7 +68,7 @@ void PageFaultManager::moveAllocationToGpuDomain(void *ptr) {
    }
 }

-void PageFaultManager::moveAllocationsWithinUMAllocsManagerToGpuDomain(SVMAllocsManager *unifiedMemoryManager) {
+void CpuPageFaultManager::moveAllocationsWithinUMAllocsManagerToGpuDomain(SVMAllocsManager *unifiedMemoryManager) {
    std::unique_lock<SpinLock> lock{mtx};
    for (auto allocPtr : unifiedMemoryManager->nonGpuDomainAllocs) {
        auto &pageFaultData = this->memoryData[allocPtr];
@@ -70,7 +77,7 @@ void PageFaultManager::moveAllocationsWithinUMAllocsManagerToGpuDomain(SVMAllocs
    unifiedMemoryManager->nonGpuDomainAllocs.clear();
 }

-inline void PageFaultManager::migrateStorageToGpuDomain(void *ptr, PageFaultData &pageFaultData) {
+inline void CpuPageFaultManager::migrateStorageToGpuDomain(void *ptr, PageFaultData &pageFaultData) {
    if (pageFaultData.domain == AllocationDomain::cpu) {
        this->setCpuAllocEvictable(false, ptr, pageFaultData.unifiedMemoryManager);
        this->allowCPUMemoryEviction(false, ptr, pageFaultData);
@@ -96,15 +103,19 @@ inline void PageFaultManager::migrateStorageToGpuDomain(void *ptr, PageFaultData
    pageFaultData.domain = AllocationDomain::gpu;
 }

-bool PageFaultManager::verifyAndHandlePageFault(void *ptr, bool handlePageFault) {
+void CpuPageFaultManager::handlePageFault(void *ptr, PageFaultData &faultData) {
+    this->setAubWritable(true, ptr, faultData.unifiedMemoryManager);
+    gpuDomainHandler(this, ptr, faultData);
+}
+
+bool CpuPageFaultManager::verifyAndHandlePageFault(void *ptr, bool handleFault) {
    std::unique_lock<SpinLock> lock{mtx};
-    for (auto &alloc : this->memoryData) {
+    for (auto &alloc : memoryData) {
        auto allocPtr = alloc.first;
        auto &pageFaultData = alloc.second;
        if (ptr >= allocPtr && ptr < ptrOffset(allocPtr, pageFaultData.size)) {
-            if (handlePageFault) {
-                this->setAubWritable(true, allocPtr, pageFaultData.unifiedMemoryManager);
-                gpuDomainHandler(this, allocPtr, pageFaultData);
+            if (handleFault) {
+                handlePageFault(allocPtr, pageFaultData);
            }
            return true;
        }
@@ -112,23 +123,23 @@ bool PageFaultManager::verifyAndHandlePageFault(void *ptr, bool handlePageFault)
    return false;
 }

-void PageFaultManager::setGpuDomainHandler(gpuDomainHandlerFunc gpuHandlerFuncPtr) {
+void CpuPageFaultManager::setGpuDomainHandler(gpuDomainHandlerFunc gpuHandlerFuncPtr) {
    this->gpuDomainHandler = gpuHandlerFuncPtr;
 }

-void PageFaultManager::transferAndUnprotectMemory(PageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
+void CpuPageFaultManager::transferAndUnprotectMemory(CpuPageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
    pageFaultHandler->migrateStorageToCpuDomain(allocPtr, pageFaultData);
    pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
    pageFaultHandler->setCpuAllocEvictable(true, allocPtr, pageFaultData.unifiedMemoryManager);
    pageFaultHandler->allowCPUMemoryEviction(true, allocPtr, pageFaultData);
 }

-void PageFaultManager::unprotectAndTransferMemory(PageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
+void CpuPageFaultManager::unprotectAndTransferMemory(CpuPageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
    pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
    pageFaultHandler->migrateStorageToCpuDomain(allocPtr, pageFaultData);
 }

-inline void PageFaultManager::migrateStorageToCpuDomain(void *ptr, PageFaultData &pageFaultData) {
+inline void CpuPageFaultManager::migrateStorageToCpuDomain(void *ptr, PageFaultData &pageFaultData) {
    if (pageFaultData.domain == AllocationDomain::gpu) {
        std::chrono::steady_clock::time_point start;
        std::chrono::steady_clock::time_point end;
@@ -144,19 +155,19 @@ inline void PageFaultManager::migrateStorageToCpuDomain(void *ptr, PageFaultData
    pageFaultData.domain = AllocationDomain::cpu;
 }

-void PageFaultManager::selectGpuDomainHandler() {
+void CpuPageFaultManager::selectGpuDomainHandler() {
    if (debugManager.flags.SetCommandStreamReceiver.get() > static_cast<int32_t>(CommandStreamReceiverType::hardware) || debugManager.flags.NEO_CAL_ENABLED.get()) {
-        this->gpuDomainHandler = &PageFaultManager::unprotectAndTransferMemory;
+        this->gpuDomainHandler = &CpuPageFaultManager::unprotectAndTransferMemory;
    }
 }

-void PageFaultManager::setAubWritable(bool writable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
+void CpuPageFaultManager::setAubWritable(bool writable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
    UNRECOVERABLE_IF(ptr == nullptr);
    auto gpuAlloc = unifiedMemoryManager->getSVMAlloc(ptr)->gpuAllocations.getDefaultGraphicsAllocation();
    gpuAlloc->setAubWritable(writable, GraphicsAllocation::allBanks);
 }

-void PageFaultManager::setCpuAllocEvictable(bool evictable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
+void CpuPageFaultManager::setCpuAllocEvictable(bool evictable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
    UNRECOVERABLE_IF(ptr == nullptr);
    auto cpuAlloc = unifiedMemoryManager->getSVMAlloc(ptr)->cpuAllocation;
    cpuAlloc->setEvictable(evictable);
--- a/shared/source/page_fault_manager/cpu_page_fault_manager.h
+++ b/shared/source/page_fault_manager/cpu_page_fault_manager.h
@@ -21,16 +21,26 @@ class Device;
 class SVMAllocsManager;
 class OSInterface;

-class PageFaultManager : public NonCopyableOrMovableClass {
+class CpuPageFaultManager : public NonCopyableClass {
  public:
-    static std::unique_ptr<PageFaultManager> create();
+    static std::unique_ptr<CpuPageFaultManager> create();

-    virtual ~PageFaultManager() = default;
+    virtual ~CpuPageFaultManager() = default;
+
+    virtual void allowCPUMemoryAccess(void *ptr, size_t size) = 0;
+    virtual void protectCPUMemoryAccess(void *ptr, size_t size) = 0;

    MOCKABLE_VIRTUAL void moveAllocationToGpuDomain(void *ptr);
    MOCKABLE_VIRTUAL void moveAllocationsWithinUMAllocsManagerToGpuDomain(SVMAllocsManager *unifiedMemoryManager);
    void insertAllocation(void *ptr, size_t size, SVMAllocsManager *unifiedMemoryManager, void *cmdQ, const MemoryProperties &memoryProperties);
    void removeAllocation(void *ptr);
+    virtual void insertAllocation(CommandStreamReceiver *csr, GraphicsAllocation *alloc, void *ptr, size_t size) {}
+    virtual void removeAllocation(GraphicsAllocation *alloc) {}
+
+    enum class FaultMode {
+        cpu,
+        tbx
+    };

    enum class AllocationDomain {
        none,
@@ -39,39 +49,49 @@ class PageFaultManager : public NonCopyableOrMovableClass {
    };

    struct PageFaultData {
-        size_t size;
-        SVMAllocsManager *unifiedMemoryManager;
-        void *cmdQ;
-        AllocationDomain domain;
+        FaultMode faultType = FaultMode::cpu;
+        size_t size = 0;
+
+        // cpu fault data
+        SVMAllocsManager *unifiedMemoryManager = nullptr;
+        void *cmdQ = nullptr;
+        AllocationDomain domain = AllocationDomain::none;
+
+        // tbx fault data
+        GraphicsAllocation *gfxAllocation = nullptr;
+        CommandStreamReceiver *csr = nullptr;
    };

-    typedef void (*gpuDomainHandlerFunc)(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
+    typedef void (*gpuDomainHandlerFunc)(CpuPageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);

    void setGpuDomainHandler(gpuDomainHandlerFunc gpuHandlerFuncPtr);

-    virtual void allowCPUMemoryAccess(void *ptr, size_t size) = 0;
-    virtual void protectCPUMemoryAccess(void *ptr, size_t size) = 0;
    MOCKABLE_VIRTUAL void transferToCpu(void *ptr, size_t size, void *cmdQ);

  protected:
-    virtual bool checkFaultHandlerFromPageFaultManager() = 0;
-    virtual void registerFaultHandler() = 0;
    virtual void evictMemoryAfterImplCopy(GraphicsAllocation *allocation, Device *device) = 0;
    virtual void allowCPUMemoryEvictionImpl(bool evict, void *ptr, CommandStreamReceiver &csr, OSInterface *osInterface) = 0;

-    MOCKABLE_VIRTUAL bool verifyAndHandlePageFault(void *ptr, bool handlePageFault);
+    virtual bool checkFaultHandlerFromPageFaultManager() = 0;
+    virtual void registerFaultHandler() = 0;
+
+    virtual bool verifyAndHandlePageFault(void *ptr, bool handlePageFault);
+
+    virtual void handlePageFault(void *ptr, PageFaultData &faultData);
+
    MOCKABLE_VIRTUAL void transferToGpu(void *ptr, void *cmdQ);
    MOCKABLE_VIRTUAL void setAubWritable(bool writable, void *ptr, SVMAllocsManager *unifiedMemoryManager);
    MOCKABLE_VIRTUAL void setCpuAllocEvictable(bool evictable, void *ptr, SVMAllocsManager *unifiedMemoryManager);
    MOCKABLE_VIRTUAL void allowCPUMemoryEviction(bool evict, void *ptr, PageFaultData &pageFaultData);

-    static void transferAndUnprotectMemory(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
-    static void unprotectAndTransferMemory(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
+    static void transferAndUnprotectMemory(CpuPageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
+    static void unprotectAndTransferMemory(CpuPageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
    void selectGpuDomainHandler();
    inline void migrateStorageToGpuDomain(void *ptr, PageFaultData &pageFaultData);
    inline void migrateStorageToCpuDomain(void *ptr, PageFaultData &pageFaultData);

-    decltype(&transferAndUnprotectMemory) gpuDomainHandler = &transferAndUnprotectMemory;
+    using gpuDomainHandlerType = decltype(&transferAndUnprotectMemory);
+    gpuDomainHandlerType gpuDomainHandler = &transferAndUnprotectMemory;

    std::unordered_map<void *, PageFaultData> memoryData;
    SpinLock mtx;
--- a/shared/source/page_fault_manager/linux/CMakeLists.txt
+++ b/shared/source/page_fault_manager/linux/CMakeLists.txt
@@ -1,5 +1,5 @@
 #
-# Copyright (C) 2019-2020 Intel Corporation
+# Copyright (C) 2019-2024 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
 #
@@ -8,6 +8,8 @@ set(NEO_CORE_PAGE_FAULT_MANAGER_LINUX
    ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
    ${CMAKE_CURRENT_SOURCE_DIR}/cpu_page_fault_manager_linux.cpp
    ${CMAKE_CURRENT_SOURCE_DIR}/cpu_page_fault_manager_linux.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/tbx_page_fault_manager_linux.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/tbx_page_fault_manager_linux.h
 )

 set_property(GLOBAL PROPERTY NEO_CORE_PAGE_FAULT_MANAGER_LINUX ${NEO_CORE_PAGE_FAULT_MANAGER_LINUX})
--- a/shared/source/page_fault_manager/linux/cpu_page_fault_manager_linux.cpp
+++ b/shared/source/page_fault_manager/linux/cpu_page_fault_manager_linux.cpp
@@ -10,15 +10,20 @@
 #include "shared/source/debug_settings/debug_settings_manager.h"
 #include "shared/source/device/device.h"
 #include "shared/source/execution_environment/root_device_environment.h"
-#include "shared/source/helpers/debug_helpers.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
+#include "shared/source/page_fault_manager/linux/tbx_page_fault_manager_linux.h"

 #include <algorithm>
 #include <sys/mman.h>

 namespace NEO {
-std::unique_ptr<PageFaultManager> PageFaultManager::create() {
-    auto pageFaultManager = std::make_unique<PageFaultManagerLinux>();
+std::unique_ptr<CpuPageFaultManager> CpuPageFaultManager::create() {
+    auto pageFaultManager = [&]() -> std::unique_ptr<CpuPageFaultManager> {
+        if (debugManager.isTbxMode()) {
+            return TbxPageFaultManager::create();
+        }
+        return std::make_unique<PageFaultManagerLinux>();
+    }();

    pageFaultManager->selectGpuDomainHandler();
    return pageFaultManager;
--- a/shared/source/page_fault_manager/linux/cpu_page_fault_manager_linux.h
+++ b/shared/source/page_fault_manager/linux/cpu_page_fault_manager_linux.h
@@ -14,7 +14,7 @@
 #include <vector>

 namespace NEO {
-class PageFaultManagerLinux : public PageFaultManager {
+class PageFaultManagerLinux : public virtual CpuPageFaultManager {
  public:
    PageFaultManagerLinux();
    ~PageFaultManagerLinux() override;
@@ -41,4 +41,7 @@ class PageFaultManagerLinux : public PageFaultManager {
    bool evictMemoryAfterCopy = false;
    int handlerIndex = 0;
 };
+
+class CpuPageFaultManagerLinux final : public PageFaultManagerLinux {};
+
 } // namespace NEO
--- a/shared/source/page_fault_manager/linux/tbx_page_fault_manager_linux.cpp
+++ b/shared/source/page_fault_manager/linux/tbx_page_fault_manager_linux.cpp
@@ -0,0 +1,14 @@
+/*
+ * Copyright (C) 2019-2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/page_fault_manager/linux/tbx_page_fault_manager_linux.h"
+
+namespace NEO {
+std::unique_ptr<TbxPageFaultManager> TbxPageFaultManager::create() {
+    return std::make_unique<TbxPageFaultManagerLinux>();
+}
+} // namespace NEO
--- a/shared/source/page_fault_manager/linux/tbx_page_fault_manager_linux.h
+++ b/shared/source/page_fault_manager/linux/tbx_page_fault_manager_linux.h
@@ -0,0 +1,17 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/page_fault_manager/linux/cpu_page_fault_manager_linux.h"
+#include "shared/source/page_fault_manager/tbx_page_fault_manager.h"
+
+namespace NEO {
+
+class TbxPageFaultManagerLinux final : public TbxPageFaultManager, public PageFaultManagerLinux {};
+
+} // namespace NEO
--- a/shared/source/page_fault_manager/tbx_page_fault_manager.cpp
+++ b/shared/source/page_fault_manager/tbx_page_fault_manager.cpp
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/page_fault_manager/tbx_page_fault_manager.h"
+
+#include "shared/source/command_stream/command_stream_receiver.h"
+#include "shared/source/memory_manager/graphics_allocation.h"
+#include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
+
+namespace NEO {
+
+void TbxPageFaultManager::handlePageFault(void *ptr, PageFaultData &faultData) {
+    if (faultData.faultType == FaultMode::cpu) {
+        CpuPageFaultManager::handlePageFault(ptr, faultData);
+        return;
+    }
+    unprotectAndTransferMemoryTbx(this, ptr, faultData);
+}
+
+void TbxPageFaultManager::unprotectAndTransferMemoryTbx(TbxPageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
+    pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
+    pageFaultData.csr->downloadAllocation(*pageFaultData.gfxAllocation);
+    pageFaultData.gfxAllocation->setTbxWritable(true, GraphicsAllocation::allBanks);
+}
+
+void TbxPageFaultManager::removeAllocation(GraphicsAllocation *alloc) {
+    std::unique_lock<SpinLock> lock{mtx};
+    for (auto &data : memoryData) {
+        auto allocPtr = data.first;
+        auto faultData = data.second;
+        if (faultData.gfxAllocation == alloc) {
+            memoryData.erase(allocPtr);
+            this->allowCPUMemoryAccess(allocPtr, faultData.size);
+            return;
+        }
+    }
+}
+
+void TbxPageFaultManager::insertAllocation(CommandStreamReceiver *csr, GraphicsAllocation *alloc, void *ptr, size_t size) {
+    std::unique_lock<SpinLock> lock{mtx};
+
+    if (memoryData.find(ptr) == memoryData.end()) {
+        PageFaultData pageFaultData{};
+        pageFaultData.faultType = FaultMode::tbx;
+        pageFaultData.size = size;
+        pageFaultData.gfxAllocation = alloc;
+        pageFaultData.csr = csr;
+        memoryData[ptr] = pageFaultData;
+    }
+    this->protectCPUMemoryAccess(ptr, size);
+}
+
+} // namespace NEO
--- a/shared/source/page_fault_manager/tbx_page_fault_manager.h
+++ b/shared/source/page_fault_manager/tbx_page_fault_manager.h
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
+
+namespace NEO {
+class CommandStreamReceiver;
+class GraphicsAllocation;
+
+class TbxPageFaultManager : public virtual CpuPageFaultManager {
+  public:
+    static std::unique_ptr<TbxPageFaultManager> create();
+
+    using CpuPageFaultManager::insertAllocation;
+    using CpuPageFaultManager::removeAllocation;
+    void insertAllocation(CommandStreamReceiver *csr, GraphicsAllocation *alloc, void *ptr, size_t size) override;
+    void removeAllocation(GraphicsAllocation *alloc) override;
+
+    using CpuPageFaultManager::checkFaultHandlerFromPageFaultManager;
+    using CpuPageFaultManager::verifyAndHandlePageFault;
+
+  protected:
+    void handlePageFault(void *ptr, PageFaultData &faultData) override;
+    static void unprotectAndTransferMemoryTbx(TbxPageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
+};
+
+} // namespace NEO
--- a/shared/source/page_fault_manager/windows/CMakeLists.txt
+++ b/shared/source/page_fault_manager/windows/CMakeLists.txt
@@ -1,5 +1,5 @@
 #
-# Copyright (C) 2019-2020 Intel Corporation
+# Copyright (C) 2019-2024 Intel Corporation
 #
 # SPDX-License-Identifier: MIT
 #
@@ -8,6 +8,8 @@ set(NEO_CORE_PAGE_FAULT_MANAGER_WINDOWS
    ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
    ${CMAKE_CURRENT_SOURCE_DIR}/cpu_page_fault_manager_windows.cpp
    ${CMAKE_CURRENT_SOURCE_DIR}/cpu_page_fault_manager_windows.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/tbx_page_fault_manager_windows.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/tbx_page_fault_manager_windows.h
 )

 if(WIN32)
--- a/shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.cpp
+++ b/shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.cpp
@@ -8,21 +8,27 @@
 #include "shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.h"

 #include "shared/source/command_stream/command_stream_receiver.h"
+#include "shared/source/debug_settings/debug_settings_manager.h"
 #include "shared/source/device/device.h"
-#include "shared/source/helpers/debug_helpers.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
 #include "shared/source/os_interface/os_interface.h"
 #include "shared/source/os_interface/windows/os_context_win.h"
+#include "shared/source/page_fault_manager/windows/tbx_page_fault_manager_windows.h"

 namespace NEO {
-std::unique_ptr<PageFaultManager> PageFaultManager::create() {
-    auto pageFaultManager = std::make_unique<PageFaultManagerWindows>();
+std::unique_ptr<CpuPageFaultManager> CpuPageFaultManager::create() {
+    auto pageFaultManager = [&]() -> std::unique_ptr<CpuPageFaultManager> {
+        if (debugManager.isTbxMode()) {
+            return TbxPageFaultManager::create();
+        }
+        return std::make_unique<CpuPageFaultManagerWindows>();
+    }();

    pageFaultManager->selectGpuDomainHandler();
    return pageFaultManager;
 }

-std::function<LONG(struct _EXCEPTION_POINTERS *exceptionInfo)> PageFaultManagerWindows::pageFaultHandler;
+std::function<LONG(struct _EXCEPTION_POINTERS *exceptionInfo)> PageFaultManagerWindows::pageFaultHandler = nullptr;

 PageFaultManagerWindows::PageFaultManagerWindows() {
    PageFaultManagerWindows::registerFaultHandler();
@@ -70,7 +76,7 @@ void PageFaultManagerWindows::protectCPUMemoryAccess(void *ptr, size_t size) {
 void PageFaultManagerWindows::evictMemoryAfterImplCopy(GraphicsAllocation *allocation, Device *device) {}

 void PageFaultManagerWindows::allowCPUMemoryEvictionImpl(bool evict, void *ptr, CommandStreamReceiver &csr, OSInterface *osInterface) {
-    NEO::SvmAllocationData *allocData = memoryData[ptr].unifiedMemoryManager->getSVMAlloc(ptr);
+    NEO::SvmAllocationData *allocData = this->memoryData[ptr].unifiedMemoryManager->getSVMAlloc(ptr);
    UNRECOVERABLE_IF(allocData == nullptr);

    if (osInterface) {
--- a/shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.h
+++ b/shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.h
@@ -13,7 +13,8 @@
 #include <functional>

 namespace NEO {
-class PageFaultManagerWindows : public PageFaultManager {
+
+class PageFaultManagerWindows : public virtual CpuPageFaultManager {
  public:
    PageFaultManagerWindows();
    ~PageFaultManagerWindows() override;
@@ -34,4 +35,6 @@ class PageFaultManagerWindows : public PageFaultManager {
    PVOID previousHandler;
 };

+class CpuPageFaultManagerWindows final : public PageFaultManagerWindows {};
+
 } // namespace NEO
--- a/shared/source/page_fault_manager/windows/tbx_page_fault_manager_windows.cpp
+++ b/shared/source/page_fault_manager/windows/tbx_page_fault_manager_windows.cpp
@@ -0,0 +1,14 @@
+/*
+ * Copyright (C) 2019-2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/page_fault_manager/windows/tbx_page_fault_manager_windows.h"
+
+namespace NEO {
+std::unique_ptr<TbxPageFaultManager> TbxPageFaultManager::create() {
+    return std::make_unique<TbxPageFaultManagerWindows>();
+}
+} // namespace NEO
--- a/shared/source/page_fault_manager/windows/tbx_page_fault_manager_windows.h
+++ b/shared/source/page_fault_manager/windows/tbx_page_fault_manager_windows.h
@@ -0,0 +1,20 @@
+/*
+ * Copyright (C) 2024 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/page_fault_manager/tbx_page_fault_manager.h"
+#include "shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.h"
+
+// Known false positive for valid virtual inheritance
+#pragma warning(disable : 4250)
+
+namespace NEO {
+
+class TbxPageFaultManagerWindows final : public TbxPageFaultManager, public PageFaultManagerWindows {};
+
+} // namespace NEO
--- a/shared/test/common/fixtures/tbx_command_stream_fixture.cpp
+++ b/shared/test/common/fixtures/tbx_command_stream_fixture.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -10,6 +10,7 @@
 #include "shared/source/command_stream/command_stream_receiver.h"
 #include "shared/source/command_stream/tbx_command_stream_receiver_hw.h"
 #include "shared/source/memory_manager/os_agnostic_memory_manager.h"
+#include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/mocks/mock_device.h"

 #include "gtest/gtest.h"
@@ -18,6 +19,8 @@ namespace NEO {

 void TbxCommandStreamFixture::setUp(MockDevice *pDevice) {
    // Create our TBX command stream receiver based on HW type
+    DebugManagerStateRestore dbgRestore;
+    debugManager.flags.SetCommandStreamReceiver.set(static_cast<int32_t>(CommandStreamReceiverType::tbx));
    pCommandStreamReceiver = TbxCommandStreamReceiver::create("", false, *pDevice->executionEnvironment, pDevice->getRootDeviceIndex(), pDevice->getDeviceBitfield());
    ASSERT_NE(nullptr, pCommandStreamReceiver);
    memoryManager = new OsAgnosticMemoryManager(*pDevice->executionEnvironment);
--- a/shared/test/common/fixtures/tbx_command_stream_fixture.h
+++ b/shared/test/common/fixtures/tbx_command_stream_fixture.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2022 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -7,6 +7,9 @@

 #pragma once

+#include "shared/source/memory_manager/os_agnostic_memory_manager.h"
+#include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
+
 namespace NEO {

 class CommandStreamReceiver;
--- a/shared/test/common/mocks/mock_allocation_properties.h
+++ b/shared/test/common/mocks/mock_allocation_properties.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
@@ -19,5 +19,6 @@ struct MockAllocationProperties : public AllocationProperties {
    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, DeviceBitfield deviceBitfield) : AllocationProperties(rootDeviceIndex, allocateMemory, size, AllocationType::internalHostMemory, false, deviceBitfield) {}
    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, AllocationType allocationType) : AllocationProperties(rootDeviceIndex, allocateMemory, size, allocationType, false, mockDeviceBitfield) {}
    MockAllocationProperties(uint32_t rootDeviceIndex, bool allocateMemory, size_t size, AllocationType allocationType, DeviceBitfield deviceBitfield) : AllocationProperties(rootDeviceIndex, allocateMemory, size, allocationType, false, deviceBitfield) {}
+    MockAllocationProperties(uint32_t rootDeviceIndex, AllocationType allocType, size_t size) : AllocationProperties(rootDeviceIndex, true, size, allocType, false, mockDeviceBitfield) {}
 };
 } // namespace NEO
--- a/shared/test/common/mocks/mock_cpu_page_fault_manager.h
+++ b/shared/test/common/mocks/mock_cpu_page_fault_manager.h
@@ -14,16 +14,17 @@

 using namespace NEO;

-class MockPageFaultManager : public PageFaultManager {
+template <class BaseFaultManager>
+class MockPageFaultManagerImpl : public BaseFaultManager {
  public:
-    using PageFaultManager::gpuDomainHandler;
-    using PageFaultManager::memoryData;
-    using PageFaultManager::PageFaultData;
-    using PageFaultManager::PageFaultManager;
-    using PageFaultManager::selectGpuDomainHandler;
-    using PageFaultManager::transferAndUnprotectMemory;
-    using PageFaultManager::unprotectAndTransferMemory;
-    using PageFaultManager::verifyAndHandlePageFault;
+    using BaseFaultManager::BaseFaultManager;
+    using BaseFaultManager::gpuDomainHandler;
+    using BaseFaultManager::memoryData;
+    using PageFaultData = typename BaseFaultManager::PageFaultData;
+    using BaseFaultManager::selectGpuDomainHandler;
+    using BaseFaultManager::transferAndUnprotectMemory;
+    using BaseFaultManager::unprotectAndTransferMemory;
+    using BaseFaultManager::verifyAndHandlePageFault;

    bool checkFaultHandlerFromPageFaultManager() override {
        checkFaultHandlerCalled++;
@@ -62,19 +63,19 @@ class MockPageFaultManager : public PageFaultManager {
        allowCPUMemoryEvictionCalled++;
    }
    void baseAubWritable(bool writable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
-        PageFaultManager::setAubWritable(writable, ptr, unifiedMemoryManager);
+        BaseFaultManager::setAubWritable(writable, ptr, unifiedMemoryManager);
    }
    void baseCpuTransfer(void *ptr, size_t size, void *cmdQ) {
-        PageFaultManager::transferToCpu(ptr, size, cmdQ);
+        BaseFaultManager::transferToCpu(ptr, size, cmdQ);
    }
    void baseGpuTransfer(void *ptr, void *cmdQ) {
-        PageFaultManager::transferToGpu(ptr, cmdQ);
+        BaseFaultManager::transferToGpu(ptr, cmdQ);
    }
    void baseCpuAllocEvictable(bool evictable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
-        PageFaultManager::setCpuAllocEvictable(evictable, ptr, unifiedMemoryManager);
+        BaseFaultManager::setCpuAllocEvictable(evictable, ptr, unifiedMemoryManager);
    }
    void baseAllowCPUMemoryEviction(bool evict, void *ptr, PageFaultData &pageFaultData) {
-        PageFaultManager::allowCPUMemoryEviction(evict, ptr, pageFaultData);
+        BaseFaultManager::allowCPUMemoryEviction(evict, ptr, pageFaultData);
    }
    void evictMemoryAfterImplCopy(GraphicsAllocation *allocation, Device *device) override {}

@@ -85,15 +86,15 @@ class MockPageFaultManager : public PageFaultManager {
    }

    void *getHwHandlerAddress() {
-        return reinterpret_cast<void *>(PageFaultManager::transferAndUnprotectMemory);
+        return reinterpret_cast<void *>(BaseFaultManager::transferAndUnprotectMemory);
    }

    void *getAubAndTbxHandlerAddress() {
-        return reinterpret_cast<void *>(PageFaultManager::unprotectAndTransferMemory);
+        return reinterpret_cast<void *>(BaseFaultManager::unprotectAndTransferMemory);
    }
    void moveAllocationToGpuDomain(void *ptr) override {
        moveAllocationToGpuDomainCalled++;
-        PageFaultManager::moveAllocationToGpuDomain(ptr);
+        BaseFaultManager::moveAllocationToGpuDomain(ptr);
    }

    int checkFaultHandlerCalled = 0;
@@ -120,6 +121,8 @@ class MockPageFaultManager : public PageFaultManager {
    EngineUsage engineUsage = EngineUsage::engineUsageCount;
 };

+class MockPageFaultManager : public MockPageFaultManagerImpl<CpuPageFaultManager> {};
+
 template <class T>
 class MockPageFaultManagerHandlerInvoke : public T {
  public:
--- a/shared/test/common/mocks/mock_tbx_csr.h
+++ b/shared/test/common/mocks/mock_tbx_csr.h
@@ -10,6 +10,8 @@
 #include "shared/source/command_stream/submission_status.h"
 #include "shared/source/command_stream/tbx_command_stream_receiver_hw.h"
 #include "shared/source/execution_environment/execution_environment.h"
+#include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
+#include "shared/source/page_fault_manager/tbx_page_fault_manager.h"

 namespace NEO {

@@ -18,6 +20,8 @@ class MockTbxCsr : public TbxCommandStreamReceiverHw<GfxFamily> {
  public:
    using TbxCommandStreamReceiverHw<GfxFamily>::writeMemory;
    using TbxCommandStreamReceiverHw<GfxFamily>::allocationsForDownload;
+    using TbxCommandStreamReceiverHw<GfxFamily>::getParametersForMemory;
+    using TbxCommandStreamReceiverHw<GfxFamily>::getTbxPageFaultManager;
    MockTbxCsr(ExecutionEnvironment &executionEnvironment, const DeviceBitfield deviceBitfield)
        : TbxCommandStreamReceiverHw<GfxFamily>(executionEnvironment, 0, deviceBitfield) {
        this->downloadAllocationImpl = [this](GraphicsAllocation &gfxAllocation) {
@@ -64,6 +68,7 @@ class MockTbxCsr : public TbxCommandStreamReceiverHw<GfxFamily> {
        TbxCommandStreamReceiverHw<GfxFamily>::dumpAllocation(gfxAllocation);
        dumpAllocationCalled = true;
    }
+
    bool initializeEngineCalled = false;
    bool writeMemoryWithAubManagerCalled = false;
    bool writeMemoryCalled = false;
--- a/shared/test/unit_test/command_stream/tbx_command_stream_tests.cpp
+++ b/shared/test/unit_test/command_stream/tbx_command_stream_tests.cpp
@@ -13,9 +13,11 @@
 #include "shared/source/helpers/engine_node_helper.h"
 #include "shared/source/helpers/gfx_core_helper.h"
 #include "shared/source/helpers/hardware_context_controller.h"
+#include "shared/source/helpers/options.h"
 #include "shared/source/helpers/ptr_math.h"
 #include "shared/source/memory_manager/memory_banks.h"
 #include "shared/source/memory_manager/os_agnostic_memory_manager.h"
+#include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
 #include "shared/test/common/fixtures/device_fixture.h"
 #include "shared/test/common/fixtures/mock_aub_center_fixture.h"
 #include "shared/test/common/fixtures/tbx_command_stream_fixture.h"
@@ -36,6 +38,7 @@
 #include "shared/test/common/test_macros/hw_test.h"

 #include <cstdint>
+#include <memory>

 using namespace NEO;

@@ -1281,3 +1284,111 @@ HWTEST_F(TbxCommandStreamTests, givenTimestampBufferAllocationWhenTbxWriteMemory

    memoryManager->freeGraphicsMemory(timestampAllocation);
 }
+
+template <typename FamilyType, CommandStreamReceiverType csrType>
+struct TbxPageFaultTestFixture {
+
+    class MockTbxCsrForPageFaultTests : public MockTbxCsr<FamilyType> {
+      public:
+        using MockTbxCsr<FamilyType>::MockTbxCsr;
+
+        CpuPageFaultManager *getTbxPageFaultManager() override {
+            return this->tbxFaultManager.get();
+        }
+
+        std::unique_ptr<TbxPageFaultManager> tbxFaultManager = TbxPageFaultManager::create();
+    };
+
+    static void runTest1(MockDevice *pDevice) {
+        DebugManagerStateRestore dbgRestore;
+        debugManager.flags.SetCommandStreamReceiver.set(static_cast<int32_t>(csrType));
+        std::unique_ptr<MockTbxCsrForPageFaultTests> tbxCsr(new MockTbxCsrForPageFaultTests(*pDevice->executionEnvironment, pDevice->getDeviceBitfield()));
+        tbxCsr->setupContext(*pDevice->getDefaultEngine().osContext);
+
+        EXPECT_TRUE(tbxCsr->tbxFaultManager->checkFaultHandlerFromPageFaultManager());
+
+        auto memoryManager = pDevice->getMemoryManager();
+
+        NEO::GraphicsAllocation *gfxAlloc1 = memoryManager->allocateGraphicsMemoryWithProperties(
+            {pDevice->getRootDeviceIndex(),
+             MemoryConstants::pageSize,
+             AllocationType::bufferHostMemory,
+             pDevice->getDeviceBitfield()});
+
+        uint64_t gpuAddress;
+        void *cpuAddress;
+        size_t size;
+
+        EXPECT_TRUE(tbxCsr->getParametersForMemory(*gfxAlloc1, gpuAddress, cpuAddress, size));
+
+        tbxCsr->writeMemory(*gfxAlloc1);
+        EXPECT_FALSE(tbxCsr->isTbxWritable(*gfxAlloc1));
+
+        // accessing outside address range does not affect inserted host allocs
+        auto ptrBelow = (void *)0x0;
+        EXPECT_FALSE(tbxCsr->tbxFaultManager->verifyAndHandlePageFault(ptrBelow, true));
+        auto ptrAbove = ptrOffset(cpuAddress, size + 1);
+        EXPECT_FALSE(tbxCsr->tbxFaultManager->verifyAndHandlePageFault(ptrAbove, true));
+        EXPECT_FALSE(tbxCsr->isTbxWritable(*gfxAlloc1));
+
+        *reinterpret_cast<char *>(cpuAddress) = 1;
+        EXPECT_TRUE(tbxCsr->isTbxWritable(*gfxAlloc1));
+
+        tbxCsr->writeMemory(*gfxAlloc1);
+        EXPECT_FALSE(tbxCsr->isTbxWritable(*gfxAlloc1));
+
+        // accessing address with offset that is still in alloc range should
+        // also make writable and download
+        reinterpret_cast<char *>(cpuAddress)[1] = 1;
+        EXPECT_TRUE(tbxCsr->isTbxWritable(*gfxAlloc1));
+
+        // for coverage
+        tbxCsr->tbxFaultManager->removeAllocation(static_cast<GraphicsAllocation *>(nullptr));
+        tbxCsr->tbxFaultManager->removeAllocation(gfxAlloc1);
+
+        memoryManager->freeGraphicsMemory(gfxAlloc1);
+    }
+
+    static void runtTest2(MockDevice *pDevice) {
+        DebugManagerStateRestore dbgRestore;
+        debugManager.flags.SetCommandStreamReceiver.set(static_cast<int32_t>(csrType));
+        std::unique_ptr<MockTbxCsrForPageFaultTests> tbxCsr(new MockTbxCsrForPageFaultTests(*pDevice->executionEnvironment, pDevice->getDeviceBitfield()));
+        tbxCsr->setupContext(*pDevice->getDefaultEngine().osContext);
+
+        EXPECT_TRUE(tbxCsr->tbxFaultManager->checkFaultHandlerFromPageFaultManager());
+
+        auto memoryManager = pDevice->getMemoryManager();
+
+        NEO::GraphicsAllocation *gfxAlloc1 = memoryManager->allocateGraphicsMemoryWithProperties(
+            {pDevice->getRootDeviceIndex(),
+             MemoryConstants::pageSize,
+             AllocationType::bufferHostMemory,
+             pDevice->getDeviceBitfield()});
+
+        uint64_t gpuAddress;
+        void *cpuAddress;
+        size_t size;
+
+        EXPECT_TRUE(tbxCsr->getParametersForMemory(*gfxAlloc1, gpuAddress, cpuAddress, size));
+
+        tbxCsr->writeMemory(*gfxAlloc1);
+        tbxCsr->downloadAllocationTbx(*gfxAlloc1);
+        EXPECT_TRUE(!tbxCsr->isTbxWritable(*gfxAlloc1));
+
+        static_cast<float *>(cpuAddress)[0] = 1.0f;
+
+        memoryManager->freeGraphicsMemory(gfxAlloc1);
+    }
+};
+
+HWTEST_F(TbxCommandStreamTests, givenTbxModeWhenHostAccessesHostAllocThenAllocShouldBeDownloadedAndWritable) {
+    TbxPageFaultTestFixture<FamilyType, CommandStreamReceiverType::tbx>::runTest1(pDevice);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxWithAubModeWhenHostAccessesHostAllocThenAllocShouldBeDownloadedAndWritable) {
+    TbxPageFaultTestFixture<FamilyType, CommandStreamReceiverType::tbxWithAub>::runTest1(pDevice);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxWithModeWhenHostBufferNotWritableAndProtectedThenDownloadShouldNotCrash) {
+    TbxPageFaultTestFixture<FamilyType, CommandStreamReceiverType::tbx>::runtTest2(pDevice);
+}
--- a/shared/test/unit_test/debug_settings/debug_settings_manager_tests.cpp
+++ b/shared/test/unit_test/debug_settings/debug_settings_manager_tests.cpp
@@ -452,3 +452,21 @@ TEST(DurationLogTest, givenDurationGetTimeStringThenTimeStringIsCorrect) {
        EXPECT_TRUE(std::isdigit(c) || c == '[' || c == ']' || c == '.' || c == ' ');
    }
 }
+
+TEST(DebugSettingsManager, GivenTbxOrTbxWithAubCsrTypeWhenCallingIsTbxModeThenReturnTrue) {
+    DebugManagerStateRestore restorer;
+    NEO::debugManager.flags.SetCommandStreamReceiver.set(2);
+    EXPECT_TRUE(NEO::debugManager.isTbxMode());
+
+    NEO::debugManager.flags.SetCommandStreamReceiver.set(4);
+    EXPECT_TRUE(NEO::debugManager.isTbxMode());
+}
+
+TEST(DebugSettingsManager, GivenHardwareOrHardwareWithAubCsrTypeWhenCallingIsTbxModeThenReturnFalse) {
+    DebugManagerStateRestore restorer;
+    NEO::debugManager.flags.SetCommandStreamReceiver.set(1);
+    EXPECT_FALSE(NEO::debugManager.isTbxMode());
+
+    NEO::debugManager.flags.SetCommandStreamReceiver.set(3);
+    EXPECT_FALSE(NEO::debugManager.isTbxMode());
+}
--- a/shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp
+++ b/shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp
@@ -5,6 +5,8 @@
 *
 */

+#include "shared/source/helpers/options.h"
+#include "shared/source/page_fault_manager/tbx_page_fault_manager.h"
 #include "shared/test/common/fixtures/cpu_page_fault_manager_tests_fixture.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
@@ -21,7 +23,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocsWhenInsertingAllocsThenAllo
    EXPECT_EQ(pageFaultManager->memoryData.size(), 1u);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].size, 10u);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].cmdQ, cmdQ);
-    EXPECT_EQ(pageFaultManager->memoryData[alloc1].domain, PageFaultManager::AllocationDomain::cpu);
+    EXPECT_EQ(pageFaultManager->memoryData[alloc1].domain, CpuPageFaultManager::AllocationDomain::cpu);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].unifiedMemoryManager, unifiedMemoryManager.get());
    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 0);
    EXPECT_EQ(pageFaultManager->protectMemoryCalled, 0);
@@ -33,11 +35,11 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocsWhenInsertingAllocsThenAllo
    EXPECT_EQ(pageFaultManager->memoryData.size(), 2u);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].size, 10u);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].cmdQ, cmdQ);
-    EXPECT_EQ(pageFaultManager->memoryData[alloc1].domain, PageFaultManager::AllocationDomain::cpu);
+    EXPECT_EQ(pageFaultManager->memoryData[alloc1].domain, CpuPageFaultManager::AllocationDomain::cpu);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].unifiedMemoryManager, unifiedMemoryManager.get());
    EXPECT_EQ(pageFaultManager->memoryData[alloc2].size, 20u);
    EXPECT_EQ(pageFaultManager->memoryData[alloc2].cmdQ, cmdQ);
-    EXPECT_EQ(pageFaultManager->memoryData[alloc2].domain, PageFaultManager::AllocationDomain::cpu);
+    EXPECT_EQ(pageFaultManager->memoryData[alloc2].domain, CpuPageFaultManager::AllocationDomain::cpu);
    EXPECT_EQ(pageFaultManager->memoryData[alloc2].unifiedMemoryManager, unifiedMemoryManager.get());
    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 0);
    EXPECT_EQ(pageFaultManager->protectMemoryCalled, 0);
@@ -62,7 +64,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenRemovingProtectedAllocTh
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].size, 10u);
    EXPECT_EQ(pageFaultManager->memoryData[alloc1].unifiedMemoryManager, unifiedMemoryManager.get());

-    pageFaultManager->memoryData[alloc1].domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData[alloc1].domain = CpuPageFaultManager::AllocationDomain::gpu;

    pageFaultManager->removeAllocation(alloc1);
    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 1);
@@ -91,7 +93,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenRemovingAllocsThenNonGpu
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 1u);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs[0], alloc2);

-    pageFaultManager->memoryData.at(alloc2).domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData.at(alloc2).domain = CpuPageFaultManager::AllocationDomain::gpu;
    pageFaultManager->removeAllocation(alloc2);
    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 1);
    EXPECT_EQ(pageFaultManager->allowedMemoryAccessAddress, alloc2);
@@ -109,7 +111,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenAllocNotPresentInNonGpuD

    pageFaultManager->moveAllocationToGpuDomain(alloc1);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 0u);
-    pageFaultManager->memoryData.at(alloc1).domain = PageFaultManager::AllocationDomain::cpu;
+    pageFaultManager->memoryData.at(alloc1).domain = CpuPageFaultManager::AllocationDomain::cpu;
    pageFaultManager->removeAllocation(alloc1);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 0u);
 }
@@ -127,7 +129,7 @@ TEST_F(PageFaultManagerTest, givenNonGpuAllocsContainerWhenMovingToGpuDomainMult
    pageFaultManager->moveAllocationToGpuDomain(alloc1);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 1u);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs[0], alloc2);
-    pageFaultManager->memoryData.at(alloc1).domain = PageFaultManager::AllocationDomain::none;
+    pageFaultManager->memoryData.at(alloc1).domain = CpuPageFaultManager::AllocationDomain::none;
    pageFaultManager->moveAllocationToGpuDomain(alloc1);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 1u);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs[0], alloc2);
@@ -146,7 +148,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocsWhenMovingToGpuDomainAllocs
    pageFaultManager->insertAllocation(alloc3, 30, unifiedMemoryManager.get(), cmdQ, {});
    EXPECT_EQ(pageFaultManager->transferToCpuCalled, 0);
    EXPECT_EQ(pageFaultManager->memoryData.size(), 3u);
-    pageFaultManager->memoryData.at(alloc3).domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData.at(alloc3).domain = CpuPageFaultManager::AllocationDomain::gpu;

    pageFaultManager->moveAllocationsWithinUMAllocsManagerToGpuDomain(unifiedMemoryManager.get());

@@ -166,9 +168,9 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenMovingAllocsOfGivenUMMan
    pageFaultManager->insertAllocation(alloc1, 10u, unifiedMemoryManager.get(), nullptr, {});
    pageFaultManager->insertAllocation(alloc2, 20u, unifiedMemoryManager.get(), nullptr, {});

-    pageFaultManager->memoryData.at(alloc2).domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData.at(alloc2).domain = CpuPageFaultManager::AllocationDomain::gpu;
    pageFaultManager->moveAllocationsWithinUMAllocsManagerToGpuDomain(unifiedMemoryManager.get());
-    EXPECT_EQ(pageFaultManager->memoryData.at(alloc1).domain, PageFaultManager::AllocationDomain::gpu);
+    EXPECT_EQ(pageFaultManager->memoryData.at(alloc1).domain, CpuPageFaultManager::AllocationDomain::gpu);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 0u);
 }

@@ -187,7 +189,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocsWhenMovingToGpuDomainWithPr
    pageFaultManager->insertAllocation(alloc3, 30, unifiedMemoryManager.get(), cmdQ, {});
    EXPECT_EQ(pageFaultManager->transferToCpuCalled, 0);
    EXPECT_EQ(pageFaultManager->memoryData.size(), 3u);
-    pageFaultManager->memoryData.at(alloc3).domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData.at(alloc3).domain = CpuPageFaultManager::AllocationDomain::gpu;

    testing::internal::CaptureStdout(); // start capturing

@@ -219,8 +221,8 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocsWhenMovingToGpuDomainAllocs
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 2u);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs[0], alloc1);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs[1], alloc2);
-    pageFaultManager->memoryData.at(alloc1).domain = PageFaultManager::AllocationDomain::cpu;
-    pageFaultManager->memoryData.at(alloc2).domain = PageFaultManager::AllocationDomain::none;
+    pageFaultManager->memoryData.at(alloc1).domain = CpuPageFaultManager::AllocationDomain::cpu;
+    pageFaultManager->memoryData.at(alloc2).domain = CpuPageFaultManager::AllocationDomain::none;

    pageFaultManager->moveAllocationsWithinUMAllocsManagerToGpuDomain(unifiedMemoryManager.get());

@@ -366,13 +368,13 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenMovingToGpuDomainThenUpd
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 2u);

    pageFaultManager->moveAllocationToGpuDomain(alloc1);
-    EXPECT_EQ(pageFaultManager->memoryData.at(alloc1).domain, PageFaultManager::AllocationDomain::gpu);
+    EXPECT_EQ(pageFaultManager->memoryData.at(alloc1).domain, CpuPageFaultManager::AllocationDomain::gpu);
    const auto &allocsVec = unifiedMemoryManager->nonGpuDomainAllocs;
    EXPECT_EQ(std::find(allocsVec.cbegin(), allocsVec.cend(), alloc1), allocsVec.cend());
    EXPECT_EQ(allocsVec.size(), 1u);
    EXPECT_EQ(allocsVec[0], alloc2);

-    pageFaultManager->memoryData.at(alloc2).domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData.at(alloc2).domain = CpuPageFaultManager::AllocationDomain::gpu;
    pageFaultManager->moveAllocationToGpuDomain(alloc2);
    EXPECT_NE(std::find(allocsVec.cbegin(), allocsVec.cend(), alloc2), allocsVec.cend());
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 1u);
@@ -413,7 +415,7 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocInGpuDomainWhenMovingToGpuDo
    pageFaultManager->insertAllocation(alloc, 10, unifiedMemoryManager.get(), cmdQ, {});
    EXPECT_EQ(pageFaultManager->transferToCpuCalled, 0);
    EXPECT_EQ(pageFaultManager->memoryData.size(), 1u);
-    pageFaultManager->memoryData.at(alloc).domain = PageFaultManager::AllocationDomain::gpu;
+    pageFaultManager->memoryData.at(alloc).domain = CpuPageFaultManager::AllocationDomain::gpu;

    pageFaultManager->moveAllocationToGpuDomain(alloc);
    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 0);
@@ -539,15 +541,15 @@ TEST_F(PageFaultManagerTest, givenAllocsFromCpuDomainWhenVerifyingPageFaultThenD
    pageFaultManager->moveAllocationsWithinUMAllocsManagerToGpuDomain(unifiedMemoryManager.get());
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 0u);

-    pageFaultManager->memoryData.at(alloc2).domain = PageFaultManager::AllocationDomain::none;
+    pageFaultManager->memoryData.at(alloc2).domain = CpuPageFaultManager::AllocationDomain::none;
    pageFaultManager->verifyAndHandlePageFault(alloc2, true);
-    EXPECT_EQ(pageFaultManager->memoryData.at(alloc2).domain, PageFaultManager::AllocationDomain::cpu);
+    EXPECT_EQ(pageFaultManager->memoryData.at(alloc2).domain, CpuPageFaultManager::AllocationDomain::cpu);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 0u);

    pageFaultManager->gpuDomainHandler = &MockPageFaultManager::unprotectAndTransferMemory;
-    pageFaultManager->memoryData.at(alloc1).domain = PageFaultManager::AllocationDomain::none;
+    pageFaultManager->memoryData.at(alloc1).domain = CpuPageFaultManager::AllocationDomain::none;
    pageFaultManager->verifyAndHandlePageFault(alloc1, true);
-    EXPECT_EQ(pageFaultManager->memoryData.at(alloc1).domain, PageFaultManager::AllocationDomain::cpu);
+    EXPECT_EQ(pageFaultManager->memoryData.at(alloc1).domain, CpuPageFaultManager::AllocationDomain::cpu);
    EXPECT_EQ(unifiedMemoryManager->nonGpuDomainAllocs.size(), 0u);
 }

@@ -902,7 +904,7 @@ struct PageFaultManagerTestWithDebugFlag : public ::testing::TestWithParam<uint3
 TEST_P(PageFaultManagerTestWithDebugFlag, givenDebugFlagWhenInsertingAllocationThenItOverridesHints) {
    DebugManagerStateRestore restore;
    debugManager.flags.UsmInitialPlacement.set(GetParam()); // Should be ignored by the driver, when passing hints
-    const auto expectedDomain = GetParam() == 1 ? PageFaultManager::AllocationDomain::none : PageFaultManager::AllocationDomain::cpu;
+    const auto expectedDomain = GetParam() == 1 ? CpuPageFaultManager::AllocationDomain::none : CpuPageFaultManager::AllocationDomain::cpu;

    void *allocs[] = {
        reinterpret_cast<void *>(0x1),
@@ -953,24 +955,58 @@ TEST_F(PageFaultManagerTest, givenNoUsmInitialPlacementFlagsWHenInsertingUsmAllo
    memoryProperties.allocFlags.usmInitialPlacementCpu = 0;
    memoryProperties.allocFlags.usmInitialPlacementGpu = 0;
    pageFaultManager->insertAllocation(allocs[0], 10, unifiedMemoryManager.get(), cmdQ, memoryProperties);
-    EXPECT_EQ(PageFaultManager::AllocationDomain::cpu, pageFaultManager->memoryData.at(allocs[0]).domain);
+    EXPECT_EQ(CpuPageFaultManager::AllocationDomain::cpu, pageFaultManager->memoryData.at(allocs[0]).domain);
    EXPECT_EQ(allocs[0], unifiedMemoryManager->nonGpuDomainAllocs[0]);

    memoryProperties.allocFlags.usmInitialPlacementCpu = 0;
    memoryProperties.allocFlags.usmInitialPlacementGpu = 1;
    pageFaultManager->insertAllocation(allocs[1], 10, unifiedMemoryManager.get(), cmdQ, memoryProperties);
-    EXPECT_EQ(PageFaultManager::AllocationDomain::none, pageFaultManager->memoryData.at(allocs[1]).domain);
+    EXPECT_EQ(CpuPageFaultManager::AllocationDomain::none, pageFaultManager->memoryData.at(allocs[1]).domain);
    EXPECT_EQ(allocs[1], unifiedMemoryManager->nonGpuDomainAllocs[1]);

    memoryProperties.allocFlags.usmInitialPlacementCpu = 1;
    memoryProperties.allocFlags.usmInitialPlacementGpu = 0;
    pageFaultManager->insertAllocation(allocs[2], 10, unifiedMemoryManager.get(), cmdQ, memoryProperties);
-    EXPECT_EQ(PageFaultManager::AllocationDomain::cpu, pageFaultManager->memoryData.at(allocs[2]).domain);
+    EXPECT_EQ(CpuPageFaultManager::AllocationDomain::cpu, pageFaultManager->memoryData.at(allocs[2]).domain);
    EXPECT_EQ(allocs[2], unifiedMemoryManager->nonGpuDomainAllocs[2]);

    memoryProperties.allocFlags.usmInitialPlacementCpu = 1;
    memoryProperties.allocFlags.usmInitialPlacementGpu = 1;
    pageFaultManager->insertAllocation(allocs[3], 10, unifiedMemoryManager.get(), cmdQ, memoryProperties);
-    EXPECT_EQ(PageFaultManager::AllocationDomain::cpu, pageFaultManager->memoryData.at(allocs[3]).domain);
+    EXPECT_EQ(CpuPageFaultManager::AllocationDomain::cpu, pageFaultManager->memoryData.at(allocs[3]).domain);
    EXPECT_EQ(allocs[3], unifiedMemoryManager->nonGpuDomainAllocs[3]);
 }
+
+TEST_F(PageFaultManagerTest, givenTbxModeWhenAllocateSharedMemoryThenTbxFaultManagerShouldBehaveLikeCpuFaultManager) {
+    auto memoryManager2 = std::make_unique<MockMemoryManager>(executionEnvironment);
+    auto unifiedMemoryManager2 = std::make_unique<SVMAllocsManager>(memoryManager2.get(), false);
+    auto pageFaultManager2 = std::make_unique<MockPageFaultManagerImpl<TbxPageFaultManager>>();
+    void *cmdQ = reinterpret_cast<void *>(0xFFFF);
+
+    RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
+    std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
+
+    auto properties = SVMAllocsManager::UnifiedMemoryProperties(InternalMemoryType::sharedUnifiedMemory, 1, rootDeviceIndices, deviceBitfields);
+    void *ptr = unifiedMemoryManager2->createSharedUnifiedMemoryAllocation(10, properties, cmdQ);
+
+    pageFaultManager2->insertAllocation(ptr, 10, unifiedMemoryManager2.get(), cmdQ, {});
+
+    EXPECT_TRUE(pageFaultManager2->verifyAndHandlePageFault(ptr, true));
+
+    unifiedMemoryManager2->freeSVMAlloc(ptr);
+}
+
+TEST_F(PageFaultManagerTest, givenHardwareModeWhenCallTbxInsertOrRemoveApiThenNothing) {
+    auto pageFaultManager2 = std::make_unique<MockPageFaultManagerImpl<CpuPageFaultManager>>();
+
+    auto ptr = reinterpret_cast<void *>(0XFFFF);
+    auto gfxAlloc = reinterpret_cast<GraphicsAllocation *>(0xFFFF);
+    auto csr = reinterpret_cast<CommandStreamReceiver *>(0xFFFF);
+    pageFaultManager2->insertAllocation(csr, gfxAlloc, ptr, 0);
+    EXPECT_EQ(pageFaultManager2->memoryData.find(ptr), pageFaultManager2->memoryData.end());
+
+    pageFaultManager2->memoryData[ptr] = {};
+    pageFaultManager2->removeAllocation(gfxAlloc);
+    EXPECT_FALSE(pageFaultManager2->memoryData.find(ptr) == pageFaultManager2->memoryData.end());
+    pageFaultManager2->memoryData.erase(ptr);
+}
--- a/shared/test/unit_test/page_fault_manager/windows/cpu_page_fault_manager_windows_tests.cpp
+++ b/shared/test/unit_test/page_fault_manager/windows/cpu_page_fault_manager_windows_tests.cpp
@@ -156,4 +156,4 @@ TEST_F(PageFaultManagerTest,
    EXPECT_EQ(0u, csr->getEvictionAllocations().size());

    unifiedMemoryManager->freeSVMAlloc(ptr);
-}
+}
--- a/shared/test/unit_test/ult_specific_config.cpp
+++ b/shared/test/unit_test/ult_specific_config.cpp
@@ -29,11 +29,11 @@ const char *neoMockSettingsFileName = "neo_mock.config";
 bool CompressionSelector::preferCompressedAllocation(const AllocationProperties &properties) {
    return false;
 }
-void PageFaultManager::transferToCpu(void *ptr, size_t size, void *cmdQ) {
+void CpuPageFaultManager::transferToCpu(void *ptr, size_t size, void *cmdQ) {
 }
-void PageFaultManager::transferToGpu(void *ptr, void *cmdQ) {
+void CpuPageFaultManager::transferToGpu(void *ptr, void *cmdQ) {
 }
-void PageFaultManager::allowCPUMemoryEviction(bool evict, void *ptr, PageFaultData &pageFaultData) {
+void CpuPageFaultManager::allowCPUMemoryEviction(bool evict, void *ptr, PageFaultData &pageFaultData) {
 }

 void RootDeviceEnvironment::initApiGfxCoreHelper() {