Non-Local Residency 1: Get Non-Local Memory Info

Adds functionality to query VideoMemoryInfo for the NON_LOCAL memory
segment.

Bug: dawn:193
Change-Id: I63c2f5a649c37617e7b39a60faa2d3b5b5077156
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/19900
Commit-Queue: Brandon Jones <brandon1.jones@intel.com>
Reviewed-by: Austin Eng <enga@chromium.org>

src/dawn_native/d3d12/D3D12Backend.cpp

@@ -51,11 +51,13 @@ namespace dawn_native { namespace d3d12 {
         : ExternalImageDescriptor(ExternalImageDescriptorType::DXGISharedHandle) {
     }
 
-    uint64_t SetExternalMemoryReservation(WGPUDevice device, uint64_t requestedReservationSize) {
+    uint64_t SetExternalMemoryReservation(WGPUDevice device,
+                                          uint64_t requestedReservationSize,
+                                          MemorySegment memorySegment) {
         Device* backendDevice = reinterpret_cast<Device*>(device);
 
         return backendDevice->GetResidencyManager()->SetExternalMemoryReservation(
-            requestedReservationSize);
+            memorySegment, requestedReservationSize);
     }
 
     WGPUTexture WrapSharedHandle(WGPUDevice device,

src/dawn_native/d3d12/ResidencyManagerD3D12.cpp

@@ -1,3 +1,4 @@
+#include "ResidencyManagerD3D12.h"
 // Copyright 2020 The Dawn Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,8 +21,6 @@
 #include "dawn_native/d3d12/Forward.h"
 #include "dawn_native/d3d12/HeapD3D12.h"
 
-#include "dawn_native/d3d12/d3d12_platform.h"
-
 namespace dawn_native { namespace d3d12 {
 
     ResidencyManager::ResidencyManager(Device* device)
@@ -84,24 +83,43 @@ namespace dawn_native { namespace d3d12 {
         }
     }
 
-    // Allows an application component external to Dawn to cap Dawn's residency budget to prevent
+    // Allows an application component external to Dawn to cap Dawn's residency budgets to prevent
-    // competition for device local memory. Returns the amount of memory reserved, which may be less
+    // competition for device memory. Returns the amount of memory reserved, which may be less
     // that the requested reservation when under pressure.
-    uint64_t ResidencyManager::SetExternalMemoryReservation(uint64_t requestedReservationSize) {
+    uint64_t ResidencyManager::SetExternalMemoryReservation(MemorySegment segment,
-        mVideoMemoryInfo.externalRequest = requestedReservationSize;
+                                                            uint64_t requestedReservationSize) {
-        UpdateVideoMemoryInfo();
+        MemorySegmentInfo* segmentInfo = nullptr;
-        return mVideoMemoryInfo.externalReservation;
+        switch (segment) {
+            case MemorySegment::Local:
+                segmentInfo = &mVideoMemoryInfo.local;
+                break;
+            case MemorySegment::NonLocal:
+                segmentInfo = &mVideoMemoryInfo.nonLocal;
+                break;
+            default:
+                UNREACHABLE();
+        }
+
+        segmentInfo->externalRequest = requestedReservationSize;
+
+        UpdateMemorySegmentInfo(segmentInfo);
+
+        return segmentInfo->externalReservation;
     }
 
     void ResidencyManager::UpdateVideoMemoryInfo() {
-        if (!mResidencyManagementEnabled) {
+        UpdateMemorySegmentInfo(&mVideoMemoryInfo.local);
-            return;
+        if (!mDevice->GetDeviceInfo().isUMA) {
+            UpdateMemorySegmentInfo(&mVideoMemoryInfo.nonLocal);
         }
+    }
 
+    void ResidencyManager::UpdateMemorySegmentInfo(MemorySegmentInfo* segmentInfo) {
         DXGI_QUERY_VIDEO_MEMORY_INFO queryVideoMemoryInfo;
+
         ToBackend(mDevice->GetAdapter())
             ->GetHardwareAdapter()
-            ->QueryVideoMemoryInfo(0, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &queryVideoMemoryInfo);
+            ->QueryVideoMemoryInfo(0, segmentInfo->dxgiSegment, &queryVideoMemoryInfo);
 
         // The video memory budget provided by QueryVideoMemoryInfo is defined by the operating
         // system, and may be lower than expected in certain scenarios. Under memory pressure, we
@@ -109,11 +127,10 @@ namespace dawn_native { namespace d3d12 {
         // component from consuming a disproportionate share of memory and ensures that Dawn can
         // continue to make forward progress. Note the choice to halve memory is arbitrarily chosen
         // and subject to future experimentation.
-        mVideoMemoryInfo.externalReservation =
+        segmentInfo->externalReservation =
-            std::min(queryVideoMemoryInfo.Budget / 2, mVideoMemoryInfo.externalRequest);
+            std::min(queryVideoMemoryInfo.Budget / 2, segmentInfo->externalRequest);
 
-        mVideoMemoryInfo.dawnUsage =
+        segmentInfo->usage = queryVideoMemoryInfo.CurrentUsage - segmentInfo->externalReservation;
-            queryVideoMemoryInfo.CurrentUsage - mVideoMemoryInfo.externalReservation;
 
         // If we're restricting the budget for testing, leave the budget as is.
         if (mRestrictBudgetForTesting) {
@@ -125,8 +142,8 @@ namespace dawn_native { namespace d3d12 {
         // for both Dawn and other applications on the system. Note the value of 95% is arbitrarily
         // chosen and subject to future experimentation.
         static constexpr float kBudgetCap = 0.95;
-        mVideoMemoryInfo.dawnBudget =
+        segmentInfo->budget =
-            (queryVideoMemoryInfo.Budget - mVideoMemoryInfo.externalReservation) * kBudgetCap;
+            (queryVideoMemoryInfo.Budget - segmentInfo->externalReservation) * kBudgetCap;
     }
 
     // Removes from the LRU and returns the least recently used heap when possible. Returns nullptr
@@ -161,18 +178,18 @@ namespace dawn_native { namespace d3d12 {
             return {};
         }
 
-        UpdateVideoMemoryInfo();
+        UpdateMemorySegmentInfo(&mVideoMemoryInfo.local);
 
-        uint64_t memoryUsageAfterMakeResident = sizeToMakeResident + mVideoMemoryInfo.dawnUsage;
+        uint64_t memoryUsageAfterMakeResident = sizeToMakeResident + mVideoMemoryInfo.local.usage;
 
         // Return when we can call MakeResident and remain under budget.
-        if (memoryUsageAfterMakeResident < mVideoMemoryInfo.dawnBudget) {
+        if (memoryUsageAfterMakeResident < mVideoMemoryInfo.local.budget) {
             return {};
         }
 
         std::vector<ID3D12Pageable*> resourcesToEvict;
         uint64_t sizeNeededToBeUnderBudget =
-            memoryUsageAfterMakeResident - mVideoMemoryInfo.dawnBudget;
+            memoryUsageAfterMakeResident - mVideoMemoryInfo.local.budget;
         uint64_t sizeEvicted = 0;
         while (sizeEvicted < sizeNeededToBeUnderBudget) {
             Heap* heap;
@@ -301,7 +318,11 @@ namespace dawn_native { namespace d3d12 {
         // value can vary depending on the environment Dawn is running in. By adding this in
         // addition to the artificial budget cap, we can create a predictable and reproducible
         // budget for testing.
-        mVideoMemoryInfo.dawnBudget = mVideoMemoryInfo.dawnUsage + artificialBudgetCap;
+        mVideoMemoryInfo.local.budget = mVideoMemoryInfo.local.usage + artificialBudgetCap;
+        if (!mDevice->GetDeviceInfo().isUMA) {
+            mVideoMemoryInfo.nonLocal.budget =
+                mVideoMemoryInfo.nonLocal.usage + artificialBudgetCap;
+        }
     }
 
 }}  // namespace dawn_native::d3d12

src/dawn_native/d3d12/ResidencyManagerD3D12.h

@@ -17,9 +17,12 @@
 
 #include "common/LinkedList.h"
 #include "common/Serial.h"
+#include "dawn_native/D3D12Backend.h"
 #include "dawn_native/Error.h"
 #include "dawn_native/dawn_platform.h"
 
+#include "dawn_native/d3d12/d3d12_platform.h"
+
 namespace dawn_native { namespace d3d12 {
 
     class Device;
@@ -34,22 +37,31 @@ namespace dawn_native { namespace d3d12 {
         MaybeError EnsureCanMakeResident(uint64_t allocationSize);
         MaybeError EnsureHeapsAreResident(Heap** heaps, size_t heapCount);
 
-        uint64_t SetExternalMemoryReservation(uint64_t requestedReservationSize);
+        uint64_t SetExternalMemoryReservation(MemorySegment segment,
+                                              uint64_t requestedReservationSize);
 
         void TrackResidentAllocation(Heap* heap);
 
-        void RestrictBudgetForTesting(uint64_t);
+        void RestrictBudgetForTesting(uint64_t artificialBudgetCap);
 
       private:
-        struct VideoMemoryInfo {
+        struct MemorySegmentInfo {
-            uint64_t dawnBudget;
+            const DXGI_MEMORY_SEGMENT_GROUP dxgiSegment;
-            uint64_t dawnUsage;
+            uint64_t budget;
+            uint64_t usage;
             uint64_t externalReservation;
             uint64_t externalRequest;
         };
+
+        struct VideoMemoryInfo {
+            MemorySegmentInfo local = {DXGI_MEMORY_SEGMENT_GROUP_LOCAL, 0, 0, 0, 0};
+            MemorySegmentInfo nonLocal = {DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, 0, 0, 0, 0};
+        };
+
         ResultOrError<Heap*> RemoveSingleEntryFromLRU();
         bool ShouldTrackHeap(Heap* heap) const;
         void UpdateVideoMemoryInfo();
+        void UpdateMemorySegmentInfo(MemorySegmentInfo* segmentInfo);
 
         Device* mDevice;
         LinkedList<Heap> mLRUCache;

src/include/dawn_native/D3D12Backend.h

@@ -30,6 +30,15 @@ namespace dawn_native { namespace d3d12 {
     DAWN_NATIVE_EXPORT WGPUTextureFormat
     GetNativeSwapChainPreferredFormat(const DawnSwapChainImplementation* swapChain);
 
+    enum MemorySegment {
+        Local,
+        NonLocal,
+    };
+
+    DAWN_NATIVE_EXPORT uint64_t SetExternalMemoryReservation(WGPUDevice device,
+                                                             uint64_t requestedReservationSize,
+                                                             MemorySegment memorySegment);
+
     struct DAWN_NATIVE_EXPORT ExternalImageDescriptorDXGISharedHandle : ExternalImageDescriptor {
       public:
         ExternalImageDescriptorDXGISharedHandle();
@@ -39,9 +48,6 @@ namespace dawn_native { namespace d3d12 {
         bool isSwapChainTexture = false;
     };
 
-    DAWN_NATIVE_EXPORT uint64_t SetExternalMemoryReservation(WGPUDevice device,
-                                                             uint64_t requestedReservationSize);
-
     // Note: SharedHandle must be a handle to a texture object.
     DAWN_NATIVE_EXPORT WGPUTexture
     WrapSharedHandle(WGPUDevice device, const ExternalImageDescriptorDXGISharedHandle* descriptor);

src/tests/white_box/D3D12ResidencyTests.cpp

@@ -12,11 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "dawn_native/ResourceMemoryAllocation.h"
+#include "dawn_native/D3D12Backend.h"
 #include "dawn_native/d3d12/BufferD3D12.h"
 #include "dawn_native/d3d12/DeviceD3D12.h"
 #include "dawn_native/d3d12/ResidencyManagerD3D12.h"
-#include "dawn_native/d3d12/ResourceHeapAllocationD3D12.h"
 #include "tests/DawnTest.h"
 #include "utils/WGPUHelpers.h"
 
@@ -301,4 +300,19 @@ TEST_P(D3D12ResidencyTests, OvercommitInASingleSubmit) {
     }
 }
 
+TEST_P(D3D12ResidencyTests, SetExternalReservation) {
+    // Set an external reservation of 20% the budget. We should succesfully reserve the amount we
+    // request.
+    uint64_t amountReserved = dawn_native::d3d12::SetExternalMemoryReservation(
+        device.Get(), kRestrictedBudgetSize * .2, dawn_native::d3d12::MemorySegment::Local);
+    EXPECT_EQ(amountReserved, kRestrictedBudgetSize * .2);
+
+    // If we're on a non-UMA device, we should also check the NON_LOCAL memory segment.
+    if (!IsUMA()) {
+        amountReserved = dawn_native::d3d12::SetExternalMemoryReservation(
+            device.Get(), kRestrictedBudgetSize * .2, dawn_native::d3d12::MemorySegment::NonLocal);
+        EXPECT_EQ(amountReserved, kRestrictedBudgetSize * .2);
+    }
+}
+
 DAWN_INSTANTIATE_TEST(D3D12ResidencyTests, D3D12Backend());