PageFaultHandler: speedup UM allocations lookup

Add a per-instance SVMAllocsManager::nonGpuDomainAllocs container for
all allocations to be removed in
moveAllocationsWithinUMAllocsManagerToGpuDomain. This approach replaces
the current iterative search and performs the task faster.

Add 7 new unit-tests to verify the functionality related to
nonGpuDomainAllocs container, both in expected and unexpected/synthetic
scenarios.

For UTs replace a dummy unifiedMemoryManager pointer with a pointer to
an instace of SVMAllocsManager, otherwise a SegFault error is thrown at
the end of tests.

Perform overall cleanup in related tests implementation, includes but
not limited to removal of:

- givenInitialPlacementGpu\
WhenMovingToGpuDomainThenFirstAccessDoesNotInvokeTransfer

As it is fully covered by:

givenAllocationMovedToGpuDomain\
WhenVerifyingPagefaultThenAllocationIsMovedToCpuDomain

- givenInitialPlacementGpu\
WhenVerifyingPagefaultThenFirstAccessDoesNotInvokeTransfer

As it is fully covered by:

givenTbxAndnitialPlacementGpu\
WhenVerifyingPagefaultThenMemoryIsUnprotectedOnly

Finally, reduce code duplication where possible.

Related-To: NEO-6658
Signed-off-by: Maciej Bielski <[email protected]>

Author: MaciejBielski

shared/source/memory_manager/unified_memory_manager.h

