metal: Create a global residency set, holding all allocated heaps

`gpu-allocator` creates heaps on which callers can allocate ranges
and create "placed" resources like textures, buffers and acceleration
structures.  These individual resources, or the heaps as a whole,
need to be made resident on the command buffer or even globally on an
entire queue.

In the previous API those heaps had to be made resident on individual
command *encoders* with `useHeap(s):` (making an entire heap resident
perfectly matches a bindless design, as opposed to making every
individual resource -either placed on the heap or allocated separately-
resident with `useResource(s):`).  Worse, this API only applies
`MTLResourceUsageRead` (exluding `RenderTarget` and `ShaderWrite`
textures) which would disallow any resources on the heap to be written.

Now with `MTLResidencySet` multiple heaps can be made resident with one
call, defeating the performance overhead of individually "using" all
heaps on *every* command *encoder*.  But without tracking this inside
`gpu-allocator`, users of our crate still have to manually rebuild
this `MTLResidencySet` each time they change their allocations, without
knowing when `gpu-allocator` created or destroyed a heap.

By managing a single updated `MTLResidencySet` in `gpu-allocator`,
callers can simply call `.commit()` on this object right before they
submit command buffers referencing resources on these heaps, as long as
they have the residency set attached to the queue in question or "used"
on the command buffer that is being submitted.  This removes all the
performance overhead of repeatedly creating `MTLResidencySet`s, which
otherwise defeats the purpose of it over plain `useHeap(s):` call(s).

Author: MarijnS95

Cargo.toml

