Skip to main content

antecedent_core/
execution.rs

1//! Execution context: parallelism, determinism, RNG, budgets, kernel policy.
2//!
3//! No core algorithm creates a global thread pool, uses an implicit global RNG,
4//! or selects architecture-specific behavior outside [`KernelPolicy`]
5//!
6//! SPDX-License-Identifier: MIT OR Apache-2.0
7
8use core::sync::atomic::{AtomicBool, Ordering};
9use std::sync::Arc;
10
11/// Parallel execution budget.
12#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
13pub struct Parallelism {
14    /// Maximum worker threads (1 = serial).
15    pub max_threads: NonZeroThreadCount,
16}
17
18/// Thread count that is at least one.
19#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
20pub struct NonZeroThreadCount(u32);
21
22impl NonZeroThreadCount {
23    /// Create from a positive thread count.
24    #[must_use]
25    pub const fn new(n: u32) -> Option<Self> {
26        if n == 0 { None } else { Some(Self(n)) }
27    }
28
29    /// Single-threaded execution.
30    #[must_use]
31    pub const fn one() -> Self {
32        Self(1)
33    }
34
35    /// Underlying count.
36    #[must_use]
37    pub const fn get(self) -> u32 {
38        self.0
39    }
40}
41
42impl Parallelism {
43    /// Serial execution.
44    #[must_use]
45    pub const fn serial() -> Self {
46        Self { max_threads: NonZeroThreadCount::one() }
47    }
48
49    /// Bounded parallelism.
50    #[must_use]
51    pub const fn bounded(max_threads: NonZeroThreadCount) -> Self {
52        Self { max_threads }
53    }
54}
55
56/// Determinism requirements for reductions and scheduling.
57#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
58pub enum Determinism {
59    /// Prefer fastest path; reductions may be nondeterministic.
60    PreferFast,
61    /// Require bitwise-reproducible results for a fixed seed and thread count.
62    Strict,
63}
64
65/// Memory budget for planned allocations.
66#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
67pub struct MemoryBudget {
68    /// Soft limit in bytes; planners should refuse or stream above this.
69    pub soft_limit_bytes: Option<u64>,
70    /// Hard limit in bytes; exceeding is a resource error.
71    pub hard_limit_bytes: Option<u64>,
72}
73
74impl MemoryBudget {
75    /// Unlimited budget (still subject to OS limits).
76    #[must_use]
77    pub const fn unlimited() -> Self {
78        Self { soft_limit_bytes: None, hard_limit_bytes: None }
79    }
80}
81
82/// Kernel selection policy.
83#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
84pub struct KernelPolicy {
85    /// Allow portable optimized kernels.
86    pub allow_portable_optimized: bool,
87    /// Allow architecture-specific SIMD after feature detection.
88    pub allow_arch_simd: bool,
89    /// Force scalar reference path (for tests / debugging).
90    pub force_scalar: bool,
91}
92
93impl KernelPolicy {
94    /// Default: optimized allowed, SIMD allowed, scalar not forced.
95    #[must_use]
96    pub const fn default_policy() -> Self {
97        Self { allow_portable_optimized: true, allow_arch_simd: true, force_scalar: false }
98    }
99
100    /// Force scalar kernels only.
101    #[must_use]
102    pub const fn scalar_only() -> Self {
103        Self { allow_portable_optimized: false, allow_arch_simd: false, force_scalar: true }
104    }
105}
106
107/// Cache usage policy.
108#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
109pub struct CachePolicy {
110    /// Whether semantic caches may be used.
111    pub enabled: bool,
112    /// Maximum cache bytes, if bounded.
113    pub max_bytes: Option<u64>,
114}
115
116impl CachePolicy {
117    /// Caching disabled.
118    #[must_use]
119    pub const fn disabled() -> Self {
120        Self { enabled: false, max_bytes: None }
121    }
122
123    /// Caching enabled with optional byte cap.
124    #[must_use]
125    pub const fn enabled(max_bytes: Option<u64>) -> Self {
126        Self { enabled: true, max_bytes }
127    }
128}
129
130/// Bounded cache budget for incremental causal state.
131#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
132pub struct CacheBudget {
133    /// Maximum retained cache bytes.
134    pub max_bytes: u64,
135    /// Bytes currently retained (updated by the state crate).
136    pub used_bytes: u64,
137}
138
139impl CacheBudget {
140    /// Fresh budget with `max_bytes` capacity and zero usage.
141    #[must_use]
142    pub const fn new(max_bytes: u64) -> Self {
143        Self { max_bytes, used_bytes: 0 }
144    }
145
146    /// Unlimited soft budget (still subject to OS limits).
147    #[must_use]
148    pub const fn unlimited() -> Self {
149        Self { max_bytes: u64::MAX, used_bytes: 0 }
150    }
151
152    /// Remaining capacity in bytes.
153    #[must_use]
154    pub const fn remaining(self) -> u64 {
155        self.max_bytes.saturating_sub(self.used_bytes)
156    }
157
158    /// Whether `additional` bytes would fit under the budget.
159    #[must_use]
160    pub const fn can_admit(self, additional: u64) -> bool {
161        self.used_bytes.saturating_add(additional) <= self.max_bytes
162    }
163}
164
165/// Shared Monte Carlo / approximate-compute budget report.
166#[derive(Clone, Debug, Default, PartialEq)]
167pub struct MonteCarloBudget {
168    /// Scalar / batch evaluations performed.
169    pub evaluations: u64,
170    /// Monte Carlo samples drawn.
171    pub samples: u64,
172    /// Exact enumerations performed, if any.
173    pub exact_enumerations: u64,
174}
175
176/// Per-estimate Monte Carlo uncertainty summary.
177#[derive(Clone, Copy, Debug, Default, PartialEq)]
178pub struct MonteCarloError {
179    /// Estimated standard error of the reported score.
180    pub stderr: f64,
181    /// Samples contributing to this estimate.
182    pub samples: u64,
183}
184
185/// Adaptive bootstrap early-stop budget (SE relative-change criterion).
186///
187/// After at least [`Self::min_replicates`] successful replicates, stop when
188/// `|SE_t − SE_{t−1}| / max(|SE_{t−1}|, ε₀) < se_rel_epsilon`. Cap remains the
189/// requested replicate count. Disabled runs always evaluate the full request.
190#[derive(Clone, Copy, Debug, PartialEq)]
191pub struct AdaptiveBootstrapBudget {
192    /// When false, evaluate all requested replicates (no early-stop).
193    pub enabled: bool,
194    /// Minimum successful replicates before early-stop is eligible.
195    pub min_replicates: u32,
196    /// Relative SE change threshold (default `0.01`).
197    pub se_rel_epsilon: f64,
198}
199
200impl AdaptiveBootstrapBudget {
201    /// Enabled defaults: min 10 replicates, 1% relative SE change.
202    #[must_use]
203    pub const fn enabled_default() -> Self {
204        Self { enabled: true, min_replicates: 10, se_rel_epsilon: 0.01 }
205    }
206
207    /// Force full requested replicate count (tests / exact-N pins).
208    #[must_use]
209    pub const fn disabled() -> Self {
210        Self { enabled: false, min_replicates: 0, se_rel_epsilon: 0.0 }
211    }
212}
213
214impl Default for AdaptiveBootstrapBudget {
215    fn default() -> Self {
216        Self::enabled_default()
217    }
218}
219
220/// Adaptive Bayesian draw budget (Laplace / conjugate path).
221///
222/// Cap by quantile-width relative change and/or ESS target under the latency
223/// `n_draws` maximum. HMC ignores this and always draws the full request.
224#[derive(Clone, Copy, Debug, PartialEq)]
225pub struct AdaptiveDrawBudget {
226    /// When false, materialize the full requested draw count.
227    pub enabled: bool,
228    /// Minimum draws before early-stop is eligible.
229    pub min_draws: usize,
230    /// Relative 95% quantile-width change threshold (default `0.01`).
231    pub quantile_width_rel_epsilon: f64,
232    /// Stop when ESS of the effect draws reaches this target (default large).
233    pub ess_target: f64,
234}
235
236impl AdaptiveDrawBudget {
237    /// Enabled defaults: min 32 draws, 1% quantile-width change, high ESS target.
238    #[must_use]
239    pub const fn enabled_default() -> Self {
240        Self {
241            enabled: true,
242            min_draws: 32,
243            quantile_width_rel_epsilon: 0.01,
244            ess_target: 10_000.0,
245        }
246    }
247
248    /// Force full requested draw count.
249    #[must_use]
250    pub const fn disabled() -> Self {
251        Self { enabled: false, min_draws: 0, quantile_width_rel_epsilon: 0.0, ess_target: 0.0 }
252    }
253}
254
255impl Default for AdaptiveDrawBudget {
256    fn default() -> Self {
257        Self::enabled_default()
258    }
259}
260
261/// Cooperative cancellation token.
262#[derive(Clone, Debug, Default)]
263pub struct CancellationToken {
264    cancelled: Arc<AtomicBool>,
265}
266
267impl CancellationToken {
268    /// Create a fresh token.
269    #[must_use]
270    pub fn new() -> Self {
271        Self::default()
272    }
273
274    /// Request cancellation.
275    pub fn cancel(&self) {
276        self.cancelled.store(true, Ordering::SeqCst);
277    }
278
279    /// Whether cancellation was requested.
280    #[must_use]
281    pub fn is_cancelled(&self) -> bool {
282        self.cancelled.load(Ordering::SeqCst)
283    }
284}
285
286/// Optional progress reporting sink.
287pub trait ProgressSink: Send + Sync {
288    /// Report progress in `[0.0, 1.0]` with an optional stage label.
289    fn report(&self, fraction: f64, stage: &str);
290}
291
292/// Factory for deterministic, independently seeded RNG streams.
293///
294/// Streams are derived from a master seed and a stream id so algorithms can
295/// request reproducible substreams without a global RNG.
296#[derive(Clone, Debug, Eq, PartialEq)]
297pub struct RngFactory {
298    master_seed: u64,
299}
300
301impl RngFactory {
302    /// Create a factory from a master seed.
303    #[must_use]
304    pub const fn from_seed(master_seed: u64) -> Self {
305        Self { master_seed }
306    }
307
308    /// Master seed.
309    #[must_use]
310    pub const fn master_seed(&self) -> u64 {
311        self.master_seed
312    }
313
314    /// Derive an independent stream for `stream_id`.
315    #[must_use]
316    pub fn stream(&self, stream_id: u64) -> CausalRng {
317        let seed = mix_seed(self.master_seed, stream_id);
318        CausalRng::from_seed(seed)
319    }
320}
321
322/// Deterministic SplitMix64-based RNG for library algorithms.
323#[derive(Clone, Debug)]
324pub struct CausalRng {
325    state: u64,
326}
327
328impl CausalRng {
329    /// Create from a 64-bit seed.
330    #[must_use]
331    pub const fn from_seed(seed: u64) -> Self {
332        // Avoid the all-zero fixed point of SplitMix by mixing once.
333        Self { state: seed ^ 0x9E37_79B9_7F4A_7C15 }
334    }
335
336    /// Restore from a previously exported [`Self::state`].
337    #[must_use]
338    pub const fn from_state(state: u64) -> Self {
339        Self { state }
340    }
341
342    /// Opaque stream state for checkpoint / CRN continuation.
343    #[must_use]
344    pub const fn state(&self) -> u64 {
345        self.state
346    }
347
348    /// Next `u64` from the stream.
349    #[must_use]
350    pub fn next_u64(&mut self) -> u64 {
351        self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
352        let mut z = self.state;
353        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
354        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
355        z ^ (z >> 31)
356    }
357
358    /// Next `f64` in `[0, 1)`.
359    #[must_use]
360    pub fn next_f64(&mut self) -> f64 {
361        // 53-bit mantissa extraction; precision loss vs full u64 is intentional.
362        let bits = self.next_u64() >> 11;
363        #[allow(clippy::cast_precision_loss)]
364        {
365            bits as f64 * (1.0 / ((1u64 << 53) as f64))
366        }
367    }
368}
369
370fn mix_seed(master: u64, stream_id: u64) -> u64 {
371    let mut z = master.wrapping_add(stream_id).wrapping_mul(0xD6E8_FEB8_6659_FD93);
372    z = (z ^ (z >> 32)).wrapping_mul(0xD6E8_FEB8_6659_FD93);
373    z ^ (z >> 32)
374}
375
376/// Full execution context passed into algorithms.
377pub struct ExecutionContext {
378    /// Parallelism budget.
379    pub parallelism: Parallelism,
380    /// Determinism policy.
381    pub determinism: Determinism,
382    /// RNG factory (no global RNG).
383    pub rng: RngFactory,
384    /// Memory budget.
385    pub memory: MemoryBudget,
386    /// Cancellation token.
387    pub cancellation: CancellationToken,
388    /// Optional progress sink.
389    pub progress: Option<Arc<dyn ProgressSink>>,
390    /// Kernel selection policy.
391    pub kernel_policy: KernelPolicy,
392    /// Cache policy.
393    pub cache_policy: CachePolicy,
394    /// Adaptive bootstrap early-stop (estimate SE path).
395    pub adaptive_bootstrap: AdaptiveBootstrapBudget,
396    /// Adaptive Bayesian draw early-stop (Laplace / conjugate).
397    pub adaptive_draws: AdaptiveDrawBudget,
398}
399
400impl ExecutionContext {
401    /// Construct a serial, strict, scalar-friendly context for tests.
402    ///
403    /// # Examples
404    ///
405    /// ```
406    /// use antecedent_core::ExecutionContext;
407    ///
408    /// let ctx = ExecutionContext::for_tests(42);
409    /// assert!(!ctx.cancellation.is_cancelled());
410    /// ```
411    #[must_use]
412    pub fn for_tests(seed: u64) -> Self {
413        Self {
414            parallelism: Parallelism::serial(),
415            determinism: Determinism::Strict,
416            rng: RngFactory::from_seed(seed),
417            memory: MemoryBudget::unlimited(),
418            cancellation: CancellationToken::new(),
419            progress: None,
420            kernel_policy: KernelPolicy::scalar_only(),
421            cache_policy: CachePolicy::disabled(),
422            // Exact-N pins in unit tests; enable explicitly for adaptive MC tests.
423            adaptive_bootstrap: AdaptiveBootstrapBudget::disabled(),
424            adaptive_draws: AdaptiveDrawBudget::disabled(),
425        }
426    }
427
428    /// Production context: optimized kernels allowed, cache enabled, bounded threads.
429    #[must_use]
430    pub fn production(seed: u64, max_threads: u32) -> Self {
431        let threads =
432            NonZeroThreadCount::new(max_threads.max(1)).unwrap_or_else(NonZeroThreadCount::one);
433        Self {
434            parallelism: Parallelism::bounded(threads),
435            determinism: Determinism::Strict,
436            rng: RngFactory::from_seed(seed),
437            memory: MemoryBudget::unlimited(),
438            cancellation: CancellationToken::new(),
439            progress: None,
440            kernel_policy: KernelPolicy::default_policy(),
441            cache_policy: CachePolicy::enabled(None),
442            adaptive_bootstrap: AdaptiveBootstrapBudget::enabled_default(),
443            adaptive_draws: AdaptiveDrawBudget::enabled_default(),
444        }
445    }
446}
447
448impl core::fmt::Debug for ExecutionContext {
449    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
450        f.debug_struct("ExecutionContext")
451            .field("parallelism", &self.parallelism)
452            .field("determinism", &self.determinism)
453            .field("rng", &self.rng)
454            .field("memory", &self.memory)
455            .field("cancellation_cancelled", &self.cancellation.is_cancelled())
456            .field("progress_is_some", &self.progress.is_some())
457            .field("kernel_policy", &self.kernel_policy)
458            .field("cache_policy", &self.cache_policy)
459            .field("adaptive_bootstrap", &self.adaptive_bootstrap)
460            .field("adaptive_draws", &self.adaptive_draws)
461            .finish()
462    }
463}
464
465#[cfg(test)]
466mod tests {
467    use super::*;
468
469    #[test]
470    fn rng_streams_are_deterministic() {
471        let factory = RngFactory::from_seed(42);
472        let mut a1 = factory.stream(0);
473        let mut a2 = factory.stream(0);
474        let mut b = factory.stream(1);
475        let seq_a1: Vec<u64> = (0..8).map(|_| a1.next_u64()).collect();
476        let seq_a2: Vec<u64> = (0..8).map(|_| a2.next_u64()).collect();
477        let seq_b: Vec<u64> = (0..8).map(|_| b.next_u64()).collect();
478        assert_eq!(seq_a1, seq_a2);
479        assert_ne!(seq_a1, seq_b);
480    }
481
482    #[test]
483    fn independent_factories_same_seed_match() {
484        let f1 = RngFactory::from_seed(7);
485        let f2 = RngFactory::from_seed(7);
486        let mut s1 = f1.stream(99);
487        let mut s2 = f2.stream(99);
488        for _ in 0..32 {
489            assert_eq!(s1.next_u64(), s2.next_u64());
490        }
491    }
492
493    #[test]
494    fn cancellation_is_shared_across_clones() {
495        let token = CancellationToken::new();
496        let clone = token.clone();
497        assert!(!token.is_cancelled());
498        clone.cancel();
499        assert!(token.is_cancelled());
500    }
501
502    #[test]
503    fn f64_draws_are_in_unit_interval() {
504        let mut rng = CausalRng::from_seed(123);
505        for _ in 0..1000 {
506            let x = rng.next_f64();
507            assert!((0.0..1.0).contains(&x));
508        }
509    }
510}