Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MarijnS95 merged 1 commit into from
Jun 20, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ objc2-metal = { version = "0.3", default-features = false, features = [
"MTLBuffer",
"MTLDevice",
"MTLHeap",
"MTLResidencySet",
"MTLResource",
"MTLTexture",
"std",
Expand Down
1 change: 1 addition & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ let mut allocator = Allocator::new(&AllocatorCreateDesc {
device: device.clone(),
debug_settings: Default::default(),
allocation_sizes: Default::default(),
create_residency_set: false,
});
```

Expand Down
1 change: 1 addition & 0 deletions examples/metal-buffer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ fn main() {
device: device.clone(),
debug_settings: Default::default(),
allocation_sizes: Default::default(),
create_residency_set: false,
})
.unwrap();

Expand Down
2 changes: 2 additions & 0 deletions src/allocator/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,8 @@ pub(crate) trait SubAllocator: SubAllocatorBase + fmt::Debug + Sync + Send {

fn report_allocations(&self) -> Vec<AllocationReport>;

/// Returns [`true`] if this allocator allows sub-allocating multiple allocations, [`false`] if
/// it is designed to only represent dedicated allocations.
#[must_use]
fn supports_general_allocations(&self) -> bool;
#[must_use]
Expand Down
36 changes: 15 additions & 21 deletions src/d3d12/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -475,28 +475,22 @@ 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

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
// We only want to destroy this now-empty block if it is either a dedicated/personal
// allocation, or a block supporting sub-allocations that is not the last one (ensuring
// there's always at least one block/allocator readily available).
let is_dedicated_or_not_last_general_block =
!mem_block.sub_allocator.supports_general_allocations()
|| self.active_general_blocks > 1;
if mem_block.sub_allocator.is_empty() && is_dedicated_or_not_last_general_block {
let block = self.memory_blocks[block_idx]
.take()
.ok_or_else(|| AllocationError::Internal("Memory block must be Some.".into()))?;

if block.sub_allocator.supports_general_allocations() {
self.active_general_blocks -= 1;
Comment thread
Athosvk marked this conversation as resolved.
}

// Note that `block` will be destroyed on `drop` here
}

Ok(())
Expand Down
2 changes: 2 additions & 0 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,7 @@
//! device: device.clone(),
//! debug_settings: Default::default(),
//! allocation_sizes: Default::default(),
//! create_residency_set: false,
//! });
//! # }
//! # #[cfg(not(feature = "metal"))]
Expand All @@ -177,6 +178,7 @@
//! # device: device.clone(),
//! # debug_settings: Default::default(),
//! # allocation_sizes: Default::default(),
//! # create_residency_set: false,
//! # })
//! # .unwrap();
//! let allocation_desc = AllocationCreateDesc::buffer(
Expand Down
106 changes: 80 additions & 26 deletions src/metal/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@ 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};
#[cfg(doc)]
use objc2_metal::{MTLAllocation, MTLResource};
use objc2_metal::{
MTLCPUCacheMode, MTLDevice, MTLHeap, MTLHeapDescriptor, MTLHeapType, MTLResourceOptions,
MTLStorageMode, MTLTextureDescriptor,
MTLCPUCacheMode, MTLDevice, MTLHeap, MTLHeapDescriptor, MTLHeapType, MTLResidencySet,
MTLResourceOptions, MTLStorageMode, MTLTextureDescriptor,
};

#[cfg(feature = "visualizer")]
Expand Down Expand Up @@ -150,6 +152,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,
Expand All @@ -166,6 +169,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,
Copy link
Copy Markdown
Member

@Jasper-Bekkers Jasper-Bekkers Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe for now its a bit early, but if we do more of these platform specific settings I don't think they belong in AllocatorCreateDesc, maybe we need some platform specific traits.

Copy link
Copy Markdown
Member Author

@MarijnS95 MarijnS95 Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AllocatorCreateDesc already is platform-specific (defined in metal/mod.rs), gpu-allocator has relatively few platform-agnostic types and definitions (just the error, debug settings, allocation reports, and the allocator algorithms themselves). Everything else is platform-specific because it integrates with platform-specific primitives (i.e. having to pass the Device of the given API).

}

#[derive(Debug)]
Expand Down Expand Up @@ -215,6 +221,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,
Expand Down Expand Up @@ -249,6 +256,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) => {
Expand Down Expand Up @@ -317,19 +328,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
};

Expand Down Expand Up @@ -373,28 +388,26 @@ 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
// We only want to destroy this now-empty block if it is either a dedicated/personal
// allocation, or a block supporting sub-allocations that is not the last one (ensuring
// there's always at least one block/allocator readily available).
let is_dedicated_or_not_last_general_block =
!mem_block.sub_allocator.supports_general_allocations()
|| self.active_general_blocks > 1;
if mem_block.sub_allocator.is_empty() && is_dedicated_or_not_last_general_block {
let block = self.memory_blocks[block_idx]
.take()
.ok_or_else(|| AllocationError::Internal("Memory block must be Some.".into()))?;

if block.sub_allocator.supports_general_allocations() {
self.active_general_blocks -= 1;
}

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 let Some(rs) = &self.global_residency_set {
unsafe { rs.removeAllocation(block.heap.as_ref()) }
}

// Note that `block` will be destroyed on `drop` here
}

Ok(())
Expand Down Expand Up @@ -427,10 +440,23 @@ impl Allocator {
}),
];

let global_residency_set = if desc.create_residency_set {
Some(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. Unsupported MacOS/iOS version?")
})
} else {
None
};

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,
Expand All @@ -448,6 +474,7 @@ impl Allocator {
debug_settings: desc.debug_settings,
memory_types,
allocation_sizes: desc.allocation_sizes,
global_residency_set,
})
}

Expand Down Expand Up @@ -557,4 +584,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
/// [`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 will
// keep using and updating the same object going forward.
self.global_residency_set.as_ref()
}
}
32 changes: 15 additions & 17 deletions src/vulkan/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -657,24 +657,22 @@ 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();
let block = block.ok_or_else(|| {
AllocationError::Internal("Memory block must be Some.".into())
})?;
block.destroy(device);

self.active_general_blocks -= 1;
}
} else {
let block = self.memory_blocks[block_idx].take();
let block = block.ok_or_else(|| {
AllocationError::Internal("Memory block must be Some.".into())
})?;
block.destroy(device);
// We only want to destroy this now-empty block if it is either a dedicated/personal
// allocation, or a block supporting sub-allocations that is not the last one (ensuring
// there's always at least one block/allocator readily available).
let is_dedicated_or_not_last_general_block =
!mem_block.sub_allocator.supports_general_allocations()
|| self.active_general_blocks > 1;
if mem_block.sub_allocator.is_empty() && is_dedicated_or_not_last_general_block {
let block = self.memory_blocks[block_idx]
.take()
.ok_or_else(|| AllocationError::Internal("Memory block must be Some.".into()))?;

if block.sub_allocator.supports_general_allocations() {
self.active_general_blocks -= 1;
}

block.destroy(device);
}

Ok(())
Expand Down