Skip to main content

tensor4all_tensorbackend/
memory.rs

1//! Process-level memory pressure helpers.
2
3/// Result reported by [`release_process_allocator_cached_memory`].
4///
5/// The fields are platform-specific because allocator pressure-relief APIs do
6/// not expose a uniform contract. On macOS, `released_bytes` is the value
7/// returned by `malloc_zone_pressure_relief`. On Linux, `success` is the boolean
8/// result returned by `malloc_trim(0)`.
9#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
10pub struct AllocatorPressureRelief {
11    /// Whether tensor4all-rs has a platform hook for the current target.
12    pub supported: bool,
13    /// Platform-reported released bytes when the allocator exposes that value.
14    pub released_bytes: Option<usize>,
15    /// Platform-reported success when the allocator exposes only a status bit.
16    pub success: Option<bool>,
17}
18
19/// Ask the process allocator to return cached/free memory to the operating system.
20///
21/// This is a diagnostic and memory-pressure hook for the platform system
22/// allocator only. It does not release memory that is still owned by live
23/// tensors or explicit buffer pools, and it may have no effect if the program is
24/// built with a custom global allocator.
25///
26/// # Examples
27///
28/// ```
29/// use tensor4all_tensorbackend::release_process_allocator_cached_memory;
30///
31/// let report = release_process_allocator_cached_memory();
32/// assert_eq!(
33///     report.supported,
34///     cfg!(any(target_os = "macos", target_os = "linux"))
35/// );
36/// ```
37pub fn release_process_allocator_cached_memory() -> AllocatorPressureRelief {
38    allocator_pressure_relief()
39}
40
41#[cfg(target_os = "macos")]
42fn allocator_pressure_relief() -> AllocatorPressureRelief {
43    use std::ffi::c_void;
44
45    extern "C" {
46        fn malloc_default_zone() -> *mut c_void;
47        fn malloc_zone_pressure_relief(zone: *mut c_void, goal: usize) -> usize;
48    }
49
50    unsafe {
51        let zone = malloc_default_zone();
52        if zone.is_null() {
53            AllocatorPressureRelief {
54                supported: true,
55                released_bytes: Some(0),
56                success: Some(false),
57            }
58        } else {
59            let released_bytes = malloc_zone_pressure_relief(zone, 0);
60            AllocatorPressureRelief {
61                supported: true,
62                released_bytes: Some(released_bytes),
63                success: Some(released_bytes > 0),
64            }
65        }
66    }
67}
68
69#[cfg(target_os = "linux")]
70fn allocator_pressure_relief() -> AllocatorPressureRelief {
71    extern "C" {
72        fn malloc_trim(pad: usize) -> i32;
73    }
74
75    let success = unsafe { malloc_trim(0) != 0 };
76    AllocatorPressureRelief {
77        supported: true,
78        released_bytes: None,
79        success: Some(success),
80    }
81}
82
83#[cfg(not(any(target_os = "macos", target_os = "linux")))]
84fn allocator_pressure_relief() -> AllocatorPressureRelief {
85    AllocatorPressureRelief {
86        supported: false,
87        released_bytes: None,
88        success: None,
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::release_process_allocator_cached_memory;
95
96    #[test]
97    fn reports_platform_support() {
98        let report = release_process_allocator_cached_memory();
99        assert_eq!(
100            report.supported,
101            cfg!(any(target_os = "macos", target_os = "linux"))
102        );
103    }
104}