iree_embedded/
arena.rs

1//! A fixed-buffer allocator exposed to IREE through its `iree_allocator_t`
2//! vtable. The buffer is supplied by the caller (a `Vec` on the host, a
3//! `static mut` on the board), so there is no global heap and memory use is
4//! bounded and known.
5
6use core::alloc::Layout;
7use core::ffi::c_void;
8use core::ptr::NonNull;
9use core::sync::atomic::{AtomicUsize, Ordering};
10
11use iree_embedded_sys as sys;
12use spin::Mutex;
13use talc::{ClaimOnOom, Span, Talc};
14
15/// Byte length of the most recent allocation the arena could not satisfy (0 if
16/// none). Useful for diagnosing on-device out-of-memory failures.
17pub static LAST_ALLOC_FAIL_SIZE: AtomicUsize = AtomicUsize::new(0);
18
19/// Bytes reserved before each allocation to store its size (IREE's `free` does
20/// not pass the size back, so we record it). Also keeps user data 16-aligned.
21const HEADER: usize = 16;
22/// IREE's default allocation alignment (`iree_max_align_t` on 64-bit targets).
23const ALIGN: usize = 16;
24
25/// A fixed-size memory pool backing every IREE allocation. Built over a static
26/// byte buffer the caller provides; IREE's runtime allocates and frees objects
27/// inside it through a `talc` allocator, so it behaves like a heap of a
28/// compile-time-constant size.
29pub struct Arena {
30    talc: Mutex<Talc<ClaimOnOom>>,
31}
32
33impl Arena {
34    /// Build an arena over `buffer`. The arena (and every IREE object created
35    /// with it) must not outlive the buffer; `'static` enforces that here.
36    pub fn new(buffer: &'static mut [u8]) -> Self {
37        let span = Span::from_base_size(buffer.as_mut_ptr(), buffer.len());
38        // SAFETY: the buffer is exclusively owned by this Talc for its life.
39        let talc = Talc::new(unsafe { ClaimOnOom::new(span) });
40        Arena {
41            talc: Mutex::new(talc),
42        }
43    }
44
45    /// An `iree_allocator_t` backed by this arena. The arena must outlive it.
46    pub fn as_iree_allocator(&self) -> sys::iree_allocator_t {
47        sys::iree_allocator_t {
48            self_: self as *const Arena as *mut c_void,
49            ctl: Some(arena_ctl),
50        }
51    }
52
53    fn alloc(&self, byte_length: usize, zero: bool) -> *mut c_void {
54        let Ok(layout) = Layout::from_size_align(byte_length + HEADER, ALIGN) else {
55            return core::ptr::null_mut();
56        };
57        let mut talc = self.talc.lock();
58        // SAFETY: layout has non-zero size (HEADER > 0).
59        let base = match unsafe { talc.malloc(layout) } {
60            Ok(p) => p.as_ptr(),
61            Err(_) => {
62                LAST_ALLOC_FAIL_SIZE.store(byte_length, Ordering::Relaxed);
63                return core::ptr::null_mut();
64            }
65        };
66        // SAFETY: base points to a fresh block of `byte_length + HEADER` bytes.
67        unsafe {
68            (base as *mut usize).write(byte_length);
69            let user = base.add(HEADER);
70            if zero {
71                core::ptr::write_bytes(user, 0, byte_length);
72            }
73            user as *mut c_void
74        }
75    }
76
77    /// SAFETY: `user` must be null or a pointer previously returned by `alloc`.
78    unsafe fn free(&self, user: *mut c_void) {
79        if user.is_null() {
80            return;
81        }
82        // SAFETY: per the caller contract, a length header written by `alloc`
83        // sits HEADER bytes below `user`.
84        unsafe {
85            let base = (user as *mut u8).sub(HEADER);
86            let byte_length = (base as *const usize).read();
87            let layout = Layout::from_size_align_unchecked(byte_length + HEADER, ALIGN);
88            let mut talc = self.talc.lock();
89            talc.free(NonNull::new_unchecked(base), layout);
90        }
91    }
92
93    /// SAFETY: `user` must be null or a pointer previously returned by `alloc`.
94    unsafe fn realloc(&self, user: *mut c_void, new_len: usize) -> *mut c_void {
95        if user.is_null() {
96            return self.alloc(new_len, false);
97        }
98        // SAFETY: per the caller contract, a length header written by `alloc`
99        // sits HEADER bytes below `user`; both blocks are at least
100        // `old_len.min(new_len)` bytes.
101        unsafe {
102            let base = (user as *mut u8).sub(HEADER);
103            let old_len = (base as *const usize).read();
104            let new_ptr = self.alloc(new_len, false);
105            if new_ptr.is_null() {
106                return core::ptr::null_mut();
107            }
108            core::ptr::copy_nonoverlapping(
109                user as *const u8,
110                new_ptr as *mut u8,
111                old_len.min(new_len),
112            );
113            self.free(user);
114            new_ptr
115        }
116    }
117}
118
119// SAFETY: all interior mutation goes through the Mutex.
120unsafe impl Send for Arena {}
121unsafe impl Sync for Arena {}
122
123/// IREE routes malloc/calloc/realloc/free through this single control function.
124/// SAFETY: invoked by IREE with a `self` pointer to a live `Arena`.
125unsafe extern "C" fn arena_ctl(
126    self_: *mut c_void,
127    command: sys::iree_allocator_command_t,
128    params: *const c_void,
129    inout_ptr: *mut *mut c_void,
130) -> sys::iree_status_t {
131    // SAFETY: IREE invokes this with `self_` pointing at a live `Arena` and
132    // `params`/`inout_ptr` valid for the given command, per the
133    // `iree_allocator_ctl_fn_t` contract.
134    unsafe {
135        let arena = &*(self_ as *const Arena);
136        let cmd = command;
137
138        if cmd == sys::IREE_ALLOCATOR_COMMAND_FREE {
139            arena.free(*inout_ptr);
140            *inout_ptr = core::ptr::null_mut();
141            return ok();
142        }
143        if cmd == sys::IREE_ALLOCATOR_COMMAND_MALLOC || cmd == sys::IREE_ALLOCATOR_COMMAND_CALLOC {
144            let byte_length = (*(params as *const sys::iree_allocator_alloc_params_t)).byte_length;
145            let zero = cmd == sys::IREE_ALLOCATOR_COMMAND_CALLOC;
146            let p = arena.alloc(byte_length, zero);
147            if p.is_null() {
148                return oom();
149            }
150            *inout_ptr = p;
151            return ok();
152        }
153        if cmd == sys::IREE_ALLOCATOR_COMMAND_REALLOC {
154            let new_len = (*(params as *const sys::iree_allocator_alloc_params_t)).byte_length;
155            let p = arena.realloc(*inout_ptr, new_len);
156            if p.is_null() {
157                return oom();
158            }
159            *inout_ptr = p;
160            return ok();
161        }
162        unimplemented_status()
163    }
164}
165
166#[inline]
167fn ok() -> sys::iree_status_t {
168    core::ptr::null_mut() // IREE_STATUS_OK
169}
170#[inline]
171fn oom() -> sys::iree_status_t {
172    sys::IREE_STATUS_RESOURCE_EXHAUSTED as usize as sys::iree_status_t
173}
174#[inline]
175fn unimplemented_status() -> sys::iree_status_t {
176    sys::IREE_STATUS_UNIMPLEMENTED as usize as sys::iree_status_t
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn alloc_and_free_roundtrip() {
185        static mut BUF: [u8; 64 * 1024] = [0; 64 * 1024];
186        // SAFETY: single-threaded test with exclusive access to BUF.
187        let arena = unsafe { Arena::new(&mut *core::ptr::addr_of_mut!(BUF)) };
188        let allocator = arena.as_iree_allocator();
189        unsafe {
190            let mut p: *mut c_void = core::ptr::null_mut();
191            let st = sys::iree_allocator_malloc(allocator, 128, &mut p);
192            assert!(st.is_null(), "malloc returned non-OK status");
193            assert!(!p.is_null());
194            // Write through the whole block to prove it is usable.
195            core::ptr::write_bytes(p as *mut u8, 0xAB, 128);
196            sys::iree_allocator_free(allocator, p);
197        }
198    }
199
200    #[test]
201    fn many_allocs_do_not_leak_across_reuse() {
202        static mut BUF: [u8; 64 * 1024] = [0; 64 * 1024];
203        // SAFETY: single-threaded test with exclusive access to BUF.
204        let arena = unsafe { Arena::new(&mut *core::ptr::addr_of_mut!(BUF)) };
205        let allocator = arena.as_iree_allocator();
206        // Allocate and free repeatedly; a leaking allocator would exhaust the
207        // 64 KiB arena well before 10_000 iterations of 256 bytes.
208        for _ in 0..10_000 {
209            unsafe {
210                let mut p: *mut c_void = core::ptr::null_mut();
211                let st = sys::iree_allocator_malloc(allocator, 256, &mut p);
212                assert!(st.is_null());
213                assert!(!p.is_null());
214                sys::iree_allocator_free(allocator, p);
215            }
216        }
217    }
218}
iree_embedded/arena.rs

iree_embedded/
arena.rs