@@ -164,6 +164,9 @@ class SVMAllocsManager {
 
     std::map<CommandStreamReceiver *, InternalAllocationsTracker> indirectAllocationsResidency;
 
+    using NonGpuDomainAllocsContainer = std::vector<void *>;
+    NonGpuDomainAllocsContainer nonGpuDomainAllocs;
+
   protected:
     void *createZeroCopySvmAllocation(size_t size, const SvmAllocationProperties &svmProperties,
                                       const std::set<uint32_t> &rootDeviceIndices,

shared/source/page_fault_manager/cpu_page_fault_manager.cpp

@@ -7,13 +7,10 @@
 
 #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/memory_properties_helpers.h"
-#include "shared/source/helpers/options.h"
-#include "shared/source/helpers/ptr_math.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
 
+#include <algorithm>
 #include <mutex>
 
 namespace NEO {
@@ -26,6 +23,7 @@ void PageFaultManager::insertAllocation(void *ptr, size_t size, SVMAllocsManager
     if (initialPlacement != GraphicsAllocation::UsmInitialPlacement::CPU) {
         this->protectCPUMemoryAccess(ptr, size);
     }
+    unifiedMemoryManager->nonGpuDomainAllocs.push_back(ptr);
 }
 
 void PageFaultManager::removeAllocation(void *ptr) {
@@ -35,6 +33,11 @@ void PageFaultManager::removeAllocation(void *ptr) {
         auto &pageFaultData = alloc->second;
         if (pageFaultData.domain == AllocationDomain::Gpu) {
             allowCPUMemoryAccess(ptr, pageFaultData.size);
+        } else {
+            auto &cpuAllocs = pageFaultData.unifiedMemoryManager->nonGpuDomainAllocs;
+            if (auto it = std::find(cpuAllocs.begin(), cpuAllocs.end(), ptr); it != cpuAllocs.end()) {
+                cpuAllocs.erase(it);
+            }
         }
         this->memoryData.erase(ptr);
     }
@@ -46,34 +49,34 @@ void PageFaultManager::moveAllocationToGpuDomain(void *ptr) {
     if (alloc != memoryData.end()) {
         auto &pageFaultData = alloc->second;
         if (pageFaultData.domain != AllocationDomain::Gpu) {
-            if (pageFaultData.domain == AllocationDomain::Cpu) {
-                if (DebugManager.flags.PrintUmdSharedMigration.get()) {
-                    printf("UMD transferring shared allocation %llx from CPU to GPU\n", reinterpret_cast<unsigned long long int>(ptr));
-                }
-                this->transferToGpu(ptr, pageFaultData.cmdQ);
-                this->protectCPUMemoryAccess(ptr, pageFaultData.size);
+            this->migrateStorageToGpuDomain(ptr, pageFaultData);
+
+            auto &cpuAllocs = pageFaultData.unifiedMemoryManager->nonGpuDomainAllocs;
+            if (auto it = std::find(cpuAllocs.begin(), cpuAllocs.end(), ptr); it != cpuAllocs.end()) {
+                cpuAllocs.erase(it);
             }
-            pageFaultData.domain = AllocationDomain::Gpu;
         }
     }
 }
 
 void PageFaultManager::moveAllocationsWithinUMAllocsManagerToGpuDomain(SVMAllocsManager *unifiedMemoryManager) {
     std::unique_lock<SpinLock> lock{mtx};
-    for (auto &alloc : this->memoryData) {
-        auto allocPtr = alloc.first;
-        auto &pageFaultData = alloc.second;
-        if (pageFaultData.unifiedMemoryManager == unifiedMemoryManager && pageFaultData.domain != AllocationDomain::Gpu) {
-            if (pageFaultData.domain == AllocationDomain::Cpu) {
-                if (DebugManager.flags.PrintUmdSharedMigration.get()) {
-                    printf("UMD transferring shared allocation %llx from CPU to GPU\n", reinterpret_cast<unsigned long long int>(allocPtr));
-                }
-                this->transferToGpu(allocPtr, pageFaultData.cmdQ);
-                this->protectCPUMemoryAccess(allocPtr, pageFaultData.size);
-            }
-            pageFaultData.domain = AllocationDomain::Gpu;
+    for (auto allocPtr : unifiedMemoryManager->nonGpuDomainAllocs) {
+        auto &pageFaultData = this->memoryData[allocPtr];
+        this->migrateStorageToGpuDomain(allocPtr, pageFaultData);
+    }
+    unifiedMemoryManager->nonGpuDomainAllocs.clear();
+}
+
+inline void PageFaultManager::migrateStorageToGpuDomain(void *ptr, PageFaultData &pageFaultData) {
+    if (pageFaultData.domain == AllocationDomain::Cpu) {
+        if (DebugManager.flags.PrintUmdSharedMigration.get()) {
+            printf("UMD transferring shared allocation %llx from CPU to GPU\n", reinterpret_cast<unsigned long long int>(ptr));
         }
+        this->transferToGpu(ptr, pageFaultData.cmdQ);
+        this->protectCPUMemoryAccess(ptr, pageFaultData.size);
     }
+    pageFaultData.domain = AllocationDomain::Gpu;
 }
 
 bool PageFaultManager::verifyPageFault(void *ptr) {
@@ -95,24 +98,22 @@ void PageFaultManager::setGpuDomainHandler(gpuDomainHandlerFunc gpuHandlerFuncPt
 }
 
 void PageFaultManager::handleGpuDomainTransferForHw(PageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
-    if (pageFaultData.domain == AllocationDomain::Gpu) {
-        if (DebugManager.flags.PrintUmdSharedMigration.get()) {
-            printf("UMD transferring shared allocation %llx from GPU to CPU\n", reinterpret_cast<unsigned long long int>(allocPtr));
-        }
-        pageFaultHandler->transferToCpu(allocPtr, pageFaultData.size, pageFaultData.cmdQ);
-    }
-    pageFaultData.domain = AllocationDomain::Cpu;
+    pageFaultHandler->migrateStorageToCpuDomain(allocPtr, pageFaultData);
     pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
 }
 
 void PageFaultManager::handleGpuDomainTransferForAubAndTbx(PageFaultManager *pageFaultHandler, void *allocPtr, PageFaultData &pageFaultData) {
     pageFaultHandler->allowCPUMemoryAccess(allocPtr, pageFaultData.size);
+    pageFaultHandler->migrateStorageToCpuDomain(allocPtr, pageFaultData);
+}
 
+inline void PageFaultManager::migrateStorageToCpuDomain(void *ptr, PageFaultData &pageFaultData) {
     if (pageFaultData.domain == AllocationDomain::Gpu) {
         if (DebugManager.flags.PrintUmdSharedMigration.get()) {
-            printf("UMD transferring shared allocation %llx from GPU to CPU\n", reinterpret_cast<unsigned long long int>(allocPtr));
+            printf("UMD transferring shared allocation %llx from GPU to CPU\n", reinterpret_cast<unsigned long long int>(ptr));
         }
-        pageFaultHandler->transferToCpu(allocPtr, pageFaultData.size, pageFaultData.cmdQ);
+        this->transferToCpu(ptr, pageFaultData.size, pageFaultData.cmdQ);
+        pageFaultData.unifiedMemoryManager->nonGpuDomainAllocs.push_back(ptr);
     }
     pageFaultData.domain = AllocationDomain::Cpu;
 }

shared/source/page_fault_manager/cpu_page_fault_manager.h

@@ -62,6 +62,8 @@ class PageFaultManager : public NonCopyableOrMovableClass {
     static void handleGpuDomainTransferForHw(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
     static void handleGpuDomainTransferForAubAndTbx(PageFaultManager *pageFaultHandler, void *alloc, PageFaultData &pageFaultData);
     void selectGpuDomainHandler();
+    inline void migrateStorageToGpuDomain(void *ptr, PageFaultData &pageFaultData);
+    inline void migrateStorageToCpuDomain(void *ptr, PageFaultData &pageFaultData);
 
     decltype(&handleGpuDomainTransferForHw) gpuDomainHandler = &handleGpuDomainTransferForHw;