iree_embedded/
device.rs

1//! A synchronous, single-threaded CPU device.
2//!
3//! Two loaders are offered: the embedded-ELF loader (host milestone; the
4//! `.vmfb` carries position-independent ELF kernels that are mapped into RAM)
5//! and the static-library loader (MCU; the kernels are compiled into the
6//! firmware image and execute in place from flash, costing no RAM).
7
8use crate::{Arena, Result, check};
9use iree_embedded_sys as sys;
10
11/// The `*_library_query` entry point emitted by
12/// `iree-compile --iree-llvmcpu-static-library-output-path=` (declared in the
13/// generated header alongside the `.o`). Declared with opaque pointers so
14/// firmware can `extern "C"` it without naming generated binding types; the
15/// shape is ABI-identical to `iree_hal_executable_library_query_fn_t`
16/// (`uint32_t`, pointer) -> pointer.
17pub type LibraryQueryFn = unsafe extern "C" fn(
18    max_version: u32,
19    environment: *const core::ffi::c_void,
20) -> *const core::ffi::c_void;
21
22/// A synchronous, single-threaded local CPU device that executes a model's
23/// kernels. Created with either the static-library loader
24/// ([`local_sync_static`](Device::local_sync_static)) or the embedded-ELF
25/// loader ([`local_sync`](Device::local_sync)).
26pub struct Device {
27    raw: *mut sys::iree_hal_device_t,
28}
29
30impl Device {
31    /// Device whose executables are statically linked into this binary.
32    ///
33    /// `libraries` are the query functions of every model linked into the
34    /// firmware; the loader resolves a `.vmfb`'s library reference by name.
35    pub fn local_sync_static(arena: &Arena, libraries: &[LibraryQueryFn]) -> Result<Self> {
36        let alloc = arena.as_iree_allocator();
37        // SAFETY: every out-pointer is valid; the arena outlives the device.
38        unsafe {
39            let mut loader = core::ptr::null_mut();
40            check(sys::iree_hal_static_library_loader_create(
41                libraries.len() as _,
42                // A non-null `LibraryQueryFn` has the same layout as the
43                // `Option`-wrapped bindgen fn pointer (niche optimization).
44                libraries.as_ptr() as *const sys::iree_hal_executable_library_query_fn_t,
45                sys::iree_hal_executable_import_provider_null(),
46                alloc,
47                &mut loader,
48            ))?;
49            Self::from_loader(loader, alloc)
50        }
51    }
52
53    /// Device using the embedded-ELF loader, for a `.vmfb` whose kernels are
54    /// position-independent ELF mapped into RAM (the host-test path; the MCU
55    /// uses [`local_sync_static`](Device::local_sync_static)).
56    pub fn local_sync(arena: &Arena) -> Result<Self> {
57        let alloc = arena.as_iree_allocator();
58        // SAFETY: every out-pointer is valid; the arena outlives the device.
59        unsafe {
60            let mut loader = core::ptr::null_mut();
61            check(sys::iree_hal_embedded_elf_loader_create(
62                core::ptr::null_mut(), // plugin_manager
63                alloc,
64                &mut loader,
65            ))?;
66            Self::from_loader(loader, alloc)
67        }
68    }
69
70    /// Build the local-sync device around `loader`, consuming one reference to
71    /// it (released here whether or not creation succeeds).
72    ///
73    /// # Safety
74    /// `loader` must be a valid executable loader and `alloc` must outlive the
75    /// returned device.
76    unsafe fn from_loader(
77        mut loader: *mut sys::iree_hal_executable_loader_t,
78        alloc: sys::iree_allocator_t,
79    ) -> Result<Self> {
80        unsafe {
81            let id = sys::iree_make_cstring_view(c"local-sync".as_ptr());
82
83            let mut device_allocator = core::ptr::null_mut();
84            let status =
85                sys::iree_hal_allocator_create_heap(id, alloc, alloc, &mut device_allocator);
86
87            let mut params: sys::iree_hal_sync_device_params_t = core::mem::zeroed();
88            sys::iree_hal_sync_device_params_initialize(&mut params);
89            // The default 32 KiB transient block is host-sized; on an MCU it
90            // starves the arena. 4 KiB is the device's documented minimum and
91            // blocks chain on demand.
92            params.arena_block_size = 4096;
93
94            let mut raw = core::ptr::null_mut();
95            let status = if status.is_null() {
96                sys::iree_hal_sync_device_create(
97                    id,
98                    &params,
99                    1, // loader_count
100                    &mut loader,
101                    device_allocator,
102                    alloc,
103                    &mut raw,
104                )
105            } else {
106                status
107            };
108
109            sys::iree_hal_allocator_release(device_allocator);
110            sys::iree_hal_executable_loader_release(loader);
111            check(status)?;
112            Ok(Device { raw })
113        }
114    }
115}
116
117impl Device {
118    pub(crate) fn raw(&self) -> *mut sys::iree_hal_device_t {
119        self.raw
120    }
121}
122
123impl Drop for Device {
124    fn drop(&mut self) {
125        // SAFETY: raw was created by iree_hal_sync_device_create.
126        unsafe { sys::iree_hal_device_release(self.raw) };
127    }
128}
iree_embedded/device.rs

iree_embedded/
device.rs
