From 13f9f3a929c0b548d95e56a7267cf46edf0afc6b Mon Sep 17 00:00:00 2001
From: Mateusz Jablonski <mateusz.jablonski@intel.com>
Date: Thu, 14 Mar 2019 08:32:56 +0100
Subject: [PATCH] Prepare for refactor map methods of wddm

Change-Id: I33249d9a097d717ab2eb5801eaa17154f68ffdce
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>
---
 .../windows/wddm_memory_manager.cpp           | 38 +++++++++++++------
 .../windows/wddm_memory_manager.h             |  2 +
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/runtime/os_interface/windows/wddm_memory_manager.cpp b/runtime/os_interface/windows/wddm_memory_manager.cpp
index e8bd23447d..cec43480ed 100644
--- a/runtime/os_interface/windows/wddm_memory_manager.cpp
+++ b/runtime/os_interface/windows/wddm_memory_manager.cpp
@@ -76,7 +76,7 @@ GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemory64kb(const Allocati
     auto cpuPtr = lockResource(wddmAllocation.get());
 
     // 64kb map is not needed
-    auto status = wddm->mapGpuVirtualAddress(wddmAllocation.get(), cpuPtr);
+    auto status = mapGpuVirtualAddressWithRetry(wddmAllocation.get(), cpuPtr);
     DEBUG_BREAK_IF(!status);
     wddmAllocation->setCpuAddress(cpuPtr);
 
@@ -225,8 +225,8 @@ GraphicsAllocation *WddmMemoryManager::createAllocationFromHandle(osHandle handl
     size_t size = allocation->getDefaultGmm()->gmmResourceInfo->getSizeAllocation();
     allocation->setSize(size);
 
-    void *ptr = nullptr;
     if (is32bit) {
+        void *ptr = nullptr;
         if (!wddm->reserveValidAddressRange(size, ptr)) {
             return nullptr;
         }
@@ -235,7 +235,7 @@ GraphicsAllocation *WddmMemoryManager::createAllocationFromHandle(osHandle handl
         allocation->set32BitAllocation(true);
         allocation->setGpuBaseAddress(GmmHelper::canonize(allocator32Bit->getBase()));
     }
-    status = wddm->mapGpuVirtualAddress(allocation.get(), ptr);
+    status = mapGpuVirtualAddressWithRetry(allocation.get(), allocation->getReservedAddressPtr());
     DEBUG_BREAK_IF(!status);
 
     DebugManager.logAllocation(allocation.get());
@@ -400,7 +400,7 @@ MemoryManager::AllocationStatus WddmMemoryManager::populateOsHandles(OsHandleSto
             allocatedFragmentsCounter++;
         }
     }
-    NTSTATUS result = wddm->createAllocationsAndMapGpuVa(handleStorage);
+    auto result = wddm->createAllocationsAndMapGpuVa(handleStorage);
 
     if (result == STATUS_GRAPHICS_NO_VIDEO_MEMORY) {
         return AllocationStatus::InvalidHostPointer;
@@ -491,17 +491,33 @@ bool WddmMemoryManager::createWddmAllocation(WddmAllocation *allocation, void *r
     }
 
     if (wddmSuccess == STATUS_SUCCESS) {
-        bool mapSuccess = wddm->mapGpuVirtualAddress(allocation, requiredGpuPtr);
-        if (!mapSuccess && deferredDeleter) {
-            deferredDeleter->drain(true);
-            mapSuccess = wddm->mapGpuVirtualAddress(allocation, requiredGpuPtr);
-        }
+        bool mapSuccess = mapGpuVirtualAddressWithRetry(allocation, requiredGpuPtr);
         if (!mapSuccess) {
             wddm->destroyAllocations(allocation->getHandles().data(), allocation->getNumHandles(), allocation->resourceHandle);
-            wddmSuccess = STATUS_UNSUCCESSFUL;
+            return false;
         }
+        return true;
     }
-    return (wddmSuccess == STATUS_SUCCESS);
+    return false;
+}
+
+bool WddmMemoryManager::mapGpuVirtualAddressWithRetry(WddmAllocation *graphicsAllocation, const void *preferredGpuVirtualAddress) {
+    uint32_t numMappedAllocations = mapGpuVirtualAddress(graphicsAllocation, preferredGpuVirtualAddress, 0u);
+    if (numMappedAllocations < graphicsAllocation->getNumHandles() && deferredDeleter) {
+        deferredDeleter->drain(true);
+        numMappedAllocations += mapGpuVirtualAddress(graphicsAllocation, preferredGpuVirtualAddress, numMappedAllocations);
+    }
+    return numMappedAllocations == graphicsAllocation->getNumHandles();
+}
+uint32_t WddmMemoryManager::mapGpuVirtualAddress(WddmAllocation *graphicsAllocation, const void *preferredGpuVirtualAddress, uint32_t startingIndex) {
+    auto numMappedAllocations = 0;
+    for (auto i = startingIndex; i < graphicsAllocation->getNumHandles(); i++) {
+        if (!wddm->mapGpuVirtualAddress(graphicsAllocation, const_cast<void *>(preferredGpuVirtualAddress))) {
+            return numMappedAllocations;
+        }
+        numMappedAllocations++;
+    }
+    return numMappedAllocations;
 }
 
 void *WddmMemoryManager::reserveCpuAddressRange(size_t size) {
diff --git a/runtime/os_interface/windows/wddm_memory_manager.h b/runtime/os_interface/windows/wddm_memory_manager.h
index 2ac13e4bdb..e4ff575f4c 100644
--- a/runtime/os_interface/windows/wddm_memory_manager.h
+++ b/runtime/os_interface/windows/wddm_memory_manager.h
@@ -83,6 +83,8 @@ class WddmMemoryManager : public MemoryManager {
     GraphicsAllocation *createAllocationFromHandle(osHandle handle, bool requireSpecificBitness, bool ntHandle);
     static bool validateAllocation(WddmAllocation *alloc);
     bool createWddmAllocation(WddmAllocation *allocation, void *requiredGpuPtr);
+    bool mapGpuVirtualAddressWithRetry(WddmAllocation *graphicsAllocation, const void *preferredGpuVirtualAddress);
+    uint32_t mapGpuVirtualAddress(WddmAllocation *graphicsAllocation, const void *preferredGpuVirtualAddress, uint32_t startingIndex);
     AlignedMallocRestrictions mallocRestrictions;
 
   private: