weir/
fixed_point_scratch.rs

1//! Reusable host scratch for GPU-driven fixed-point dataflow loops.
2//!
3//! This module owns the per-session buffers that are safe to reuse across
4//! reaching, liveness, points-to, and future fixed-point analyses. The graph
5//! and transfer semantics stay in each analysis module; this file only owns
6//! scratch lifetime and reuse.
7
8mod buffers;
9mod telemetry;
10
11#[cfg(test)]
12mod tests;
13
14pub(crate) use buffers::FixedPointForwardChangedProgramCache;
15pub use buffers::{count_changed_domain_bits, count_set_domain_bits, FixedPointScratch};
16#[cfg(feature = "gpu-telemetry")]
17pub use buffers::{count_changed_domain_bits_gpu, count_set_domain_bits_gpu};
18pub use telemetry::{FrontierDensityTelemetry, FrontierExecutionMode};
19
20/// Generic retention pool for reusable host-side scratch vectors.
21///
22/// Clears logical contents while preserving allocation capacity across
23/// dispatches. Derefs to [`Vec<T>`] so existing code can push, extend, and
24/// index without extra boilerplate.
25#[derive(Clone, Debug, PartialEq)]
26pub struct ScratchPool<T> {
27    vec: Vec<T>,
28}
29
30impl<T> Default for ScratchPool<T> {
31    fn default() -> Self {
32        Self { vec: Vec::new() }
33    }
34}
35
36impl<T> ScratchPool<T> {
37    /// Create an empty scratch pool.
38    #[inline]
39    #[must_use]
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    /// Create a scratch pool with pre-reserved capacity.
45    #[inline]
46    #[must_use]
47    pub fn with_capacity(capacity: usize) -> Self {
48        Self {
49            vec: Vec::with_capacity(capacity),
50        }
51    }
52
53    /// Drop logical contents while preserving allocation capacity.
54    #[inline]
55    pub fn clear(&mut self) {
56        self.vec.clear();
57    }
58
59    /// Current retained capacity.
60    #[inline]
61    #[must_use]
62    pub fn capacity(&self) -> usize {
63        self.vec.capacity()
64    }
65}
66
67impl<T> std::ops::Deref for ScratchPool<T> {
68    type Target = Vec<T>;
69
70    fn deref(&self) -> &Self::Target {
71        &self.vec
72    }
73}
74
75impl<T> std::ops::DerefMut for ScratchPool<T> {
76    fn deref_mut(&mut self) -> &mut Self::Target {
77        &mut self.vec
78    }
79}
80
81/// Build a `vyre_primitives::bitset::equal` program for `words`.
82#[inline]
83pub(crate) fn bitset_equal_program(words: usize) -> vyre::ir::Program {
84    vyre_primitives::bitset::equal::bitset_equal("lhs", "rhs", "out", words as u32)
85}
86
87/// GPU convergence check using a generic dispatch closure.
88///
89/// Dispatches `bitset_equal(next_bytes, current_bytes, out_scalar)` and
90/// returns `true` iff `out_scalar[0] == 1`.
91///
92/// `next_bytes` and `current_bytes` must each be exactly `words * 4`
93/// little-endian bytes.
94#[allow(dead_code)]
95pub(crate) fn converged_via_gpu_into<F>(
96    next_bytes: &[u8],
97    current_bytes: &[u8],
98    words: usize,
99    dispatch: F,
100) -> Result<bool, String>
101where
102    F: FnOnce(
103        &vyre::ir::Program,
104        &[&[u8]],
105        Option<[u32; 3]>,
106        &mut Vec<Vec<u8>>,
107    ) -> Result<(), String>,
108{
109    if next_bytes.len() != words * 4 || current_bytes.len() != words * 4 {
110        return Ok(false);
111    }
112    let program = bitset_equal_program(words);
113    let scalar_zero = 0u32.to_le_bytes();
114    let mut outputs = Vec::new();
115    dispatch(
116        &program,
117        &[next_bytes, current_bytes, &scalar_zero],
118        Some([1, 1, 1]),
119        &mut outputs,
120    )?;
121    let scalar = crate::dispatch_decode::unpack_exact_u32_scalar(
122        crate::dispatch_decode::only_output(&outputs, "bitset_equal", "scalar")?,
123        "bitset_equal",
124    )?;
125    Ok(scalar == 1)
126}
127
128/// GPU convergence check using a [`vyre::VyreBackend`] directly.
129///
130/// Dispatches `bitset_equal` through `backend.dispatch_borrowed` and returns
131/// `true` iff the scalar output is `1`.
132pub(crate) fn converged_via_gpu_backend(
133    next_bytes: &[u8],
134    current_bytes: &[u8],
135    words: usize,
136    backend: &dyn vyre::VyreBackend,
137) -> Result<bool, String> {
138    if next_bytes.len() != words * 4 || current_bytes.len() != words * 4 {
139        return Ok(false);
140    }
141    let program = bitset_equal_program(words);
142    let scalar_zero = 0u32.to_le_bytes();
143    let outputs = backend
144        .dispatch_borrowed(
145            &program,
146            &[next_bytes, current_bytes, &scalar_zero],
147            &vyre::DispatchConfig::default(),
148        )
149        .map_err(|error| error.to_string())?;
150    let scalar = crate::dispatch_decode::unpack_exact_u32_scalar(
151        crate::dispatch_decode::only_output(&outputs, "bitset_equal", "scalar")?,
152        "bitset_equal",
153    )?;
154    Ok(scalar == 1)
155}
156
157#[inline]
158pub(crate) fn copy_if_converged(
159    next: &[u32],
160    current: &[u32],
161    result: &mut Vec<u32>,
162) -> Result<bool, String> {
163    if next != current {
164        return Ok(false);
165    }
166    copy_words_into(next, result)?;
167    Ok(true)
168}
169
170#[inline]
171pub(crate) fn copy_words_into(words: &[u32], result: &mut Vec<u32>) -> Result<(), String> {
172    if result.len() == words.len() {
173        result.copy_from_slice(words);
174        return Ok(());
175    }
176    // Slow allocation path is cold after scratch warms up.
177    copy_words_into_slow_path(words, result)
178}
179
180#[cold]
181fn copy_words_into_slow_path(words: &[u32], result: &mut Vec<u32>) -> Result<(), String> {
182    crate::staging_reserve::reserve_vec(result, words.len(), "fixed-point result word copy")?;
183    result.clear();
184    result.extend_from_slice(words);
185    Ok(())
186}
weir/fixed_point_scratch.rs

weir/
fixed_point_scratch.rs
