Skip to main content

weir/
fixed_point_resident_frontier.rs

1//! Reusable backend-resident frontier resources for fixed-point hot loops.
2
3/// Reusable resident frontier resources for fixed-point hot loops.
4#[derive(Clone, Debug, Default, PartialEq)]
5pub struct FixedPointResidentFrontierScratch {
6    frontier_in: Option<vyre::backend::Resource>,
7    frontier_out: Option<vyre::backend::Resource>,
8    changed: Option<vyre::backend::Resource>,
9    byte_len: usize,
10}
11
12impl FixedPointResidentFrontierScratch {
13    /// Ensure the resident frontier pair exists with at least `byte_len` bytes.
14    pub fn ensure_pair_refs(
15        &mut self,
16        backend: &dyn vyre::VyreBackend,
17        byte_len: usize,
18    ) -> Result<(&vyre::backend::Resource, &vyre::backend::Resource), String> {
19        if self.byte_len < byte_len || self.frontier_in.is_none() || self.frontier_out.is_none() {
20            self.clear(backend).map_err(|error| error.to_string())?;
21            let frontier_in = backend
22                .allocate_resident(byte_len)
23                .map_err(|error| error.to_string())?;
24            let frontier_out = match backend.allocate_resident(byte_len) {
25                Ok(resource) => resource,
26                Err(error) => {
27                    let _ = backend.free_resident(frontier_in);
28                    return Err(error.to_string());
29                }
30            };
31            self.frontier_in = Some(frontier_in);
32            self.frontier_out = Some(frontier_out);
33            self.byte_len = byte_len;
34        }
35        Ok((
36            self.frontier_in.as_ref().ok_or_else(|| {
37                "weir resident fixed-point frontier input resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
38                    .to_string()
39            })?,
40            self.frontier_out.as_ref().ok_or_else(|| {
41                "weir resident fixed-point frontier output resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
42                    .to_string()
43            })?,
44        ))
45    }
46
47    /// Ensure an in-place resident frontier and one u32 changed flag exist.
48    pub fn ensure_inplace_refs(
49        &mut self,
50        backend: &dyn vyre::VyreBackend,
51        byte_len: usize,
52    ) -> Result<(&vyre::backend::Resource, &vyre::backend::Resource), String> {
53        let mut allocated_frontier_this_call = false;
54        if self.byte_len < byte_len || self.frontier_in.is_none() {
55            self.clear(backend).map_err(|error| error.to_string())?;
56            let frontier_in = backend
57                .allocate_resident(byte_len)
58                .map_err(|error| error.to_string())?;
59            self.frontier_in = Some(frontier_in);
60            self.byte_len = byte_len;
61            allocated_frontier_this_call = true;
62        }
63        if self.changed.is_none() {
64            match backend.allocate_resident(std::mem::size_of::<u32>()) {
65                Ok(changed) => self.changed = Some(changed),
66                Err(error) => {
67                    if allocated_frontier_this_call {
68                        let _ = self.clear(backend);
69                    }
70                    return Err(error.to_string());
71                }
72            }
73        }
74        Ok((
75            self.frontier_in.as_ref().ok_or_else(|| {
76                "weir resident fixed-point in-place frontier resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
77                    .to_string()
78            })?,
79            self.changed.as_ref().ok_or_else(|| {
80                "weir resident fixed-point changed-flag resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
81                    .to_string()
82            })?,
83        ))
84    }
85
86    /// Ensure the resident frontier pair and return owned handles for tests.
87    #[cfg(test)]
88    pub fn ensure_pair(
89        &mut self,
90        backend: &dyn vyre::VyreBackend,
91        byte_len: usize,
92    ) -> Result<(vyre::backend::Resource, vyre::backend::Resource), String> {
93        let (frontier_in, frontier_out) = self.ensure_pair_refs(backend, byte_len)?;
94        Ok((frontier_in.clone(), frontier_out.clone()))
95    }
96
97    /// Current resident frontier byte capacity.
98    #[must_use]
99    pub fn byte_len(&self) -> usize {
100        self.byte_len
101    }
102
103    /// Whether both resident frontier resources are currently allocated.
104    #[must_use]
105    pub fn is_allocated(&self) -> bool {
106        self.frontier_in.is_some() && self.frontier_out.is_some()
107    }
108
109    /// Free resident frontier resources, preserving the scratch object itself.
110    pub fn clear(&mut self, backend: &dyn vyre::VyreBackend) -> Result<(), vyre::BackendError> {
111        let mut first_error = None;
112        if let Some(resource) = self.frontier_in.take() {
113            if let Err(error) = backend.free_resident(resource) {
114                first_error.get_or_insert(error);
115            }
116        }
117        if let Some(resource) = self.frontier_out.take() {
118            if let Err(error) = backend.free_resident(resource) {
119                first_error.get_or_insert(error);
120            }
121        }
122        if let Some(resource) = self.changed.take() {
123            if let Err(error) = backend.free_resident(resource) {
124                first_error.get_or_insert(error);
125            }
126        }
127        self.byte_len = 0;
128        if let Some(error) = first_error {
129            Err(error)
130        } else {
131            Ok(())
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use std::collections::HashSet;
139    use std::sync::atomic::{AtomicU64, Ordering};
140    use std::sync::Mutex;
141    use vyre::backend::{private, BackendError, DispatchConfig, Resource};
142
143    use super::*;
144
145    struct FakeBackend {
146        next: AtomicU64,
147        allocations: Mutex<Vec<usize>>,
148        freed: Mutex<Vec<u64>>,
149        fail_byte_len: Option<usize>,
150    }
151
152    impl FakeBackend {
153        fn new() -> Self {
154            Self {
155                next: AtomicU64::new(1),
156                allocations: Mutex::new(Vec::new()),
157                freed: Mutex::new(Vec::new()),
158                fail_byte_len: None,
159            }
160        }
161
162        fn failing_on(byte_len: usize) -> Self {
163            Self {
164                next: AtomicU64::new(1),
165                allocations: Mutex::new(Vec::new()),
166                freed: Mutex::new(Vec::new()),
167                fail_byte_len: Some(byte_len),
168            }
169        }
170    }
171
172    impl private::Sealed for FakeBackend {}
173
174    impl vyre::VyreBackend for FakeBackend {
175        fn id(&self) -> &'static str {
176            "weir_test_frontier_scratch"
177        }
178
179        fn supported_ops(&self) -> &HashSet<vyre::ir::OpId> {
180            static OPS: std::sync::OnceLock<HashSet<vyre::ir::OpId>> = std::sync::OnceLock::new();
181            OPS.get_or_init(HashSet::new)
182        }
183
184        fn dispatch(
185            &self,
186            _program: &vyre::ir::Program,
187            _inputs: &[Vec<u8>],
188            _config: &DispatchConfig,
189        ) -> Result<Vec<Vec<u8>>, BackendError> {
190            unreachable!("frontier scratch tests do not dispatch")
191        }
192
193        fn allocate_resident(&self, byte_len: usize) -> Result<Resource, BackendError> {
194            self.allocations.lock().unwrap().push(byte_len);
195            if self.fail_byte_len == Some(byte_len) {
196                return Err(BackendError::InvalidProgram {
197                    fix: format!("fake allocation failure for {byte_len} bytes"),
198                });
199            }
200            Ok(Resource::Resident(
201                self.next.fetch_add(1, Ordering::Relaxed),
202            ))
203        }
204
205        fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
206            let Resource::Resident(handle) = resource else {
207                panic!("fake backend only frees resident resources");
208            };
209            self.freed.lock().unwrap().push(handle);
210            Ok(())
211        }
212    }
213
214    #[test]
215    fn resident_frontier_scratch_reuses_larger_capacity_for_smaller_runs() {
216        let backend = FakeBackend::new();
217        let mut scratch = FixedPointResidentFrontierScratch::default();
218
219        let first = scratch
220            .ensure_pair(&backend, 128)
221            .expect("first frontier allocation must succeed");
222        let second = scratch
223            .ensure_pair(&backend, 32)
224            .expect("smaller frontier must reuse larger allocation");
225
226        assert_eq!(first, second);
227        assert_eq!(scratch.byte_len(), 128);
228        assert_eq!(backend.allocations.lock().unwrap().as_slice(), &[128, 128]);
229        assert!(
230            backend.freed.lock().unwrap().is_empty(),
231            "Fix: smaller resident frontier requests must not free and reallocate larger reusable slots."
232        );
233    }
234
235    #[test]
236    fn resident_inplace_frontier_rolls_back_new_frontier_when_changed_flag_fails() {
237        let backend = FakeBackend::failing_on(std::mem::size_of::<u32>());
238        let mut scratch = FixedPointResidentFrontierScratch::default();
239
240        let error = scratch
241            .ensure_inplace_refs(&backend, 128)
242            .expect_err("changed-flag allocation failure must fail the in-place frontier ensure");
243
244        assert!(error.contains("fake allocation failure"));
245        assert_eq!(
246            backend.allocations.lock().unwrap().as_slice(),
247            &[128, std::mem::size_of::<u32>()],
248            "Fix: in-place frontier ensure must attempt frontier then changed flag exactly once."
249        );
250        assert_eq!(
251            backend.freed.lock().unwrap().as_slice(),
252            &[1],
253            "Fix: newly allocated frontier must be released when changed flag allocation fails."
254        );
255        assert_eq!(scratch.byte_len(), 0);
256        assert!(!scratch.is_allocated());
257    }
258
259    #[test]
260    fn resident_frontier_scratch_production_ensure_has_no_panic_path() {
261        let source = include_str!("fixed_point_resident_frontier.rs");
262        let production = source
263            .split("#[cfg(test)]")
264            .next()
265            .expect("production frontier source must precede tests");
266
267        for forbidden in [
268            concat!(".", "expect("),
269            concat!("panic", "!("),
270            concat!("unimplemented", "!("),
271            concat!("todo", "!("),
272        ] {
273            assert!(
274                !production.contains(forbidden),
275                "resident frontier production ensure path must return typed errors instead of using {forbidden}"
276            );
277        }
278    }
279}