From 6fb34b6dccfecf6462983b1099931ab2d7e307cc Mon Sep 17 00:00:00 2001
From: Bryan Bernhart <bryan.bernhart@intel.com>
Date: Mon, 11 Jul 2022 15:35:43 -0700
Subject: [PATCH] Fix bug preventing resource heaps from being evicted.

PR #495 mistakenly forgot to specify the HEAP_DESC::AlwaysInBudget.
---
 .../d3d12/ResourceHeapAllocatorD3D12.cpp      |  1 +
 .../end2end/D3D12ResidencyManagerTests.cpp    | 37 ++++++++++++++++---
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/gpgmm/d3d12/ResourceHeapAllocatorD3D12.cpp b/src/gpgmm/d3d12/ResourceHeapAllocatorD3D12.cpp
index a80d49fa0..13d7d3c2e 100644
--- a/src/gpgmm/d3d12/ResourceHeapAllocatorD3D12.cpp
+++ b/src/gpgmm/d3d12/ResourceHeapAllocatorD3D12.cpp
@@ -62,6 +62,7 @@ namespace gpgmm::d3d12 {
         resourceHeapDesc.IsExternal = false;
         resourceHeapDesc.DebugName = "Resource heap";
         resourceHeapDesc.Alignment = request.Alignment;
+        resourceHeapDesc.AlwaysInBudget = !(mHeapFlags & D3D12_HEAP_FLAG_CREATE_NOT_RESIDENT);
         resourceHeapDesc.HeapType = mHeapType;
 
         Heap* resourceHeap = nullptr;
diff --git a/src/tests/end2end/D3D12ResidencyManagerTests.cpp b/src/tests/end2end/D3D12ResidencyManagerTests.cpp
index b3cda07c1..1837e6e90 100644
--- a/src/tests/end2end/D3D12ResidencyManagerTests.cpp
+++ b/src/tests/end2end/D3D12ResidencyManagerTests.cpp
@@ -190,20 +190,46 @@ TEST_F(D3D12ResidencyManagerTests, OverBudget) {
 
     const D3D12_RESOURCE_DESC bufferDesc = CreateBasicBufferDesc(GPGMM_MB_TO_BYTES(1));
 
-    std::vector<ComPtr<ResourceAllocation>> allocations = {};
+    // Keep allocating until we reach the budget.
+    std::vector<ComPtr<ResourceAllocation>> allocationsBelowBudget = {};
     while (!IsOverBudget(residencyManager.Get())) {
         ComPtr<ResourceAllocation> allocation;
         ASSERT_SUCCEEDED(resourceAllocator->CreateResource(
             {}, bufferDesc, D3D12_RESOURCE_STATE_COMMON, nullptr, &allocation));
+        allocationsBelowBudget.push_back(std::move(allocation));
 
-        allocations.push_back(std::move(allocation));
+        // Prevent the first created resources from being evicted once over budget.
         ASSERT_SUCCEEDED(residencyManager->UpdateVideoMemorySegments());
     }
 
-    // All allocations should be created resident.
-    for (auto& allocation : allocations) {
+    // Created allocations below the budget should be resident.
+    for (auto& allocation : allocationsBelowBudget) {
+        EXPECT_TRUE(allocation->IsResident());
+    }
+
+    // Keep allocating |kMemoryOverBudget| over the budget.
+    constexpr uint64_t kMemoryOverBudget = GPGMM_MB_TO_BYTES(10);
+
+    // Allocating the same amount over budget, where older allocations will be evicted.
+    std::vector<ComPtr<ResourceAllocation>> allocationsAboveBudget = {};
+    const uint64_t currentMemoryUsage = resourceAllocator->GetInfo().UsedMemoryUsage;
+
+    while (currentMemoryUsage + kMemoryOverBudget > resourceAllocator->GetInfo().UsedMemoryUsage) {
+        ComPtr<ResourceAllocation> allocation;
+        ASSERT_SUCCEEDED(resourceAllocator->CreateResource(
+            {}, bufferDesc, D3D12_RESOURCE_STATE_COMMON, nullptr, &allocation));
+        allocationsAboveBudget.push_back(std::move(allocation));
+    }
+
+    // Created allocations above the budget should be resident.
+    for (auto& allocation : allocationsAboveBudget) {
         EXPECT_TRUE(allocation->IsResident());
     }
+
+    // Created allocations below the budget should NOT be resident.
+    for (auto& allocation : allocationsBelowBudget) {
+        EXPECT_FALSE(allocation->IsResident());
+    }
 }
 
 // Keeps allocating until it goes over the OS limited budget.
@@ -263,10 +289,11 @@ TEST_F(D3D12ResidencyManagerTests, OverBudgetWithGrowth) {
         resourceHeaps.push_back(allocation->GetMemory());
         allocations.push_back(std::move(allocation));
 
+        // Prevent the first created resources from being evicted once over budget.
         ASSERT_SUCCEEDED(residencyManager->UpdateVideoMemorySegments());
     }
 
-    // All allocations should be created resident.
+    // Created allocations above the budget should be resident.
     for (auto& allocation : allocations) {
         EXPECT_TRUE(allocation->IsResident());
     }