Non-Local Residency 2: Implement Non-Local Management Logic

Implements logic for managing the NON_LOCAL memory segment for UPLOAD
and READBACK heaps on Non-UMA devices.

Bug: dawn:193
Change-Id: I2426bf6b5f7a7ccd4420f830f344379af9faf73c
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/19901
Reviewed-by: Rafael Cintron <rafael.cintron@microsoft.com>
Reviewed-by: Austin Eng <enga@chromium.org>
Commit-Queue: Brandon Jones <brandon1.jones@intel.com>

src/dawn_native/d3d12/HeapAllocatorD3D12.cpp

@@ -22,8 +22,12 @@ namespace dawn_native { namespace d3d12 {
 
     HeapAllocator::HeapAllocator(Device* device,
                                  D3D12_HEAP_TYPE heapType,
-                                 D3D12_HEAP_FLAGS heapFlags)
+                                 D3D12_HEAP_FLAGS heapFlags,
-        : mDevice(device), mHeapType(heapType), mHeapFlags(heapFlags) {
+                                 MemorySegment memorySegment)
+        : mDevice(device),
+          mHeapType(heapType),
+          mHeapFlags(heapFlags),
+          mMemorySegment(memorySegment) {
     }
 
     ResultOrError<std::unique_ptr<ResourceHeapBase>> HeapAllocator::AllocateResourceHeap(
@@ -44,7 +48,7 @@ namespace dawn_native { namespace d3d12 {
 
         // CreateHeap will implicitly make the created heap resident. We must ensure enough free
         // memory exists before allocating to avoid an out-of-memory error when overcommitted.
-        DAWN_TRY(mDevice->GetResidencyManager()->EnsureCanMakeResident(size));
+        DAWN_TRY(mDevice->GetResidencyManager()->EnsureCanAllocate(size, mMemorySegment));
 
         ComPtr<ID3D12Heap> d3d12Heap;
         DAWN_TRY(CheckOutOfMemoryHRESULT(
@@ -52,7 +56,7 @@ namespace dawn_native { namespace d3d12 {
             "ID3D12Device::CreateHeap"));
 
         std::unique_ptr<ResourceHeapBase> heapBase =
-            std::make_unique<Heap>(std::move(d3d12Heap), heapDesc.Properties.Type, size);
+            std::make_unique<Heap>(std::move(d3d12Heap), mMemorySegment, size);
 
         // Calling CreateHeap implicitly calls MakeResident on the new heap. We must track this to
         // avoid calling MakeResident a second time.

src/dawn_native/d3d12/HeapAllocatorD3D12.h

@@ -15,6 +15,7 @@
 #ifndef DAWNNATIVE_D3D12_HEAPALLOCATORD3D12_H_
 #define DAWNNATIVE_D3D12_HEAPALLOCATORD3D12_H_
 
+#include "dawn_native/D3D12Backend.h"
 #include "dawn_native/ResourceHeapAllocator.h"
 #include "dawn_native/d3d12/d3d12_platform.h"
 
@@ -25,7 +26,10 @@ namespace dawn_native { namespace d3d12 {
     // Wrapper to allocate a D3D12 heap.
     class HeapAllocator : public ResourceHeapAllocator {
       public:
-        HeapAllocator(Device* device, D3D12_HEAP_TYPE heapType, D3D12_HEAP_FLAGS heapFlags);
+        HeapAllocator(Device* device,
+                      D3D12_HEAP_TYPE heapType,
+                      D3D12_HEAP_FLAGS heapFlags,
+                      MemorySegment memorySegment);
         ~HeapAllocator() override = default;
 
         ResultOrError<std::unique_ptr<ResourceHeapBase>> AllocateResourceHeap(
@@ -36,6 +40,7 @@ namespace dawn_native { namespace d3d12 {
         Device* mDevice;
         D3D12_HEAP_TYPE mHeapType;
         D3D12_HEAP_FLAGS mHeapFlags;
+        MemorySegment mMemorySegment;
     };
 
 }}  // namespace dawn_native::d3d12

src/dawn_native/d3d12/HeapD3D12.cpp

@@ -15,22 +15,23 @@
 #include "dawn_native/d3d12/HeapD3D12.h"
 
 namespace dawn_native { namespace d3d12 {
-    Heap::Heap(ComPtr<ID3D12Pageable> d3d12Pageable, D3D12_HEAP_TYPE d3d12HeapType, uint64_t size)
+    Heap::Heap(ComPtr<ID3D12Pageable> d3d12Pageable, MemorySegment memorySegment, uint64_t size)
-        : mD3d12Pageable(std::move(d3d12Pageable)), mD3d12HeapType(d3d12HeapType), mSize(size) {
+        : mD3d12Pageable(std::move(d3d12Pageable)), mMemorySegment(memorySegment), mSize(size) {
     }
 
     Heap::~Heap() {
-        // When a heap is destroyed, it no longer resides in resident memory, so we must evict it
+        // When a heap is destroyed, it no longer resides in resident memory, so we must evict
-        // from the LRU cache. If this heap is not manually removed from the LRU-cache, the
+        // it from the LRU cache. If this heap is not manually removed from the LRU-cache, the
         // ResidencyManager will attempt to use it after it has been deallocated.
         if (IsInResidencyLRUCache()) {
             RemoveFromList();
         }
     }
 
-    // This function should only be used when mD3D12Pageable was initialized from a ID3D12Pageable
+    // This function should only be used when mD3D12Pageable was initialized from a
-    // that was initially created as an ID3D12Heap (i.e. SubAllocation). If the ID3D12Pageable was
+    // ID3D12Pageable that was initially created as an ID3D12Heap (i.e. SubAllocation). If the
-    // initially created as an ID3D12Resource (i.e. DirectAllocation), then use GetD3D12Pageable().
+    // ID3D12Pageable was initially created as an ID3D12Resource (i.e. DirectAllocation), then
+    // use GetD3D12Pageable().
     ComPtr<ID3D12Heap> Heap::GetD3D12Heap() const {
         ComPtr<ID3D12Heap> heap;
         HRESULT result = mD3d12Pageable.As(&heap);
@@ -42,8 +43,8 @@ namespace dawn_native { namespace d3d12 {
         return mD3d12Pageable;
     }
 
-    D3D12_HEAP_TYPE Heap::GetD3D12HeapType() const {
+    MemorySegment Heap::GetMemorySegment() const {
-        return mD3d12HeapType;
+        return mMemorySegment;
     }
 
     Serial Heap::GetLastUsage() const {

src/dawn_native/d3d12/HeapD3D12.h

@@ -17,11 +17,14 @@
 
 #include "common/LinkedList.h"
 #include "common/Serial.h"
+#include "dawn_native/D3D12Backend.h"
 #include "dawn_native/ResourceHeap.h"
 #include "dawn_native/d3d12/d3d12_platform.h"
 
 namespace dawn_native { namespace d3d12 {
 
+    class Device;
+
     // This class is used to represent heap allocations, but also serves as a node within the
     // ResidencyManager's LRU cache. This node is inserted into the LRU-cache when it is first
     // allocated, and any time it is scheduled to be used by the GPU. This node is removed from the
@@ -29,12 +32,12 @@ namespace dawn_native { namespace d3d12 {
     // is destroyed.
     class Heap : public ResourceHeapBase, public LinkNode<Heap> {
       public:
-        Heap(ComPtr<ID3D12Pageable> d3d12Pageable, D3D12_HEAP_TYPE heapType, uint64_t size);
+        Heap(ComPtr<ID3D12Pageable> d3d12Pageable, MemorySegment memorySegment, uint64_t size);
         ~Heap();
 
         ComPtr<ID3D12Heap> GetD3D12Heap() const;
         ComPtr<ID3D12Pageable> GetD3D12Pageable() const;
-        D3D12_HEAP_TYPE GetD3D12HeapType() const;
+        MemorySegment GetMemorySegment() const;
 
         // We set mLastRecordingSerial to denote the serial this heap was last recorded to be used.
         // We must check this serial against the current serial when recording heap usages to ensure
@@ -61,7 +64,7 @@ namespace dawn_native { namespace d3d12 {
 
       private:
         ComPtr<ID3D12Pageable> mD3d12Pageable;
-        D3D12_HEAP_TYPE mD3d12HeapType;
+        MemorySegment mMemorySegment;
         // mLastUsage denotes the last time this heap was recorded for use.
         Serial mLastUsage = 0;
         // mLastSubmission denotes the last time this heap was submitted to the GPU. Note that

src/dawn_native/d3d12/ResidencyManagerD3D12.cpp

@@ -1,4 +1,3 @@
-#include "ResidencyManagerD3D12.h"
 // Copyright 2020 The Dawn Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -36,18 +35,13 @@ namespace dawn_native { namespace d3d12 {
             return {};
         }
 
-        // Depending on device architecture, the heap may not need tracked.
-        if (!ShouldTrackHeap(heap)) {
-            return {};
-        }
-
         // If the heap isn't already resident, make it resident.
         if (!heap->IsInResidencyLRUCache() && !heap->IsResidencyLocked()) {
-            DAWN_TRY(EnsureCanMakeResident(heap->GetSize()));
+            DAWN_TRY(EnsureCanMakeResident(heap->GetSize(),
+                                           GetMemorySegmentInfo(heap->GetMemorySegment())));
             ID3D12Pageable* pageable = heap->GetD3D12Pageable().Get();
             DAWN_TRY(CheckHRESULT(mDevice->GetD3D12Device()->MakeResident(1, &pageable),
-                                  "Making a scheduled-to-be-used resource resident in "
+                                  "Making a scheduled-to-be-used resource resident"));
-                                  "device local memory"));
         }
 
         // Since we can't evict the heap, it's unnecessary to track the heap in the LRU Cache.
@@ -67,19 +61,31 @@ namespace dawn_native { namespace d3d12 {
             return;
         }
 
-        // Depending on device architecture, the heap may not need tracked.
-        if (!ShouldTrackHeap(heap)) {
-            return;
-        }
-
         ASSERT(heap->IsResidencyLocked());
         ASSERT(!heap->IsInResidencyLRUCache());
         heap->DecrementResidencyLock();
 
+        // If another lock still exists on the heap, nothing further should be done.
+        if (heap->IsResidencyLocked()) {
+            return;
+        }
+
         // When all locks have been removed, the resource remains resident and becomes tracked in
-        // the LRU.
+        // the corresponding LRU.
-        if (!heap->IsResidencyLocked()) {
+        TrackResidentAllocation(heap);
-            mLRUCache.Append(heap);
+    }
+
+    // Returns the appropriate MemorySegmentInfo for a given MemorySegment.
+    ResidencyManager::MemorySegmentInfo* ResidencyManager::GetMemorySegmentInfo(
+        MemorySegment memorySegment) {
+        switch (memorySegment) {
+            case MemorySegment::Local:
+                return &mVideoMemoryInfo.local;
+            case MemorySegment::NonLocal:
+                ASSERT(!mDevice->GetDeviceInfo().isUMA);
+                return &mVideoMemoryInfo.nonLocal;
+            default:
+                UNREACHABLE();
         }
     }
 
@@ -88,17 +94,7 @@ namespace dawn_native { namespace d3d12 {
     // that the requested reservation when under pressure.
     uint64_t ResidencyManager::SetExternalMemoryReservation(MemorySegment segment,
                                                             uint64_t requestedReservationSize) {
-        MemorySegmentInfo* segmentInfo = nullptr;
+        MemorySegmentInfo* segmentInfo = GetMemorySegmentInfo(segment);
-        switch (segment) {
-            case MemorySegment::Local:
-                segmentInfo = &mVideoMemoryInfo.local;
-                break;
-            case MemorySegment::NonLocal:
-                segmentInfo = &mVideoMemoryInfo.nonLocal;
-                break;
-            default:
-                UNREACHABLE();
-        }
 
         segmentInfo->externalRequest = requestedReservationSize;
 
@@ -146,11 +142,13 @@ namespace dawn_native { namespace d3d12 {
             (queryVideoMemoryInfo.Budget - segmentInfo->externalReservation) * kBudgetCap;
     }
 
-    // Removes from the LRU and returns the least recently used heap when possible. Returns nullptr
+    // Removes a heap from the LRU and returns the least recently used heap when possible. Returns
-    // when nothing further can be evicted.
+    // nullptr when nothing further can be evicted.
-    ResultOrError<Heap*> ResidencyManager::RemoveSingleEntryFromLRU() {
+    ResultOrError<Heap*> ResidencyManager::RemoveSingleEntryFromLRU(
-        ASSERT(!mLRUCache.empty());
+        MemorySegmentInfo* memorySegment) {
-        Heap* heap = mLRUCache.head()->value();
+        ASSERT(!memorySegment->lruCache.empty());
+        Heap* heap = memorySegment->lruCache.head()->value();
+
         Serial lastSubmissionSerial = heap->GetLastSubmission();
 
         // If the next candidate for eviction was inserted into the LRU during the current serial,
@@ -170,30 +168,37 @@ namespace dawn_native { namespace d3d12 {
         return heap;
     }
 
-    // Any time we need to make something resident in local memory, we must check that we have
+    MaybeError ResidencyManager::EnsureCanAllocate(uint64_t allocationSize,
-    // enough free memory to make the new object resident while also staying within our budget.
+                                                   MemorySegment memorySegment) {
-    // If there isn't enough memory, we should evict until there is.
-    MaybeError ResidencyManager::EnsureCanMakeResident(uint64_t sizeToMakeResident) {
         if (!mResidencyManagementEnabled) {
             return {};
         }
 
-        UpdateMemorySegmentInfo(&mVideoMemoryInfo.local);
+        return EnsureCanMakeResident(allocationSize, GetMemorySegmentInfo(memorySegment));
+    }
 
-        uint64_t memoryUsageAfterMakeResident = sizeToMakeResident + mVideoMemoryInfo.local.usage;
+    // Any time we need to make something resident, we must check that we have enough free memory to
+    // make the new object resident while also staying within budget. If there isn't enough
+    // memory, we should evict until there is.
+    MaybeError ResidencyManager::EnsureCanMakeResident(uint64_t sizeToMakeResident,
+                                                       MemorySegmentInfo* memorySegment) {
+        ASSERT(mResidencyManagementEnabled);
+
+        UpdateMemorySegmentInfo(memorySegment);
+
+        uint64_t memoryUsageAfterMakeResident = sizeToMakeResident + memorySegment->usage;
 
         // Return when we can call MakeResident and remain under budget.
-        if (memoryUsageAfterMakeResident < mVideoMemoryInfo.local.budget) {
+        if (memoryUsageAfterMakeResident < memorySegment->budget) {
             return {};
         }
 
         std::vector<ID3D12Pageable*> resourcesToEvict;
-        uint64_t sizeNeededToBeUnderBudget =
+        uint64_t sizeNeededToBeUnderBudget = memoryUsageAfterMakeResident - memorySegment->budget;
-            memoryUsageAfterMakeResident - mVideoMemoryInfo.local.budget;
         uint64_t sizeEvicted = 0;
         while (sizeEvicted < sizeNeededToBeUnderBudget) {
             Heap* heap;
-            DAWN_TRY_ASSIGN(heap, RemoveSingleEntryFromLRU());
+            DAWN_TRY_ASSIGN(heap, RemoveSingleEntryFromLRU(memorySegment));
 
             // If no heap was returned, then nothing more can be evicted.
             if (heap == nullptr) {
@@ -207,29 +212,12 @@ namespace dawn_native { namespace d3d12 {
         if (resourcesToEvict.size() > 0) {
             DAWN_TRY(CheckHRESULT(
                 mDevice->GetD3D12Device()->Evict(resourcesToEvict.size(), resourcesToEvict.data()),
-                "Evicting resident heaps to free device local memory"));
+                "Evicting resident heaps to free memory"));
         }
 
         return {};
     }
 
-    // Ensure that we are only tracking heaps that exist in DXGI_MEMORY_SEGMENT_LOCAL.
-    bool ResidencyManager::ShouldTrackHeap(Heap* heap) const {
-        D3D12_HEAP_PROPERTIES heapProperties =
-            mDevice->GetD3D12Device()->GetCustomHeapProperties(0, heap->GetD3D12HeapType());
-
-        if (mDevice->GetDeviceInfo().isUMA) {
-            // On UMA devices, MEMORY_POOL_L0 corresponds to MEMORY_SEGMENT_LOCAL, so we must track
-            // heaps in MEMORY_POOL_L0. For UMA, all heaps types exist in MEMORY_POOL_L0.
-            return heapProperties.MemoryPoolPreference == D3D12_MEMORY_POOL_L0;
-        }
-
-        // On non-UMA devices, MEMORY_POOL_L1 corresponds to MEMORY_SEGMENT_LOCAL, so only track the
-        // heap if it is in MEMORY_POOL_L1. For non-UMA, DEFAULT heaps exist in MEMORY_POOL_L1,
-        // while READBACK and UPLOAD heaps exist in MEMORY_POOL_L0.
-        return heapProperties.MemoryPoolPreference == D3D12_MEMORY_POOL_L1;
-    }
-
     // Given a list of heaps that are pending usage, this function will estimate memory needed,
     // evict resources until enough space is available, then make resident any heaps scheduled for
     // usage.
@@ -239,17 +227,13 @@ namespace dawn_native { namespace d3d12 {
         }
 
         std::vector<ID3D12Pageable*> heapsToMakeResident;
-        uint64_t sizeToMakeResident = 0;
+        uint64_t localSizeToMakeResident = 0;
+        uint64_t nonLocalSizeToMakeResident = 0;
 
         Serial pendingCommandSerial = mDevice->GetPendingCommandSerial();
         for (size_t i = 0; i < heapCount; i++) {
             Heap* heap = heaps[i];
 
-            // Depending on device architecture, the heap may not need tracked.
-            if (!ShouldTrackHeap(heap)) {
-                continue;
-            }
-
             // Heaps that are locked resident are not tracked in the LRU cache.
             if (heap->IsResidencyLocked()) {
                 continue;
@@ -261,7 +245,11 @@ namespace dawn_native { namespace d3d12 {
                 heap->RemoveFromList();
             } else {
                 heapsToMakeResident.push_back(heap->GetD3D12Pageable().Get());
-                sizeToMakeResident += heap->GetSize();
+                if (heap->GetMemorySegment() == MemorySegment::Local) {
+                    localSizeToMakeResident += heap->GetSize();
+                } else {
+                    nonLocalSizeToMakeResident += heap->GetSize();
+                }
             }
 
             // If we submit a command list to the GPU, we must ensure that heaps referenced by that
@@ -270,12 +258,20 @@ namespace dawn_native { namespace d3d12 {
             // eligible for eviction, even though some evictions may be possible.
             heap->SetLastSubmission(pendingCommandSerial);
 
-            mLRUCache.Append(heap);
+            // Insert the heap into the appropriate LRU.
+            TrackResidentAllocation(heap);
+        }
+
+        if (localSizeToMakeResident > 0) {
+            DAWN_TRY(EnsureCanMakeResident(localSizeToMakeResident, &mVideoMemoryInfo.local));
+        }
+
+        if (nonLocalSizeToMakeResident > 0) {
+            ASSERT(!mDevice->GetDeviceInfo().isUMA);
+            DAWN_TRY(EnsureCanMakeResident(nonLocalSizeToMakeResident, &mVideoMemoryInfo.nonLocal));
         }
 
         if (heapsToMakeResident.size() != 0) {
-            DAWN_TRY(EnsureCanMakeResident(sizeToMakeResident));
-
             // Note that MakeResident is a synchronous function and can add a significant
             // overhead to command recording. In the future, it may be possible to decrease this
             // overhead by using MakeResident on a secondary thread, or by instead making use of
@@ -283,32 +279,28 @@ namespace dawn_native { namespace d3d12 {
             // platforms).
             DAWN_TRY(CheckHRESULT(mDevice->GetD3D12Device()->MakeResident(
                                       heapsToMakeResident.size(), heapsToMakeResident.data()),
-                                  "Making scheduled-to-be-used resources resident in "
+                                  "Making scheduled-to-be-used resources resident"));
-                                  "device local memory"));
         }
 
         return {};
     }
 
-    // When a new heap is allocated, the heap will be made resident upon creation. We must track
+    // Inserts a heap at the bottom of the LRU. The passed heap must be resident or scheduled to
-    // when this happens to avoid calling MakeResident a second time.
+    // become resident within the current serial.
     void ResidencyManager::TrackResidentAllocation(Heap* heap) {
         if (!mResidencyManagementEnabled) {
             return;
         }
 
-        // Depending on device architecture and heap type, the heap may not need tracked.
+        ASSERT(heap->IsInList() == false);
-        if (!ShouldTrackHeap(heap)) {
+        GetMemorySegmentInfo(heap->GetMemorySegment())->lruCache.Append(heap);
-            return;
-        }
-
-        mLRUCache.Append(heap);
     }
 
     // Places an artifical cap on Dawn's budget so we can test in a predictable manner. If used,
     // this function must be called before any resources have been created.
     void ResidencyManager::RestrictBudgetForTesting(uint64_t artificialBudgetCap) {
-        ASSERT(mLRUCache.empty());
+        ASSERT(mVideoMemoryInfo.local.lruCache.empty());
+        ASSERT(mVideoMemoryInfo.nonLocal.lruCache.empty());
         ASSERT(!mRestrictBudgetForTesting);
 
         mRestrictBudgetForTesting = true;

src/dawn_native/d3d12/ResidencyManagerD3D12.h

@@ -34,7 +34,8 @@ namespace dawn_native { namespace d3d12 {
 
         MaybeError LockHeap(Heap* heap);
         void UnlockHeap(Heap* heap);
-        MaybeError EnsureCanMakeResident(uint64_t allocationSize);
+
+        MaybeError EnsureCanAllocate(uint64_t allocationSize, MemorySegment memorySegment);
         MaybeError EnsureHeapsAreResident(Heap** heaps, size_t heapCount);
 
         uint64_t SetExternalMemoryReservation(MemorySegment segment,
@@ -47,24 +48,25 @@ namespace dawn_native { namespace d3d12 {
       private:
         struct MemorySegmentInfo {
             const DXGI_MEMORY_SEGMENT_GROUP dxgiSegment;
-            uint64_t budget;
+            LinkedList<Heap> lruCache = {};
-            uint64_t usage;
+            uint64_t budget = 0;
-            uint64_t externalReservation;
+            uint64_t usage = 0;
-            uint64_t externalRequest;
+            uint64_t externalReservation = 0;
+            uint64_t externalRequest = 0;
         };
 
         struct VideoMemoryInfo {
-            MemorySegmentInfo local = {DXGI_MEMORY_SEGMENT_GROUP_LOCAL, 0, 0, 0, 0};
+            MemorySegmentInfo local = {DXGI_MEMORY_SEGMENT_GROUP_LOCAL};
-            MemorySegmentInfo nonLocal = {DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, 0, 0, 0, 0};
+            MemorySegmentInfo nonLocal = {DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL};
         };
 
-        ResultOrError<Heap*> RemoveSingleEntryFromLRU();
+        MemorySegmentInfo* GetMemorySegmentInfo(MemorySegment memorySegment);
-        bool ShouldTrackHeap(Heap* heap) const;
+        MaybeError EnsureCanMakeResident(uint64_t allocationSize, MemorySegmentInfo* memorySegment);
+        ResultOrError<Heap*> RemoveSingleEntryFromLRU(MemorySegmentInfo* memorySegment);
         void UpdateVideoMemoryInfo();
         void UpdateMemorySegmentInfo(MemorySegmentInfo* segmentInfo);
 
         Device* mDevice;
-        LinkedList<Heap> mLRUCache;
         bool mResidencyManagementEnabled = false;
         bool mRestrictBudgetForTesting = false;
         VideoMemoryInfo mVideoMemoryInfo = {};

src/dawn_native/d3d12/ResourceAllocatorManagerD3D12.cpp

@@ -22,6 +22,21 @@
 
 namespace dawn_native { namespace d3d12 {
     namespace {
+        MemorySegment GetMemorySegment(Device* device, D3D12_HEAP_TYPE heapType) {
+            if (device->GetDeviceInfo().isUMA) {
+                return MemorySegment::Local;
+            }
+
+            D3D12_HEAP_PROPERTIES heapProperties =
+                device->GetD3D12Device()->GetCustomHeapProperties(0, heapType);
+
+            if (heapProperties.MemoryPoolPreference == D3D12_MEMORY_POOL_L1) {
+                return MemorySegment::Local;
+            }
+
+            return MemorySegment::NonLocal;
+        }
+
         D3D12_HEAP_TYPE GetD3D12HeapType(ResourceHeapKind resourceHeapKind) {
             switch (resourceHeapKind) {
                 case Readback_OnlyBuffers:
@@ -143,7 +158,8 @@ namespace dawn_native { namespace d3d12 {
         for (uint32_t i = 0; i < ResourceHeapKind::EnumCount; i++) {
             const ResourceHeapKind resourceHeapKind = static_cast<ResourceHeapKind>(i);
             mHeapAllocators[i] = std::make_unique<HeapAllocator>(
-                mDevice, GetD3D12HeapType(resourceHeapKind), GetD3D12HeapFlags(resourceHeapKind));
+                mDevice, GetD3D12HeapType(resourceHeapKind), GetD3D12HeapFlags(resourceHeapKind),
+                GetMemorySegment(device, GetD3D12HeapType(resourceHeapKind)));
             mSubAllocatedResourceAllocators[i] = std::make_unique<BuddyMemoryAllocator>(
                 kMaxHeapSize, kMinHeapSize, mHeapAllocators[i].get());
         }
@@ -310,7 +326,8 @@ namespace dawn_native { namespace d3d12 {
         // CreateCommittedResource will implicitly make the created resource resident. We must
         // ensure enough free memory exists before allocating to avoid an out-of-memory error when
         // overcommitted.
-        DAWN_TRY(mDevice->GetResidencyManager()->EnsureCanMakeResident(resourceInfo.SizeInBytes));
+        DAWN_TRY(mDevice->GetResidencyManager()->EnsureCanAllocate(
+            resourceInfo.SizeInBytes, GetMemorySegment(mDevice, heapType)));
 
         // Note: Heap flags are inferred by the resource descriptor and do not need to be explicitly
         // provided to CreateCommittedResource.
@@ -326,7 +343,8 @@ namespace dawn_native { namespace d3d12 {
         // heap granularity, every directly allocated ResourceHeapAllocation also stores a Heap
         // object. This object is created manually, and must be deleted manually upon deallocation
         // of the committed resource.
-        Heap* heap = new Heap(committedResource, heapType, resourceInfo.SizeInBytes);
+        Heap* heap = new Heap(committedResource, GetMemorySegment(mDevice, heapType),
+                              resourceInfo.SizeInBytes);
 
         // Calling CreateCommittedResource implicitly calls MakeResident on the resource. We must
         // track this to avoid calling MakeResident a second time.

src/tests/white_box/D3D12ResidencyTests.cpp

@@ -26,6 +26,12 @@ constexpr uint32_t kDirectlyAllocatedResourceSize = 5000000;  // 5MB
 constexpr uint32_t kSuballocatedResourceSize = 1000000;       // 1MB
 constexpr uint32_t kSourceBufferSize = 4;                     // 4B
 
+constexpr wgpu::BufferUsage kMapReadBufferUsage =
+    wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead;
+constexpr wgpu::BufferUsage kMapWriteBufferUsage =
+    wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::MapWrite;
+constexpr wgpu::BufferUsage kNonMappableBufferUsage = wgpu::BufferUsage::CopyDst;
+
 class D3D12ResidencyTests : public DawnTest {
   protected:
     void TestSetUp() override {
@@ -43,11 +49,13 @@ class D3D12ResidencyTests : public DawnTest {
             utils::CreateBufferFromData(device, &one, sizeof(one), wgpu::BufferUsage::CopySrc);
     }
 
-    std::vector<wgpu::Buffer> AllocateBuffers(uint32_t bufferSize, uint32_t numberOfBuffers) {
+    std::vector<wgpu::Buffer> AllocateBuffers(uint32_t bufferSize,
+                                              uint32_t numberOfBuffers,
+                                              wgpu::BufferUsage usage) {
         std::vector<wgpu::Buffer> buffers;
 
         for (uint64_t i = 0; i < numberOfBuffers; i++) {
-            buffers.push_back(CreateBuffer(bufferSize, wgpu::BufferUsage::CopyDst));
+            buffers.push_back(CreateBuffer(bufferSize, usage));
         }
 
         return buffers;
@@ -121,7 +129,8 @@ class D3D12ResidencyTests : public DawnTest {
 TEST_P(D3D12ResidencyTests, OvercommitSmallResources) {
     // Create suballocated buffers to fill half the budget.
     std::vector<wgpu::Buffer> bufferSet1 = AllocateBuffers(
-        kSuballocatedResourceSize, ((kRestrictedBudgetSize / 2) / kSuballocatedResourceSize));
+        kSuballocatedResourceSize, ((kRestrictedBudgetSize / 2) / kSuballocatedResourceSize),
+        kNonMappableBufferUsage);
 
     // Check that all the buffers allocated are resident. Also make sure they were suballocated
     // internally.
@@ -133,7 +142,8 @@ TEST_P(D3D12ResidencyTests, OvercommitSmallResources) {
 
     // Create enough directly-allocated buffers to use the entire budget.
     std::vector<wgpu::Buffer> bufferSet2 = AllocateBuffers(
-        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize);
+        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize,
+        kNonMappableBufferUsage);
 
     // Check that everything in bufferSet1 is now evicted.
     for (uint32_t i = 0; i < bufferSet1.size(); i++) {
@@ -157,9 +167,9 @@ TEST_P(D3D12ResidencyTests, OvercommitSmallResources) {
 // correctly.
 TEST_P(D3D12ResidencyTests, OvercommitLargeResources) {
     // Create directly-allocated buffers to fill half the budget.
-    std::vector<wgpu::Buffer> bufferSet1 =
+    std::vector<wgpu::Buffer> bufferSet1 = AllocateBuffers(
-        AllocateBuffers(kDirectlyAllocatedResourceSize,
+        kDirectlyAllocatedResourceSize,
-                        ((kRestrictedBudgetSize / 2) / kDirectlyAllocatedResourceSize));
+        ((kRestrictedBudgetSize / 2) / kDirectlyAllocatedResourceSize), kNonMappableBufferUsage);
 
     // Check that all the allocated buffers are resident. Also make sure they were directly
     // allocated internally.
@@ -170,7 +180,8 @@ TEST_P(D3D12ResidencyTests, OvercommitLargeResources) {
 
     // Create enough directly-allocated buffers to use the entire budget.
     std::vector<wgpu::Buffer> bufferSet2 = AllocateBuffers(
-        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize);
+        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize,
+        kNonMappableBufferUsage);
 
     // Check that everything in bufferSet1 is now evicted.
     for (uint32_t i = 0; i < bufferSet1.size(); i++) {
@@ -191,12 +202,8 @@ TEST_P(D3D12ResidencyTests, OvercommitLargeResources) {
 
 // Check that calling MapReadAsync makes the buffer resident and keeps it locked resident.
 TEST_P(D3D12ResidencyTests, AsyncMappedBufferRead) {
-    // Dawn currently only manages LOCAL_MEMORY. Mappable buffers exist in NON_LOCAL_MEMORY on
-    // discrete devices.
-    DAWN_SKIP_TEST_IF(!IsUMA());
-
     // Create a mappable buffer.
-    wgpu::Buffer buffer = CreateBuffer(4, wgpu::BufferUsage::MapRead | wgpu::BufferUsage::CopyDst);
+    wgpu::Buffer buffer = CreateBuffer(4, kMapReadBufferUsage);
 
     uint32_t data = 12345;
     buffer.SetSubData(0, sizeof(uint32_t), &data);
@@ -206,7 +213,8 @@ TEST_P(D3D12ResidencyTests, AsyncMappedBufferRead) {
 
     // Create and touch enough buffers to use the entire budget.
     std::vector<wgpu::Buffer> bufferSet = AllocateBuffers(
-        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize);
+        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize,
+        kMapReadBufferUsage);
     TouchBuffers(0, bufferSet.size(), bufferSet);
 
     // The mappable buffer should have been evicted.
@@ -229,25 +237,23 @@ TEST_P(D3D12ResidencyTests, AsyncMappedBufferRead) {
     // This should evict the mappable buffer.
     buffer.Unmap();
     std::vector<wgpu::Buffer> bufferSet2 = AllocateBuffers(
-        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize);
+        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize,
+        kMapReadBufferUsage);
     TouchBuffers(0, bufferSet2.size(), bufferSet2);
     EXPECT_FALSE(CheckIfBufferIsResident(buffer));
 }
 
 // Check that calling MapWriteAsync makes the buffer resident and keeps it locked resident.
 TEST_P(D3D12ResidencyTests, AsyncMappedBufferWrite) {
-    // Dawn currently only manages LOCAL_MEMORY. Mappable buffers exist in NON_LOCAL_MEMORY on
-    // discrete devices.
-    DAWN_SKIP_TEST_IF(!IsUMA());
-
     // Create a mappable buffer.
-    wgpu::Buffer buffer = CreateBuffer(4, wgpu::BufferUsage::MapWrite | wgpu::BufferUsage::CopySrc);
+    wgpu::Buffer buffer = CreateBuffer(4, kMapWriteBufferUsage);
     // The mappable buffer should be resident.
     EXPECT_TRUE(CheckIfBufferIsResident(buffer));
 
     // Create and touch enough buffers to use the entire budget.
     std::vector<wgpu::Buffer> bufferSet1 = AllocateBuffers(
-        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize);
+        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize,
+        kMapReadBufferUsage);
     TouchBuffers(0, bufferSet1.size(), bufferSet1);
 
     // The mappable buffer should have been evicted.
@@ -270,7 +276,8 @@ TEST_P(D3D12ResidencyTests, AsyncMappedBufferWrite) {
     // This should evict the mappable buffer.
     buffer.Unmap();
     std::vector<wgpu::Buffer> bufferSet2 = AllocateBuffers(
-        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize);
+        kDirectlyAllocatedResourceSize, kRestrictedBudgetSize / kDirectlyAllocatedResourceSize,
+        kMapReadBufferUsage);
     TouchBuffers(0, bufferSet2.size(), bufferSet2);
     EXPECT_FALSE(CheckIfBufferIsResident(buffer));
 }
@@ -281,7 +288,8 @@ TEST_P(D3D12ResidencyTests, OvercommitInASingleSubmit) {
     constexpr uint32_t numberOfBuffersToOvercommit = 5;
     std::vector<wgpu::Buffer> bufferSet1 = AllocateBuffers(
         kDirectlyAllocatedResourceSize,
-        (kRestrictedBudgetSize / kDirectlyAllocatedResourceSize) + numberOfBuffersToOvercommit);
+        (kRestrictedBudgetSize / kDirectlyAllocatedResourceSize) + numberOfBuffersToOvercommit,
+        kNonMappableBufferUsage);
     // Touch the buffers, which creates an overcommitted command list.
     TouchBuffers(0, bufferSet1.size(), bufferSet1);
     // Ensure that all of these buffers are resident, even though we're exceeding the budget.
@@ -292,7 +300,8 @@ TEST_P(D3D12ResidencyTests, OvercommitInASingleSubmit) {
     // Allocate another set of buffers that exceeds the budget.
     std::vector<wgpu::Buffer> bufferSet2 = AllocateBuffers(
         kDirectlyAllocatedResourceSize,
-        (kRestrictedBudgetSize / kDirectlyAllocatedResourceSize) + numberOfBuffersToOvercommit);
+        (kRestrictedBudgetSize / kDirectlyAllocatedResourceSize) + numberOfBuffersToOvercommit,
+        kNonMappableBufferUsage);
     // Ensure the first <numberOfBuffersToOvercommit> buffers in the second buffer set were evicted,
     // since they shouldn't fit in the budget.
     for (uint32_t i = 0; i < numberOfBuffersToOvercommit; i++) {