Port dx12 backend to new suballocator

Author: cwfitzgerald

Cargo.lock

@@ -154,8 +154,9 @@ dx12 = [
     "dep:parking_lot",
     "dep:profiling",
     "dep:range-alloc",
+    "dep:wgpu-block-pool",
+    "dep:wgpu-offset-allocator",
     "dep:windows-core",
-    "gpu-allocator/d3d12",
     "naga/hlsl-out",
     "once_cell/std",
     "windows/Win32_Devices_DeviceAndDriverInstallation",
@@ -218,7 +219,7 @@ naga.workspace = true
 wgpu-naga-bridge.workspace = true
 naga-types = { workspace = true }
 wgpu-types = { workspace = true, default-features = false }
-# Backend: Vulkan (in-repo suballocator)
+# Backend: Vulkan and Dx12 (in-repo suballocator)
 wgpu-block-pool = { workspace = true, optional = true }
 wgpu-offset-allocator = { workspace = true, optional = true }
 

wgpu-hal/Cargo.toml

@@ -305,6 +305,7 @@ impl super::CommandEncoder {
                     &crate::dx12::QuerySet {
                         raw: query_set_raw,
                         raw_ty: Direct3D12::D3D12_QUERY_TYPE_TIMESTAMP,
+                        tracked_query_bytes: 0,
                     },
                     index,
                 );

wgpu-hal/src/dx12/command.rs

@@ -60,8 +60,12 @@ impl super::Device {
             auxil::dxgi::exception::register_exception_handler();
         }
 
-        let mem_allocator =
-            suballocation::Allocator::new(&raw, memory_hints, memory_budget_thresholds)?;
+        let mem_allocator = suballocation::Allocator::new(
+            &raw,
+            memory_hints,
+            &private_caps,
+            memory_budget_thresholds,
+        )?;
 
         let idle_fence: Direct3D12::ID3D12Fence = unsafe {
             profiling::scope!("ID3D12Device::CreateFence");
@@ -2293,26 +2297,17 @@ impl crate::Device for super::Device {
             ),
         };
 
-        if let Some(threshold) = self
+        // Query heaps allocate device memory outside the pool machinery, so we
+        // gate them through the allocator's budget probe. Assume each query is
+        // 256 bytes. On an AMD W6800 with driver version 32.0.12030.9, occlusion
+        // and pipeline statistics are 256 bytes, timestamp is 8.
+        let query_bytes = (desc.count as u64).saturating_mul(256);
+        let tracked_query_bytes = self
             .mem_allocator
-            .memory_budget_thresholds
-            .for_resource_creation
-        {
-            let info = self
-                .shared
-                .adapter
-                .query_video_memory_info(Dxgi::DXGI_MEMORY_SEGMENT_GROUP_LOCAL)?;
-
-            // Assume each query is 256 bytes.
-            // On an AMD W6800 with driver version 32.0.12030.9, occlusion and pipeline statistics are 256, timestamp is 8.
-
-            if info.CurrentUsage + desc.count as u64 * 256 >= info.Budget / 100 * threshold as u64 {
-                return Err(crate::DeviceError::OutOfMemory);
-            }
-        }
+            .add_external_allocation(&self.shared.adapter, query_bytes)?;
 
         let mut raw = None::<Direct3D12::ID3D12QueryHeap>;
-        unsafe {
+        let create_result = unsafe {
             self.raw.CreateQueryHeap(
                 &Direct3D12::D3D12_QUERY_HEAP_DESC {
                     Type: heap_ty,
@@ -2322,20 +2317,39 @@ impl crate::Device for super::Device {
                 &mut raw,
             )
         }
-        .into_device_result("Query heap creation")?;
+        .into_device_result("Query heap creation");
+        if let Err(error) = create_result {
+            self.mem_allocator
+                .remove_external_allocation(tracked_query_bytes);
+            return Err(error);
+        }
 
-        let raw = raw.ok_or(crate::DeviceError::Unexpected)?;
+        let Some(raw) = raw else {
+            self.mem_allocator
+                .remove_external_allocation(tracked_query_bytes);
+            return Err(crate::DeviceError::Unexpected);
+        };
 
         if let Some(label) = desc.label {
-            raw.set_name(label)?;
+            if let Err(error) = raw.set_name(label) {
+                self.mem_allocator
+                    .remove_external_allocation(tracked_query_bytes);
+                return Err(error);
+            }
         }
 
         self.counters.query_sets.add(1);
 
-        Ok(super::QuerySet { raw, raw_ty })
+        Ok(super::QuerySet {
+            raw,
+            raw_ty,
+            tracked_query_bytes,
+        })
     }
 
-    unsafe fn destroy_query_set(&self, _set: super::QuerySet) {
+    unsafe fn destroy_query_set(&self, set: super::QuerySet) {
+        self.mem_allocator
+            .remove_external_allocation(set.tracked_query_bytes);
         self.counters.query_sets.sub(1);
     }
 

wgpu-hal/src/dx12/device.rs

@@ -559,7 +559,6 @@ enum MemoryArchitecture {
 struct PrivateCapabilities {
     instance_flags: wgt::InstanceFlags,
     workarounds: Workarounds,
-    #[allow(unused)]
     heterogeneous_resource_heaps: bool,
     memory_architecture: MemoryArchitecture,
     heap_create_not_zeroed: bool,
@@ -1087,6 +1086,7 @@ static_assertions::assert_impl_all!(Sampler: Send, Sync);
 pub struct QuerySet {
     raw: Direct3D12::ID3D12QueryHeap,
     raw_ty: Direct3D12::D3D12_QUERY_TYPE,
+    tracked_query_bytes: u64,
 }
 
 impl crate::DynQuerySet for QuerySet {}