[d3d12] refactor: move allocation related fields into a new Allocator struct

wgpu-hal/src/dx12/device.rs

@@ -1,6 +1,6 @@
 use std::{
     borrow::Cow,
-    ffi, mem,
+    ffi,
     num::NonZeroU32,
     ptr,
     string::{String, ToString as _},
@@ -57,9 +57,8 @@ impl super::Device {
             auxil::dxgi::exception::register_exception_handler();
         }
 
-        let (mem_allocator, device_memblock_size, host_memblock_size) =
+        let mem_allocator =
-            suballocation::create_allocator(&raw, memory_hints)?;
+            suballocation::Allocator::new(&raw, memory_hints, memory_budget_thresholds)?;
-        let mem_allocator = Arc::new(mem_allocator);
 
         let idle_fence: Direct3D12::ID3D12Fence = unsafe {
             profiling::scope!("ID3D12Device::CreateFence");
@@ -158,9 +157,6 @@ impl super::Device {
             )?,
             sampler_heap: super::sampler::SamplerHeap::new(&raw, &private_caps)?,
             private_caps,
-            device_memblock_size,
-            host_memblock_size,
-            memory_budget_thresholds,
         };
 
         let mut rtv_pool =
@@ -1939,7 +1935,11 @@ impl crate::Device for super::Device {
             ),
         };
 
