Skip to main content

mtl_gpu/
residency_set.rs

1//! Metal residency set for VRAM management.
2//!
3//! Corresponds to `Metal/MTLResidencySet.hpp`.
4//!
5//! Residency sets manage which allocations are resident in GPU memory.
6
7use std::ffi::c_void;
8use std::ptr::NonNull;
9
10use mtl_foundation::{Referencing, UInteger};
11use mtl_sys::{msg_send_0, msg_send_1, msg_send_2, sel};
12
13use crate::Device;
14
15// ============================================================================
16// ResidencySetDescriptor
17// ============================================================================
18
19/// Descriptor for creating a residency set.
20///
21/// C++ equivalent: `MTL::ResidencySetDescriptor`
22#[repr(transparent)]
23pub struct ResidencySetDescriptor(pub(crate) NonNull<c_void>);
24
25impl ResidencySetDescriptor {
26    /// Create a new residency set descriptor.
27    ///
28    /// C++ equivalent: `ResidencySetDescriptor* alloc()->init()`
29    pub fn new() -> Option<Self> {
30        unsafe {
31            let class = mtl_sys::Class::get("MTLResidencySetDescriptor")?;
32            let ptr: *mut c_void = msg_send_0(class.as_ptr(), sel!(alloc));
33            if ptr.is_null() {
34                return None;
35            }
36            let ptr: *mut c_void = msg_send_0(ptr, sel!(init));
37            Self::from_raw(ptr)
38        }
39    }
40
41    /// Create from a raw pointer.
42    ///
43    /// # Safety
44    ///
45    /// The pointer must be a valid Metal ResidencySetDescriptor.
46    #[inline]
47    pub unsafe fn from_raw(ptr: *mut c_void) -> Option<Self> {
48        NonNull::new(ptr).map(Self)
49    }
50
51    /// Get the raw pointer.
52    #[inline]
53    pub fn as_raw(&self) -> *mut c_void {
54        self.0.as_ptr()
55    }
56
57    /// Get the initial capacity.
58    ///
59    /// C++ equivalent: `NS::UInteger initialCapacity() const`
60    #[inline]
61    pub fn initial_capacity(&self) -> UInteger {
62        unsafe { msg_send_0(self.as_ptr(), sel!(initialCapacity)) }
63    }
64
65    /// Set the initial capacity.
66    ///
67    /// C++ equivalent: `void setInitialCapacity(NS::UInteger)`
68    pub fn set_initial_capacity(&self, capacity: UInteger) {
69        unsafe {
70            let _: () = msg_send_1(self.as_ptr(), sel!(setInitialCapacity:), capacity);
71        }
72    }
73
74    /// Get the label.
75    ///
76    /// C++ equivalent: `NS::String* label() const`
77    pub fn label(&self) -> Option<String> {
78        unsafe {
79            let ptr: *mut c_void = msg_send_0(self.as_ptr(), sel!(label));
80            if ptr.is_null() {
81                return None;
82            }
83            let utf8_ptr: *const std::ffi::c_char =
84                mtl_sys::msg_send_0(ptr as *const c_void, sel!(UTF8String));
85            if utf8_ptr.is_null() {
86                return None;
87            }
88            let c_str = std::ffi::CStr::from_ptr(utf8_ptr);
89            Some(c_str.to_string_lossy().into_owned())
90        }
91    }
92
93    /// Set the label.
94    ///
95    /// C++ equivalent: `void setLabel(const NS::String*)`
96    pub fn set_label(&self, label: &str) {
97        if let Some(ns_label) = mtl_foundation::String::from_str(label) {
98            unsafe {
99                let _: () = msg_send_1(self.as_ptr(), sel!(setLabel:), ns_label.as_ptr());
100            }
101        }
102    }
103}
104
105impl Default for ResidencySetDescriptor {
106    fn default() -> Self {
107        Self::new().expect("failed to create ResidencySetDescriptor")
108    }
109}
110
111impl Clone for ResidencySetDescriptor {
112    fn clone(&self) -> Self {
113        unsafe {
114            let ptr: *mut c_void = msg_send_0(self.as_ptr(), sel!(copy));
115            Self::from_raw(ptr).expect("copy should succeed")
116        }
117    }
118}
119
120impl Drop for ResidencySetDescriptor {
121    fn drop(&mut self) {
122        unsafe {
123            msg_send_0::<()>(self.as_ptr(), sel!(release));
124        }
125    }
126}
127
128impl Referencing for ResidencySetDescriptor {
129    #[inline]
130    fn as_ptr(&self) -> *const c_void {
131        self.0.as_ptr()
132    }
133}
134
135unsafe impl Send for ResidencySetDescriptor {}
136unsafe impl Sync for ResidencySetDescriptor {}
137
138// ============================================================================
139// ResidencySet
140// ============================================================================
141
142/// A residency set for managing GPU memory residency.
143///
144/// C++ equivalent: `MTL::ResidencySet`
145#[repr(transparent)]
146pub struct ResidencySet(pub(crate) NonNull<c_void>);
147
148impl ResidencySet {
149    /// Create from a raw pointer.
150    ///
151    /// # Safety
152    ///
153    /// The pointer must be a valid Metal ResidencySet.
154    #[inline]
155    pub unsafe fn from_raw(ptr: *mut c_void) -> Option<Self> {
156        NonNull::new(ptr).map(Self)
157    }
158
159    /// Get the raw pointer.
160    #[inline]
161    pub fn as_raw(&self) -> *mut c_void {
162        self.0.as_ptr()
163    }
164
165    /// Get the device that created this residency set.
166    ///
167    /// C++ equivalent: `Device* device() const`
168    pub fn device(&self) -> Device {
169        unsafe {
170            let ptr: *mut c_void = msg_send_0(self.as_ptr(), sel!(device));
171            // Retain for our reference
172            msg_send_0::<*mut c_void>(ptr as *const c_void, sel!(retain));
173            Device::from_raw(ptr).expect("device should be valid")
174        }
175    }
176
177    /// Get the label.
178    ///
179    /// C++ equivalent: `NS::String* label() const`
180    pub fn label(&self) -> Option<String> {
181        unsafe {
182            let ptr: *mut c_void = msg_send_0(self.as_ptr(), sel!(label));
183            if ptr.is_null() {
184                return None;
185            }
186            let utf8_ptr: *const std::ffi::c_char =
187                mtl_sys::msg_send_0(ptr as *const c_void, sel!(UTF8String));
188            if utf8_ptr.is_null() {
189                return None;
190            }
191            let c_str = std::ffi::CStr::from_ptr(utf8_ptr);
192            Some(c_str.to_string_lossy().into_owned())
193        }
194    }
195
196    /// Get the allocated size in bytes.
197    ///
198    /// C++ equivalent: `uint64_t allocatedSize() const`
199    #[inline]
200    pub fn allocated_size(&self) -> u64 {
201        unsafe { msg_send_0(self.as_ptr(), sel!(allocatedSize)) }
202    }
203
204    /// Get the allocation count.
205    ///
206    /// C++ equivalent: `NS::UInteger allocationCount() const`
207    #[inline]
208    pub fn allocation_count(&self) -> UInteger {
209        unsafe { msg_send_0(self.as_ptr(), sel!(allocationCount)) }
210    }
211
212    /// Get all allocations as a raw NS::Array pointer.
213    ///
214    /// C++ equivalent: `NS::Array* allAllocations() const`
215    #[inline]
216    pub fn all_allocations_ptr(&self) -> *const c_void {
217        unsafe { msg_send_0(self.as_ptr(), sel!(allAllocations)) }
218    }
219
220    /// Add an allocation using a raw pointer.
221    ///
222    /// C++ equivalent: `void addAllocation(const MTL::Allocation*)`
223    pub fn add_allocation_ptr(&self, allocation: *const c_void) {
224        unsafe {
225            let _: () = msg_send_1(self.as_ptr(), sel!(addAllocation:), allocation);
226        }
227    }
228
229    /// Add multiple allocations.
230    ///
231    /// C++ equivalent: `void addAllocations(const MTL::Allocation* const[], NS::UInteger)`
232    pub fn add_allocations_ptr(&self, allocations: *const *const c_void, count: UInteger) {
233        unsafe {
234            let _: () = msg_send_2(
235                self.as_ptr(),
236                sel!(addAllocations:count:),
237                allocations,
238                count,
239            );
240        }
241    }
242
243    /// Remove an allocation using a raw pointer.
244    ///
245    /// C++ equivalent: `void removeAllocation(const MTL::Allocation*)`
246    pub fn remove_allocation_ptr(&self, allocation: *const c_void) {
247        unsafe {
248            let _: () = msg_send_1(self.as_ptr(), sel!(removeAllocation:), allocation);
249        }
250    }
251
252    /// Remove multiple allocations.
253    ///
254    /// C++ equivalent: `void removeAllocations(const MTL::Allocation* const[], NS::UInteger)`
255    pub fn remove_allocations_ptr(&self, allocations: *const *const c_void, count: UInteger) {
256        unsafe {
257            let _: () = msg_send_2(
258                self.as_ptr(),
259                sel!(removeAllocations:count:),
260                allocations,
261                count,
262            );
263        }
264    }
265
266    /// Remove all allocations.
267    ///
268    /// C++ equivalent: `void removeAllAllocations()`
269    pub fn remove_all_allocations(&self) {
270        unsafe {
271            msg_send_0::<()>(self.as_ptr(), sel!(removeAllAllocations));
272        }
273    }
274
275    /// Check if the set contains an allocation.
276    ///
277    /// C++ equivalent: `bool containsAllocation(const MTL::Allocation*)`
278    pub fn contains_allocation_ptr(&self, allocation: *const c_void) -> bool {
279        unsafe { msg_send_1(self.as_ptr(), sel!(containsAllocation:), allocation) }
280    }
281
282    /// Request residency for all allocations in the set.
283    ///
284    /// C++ equivalent: `void requestResidency()`
285    pub fn request_residency(&self) {
286        unsafe {
287            msg_send_0::<()>(self.as_ptr(), sel!(requestResidency));
288        }
289    }
290
291    /// End residency for all allocations in the set.
292    ///
293    /// C++ equivalent: `void endResidency()`
294    pub fn end_residency(&self) {
295        unsafe {
296            msg_send_0::<()>(self.as_ptr(), sel!(endResidency));
297        }
298    }
299
300    /// Commit changes to the residency set.
301    ///
302    /// C++ equivalent: `void commit()`
303    pub fn commit(&self) {
304        unsafe {
305            msg_send_0::<()>(self.as_ptr(), sel!(commit));
306        }
307    }
308}
309
310impl Clone for ResidencySet {
311    fn clone(&self) -> Self {
312        unsafe {
313            msg_send_0::<*mut c_void>(self.as_ptr(), sel!(retain));
314        }
315        Self(self.0)
316    }
317}
318
319impl Drop for ResidencySet {
320    fn drop(&mut self) {
321        unsafe {
322            msg_send_0::<()>(self.as_ptr(), sel!(release));
323        }
324    }
325}
326
327impl Referencing for ResidencySet {
328    #[inline]
329    fn as_ptr(&self) -> *const c_void {
330        self.0.as_ptr()
331    }
332}
333
334unsafe impl Send for ResidencySet {}
335unsafe impl Sync for ResidencySet {}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340
341    #[test]
342    fn test_residency_set_descriptor_creation() {
343        let descriptor = ResidencySetDescriptor::new();
344        assert!(descriptor.is_some());
345    }
346
347    #[test]
348    fn test_residency_set_descriptor_size() {
349        assert_eq!(
350            std::mem::size_of::<ResidencySetDescriptor>(),
351            std::mem::size_of::<*mut c_void>()
352        );
353    }
354
355    #[test]
356    fn test_residency_set_size() {
357        assert_eq!(
358            std::mem::size_of::<ResidencySet>(),
359            std::mem::size_of::<*mut c_void>()
360        );
361    }
362
363    #[test]
364    fn test_residency_set_descriptor_initial_capacity() {
365        let descriptor = ResidencySetDescriptor::new().unwrap();
366        descriptor.set_initial_capacity(100);
367        assert_eq!(descriptor.initial_capacity(), 100);
368    }
369}