Fix pagefault Cpu transfers in TBX mode

Related-To: NEO-5286

Signed-off-by: Mateusz Hoppe <mateusz.hoppe@intel.com>

shared/source/page_fault_manager/cpu_page_fault_manager.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2020 Intel Corporation
+ * Copyright (C) 2019-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -7,7 +7,9 @@
 
 #include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
 
+#include "shared/source/debug_settings/debug_settings_manager.h"
 #include "shared/source/helpers/debug_helpers.h"
+#include "shared/source/helpers/options.h"
 #include "shared/source/helpers/ptr_math.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
 
@@ -77,17 +79,38 @@ bool PageFaultManager::verifyPageFault(void *ptr) {
         auto &pageFaultData = alloc.second;
         if (ptr >= allocPtr && ptr < ptrOffset(allocPtr, pageFaultData.size)) {
             this->setAubWritable(true, allocPtr, pageFaultData.unifiedMemoryManager);
-            if (pageFaultData.domain == AllocationDomain::Gpu) {
+
-                this->transferToCpu(allocPtr, pageFaultData.size, pageFaultData.cmdQ);
+            gpuDomainHandler(this, allocPtr, pageFaultData);
-            }
-            pageFaultData.domain = AllocationDomain::Cpu;
-            this->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
             return true;
         }
     }
     return false;
 }
 