-        if let Some(threshold) = self.shared.memory_budget_thresholds.for_resource_creation {
+        if let Some(threshold) = self
+            .mem_allocator
+            .memory_budget_thresholds
+            .for_resource_creation
+        {
             let info = self
                 .shared
                 .adapter
@@ -2285,33 +2285,7 @@ impl crate::Device for super::Device {
     }
 
     fn generate_allocator_report(&self) -> Option<wgt::AllocatorReport> {
-        let mut upstream = self.mem_allocator.lock().generate_report();
+        Some(self.mem_allocator.generate_report())
-
-        let allocations = upstream
-            .allocations
-            .iter_mut()
-            .map(|alloc| wgt::AllocationReport {
-                name: mem::take(&mut alloc.name),
-                offset: alloc.offset,
-                size: alloc.size,
-            })
-            .collect();
-
-        let blocks = upstream
-            .blocks
-            .iter()
-            .map(|block| wgt::MemoryBlockReport {
-                size: block.size,
-                allocations: block.allocations.clone(),
-            })
-            .collect();
-
-        Some(wgt::AllocatorReport {
-            allocations,
-            blocks,
-            total_allocated_bytes: upstream.total_allocated_bytes,
-            total_reserved_bytes: upstream.total_reserved_bytes,
-        })
     }
 
     fn tlas_instance_to_bytes(&self, instance: TlasInstance) -> Vec<u8> {
@@ -2329,7 +2303,7 @@ impl crate::Device for super::Device {
     }
 
     fn check_if_oom(&self) -> Result<(), crate::DeviceError> {
-        let Some(threshold) = self.shared.memory_budget_thresholds.for_device_loss else {
+        let Some(threshold) = self.mem_allocator.memory_budget_thresholds.for_device_loss else {
             return Ok(());
         };
 

wgpu-hal/src/dx12/mod.rs

@@ -89,8 +89,8 @@ mod view;
 use std::{borrow::ToOwned as _, ffi, fmt, mem, num::NonZeroU32, ops::Deref, sync::Arc, vec::Vec};
 
 use arrayvec::ArrayVec;
-use gpu_allocator::d3d12::Allocator;
 use parking_lot::{Mutex, RwLock};
+use suballocation::Allocator;
 use windows::{
     core::{Free, Interface},
     Win32::{
@@ -635,9 +635,6 @@ struct DeviceShared {
     heap_views: descriptor::GeneralHeap,
     sampler_heap: sampler::SamplerHeap,
     private_caps: PrivateCapabilities,
-    device_memblock_size: u64,
-    host_memblock_size: u64,
-    memory_budget_thresholds: wgt::MemoryBudgetThresholds,
 }
 
 unsafe impl Send for DeviceShared {}
@@ -658,7 +655,7 @@ pub struct Device {
     #[cfg(feature = "renderdoc")]
     render_doc: auxil::renderdoc::RenderDoc,
     null_rtv_handle: descriptor::Handle,
-    mem_allocator: Arc<Mutex<Allocator>>,
+    mem_allocator: Allocator,
     dxc_container: Option<Arc<shader_compilation::DxcContainer>>,
     counters: Arc<wgt::HalCounters>,
 }
@@ -800,7 +797,7 @@ pub struct CommandEncoder {
     allocator: Direct3D12::ID3D12CommandAllocator,
     device: Direct3D12::ID3D12Device,
     shared: Arc<DeviceShared>,
-    mem_allocator: Arc<Mutex<Allocator>>,
+    mem_allocator: Allocator,
 
     null_rtv_handle: descriptor::Handle,
     list: Option<Direct3D12::ID3D12GraphicsCommandList>,

wgpu-hal/src/dx12/suballocation.rs

@@ -1,8 +1,6 @@
-use gpu_allocator::{
+use gpu_allocator::{d3d12::AllocationCreateDesc, MemoryLocation};
-    d3d12::{AllocationCreateDesc, Allocator},
-    MemoryLocation,
-};
 use parking_lot::Mutex;
+use std::sync::Arc;
 use windows::Win32::Graphics::{Direct3D12, Dxgi};
 
 use crate::{
@@ -61,10 +59,20 @@ impl Allocation {
     }
 }
 
-pub(crate) fn create_allocator(
+#[derive(Clone)]
+pub(crate) struct Allocator {
+    inner: Arc<Mutex<gpu_allocator::d3d12::Allocator>>,
+    device_memblock_size: u64,
+    host_memblock_size: u64,
+    pub memory_budget_thresholds: wgt::MemoryBudgetThresholds,
+}
+
+impl Allocator {
+    pub(crate) fn new(
         raw: &Direct3D12::ID3D12Device,
         memory_hints: &wgt::MemoryHints,
-) -> Result<(Mutex<Allocator>, u64, u64), crate::DeviceError> {
+        memory_budget_thresholds: wgt::MemoryBudgetThresholds,
+    ) -> Result<Self, crate::DeviceError> {
         // TODO: the allocator's configuration should take hardware capability into
         // account.
         const MB: u64 = 1024 * 1024;
@@ -97,15 +105,47 @@ pub(crate) fn create_allocator(
             allocation_sizes,
         };
 
-    let allocator = Allocator::new(&allocator_desc).inspect_err(|e| {
+        let allocator = gpu_allocator::d3d12::Allocator::new(&allocator_desc).inspect_err(|e| {
             log::error!("Failed to create d3d12 allocator, error: {}", e);
         })?;
 
-    Ok((
+        Ok(Self {
-        Mutex::new(allocator),
+            inner: Arc::new(Mutex::new(allocator)),
             device_memblock_size,
             host_memblock_size,
-    ))
+            memory_budget_thresholds,
+        })
+    }
+
+    pub(crate) fn generate_report(&self) -> wgt::AllocatorReport {
+        let mut upstream = self.inner.lock().generate_report();
+
+        let allocations = upstream
+            .allocations
+            .iter_mut()
+            .map(|alloc| wgt::AllocationReport {
+                name: core::mem::take(&mut alloc.name),
+                offset: alloc.offset,
+                size: alloc.size,
+            })
+            .collect();
+
+        let blocks = upstream
+            .blocks
+            .iter()
+            .map(|block| wgt::MemoryBlockReport {
+                size: block.size,
+                allocations: block.allocations.clone(),
+            })
+            .collect();
+
+        wgt::AllocatorReport {
+            allocations,
+            blocks,
+            total_allocated_bytes: upstream.total_allocated_bytes,
+            total_reserved_bytes: upstream.total_reserved_bytes,
+        }
+    }
 }
 
 /// To allow us to construct buffers from both a `Device` and `CommandEncoder`
@@ -114,7 +154,7 @@ pub(crate) fn create_allocator(
 pub(crate) struct DeviceAllocationContext<'a> {
     pub(crate) raw: &'a Direct3D12::ID3D12Device,
     pub(crate) shared: &'a super::DeviceShared,
-    pub(crate) mem_allocator: &'a Mutex<Allocator>,
+    pub(crate) mem_allocator: &'a Allocator,
     pub(crate) counters: &'a wgt::HalCounters,
 }
 
@@ -248,7 +288,7 @@ impl<'a> DeviceAllocationContext<'a> {
         counter.sub(allocation.size() as isize);
 
         if let AllocationInner::Placed { inner } = allocation.inner {
-            match self.mem_allocator.lock().free(inner) {
+            match self.mem_allocator.inner.lock().free(inner) {
                 Ok(_) => (),
                 // TODO: Don't panic here
                 Err(e) => panic!("Failed to destroy dx12 {:?}, {e}", allocation.ty),
@@ -269,7 +309,7 @@ impl<'a> DeviceAllocationContext<'a> {
     ) -> Result<(Direct3D12::ID3D12Resource, Allocation), crate::DeviceError> {
         let name = desc.label.unwrap_or("Unlabeled buffer");
 
-        let mut allocator = self.mem_allocator.lock();
+        let mut allocator = self.mem_allocator.inner.lock();
 
         let allocation_desc = AllocationCreateDesc {
             name,
@@ -308,7 +348,7 @@ impl<'a> DeviceAllocationContext<'a> {
     ) -> Result<(Direct3D12::ID3D12Resource, Allocation), crate::DeviceError> {
         let name = desc.label.unwrap_or("Unlabeled texture");
 
-        let mut allocator = self.mem_allocator.lock();
+        let mut allocator = self.mem_allocator.inner.lock();
 
         let allocation_desc = AllocationCreateDesc {
             name,
@@ -347,7 +387,7 @@ impl<'a> DeviceAllocationContext<'a> {
     ) -> Result<(Direct3D12::ID3D12Resource, Allocation), crate::DeviceError> {
         let name = desc.label.unwrap_or("Unlabeled acceleration structure");
 
-        let mut allocator = self.mem_allocator.lock();
+        let mut allocator = self.mem_allocator.inner.lock();
 
         let allocation_desc = AllocationCreateDesc {
             name,
@@ -526,7 +566,11 @@ impl<'a> DeviceAllocationContext<'a> {
                 .GetResourceAllocationInfo(0, std::slice::from_ref(desc))
         };
 
-        let Some(threshold) = self.shared.memory_budget_thresholds.for_resource_creation else {
+        let Some(threshold) = self
+            .mem_allocator
+            .memory_budget_thresholds
+            .for_resource_creation
+        else {
             return Ok(allocation_info);
         };
 
@@ -552,8 +596,10 @@ impl<'a> DeviceAllocationContext<'a> {
 
         let memblock_size = match location {
             MemoryLocation::Unknown => unreachable!(),
-            MemoryLocation::GpuOnly => self.shared.device_memblock_size,
+            MemoryLocation::GpuOnly => self.mem_allocator.device_memblock_size,
-            MemoryLocation::CpuToGpu | MemoryLocation::GpuToCpu => self.shared.host_memblock_size,
+            MemoryLocation::CpuToGpu | MemoryLocation::GpuToCpu => {
+                self.mem_allocator.host_memblock_size
+            }
         };
 
         if info.CurrentUsage + allocation_info.SizeInBytes.max(memblock_size)