Uses MTLResidencySet for TLAS dependencies.

Author: Lichtso

wgpu-hal/src/metal/command.rs

@@ -1,3 +1,4 @@
+use hashbrown::HashSet;
 use objc2::{
     rc::{autoreleasepool, Retained},
     runtime::ProtocolObject,
@@ -7,8 +8,8 @@ use objc2_metal::{
     MTLAccelerationStructureCommandEncoder, MTLBlitCommandEncoder, MTLBlitPassDescriptor,
     MTLCommandBuffer, MTLCommandEncoder, MTLCommandQueue, MTLComputeCommandEncoder,
     MTLComputePassDescriptor, MTLCounterDontSample, MTLLoadAction, MTLPrimitiveType,
-    MTLRenderCommandEncoder, MTLRenderPassDescriptor, MTLScissorRect, MTLSize, MTLStoreAction,
-    MTLTexture, MTLViewport, MTLVisibilityResultMode,
+    MTLRenderCommandEncoder, MTLRenderPassDescriptor, MTLResidencySet, MTLScissorRect, MTLSize,
+    MTLStoreAction, MTLTexture, MTLViewport, MTLVisibilityResultMode,
 };
 
 use super::{conv, TimestampQuerySupport};
@@ -37,6 +38,7 @@ impl Default for super::CommandState {
                 height: 0,
             },
             stage_infos: Default::default(),
+            residency_sets: Default::default(),
             storage_buffer_length_map: Default::default(),
             vertex_buffer_size_map: Default::default(),
             work_group_memory_sizes: Vec::new(),
@@ -284,8 +286,11 @@ impl crate::CommandEncoder for super::CommandEncoder {
         debug_assert!(self.state.compute.is_none());
         debug_assert!(self.state.pending_timer_queries.is_empty());
 
+        let mut residency_sets = HashSet::new();
+        core::mem::swap(&mut residency_sets, &mut self.state.residency_sets);
         Ok(super::CommandBuffer {
             raw: self.raw_cmd_buf.take().unwrap(),
+            residency_sets,
         })
     }
 
@@ -455,6 +460,14 @@ impl crate::CommandEncoder for super::CommandEncoder {
         dst: &super::AccelerationStructure,
         copy: wgt::AccelerationStructureCopy,
     ) {
+        self.state.residency_sets.insert(dst.residency_set.clone());
+        dst.residency_set.requestResidency();
+        dst.residency_set.removeAllAllocations();
+        let allocations = src.residency_set.allAllocations();
+        for index in 0..allocations.count() {
+            dst.residency_set
+                .addAllocation(&allocations.objectAtIndex(index));
+        }
         let command_encoder = self.enter_acceleration_structure_builder();
         match copy {
             wgt::AccelerationStructureCopy::Clone => {
@@ -1479,6 +1492,31 @@ impl crate::CommandEncoder for super::CommandEncoder {
         for descriptor in descriptors {
             let acceleration_structure_descriptor =
                 conv::map_acceleration_structure_descriptor(descriptor.entries, descriptor.flags);
+            if matches!(
+                descriptor.entries,
+                crate::AccelerationStructureEntries::Instances(_)
+            ) {
+                self.state.residency_sets.insert(
+                    descriptor
+                        .destination_acceleration_structure
+                        .residency_set
+                        .clone(),
+                );
+                descriptor
+                    .destination_acceleration_structure
+                    .residency_set
+                    .requestResidency();
+                descriptor
+                    .destination_acceleration_structure
+                    .residency_set
+                    .removeAllAllocations();
+                for dependency in descriptor.dependencies.iter() {
+                    descriptor
+                        .destination_acceleration_structure
+                        .residency_set
+                        .addAllocation(dependency.raw.as_ref());
+                }
+            }
             match descriptor.mode {
                 crate::AccelerationStructureBuildMode::Build => {
                     command_encoder

wgpu-hal/src/metal/device.rs

@@ -17,10 +17,11 @@ use objc2_metal::{
     MTLDevice, MTLFunction, MTLIndirectAccelerationStructureInstanceDescriptor, MTLLanguageVersion,
     MTLLibrary, MTLMutability, MTLPackedFloat3, MTLPackedFloat4x3,
     MTLPipelineBufferDescriptorArray, MTLPixelFormat, MTLPrimitiveTopologyClass,
-    MTLRenderPipelineDescriptor, MTLResource, MTLResourceID, MTLResourceOptions,
-    MTLSamplerAddressMode, MTLSamplerDescriptor, MTLSamplerMipFilter, MTLSamplerState, MTLSize,
-    MTLStencilDescriptor, MTLStorageMode, MTLTexture, MTLTextureDescriptor, MTLTextureType,
-    MTLTriangleFillMode, MTLVertexDescriptor, MTLVertexStepFunction,
+    MTLRenderPipelineDescriptor, MTLResidencySetDescriptor, MTLResource, MTLResourceID,
+    MTLResourceOptions, MTLSamplerAddressMode, MTLSamplerDescriptor, MTLSamplerMipFilter,
+    MTLSamplerState, MTLSize, MTLStencilDescriptor, MTLStorageMode, MTLTexture,
+    MTLTextureDescriptor, MTLTextureType, MTLTriangleFillMode, MTLVertexDescriptor,
+    MTLVertexStepFunction,
 };
 use parking_lot::Mutex;
 
@@ -1694,10 +1695,18 @@ impl crate::Device for super::Device {
         // self.counters.acceleration_structures.add(1);
         let device = self.shared.device.lock();
         autoreleasepool(|_| {
+            let residency_set_descriptor = MTLResidencySetDescriptor::new();
+            residency_set_descriptor.setInitialCapacity(1);
             Ok(super::AccelerationStructure {
                 raw: device
                     .newAccelerationStructureWithSize(descriptor.size as usize)
                     .ok_or(crate::DeviceError::OutOfMemory)?,
+                residency_set: device
+                    .newResidencySetWithDescriptor_error(&residency_set_descriptor)
+                    .map_err(|err| {
+                        log::error!("Failed to create residency set: {err:?}");
+                        crate::DeviceError::Unexpected
+                    })?,
             })
         })
     }

wgpu-hal/src/metal/mod.rs

@@ -37,7 +37,7 @@ use core::{fmt, iter, ops, ptr::NonNull, sync::atomic};
 
 use arrayvec::ArrayVec;
 use bitflags::bitflags;
-use hashbrown::HashMap;
+use hashbrown::{HashMap, HashSet};
 use naga::FastHashMap;
 use objc2::{
     rc::{autoreleasepool, Retained},
@@ -50,7 +50,7 @@ use objc2_metal::{
     MTLComputeCommandEncoder, MTLComputePipelineState, MTLCounterSampleBuffer, MTLCullMode,
     MTLDepthClipMode, MTLDepthStencilState, MTLDevice, MTLDrawable, MTLFunction, MTLIndexType,
     MTLLanguageVersion, MTLLibrary, MTLPrimitiveType, MTLReadWriteTextureTier,
-    MTLRenderCommandEncoder, MTLRenderPipelineState, MTLRenderStages, MTLResource,
+    MTLRenderCommandEncoder, MTLRenderPipelineState, MTLRenderStages, MTLResidencySet, MTLResource,
     MTLResourceUsage, MTLSamplerState, MTLSharedEvent, MTLSize, MTLTexture, MTLTextureType,
     MTLTriangleFillMode, MTLWinding,
 };
@@ -489,6 +489,9 @@ impl crate::Queue for Queue {
 
             for cmd_buffer in command_buffers {
                 cmd_buffer.raw.commit();
+                for residency_set in &cmd_buffer.residency_sets {
+                    residency_set.commit();
+                }
             }
 
             if let Some(raw) = extra_command_buffer {
@@ -929,6 +932,7 @@ struct CommandState {
     index: Option<IndexState>,
     raw_wg_size: MTLSize,
     stage_infos: MultiStageData<PipelineStageInfo>,
+    residency_sets: HashSet<Retained<ProtocolObject<dyn MTLResidencySet>>>,
 
     /// Sizes of currently bound [`wgt::BufferBindingType::Storage`] buffers.
     ///
@@ -984,6 +988,7 @@ unsafe impl Sync for CommandEncoder {}
 #[derive(Debug)]
 pub struct CommandBuffer {
     raw: Retained<ProtocolObject<dyn MTLCommandBuffer>>,
+    residency_sets: HashSet<Retained<ProtocolObject<dyn MTLResidencySet>>>,
 }
 
 impl crate::DynCommandBuffer for CommandBuffer {}
@@ -999,6 +1004,7 @@ impl crate::DynPipelineCache for PipelineCache {}
 #[derive(Debug)]
 pub struct AccelerationStructure {
     raw: Retained<ProtocolObject<dyn MTLAccelerationStructure>>,
+    residency_set: Retained<ProtocolObject<dyn MTLResidencySet>>,
 }
 
 impl AccelerationStructure {