[dawn][d3d12] Support D3D12_HEAP_TYPE_CUSTOM in SBM

Custom heap types with specific heap properties are equivalent to
the other heap types.

Bug: 441704688
Change-Id: Iae22f4d44d0ecd613a4e2525ca605567bf700503
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/259896
Commit-Queue: Enrico Galli <enrico.galli@intel.com>
Reviewed-by: Rafael Cintron <rafael.cintron@microsoft.com>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
diff --git a/src/dawn/native/d3d12/SharedBufferMemoryD3D12.cpp b/src/dawn/native/d3d12/SharedBufferMemoryD3D12.cpp
index 4dca66d..8e94eff 100644
--- a/src/dawn/native/d3d12/SharedBufferMemoryD3D12.cpp
+++ b/src/dawn/native/d3d12/SharedBufferMemoryD3D12.cpp
@@ -40,6 +40,53 @@
 
 namespace dawn::native::d3d12 {
 
+namespace {
+
+enum class HeapAccessType {
+    Upload,
+    Readback,
+    GPUQueueAccessible,
+};
+
+ResultOrError<HeapAccessType> MapToHeapAccessType(const D3D12_HEAP_PROPERTIES& heapProperties,
+                                                  const Device* device) {
+    switch (heapProperties.Type) {
+        case D3D12_HEAP_TYPE_UPLOAD:
+            return HeapAccessType::Upload;
+        case D3D12_HEAP_TYPE_READBACK:
+            return HeapAccessType::Readback;
+        case D3D12_HEAP_TYPE_DEFAULT:
+            return HeapAccessType::GPUQueueAccessible;
+        case D3D12_HEAP_TYPE_CUSTOM:
+            if (device->GetDeviceInfo().isUMA) {
+                // On UMA systems, all heaps are always GPU accessible.
+                return HeapAccessType::GPUQueueAccessible;
+            }
+
+            // Map D3D12_HEAP_TYPE_CUSTOM heap to one of the standard heap types if possible.
+            // See:
+            // https://learn.microsoft.com/en-us/windows/win32/api/d3d12/nf-d3d12-id3d12device-getcustomheapproperties(uint_d3d12_heap_type)
+            if (heapProperties.CPUPageProperty == D3D12_CPU_PAGE_PROPERTY_NOT_AVAILABLE &&
+                heapProperties.MemoryPoolPreference == D3D12_MEMORY_POOL_L1) {
+                // A CUSTOM heap with no CPU access and in L1 is equivalent to a DEFAULT heap.
+                return HeapAccessType::GPUQueueAccessible;
+            } else if (heapProperties.CPUPageProperty == D3D12_CPU_PAGE_PROPERTY_WRITE_BACK &&
+                       heapProperties.MemoryPoolPreference == D3D12_MEMORY_POOL_L0) {
+                // A CUSTOM heap with WRITE_BACK + L0 is equivalent to a READBACK heap.
+                return HeapAccessType::Readback;
+            } else if (heapProperties.CPUPageProperty == D3D12_CPU_PAGE_PROPERTY_WRITE_COMBINE &&
+                       heapProperties.MemoryPoolPreference == D3D12_MEMORY_POOL_L0) {
+                // A CUSTOM heap with WRITE_COMBINE + L0 is equivalent to a UPLOAD heap.
+                return HeapAccessType::Upload;
+            } else {
+                return DAWN_VALIDATION_ERROR("ID3D12Resources allocated on unsupported heap.");
+            }
+        default:
+            return DAWN_VALIDATION_ERROR("ID3D12Resources allocated on unsupported heap.");
+    }
+}
+}  // namespace
+
 SharedBufferMemory::SharedBufferMemory(Device* device,
                                        StringView label,
                                        SharedBufferMemoryProperties properties,
@@ -74,15 +121,17 @@
     d3d12Resource->GetHeapProperties(&heapProperties, &heapFlags);
 
     wgpu::BufferUsage usages = wgpu::BufferUsage::None;
+    HeapAccessType heapType;
+    DAWN_TRY_ASSIGN(heapType, MapToHeapAccessType(heapProperties, device));
 
-    switch (heapProperties.Type) {
-        case D3D12_HEAP_TYPE_UPLOAD:
+    switch (heapType) {
+        case HeapAccessType::Upload:
             usages |= wgpu::BufferUsage::MapWrite | wgpu::BufferUsage::CopySrc;
             break;
-        case D3D12_HEAP_TYPE_READBACK:
+        case HeapAccessType::Readback:
             usages |= wgpu::BufferUsage::MapRead | wgpu::BufferUsage::CopyDst;
             break;
-        case D3D12_HEAP_TYPE_DEFAULT:
+        case HeapAccessType::GPUQueueAccessible:
             usages |= wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst |
                       wgpu::BufferUsage::Vertex | wgpu::BufferUsage::Index |
                       wgpu::BufferUsage::Indirect | wgpu::BufferUsage::QueryResolve;
@@ -92,14 +141,26 @@
             if (IsAligned(desc.Width, D3D12_CONSTANT_BUFFER_DATA_PLACEMENT_ALIGNMENT)) {
                 usages |= wgpu::BufferUsage::Uniform;
             }
+
+            if (device->GetDeviceInfo().isUMA &&
+                device->HasFeature(Feature::BufferMapExtendedUsages)) {
+                // On UMA systems, buffers with WRITE_COMBINE or WRITE_BACK heaps can also be
+                // mapped.
+                if (heapProperties.CPUPageProperty == D3D12_CPU_PAGE_PROPERTY_WRITE_COMBINE) {
+                    usages |= wgpu::BufferUsage::MapWrite;
+                } else if (heapProperties.CPUPageProperty == D3D12_CPU_PAGE_PROPERTY_WRITE_BACK) {
+                    // On cache-coherent UMA systems, writes are immediately visible to the GPU. On
+                    // non-cache-coherent UMA systems, writes are flushed to the GPU when unmapping
+                    // or submitting work to the queue (driver dependent). Since Dawn doesn't
+                    // support submitting work to the queue while the buffer is mapped, it should be
+                    // safe to allow MapWrite on WRITE_BACK heaps. For reads, the data is guaranteed
+                    // to be available to the CPU after Map().
+                    usages |= wgpu::BufferUsage::MapRead | wgpu::BufferUsage::MapWrite;
+                }
+            }
             break;
-        case D3D12_HEAP_TYPE_CUSTOM:
-            return DAWN_VALIDATION_ERROR(
-                "ID3D12Resources allocated on D3D12_HEAP_TYPE_CUSTOM heaps are not supported by "
-                "SharedBufferMemory.");
-        default:
-            DAWN_UNREACHABLE();
     }
+
     SharedBufferMemoryProperties properties;
     properties.size = desc.Width;
     properties.usage = usages;
diff --git a/src/dawn/tests/white_box/SharedBufferMemoryTests_win.cpp b/src/dawn/tests/white_box/SharedBufferMemoryTests_win.cpp
index f70374a..8d000e7 100644
--- a/src/dawn/tests/white_box/SharedBufferMemoryTests_win.cpp
+++ b/src/dawn/tests/white_box/SharedBufferMemoryTests_win.cpp
@@ -169,9 +169,17 @@
     ComPtr<ID3D12Resource> CreateD3D12Buffer(ID3D12Device* device,
                                              D3D12_HEAP_TYPE heapType,
                                              uint32_t bufferSize = kBufferSize) {
+        D3D12_HEAP_PROPERTIES heapProperties = {heapType, D3D12_CPU_PAGE_PROPERTY_UNKNOWN,
+                                                D3D12_MEMORY_POOL_UNKNOWN, 0, 0};
+        return CreateD3D12Buffer(device, heapProperties, bufferSize);
+    }
+
+    ComPtr<ID3D12Resource> CreateD3D12Buffer(ID3D12Device* device,
+                                             D3D12_HEAP_PROPERTIES heapProperties,
+                                             uint32_t bufferSize = kBufferSize) {
         D3D12_RESOURCE_STATES initialResourceState;
         D3D12_RESOURCE_FLAGS resourceFlags = D3D12_RESOURCE_FLAG_NONE;
-        switch (heapType) {
+        switch (heapProperties.Type) {
             case D3D12_HEAP_TYPE_UPLOAD:
                 initialResourceState = D3D12_RESOURCE_STATE_GENERIC_READ;
                 break;
@@ -183,9 +191,6 @@
                 resourceFlags = D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS;
         }
 
-        D3D12_HEAP_PROPERTIES heapProperties = {heapType, D3D12_CPU_PAGE_PROPERTY_UNKNOWN,
-                                                D3D12_MEMORY_POOL_UNKNOWN, 0, 0};
-
         D3D12_RESOURCE_DESC descriptor;
         descriptor.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
         descriptor.Alignment = 0;
@@ -227,7 +232,7 @@
         static_cast<Backend*>(GetParam().mBackend)->CreateD3D12Device(device, true);
     ComPtr<ID3D12Resource> d3d12Resource =
         static_cast<Backend*>(GetParam().mBackend)
-            ->CreateD3D12Buffer(warpDevice.Get(), D3D12_HEAP_TYPE_UPLOAD, D3D12_RESOURCE_FLAG_NONE);
+            ->CreateD3D12Buffer(warpDevice.Get(), D3D12_HEAP_TYPE_UPLOAD);
     wgpu::SharedBufferMemoryDescriptor desc;
     native::d3d12::SharedBufferMemoryD3D12ResourceDescriptor sharedD3d12ResourceDesc;
     sharedD3d12ResourceDesc.resource = d3d12Resource.Get();
@@ -236,6 +241,44 @@
     ASSERT_DEVICE_ERROR(device.ImportSharedBufferMemory(&desc));
 }
 
+// Validate that importing an ID3D12Resource allocated on a CUSTOM heap that is equivalent to UPLOAD
+// works correctly.
+TEST_P(SharedBufferMemoryTests, CustomUploadHeapImport) {
+    ComPtr<ID3D12Device> d3d12Device =
+        static_cast<Backend*>(GetParam().mBackend)->CreateD3D12Device(device, false);
+    D3D12_HEAP_PROPERTIES heapProperties =
+        d3d12Device->GetCustomHeapProperties(0, D3D12_HEAP_TYPE_UPLOAD);
+    wgpu::SharedBufferMemoryDescriptor desc;
+    ComPtr<ID3D12Resource> d3d12Resource =
+        static_cast<Backend*>(GetParam().mBackend)
+            ->CreateD3D12Buffer(d3d12Device.Get(), heapProperties);
+    native::d3d12::SharedBufferMemoryD3D12ResourceDescriptor sharedD3d12ResourceDesc;
+    sharedD3d12ResourceDesc.resource = d3d12Resource.Get();
+    desc.nextInChain = &sharedD3d12ResourceDesc;
+
+    wgpu::SharedBufferMemory sharedBufferMemory = device.ImportSharedBufferMemory(&desc);
+    ASSERT_TRUE(sharedBufferMemory.CreateBuffer().Get());
+}
+
+// Validate that importing an ID3D12Resource allocated on a CUSTOM heap that is equivalent to
+// READBACK works correctly.
+TEST_P(SharedBufferMemoryTests, CustomReadbackHeapImport) {
+    ComPtr<ID3D12Device> d3d12Device =
+        static_cast<Backend*>(GetParam().mBackend)->CreateD3D12Device(device, false);
+    D3D12_HEAP_PROPERTIES heapProperties =
+        d3d12Device->GetCustomHeapProperties(0, D3D12_HEAP_TYPE_READBACK);
+    wgpu::SharedBufferMemoryDescriptor desc;
+    ComPtr<ID3D12Resource> d3d12Resource =
+        static_cast<Backend*>(GetParam().mBackend)
+            ->CreateD3D12Buffer(d3d12Device.Get(), heapProperties);
+    native::d3d12::SharedBufferMemoryD3D12ResourceDescriptor sharedD3d12ResourceDesc;
+    sharedD3d12ResourceDesc.resource = d3d12Resource.Get();
+    desc.nextInChain = &sharedD3d12ResourceDesc;
+
+    wgpu::SharedBufferMemory sharedBufferMemory = device.ImportSharedBufferMemory(&desc);
+    ASSERT_TRUE(sharedBufferMemory.CreateBuffer().Get());
+}
+
 DAWN_INSTANTIATE_PREFIXED_TEST_P(D3D12,
                                  SharedBufferMemoryTests,
                                  {D3D12Backend()},