From d086df703f141d9a8ddcd46a3068859dd88a10f2 Mon Sep 17 00:00:00 2001
From: "Bernhart, Bryan" <bryan.bernhart@intel.com>
Date: Tue, 21 Feb 2023 12:23:43 -0800
Subject: [PATCH] Do not restrict available memory if NOT_RESIDENT is used.

Also, count pending request size in usage before restricting.
---
 src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp | 45 +++++++++++++---------
 src/gpgmm/d3d12/ResourceAllocatorD3D12.h   |  2 +-
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp b/src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp
index 1ab444f38..3718b019a 100644
--- a/src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp
+++ b/src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp
@@ -575,7 +575,7 @@ namespace gpgmm::d3d12 {
           mCaps(std::move(caps)),
           mResourceHeapTier(descriptor.ResourceHeapTier),
           mIsAlwaysCommitted(descriptor.Flags & ALLOCATOR_FLAG_ALWAYS_COMMITTED),
-          mIsHeapAlwaysCreatedInBudget(descriptor.Flags & ALLOCATOR_FLAG_ALWAYS_IN_BUDGET),
+          mIsAlwaysCreatedInBudget(descriptor.Flags & ALLOCATOR_FLAG_ALWAYS_IN_BUDGET),
           mFlushEventBuffersOnDestruct(descriptor.RecordOptions.EventScope &
                                        EventRecordScope::kPerInstance),
           mUseDetailedTimingEvents(descriptor.RecordOptions.UseDetailedTimingEvents),
@@ -1059,22 +1059,31 @@ namespace gpgmm::d3d12 {
             DXGI_QUERY_VIDEO_MEMORY_INFO* currentVideoInfo =
                 residencyManager->GetVideoMemoryInfo(segment);
 
-            // If over-budget, only free memory is considered available.
-            // TODO: Consider optimizing GetStatsInternal().
-            if (currentVideoInfo->CurrentUsage > currentVideoInfo->Budget) {
-                RESOURCE_ALLOCATOR_STATS allocationStats = {};
-                ReturnIfFailed(QueryStatsInternal(&allocationStats));
-
-                request.AvailableForAllocation = allocationStats.FreeMemoryUsage;
-
-                DebugLog(MessageId::kBudgetExceeded)
-                    << "Current usage exceeded budget ("
-                    << std::to_string(currentVideoInfo->CurrentUsage) << " vs "
-                    << std::to_string(currentVideoInfo->Budget) + " bytes).";
-
-            } else {
-                request.AvailableForAllocation =
-                    currentVideoInfo->Budget - currentVideoInfo->CurrentUsage;
+            // If the allocation must be created within the budget, restrict the amount of memory
+            // to prevent OOM to free memory only or to the amount of budget left. The allocator
+            // checks this amount to determine if its appropriate to pre-allocate more memory or
+            // not.
+            if (!IsCreateHeapNotResident()) {
+                // If over-budget, only free memory is considered available.
+                // TODO: Consider optimizing GetStatsInternal().
+                if (currentVideoInfo->CurrentUsage + request.SizeInBytes >
+                    currentVideoInfo->Budget) {
+                    RESOURCE_ALLOCATOR_STATS allocationStats = {};
+                    ReturnIfFailed(QueryStatsInternal(&allocationStats));
+
+                    request.AvailableForAllocation = allocationStats.FreeMemoryUsage;
+
+                    DebugLog(MessageId::kBudgetExceeded)
+                        << "Current usage exceeded budget: "
+                        << GPGMM_BYTES_TO_MB(currentVideoInfo->CurrentUsage) << " vs "
+                        << GPGMM_BYTES_TO_MB(currentVideoInfo->Budget) << " MBs ("
+                        << GPGMM_BYTES_TO_MB(request.AvailableForAllocation) << " MBs free).";
+
+                    // Otherwise, only memory in budget is considered available.
+                } else {
+                    request.AvailableForAllocation =
+                        currentVideoInfo->Budget - currentVideoInfo->CurrentUsage;
+                }
             }
         }
 
@@ -1529,7 +1538,7 @@ namespace gpgmm::d3d12 {
     }
 
     bool ResourceAllocator::IsCreateHeapNotResident() const {
-        return IsResidencyEnabled() && !mIsHeapAlwaysCreatedInBudget;
+        return IsResidencyEnabled() && !mIsAlwaysCreatedInBudget;
     }
 
     bool ResourceAllocator::IsResidencyEnabled() const {
diff --git a/src/gpgmm/d3d12/ResourceAllocatorD3D12.h b/src/gpgmm/d3d12/ResourceAllocatorD3D12.h
index 1b8b1cd28..e001cd71a 100644
--- a/src/gpgmm/d3d12/ResourceAllocatorD3D12.h
+++ b/src/gpgmm/d3d12/ResourceAllocatorD3D12.h
@@ -177,7 +177,7 @@ namespace gpgmm::d3d12 {
 
         const D3D12_RESOURCE_HEAP_TIER mResourceHeapTier;
         const bool mIsAlwaysCommitted;
-        const bool mIsHeapAlwaysCreatedInBudget;
+        const bool mIsAlwaysCreatedInBudget;
         const bool mFlushEventBuffersOnDestruct;
         const bool mUseDetailedTimingEvents;
         const bool mIsCustomHeapsDisabled;