@@ -42,6 +42,7 @@ objc2-metal = { version = "0.3", default-features = false, features = [
     "MTLBuffer",
     "MTLDevice",
     "MTLHeap",
+    "MTLResidencySet",
     "MTLResource",
     "MTLTexture",
     "std",

src/metal/mod.rs

@@ -2,10 +2,10 @@ use std::{backtrace::Backtrace, sync::Arc};
 
 use log::debug;
 use objc2::{rc::Retained, runtime::ProtocolObject};
-use objc2_foundation::NSString;
+use objc2_foundation::{ns_string, NSString};
 use objc2_metal::{
-    MTLCPUCacheMode, MTLDevice, MTLHeap, MTLHeapDescriptor, MTLHeapType, MTLResourceOptions,
-    MTLStorageMode, MTLTextureDescriptor,
+    MTLCPUCacheMode, MTLDevice, MTLHeap, MTLHeapDescriptor, MTLHeapType, MTLResidencySet,
+    MTLResourceOptions, MTLStorageMode, MTLTextureDescriptor,
 };
 
 #[cfg(feature = "visualizer")]
@@ -150,6 +150,7 @@ impl<'a> AllocationCreateDesc<'a> {
 
 pub struct Allocator {
     device: Retained<ProtocolObject<dyn MTLDevice>>,
+    global_residency_set: Option<Retained<ProtocolObject<dyn MTLResidencySet>>>,
     debug_settings: AllocatorDebugSettings,
     memory_types: Vec<MemoryType>,
     allocation_sizes: AllocationSizes,
@@ -166,6 +167,9 @@ pub struct AllocatorCreateDesc {
     pub device: Retained<ProtocolObject<dyn MTLDevice>>,
     pub debug_settings: AllocatorDebugSettings,
     pub allocation_sizes: AllocationSizes,
+    /// Whether to create a [`MTLResidencySet`] containing all live heaps, that can be retrieved via
+    /// [`Allocator::residency_set()`].  Only supported on MacOS 15.0+ / iOS 18.0+.
+    pub create_residency_set: bool,
 }
 
 #[derive(Debug)]
@@ -215,6 +219,7 @@ impl MemoryBlock {
 
 #[derive(Debug)]
 struct MemoryType {
+    global_residency_set: Option<Retained<ProtocolObject<dyn MTLResidencySet>>>,
     memory_blocks: Vec<Option<MemoryBlock>>,
     _committed_allocations: CommittedAllocationStatistics,
     memory_location: MemoryLocation,
@@ -249,6 +254,10 @@ impl MemoryType {
                 self.memory_location,
             )?;
 
+            if let Some(rs) = &self.global_residency_set {
+                unsafe { rs.addAllocation(mem_block.heap.as_ref()) }
+            }
+
             let block_index = self.memory_blocks.iter().position(|block| block.is_none());
             let block_index = match block_index {
                 Some(i) => {
@@ -317,19 +326,23 @@ impl MemoryType {
             }
         }
 
-        let new_memory_block = MemoryBlock::new(
+        let mem_block = MemoryBlock::new(
             device,
             memblock_size,
             &self.heap_properties,
             false,
             self.memory_location,
         )?;
 
+        if let Some(rs) = &self.global_residency_set {
+            unsafe { rs.addAllocation(mem_block.heap.as_ref()) }
+        }
+
         let new_block_index = if let Some(block_index) = empty_block_index {
-            self.memory_blocks[block_index] = Some(new_memory_block);
+            self.memory_blocks[block_index] = Some(mem_block);
             block_index
         } else {
-            self.memory_blocks.push(Some(new_memory_block));
+            self.memory_blocks.push(Some(mem_block));
             self.memory_blocks.len() - 1
         };
 
@@ -373,28 +386,25 @@ impl MemoryType {
 
         mem_block.sub_allocator.free(allocation.chunk_id)?;
 
-        if mem_block.sub_allocator.is_empty() {
-            if mem_block.sub_allocator.supports_general_allocations() {
-                if self.active_general_blocks > 1 {
-                    let block = self.memory_blocks[block_idx].take();
-                    if block.is_none() {
-                        return Err(AllocationError::Internal(
-                            "Memory block must be Some.".into(),
-                        ));
-                    }
-                    // Note that `block` will be destroyed on `drop` here
+        if mem_block.sub_allocator.is_empty()
+            && (!mem_block.sub_allocator.supports_general_allocations()
+                || self.active_general_blocks > 1)
+        {
+            let Some(block) = self.memory_blocks[block_idx].take() else {
+                return Err(AllocationError::Internal(
+                    "Memory block must be Some.".into(),
+                ));
+            };
 
-                    self.active_general_blocks -= 1;
-                }
-            } else {
-                let block = self.memory_blocks[block_idx].take();
-                if block.is_none() {
-                    return Err(AllocationError::Internal(
-                        "Memory block must be Some.".into(),
-                    ));
-                }
-                // Note that `block` will be destroyed on `drop` here
+            if block.sub_allocator.supports_general_allocations() {
+                self.active_general_blocks -= 1;
+            }
+
+            if let Some(rs) = &self.global_residency_set {
+                unsafe { rs.removeAllocation(block.heap.as_ref()) }
             }
+
+            // Note that `block` will be destroyed on `drop` here
         }
 
         Ok(())
@@ -427,10 +437,19 @@ impl Allocator {
             }),
         ];
 
+        let global_residency_set = desc.create_residency_set.then(|| unsafe {
+            let rs_desc = objc2_metal::MTLResidencySetDescriptor::new();
+            rs_desc.setLabel(Some(ns_string!("gpu-allocator global residency set")));
+            desc.device
+                .newResidencySetWithDescriptor_error(&rs_desc)
+                .expect("Failed to create MTLResidencySet.  Too low MacOS/iOS version?")
+        });
+
         let memory_types = heap_types
             .into_iter()
             .enumerate()
             .map(|(i, (memory_location, heap_descriptor))| MemoryType {
+                global_residency_set: global_residency_set.clone(),
                 memory_blocks: vec![],
                 _committed_allocations: CommittedAllocationStatistics {
                     num_allocations: 0,
@@ -448,6 +467,7 @@ impl Allocator {
             debug_settings: desc.debug_settings,
             memory_types,
             allocation_sizes: desc.allocation_sizes,
+            global_residency_set,
         })
     }
 
@@ -557,4 +577,31 @@ impl Allocator {
 
         total_capacity_bytes
     }
+
+    /// Optional residency set containing all heap allocations created/owned by this allocator to
+    /// be made resident at once when its allocations are used on the GPU.  The caller _must_ invoke
+    /// [`MTLResidencySet::commit()`] whenever these resources are used to make sure the latest
+    /// changes are visible to Metal, e.g. before committing a command buffer.
+    ///
+    /// This residency set can be attached to individual command buffers or to a queue directly
+    /// since usage of allocated resources is expected to be global.
+    ///
+    /// Alternatively callers can build up their own residency set(s) based on individual
+    /// [`objc2_metal::MTLAllocation`]s [^heap-allocation] rather than making all heaps allocated
+    /// via `gpu-allocator` resident at once.
+    ///
+    /// [^heap-allocation]: Note that [`MTLHeap`]s returned by [`Allocator::heaps()`] are also
+    /// allocations.  If individual placed [`MTLResource`]s on a heap are made resident, the entire
+    /// heap will be made resident.
+    ///
+    /// Callers still need to be careful to make resources created outside of `gpu-allocator`
+    /// resident on the GPU, such as indirect command buffers.
+    ///
+    /// This residency set is only available when requested via
+    /// [`AllocatorCreateDesc::create_residency_set`], otherwise this function returns [`None`].
+    pub fn residency_set(&self) -> Option<&Retained<ProtocolObject<dyn MTLResidencySet>>> {
+        // Return the retained object so that the caller also has a way to store it,
+        // since we
+        self.global_residency_set.as_ref()
+    }
 }