Uses MTLResidencySet for TLAS dependencies.

Lichtso · Lichtso · commit aa6be884f133 · 2025-08-10T18:04:20.000+02:00
diff --git a/wgpu-hal/src/metal/command.rs b/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,16 @@ impl crate::CommandEncoder for super::CommandEncoder {
         dst: &super::AccelerationStructure,
         copy: wgt::AccelerationStructureCopy,
     ) {
+        // Store a reference of the residency_set in the encoder state to be committed later
+        self.state
+            .residency_sets
+            .insert(Retained::clone(&dst.residency_set));
+        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 => {
@@ -773,11 +788,12 @@ impl crate::CommandEncoder for super::CommandEncoder {
                             changes_sizes_buffer = true;
                         }
                     }
-                    super::BufferLikeResource::AccelerationStructure(ptr) => {
+                    super::BufferLikeResource::AccelerationStructure(ptr, residency_set) => {
                         encoder.setVertexAccelerationStructure_atBufferIndex(
                             Some(ptr.as_ref()),
                             (bg_info.base_resource_indices.vs.buffers + index) as usize,
                         );
+                        residency_set.as_ref().requestResidency();
                     }
                 }
             }
@@ -822,11 +838,12 @@ impl crate::CommandEncoder for super::CommandEncoder {
                             changes_sizes_buffer = true;
                         }
                     }
-                    super::BufferLikeResource::AccelerationStructure(ptr) => {
+                    super::BufferLikeResource::AccelerationStructure(ptr, residency_set) => {
                         encoder.setFragmentAccelerationStructure_atBufferIndex(
                             Some(ptr.as_ref()),
                             (bg_info.base_resource_indices.fs.buffers + index) as usize,
                         );
+                        residency_set.as_ref().requestResidency();
                     }
                 }
             }
@@ -914,11 +931,12 @@ impl crate::CommandEncoder for super::CommandEncoder {
                             changes_sizes_buffer = true;
                         }
                     }
-                    super::BufferLikeResource::AccelerationStructure(ptr) => {
+                    super::BufferLikeResource::AccelerationStructure(ptr, residency_set) => {
                         encoder.setAccelerationStructure_atBufferIndex(
                             Some(ptr.as_ref()),
                             (bg_info.base_resource_indices.cs.buffers + index) as usize,
                         );
+                        residency_set.as_ref().requestResidency();
                     }
                 }
             }
@@ -1479,6 +1497,25 @@ 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(_)
+            ) {
+                // Store a reference of the residency_set in the encoder state to be committed later
+                self.state.residency_sets.insert(Retained::clone(
+                    &descriptor.destination_acceleration_structure.residency_set,
+                ));
+                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
diff --git a/wgpu-hal/src/metal/device.rs b/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, MTLResidencySet, MTLResidencySetDescriptor, MTLResource,
+    MTLResourceID, MTLResourceOptions, MTLSamplerAddressMode, MTLSamplerDescriptor,
+    MTLSamplerMipFilter, MTLSamplerState, MTLSize, MTLStencilDescriptor, MTLStorageMode,
+    MTLTexture, MTLTextureDescriptor, MTLTextureType, MTLTriangleFillMode, MTLVertexDescriptor,
+    MTLVertexStepFunction,
 };
 use parking_lot::Mutex;
 
@@ -948,10 +949,15 @@ impl crate::Device for super::Device {
                                 {
                                     contents[idx] = acceleration_structure.raw.gpuResourceID();
 
-                                    let use_info = bg
-                                        .resources_to_use
-                                        .entry(acceleration_structure.as_raw().cast())
-                                        .or_default();
+                                    let super::BufferLikeResource::AccelerationStructure(
+                                        structure,
+                                        _residency_set,
+                                    ) = acceleration_structure.as_raw()
+                                    else {
+                                        unreachable!();
+                                    };
+                                    let use_info =
+                                        bg.resources_to_use.entry(structure.cast()).or_default();
                                     use_info.stages |= stages;
                                     use_info.uses |= uses;
                                     use_info.visible_in_compute |=
@@ -1045,11 +1051,7 @@ impl crate::Device for super::Device {
                                 let end = start + 1;
                                 bg.buffers.extend(
                                     desc.acceleration_structures[start..end].iter().map(
-                                        |acceleration_structure| {
-                                            super::BufferLikeResource::AccelerationStructure(
-                                                acceleration_structure.as_raw(),
-                                            )
-                                        },
+                                        |acceleration_structure| acceleration_structure.as_raw(),
                                     ),
                                 );
                                 counter.buffers += 1;
@@ -1694,19 +1696,28 @@ 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
+                    })?,
             })
         })
     }
 
     unsafe fn destroy_acceleration_structure(
         &self,
-        _acceleration_structure: super::AccelerationStructure,
+        acceleration_structure: super::AccelerationStructure,
     ) {
         // self.counters.acceleration_structures.sub(1);
+        acceleration_structure.residency_set.endResidency();
     }
 
     fn tlas_instance_to_bytes(&self, instance: TlasInstance) -> Vec<u8> {
diff --git a/wgpu-hal/src/metal/mod.rs b/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 {
@@ -718,7 +721,10 @@ enum BufferLikeResource {
 
         binding_location: u32,
     },
-    AccelerationStructure(NonNull<ProtocolObject<dyn MTLAccelerationStructure>>),
+    AccelerationStructure(
+        NonNull<ProtocolObject<dyn MTLAccelerationStructure>>,
+        NonNull<ProtocolObject<dyn MTLResidencySet>>,
+    ),
 }
 
 #[derive(Debug)]
@@ -929,6 +935,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 +991,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,11 +1007,17 @@ impl crate::DynPipelineCache for PipelineCache {}
 #[derive(Debug)]
 pub struct AccelerationStructure {
     raw: Retained<ProtocolObject<dyn MTLAccelerationStructure>>,
+    residency_set: Retained<ProtocolObject<dyn MTLResidencySet>>,
 }
 
 impl AccelerationStructure {
-    fn as_raw(&self) -> NonNull<ProtocolObject<dyn MTLAccelerationStructure>> {
-        unsafe { NonNull::new_unchecked(Retained::as_ptr(&self.raw) as *mut _) }
+    fn as_raw(&self) -> BufferLikeResource {
+        unsafe {
+            BufferLikeResource::AccelerationStructure(
+                NonNull::new_unchecked(Retained::as_ptr(&self.raw) as *mut _),
+                NonNull::new_unchecked(Retained::as_ptr(&self.residency_set) as *mut _),
+            )
+        }
     }
 }
 