+void PageFaultManager::handleGpuDomainTransferForHw(PageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
+    if (pageFaultData.domain == AllocationDomain::Gpu) {
+        pageFaultHandler->transferToCpu(allocPtr, pageFaultData.size, pageFaultData.cmdQ);
+    }
+    pageFaultData.domain = AllocationDomain::Cpu;
+    pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
+}
+
+void PageFaultManager::handleGpuDomainTransferForTbx(PageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
+    pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
+
+    if (pageFaultData.domain == AllocationDomain::Gpu) {
+        pageFaultHandler->transferToCpu(allocPtr, pageFaultData.size, pageFaultData.cmdQ);
+    }
+    pageFaultData.domain = AllocationDomain::Cpu;
+}
+
+void PageFaultManager::selectGpuDomainHandler() {
+    if (DebugManager.flags.SetCommandStreamReceiver.get() == CommandStreamReceiverType::CSR_TBX ||
+        DebugManager.flags.SetCommandStreamReceiver.get() == CommandStreamReceiverType::CSR_TBX_WITH_AUB) {
+        this->gpuDomainHandler = &PageFaultManager::handleGpuDomainTransferForTbx;
+    }
+}
+
 void PageFaultManager::setAubWritable(bool writable, void *ptr, SVMAllocsManager *unifiedMemoryManager) {
     UNRECOVERABLE_IF(ptr == nullptr);
     auto gpuAlloc = unifiedMemoryManager->getSVMAlloc(ptr)->gpuAllocations.getDefaultGraphicsAllocation();

shared/source/page_fault_manager/cpu_page_fault_manager.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2020 Intel Corporation
+ * Copyright (C) 2019-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -55,6 +55,11 @@ class PageFaultManager : public NonCopyableOrMovableClass {
     MOCKABLE_VIRTUAL void transferToGpu(void *ptr, void *cmdQ);
     MOCKABLE_VIRTUAL void setAubWritable(bool writable, void *ptr, SVMAllocsManager *unifiedMemoryManager);
 
+    static void handleGpuDomainTransferForHw(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
+    static void handleGpuDomainTransferForTbx(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
+    void selectGpuDomainHandler();
+
+    decltype(&handleGpuDomainTransferForHw) gpuDomainHandler = &handleGpuDomainTransferForHw;
     std::unordered_map<void *, PageFaultData> memoryData;
     SpinLock mtx;
 };

shared/source/page_fault_manager/linux/cpu_page_fault_manager_linux.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2020 Intel Corporation
+ * Copyright (C) 2019-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -16,7 +16,10 @@
 
 namespace NEO {
 std::unique_ptr<PageFaultManager> PageFaultManager::create() {
-    return std::make_unique<PageFaultManagerLinux>();
+    auto pageFaultManager = std::make_unique<PageFaultManagerLinux>();
+
+    pageFaultManager->selectGpuDomainHandler();
+    return pageFaultManager;
 }
 
 std::function<void(int signal, siginfo_t *info, void *context)> PageFaultManagerLinux::pageFaultHandler;

shared/source/page_fault_manager/windows/cpu_page_fault_manager_windows.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2020 Intel Corporation
+ * Copyright (C) 2019-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -11,7 +11,10 @@
 
 namespace NEO {
 std::unique_ptr<PageFaultManager> PageFaultManager::create() {
-    return std::make_unique<PageFaultManagerWindows>();
+    auto pageFaultManager = std::make_unique<PageFaultManagerWindows>();
+
+    pageFaultManager->selectGpuDomainHandler();
+    return pageFaultManager;
 }
 
 std::function<LONG(struct _EXCEPTION_POINTERS *exceptionInfo)> PageFaultManagerWindows::pageFaultHandler;

shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests.cpp

@@ -8,6 +8,7 @@
 #include "shared/source/memory_manager/graphics_allocation.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
 #include "shared/source/unified_memory/unified_memory.h"
+#include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
 #include "shared/test/common/test_macros/test_checks_shared.h"
 #include "shared/test/unit_test/page_fault_manager/cpu_page_fault_manager_tests_fixture.h"
@@ -271,6 +272,54 @@ TEST_F(PageFaultManagerTest, givenInitialPlacementGpuWhenVerifyingPagefaultThenF
     EXPECT_TRUE(pageFaultManager->isAubWritable);
 }
 
+TEST_F(PageFaultManagerTest, givenTbxWhenVerifyingPagefaultThenVerifyPagefaultUnprotectsAndTransfersMemory) {
+    void *alloc = reinterpret_cast<void *>(0x1);
+
+    pageFaultManager->gpuDomainHandler = &MockPageFaultManager::handleGpuDomainTransferForTbx;
+
+    MemoryProperties memoryProperties{};
+    memoryProperties.allocFlags.usmInitialPlacementGpu = 1;
+    pageFaultManager->insertAllocation(alloc, 10, reinterpret_cast<SVMAllocsManager *>(unifiedMemoryManager), nullptr, memoryProperties);
+
+    pageFaultManager->moveAllocationToGpuDomain(alloc);
+
+    EXPECT_EQ(pageFaultManager->protectMemoryCalled, 1);
+    EXPECT_EQ(pageFaultManager->transferToGpuCalled, 0);
+    EXPECT_EQ(pageFaultManager->protectedMemoryAccessAddress, alloc);
+    EXPECT_EQ(pageFaultManager->protectedSize, 10u);
+
+    pageFaultManager->verifyPageFault(alloc);
+
+    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 1);
+    EXPECT_EQ(pageFaultManager->transferToCpuCalled, 1);
+    EXPECT_EQ(pageFaultManager->allowedMemoryAccessAddress, alloc);
+    EXPECT_EQ(pageFaultManager->accessAllowedSize, 10u);
+    EXPECT_TRUE(pageFaultManager->isAubWritable);
+}
+
+TEST_F(PageFaultManagerTest, givenTbxAndInitialPlacementGpuWhenVerifyingPagefaultThenMemoryIsUnprotectedOnly) {
+    void *alloc = reinterpret_cast<void *>(0x1);
+
+    pageFaultManager->gpuDomainHandler = &MockPageFaultManager::handleGpuDomainTransferForTbx;
+
+    MemoryProperties memoryProperties{};
+    memoryProperties.allocFlags.usmInitialPlacementGpu = 1;
+    pageFaultManager->insertAllocation(alloc, 10, reinterpret_cast<SVMAllocsManager *>(unifiedMemoryManager), nullptr, memoryProperties);
+
+    EXPECT_EQ(pageFaultManager->protectMemoryCalled, 1);
+    EXPECT_EQ(pageFaultManager->transferToGpuCalled, 0);
+    EXPECT_EQ(pageFaultManager->protectedMemoryAccessAddress, alloc);
+    EXPECT_EQ(pageFaultManager->protectedSize, 10u);
+
+    pageFaultManager->verifyPageFault(alloc);
+
+    EXPECT_EQ(pageFaultManager->allowMemoryAccessCalled, 1);
+    EXPECT_EQ(pageFaultManager->transferToCpuCalled, 0);
+    EXPECT_EQ(pageFaultManager->allowedMemoryAccessAddress, alloc);
+    EXPECT_EQ(pageFaultManager->accessAllowedSize, 10u);
+    EXPECT_TRUE(pageFaultManager->isAubWritable);
+}
+
 TEST_F(PageFaultManagerTest, givenInitialPlacementCpuWhenMovingToGpuDomainThenFirstAccessInvokesTransfer) {
     void *cmdQ = reinterpret_cast<void *>(0xFFFF);
 
@@ -369,3 +418,42 @@ TEST_F(PageFaultManagerTest, givenUnifiedMemoryAllocWhenSetAubWritableIsCalledTh
 
     unifiedMemoryManager->freeSVMAlloc(alloc1);
 }
+
+TEST(PageFaultManager, givenTbxCsrWhenSelectingHandlerThenTbxGpuDomainHandlerIsSet) {
+    DebugManagerStateRestore restorer;
+
+    auto pageFaultManager = std::make_unique<MockPageFaultManager>();
+    DebugManager.flags.SetCommandStreamReceiver.set(2);
+
+    pageFaultManager = std::make_unique<MockPageFaultManager>();
+    pageFaultManager->selectGpuDomainHandler();
+
+    EXPECT_EQ(pageFaultManager->getTbxHandlerAddress(), reinterpret_cast<void *>(pageFaultManager->gpuDomainHandler));
+
+    DebugManager.flags.SetCommandStreamReceiver.set(4);
+
+    pageFaultManager->selectGpuDomainHandler();
+
+    EXPECT_EQ(pageFaultManager->getTbxHandlerAddress(), reinterpret_cast<void *>(pageFaultManager->gpuDomainHandler));
+}
+
+TEST(PageFaultManager, givenNonTbxCsrWhenSelectingHandlerThenHwGpuDomainHandlerIsSet) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.SetCommandStreamReceiver.set(0);
+
+    auto pageFaultManager = std::make_unique<MockPageFaultManager>();
+    auto defaultHandler = pageFaultManager->gpuDomainHandler;
+
+    EXPECT_EQ(pageFaultManager->getHwHandlerAddress(), reinterpret_cast<void *>(pageFaultManager->gpuDomainHandler));
+
+    pageFaultManager->selectGpuDomainHandler();
+
+    EXPECT_EQ(defaultHandler, pageFaultManager->gpuDomainHandler);
+
+    DebugManager.flags.SetCommandStreamReceiver.set(3);
+
+    pageFaultManager = std::make_unique<MockPageFaultManager>();
+    pageFaultManager->selectGpuDomainHandler();
+
+    EXPECT_EQ(defaultHandler, pageFaultManager->gpuDomainHandler);
+}

shared/test/unit_test/page_fault_manager/mock_cpu_page_fault_manager.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2020 Intel Corporation
+ * Copyright (C) 2019-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -13,9 +13,13 @@ using namespace NEO;
 
 class MockPageFaultManager : public PageFaultManager {
   public:
+    using PageFaultManager::gpuDomainHandler;
+    using PageFaultManager::handleGpuDomainTransferForHw;
+    using PageFaultManager::handleGpuDomainTransferForTbx;
     using PageFaultManager::memoryData;
     using PageFaultManager::PageFaultData;
     using PageFaultManager::PageFaultManager;
+    using PageFaultManager::selectGpuDomainHandler;
     using PageFaultManager::verifyPageFault;
 
     void allowCPUMemoryAccess(void *ptr, size_t size) override {
@@ -51,6 +55,14 @@ class MockPageFaultManager : public PageFaultManager {
     }
     void evictMemoryAfterImplCopy(GraphicsAllocation *allocation, Device *device) override {}
 
+    void *getHwHandlerAddress() {
+        return reinterpret_cast<void *>(PageFaultManager::handleGpuDomainTransferForHw);
+    }
+
+    void *getTbxHandlerAddress() {
+        return reinterpret_cast<void *>(PageFaultManager::handleGpuDomainTransferForTbx);
+    }
+
     int allowMemoryAccessCalled = 0;
     int protectMemoryCalled = 0;
     int transferToCpuCalled = 0;