handwired-rs 0.6.0

//! Parameters, return values, and the compiled program a `#[handwired]`
//! function produces.

use std::cell::Cell;
use std::fmt::Debug;
use std::panic::{self, AssertUnwindSafe};
use std::sync::Once;

use super::{CONTEXT, Kind, Mux, Scalar, Val, is_wide, with_circuit};
use crate::Error;
use crate::batch::BatchExecutor;
use crate::circuit::{BackpropOptions, Circuit};
use crate::tensor::{TensorExecutor, TensorProgram};
use crate::wire::Wire;

/// What a parameter declares with `#[range(lo, hi)]`, `#[ranges([...])]` and `#[grid(step)]`. A parameter of a
/// compound type, e.g. `[u8; 4]`, applies `range` and `grid` to each scalar; `ranges` gives one range per scalar, in
/// [`HwType::encode`] order.
#[derive(Copy, Clone, Debug, Default)]
pub struct Attrs {
    pub range: Option<(f64, f64)>,
    pub grid: Option<f64>,
    pub ranges: Option<&'static [(f64, f64)]>,
}

impl Attrs {
    pub const NONE: Attrs = Attrs { range: None, grid: None, ranges: None };

    /// The attributes of the `cells` scalars starting at `offset` within this value.
    pub fn part(&self, offset: usize, cells: usize) -> Attrs {
        match self.ranges {
            Some(ranges) => Attrs { ranges: Some(&ranges[offset..offset + cells]), ..*self },
            None => *self,
        }
    }

    /// A struct field's attributes: the parameter's per-scalar `ranges` when given, else the field's own
    /// declarations, else the parameter's `range` and `grid`.
    pub fn field(&self, own: &Attrs, offset: usize, cells: usize) -> Attrs {
        if self.ranges.is_some() {
            self.part(offset, cells)
        } else if own.range.is_some() || own.ranges.is_some() || own.grid.is_some() {
            *own
        } else {
            *self
        }
    }

    /// Check that `ranges`, if given, has one range per scalar of a value of `cells` scalars.
    pub fn expect_cells(&self, cells: usize) {
        if let Some(ranges) = self.ranges {
            assert_eq!(ranges.len(), cells, "#[ranges] gives {} ranges for {cells} scalars", ranges.len());
        }
    }

    fn scalar_range(&self) -> Option<(f64, f64)> {
        match self.ranges {
            Some(ranges) => {
                self.expect_cells(1);
                Some(ranges[0])
            }
            None => self.range,
        }
    }
}

/// Hands out source cells while a function's parameters are declared.
pub struct Inputs {
    next: usize,
}

/// A deterministic random source for sampling inputs.
pub struct Rng(u64);

impl Rng {
    pub fn new(seed: u64) -> Rng {
        Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1)
    }

    pub fn next_u64(&mut self) -> u64 {
        self.0 ^= self.0 >> 12;
        self.0 ^= self.0 << 25;
        self.0 ^= self.0 >> 27;
        self.0.wrapping_mul(0x2545_F491_4F6C_DD1D)
    }

    /// Uniform in 0..1.
    pub fn unit(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
    }
}

/// A type a `#[handwired]` function can take or return.
pub trait HwType: Sized {
    /// The circuit form of the type, e.g. `Val<u8>` for `u8`.
    type Val: Mux;
    const CELLS: usize;
    fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Self::Val;
    fn encode(&self, out: &mut Vec<f32>);
    fn decode(values: &mut dyn Iterator<Item = f32>) -> Self;
    fn outputs(value: &Self::Val, out: &mut Vec<Wire>);
    fn sample(rng: &mut Rng, attrs: &Attrs) -> Self;
    /// A start value for a binding declared without one.
    fn zero() -> Self::Val;
}

fn scalar_range<T: Scalar>(attrs: &Attrs) -> (f64, f64) {
    match (T::KIND, attrs.scalar_range()) {
        (Kind::Bool, _) => (0.0, 1.0),
        (_, Some((low, high))) => {
            assert!(
                low >= T::MIN && high <= T::MAX && low <= high,
                "#[range({low}, {high})] doesn't fit in {}",
                T::NAME
            );
            if T::KIND == Kind::Float { (f32_at_most(low), f32_at_least(high)) } else { (low, high) }
        }
        (Kind::Int, None) if !is_wide::<T>() => (T::MIN, T::MAX),
        (Kind::Int, None) => panic!("a {} parameter needs #[range(lo, hi)] within ±2^24", T::NAME),
        (Kind::Float, None) => panic!("an f32 parameter needs #[range(lo, hi)]"),
    }
}

/// The largest f32 at or below `value`, so a declared range like 1e-3
/// covers every f32 a caller can pass for it.
fn f32_at_most(value: f64) -> f64 {
    let rounded = value as f32;
    (if rounded as f64 > value { rounded.next_down() } else { rounded }) as f64
}

fn f32_at_least(value: f64) -> f64 {
    let rounded = value as f32;
    (if (rounded as f64) < value { rounded.next_up() } else { rounded }) as f64
}

/// Declare one scalar parameter cell: its range, and for integers and enums a grid of 1.
pub fn declare_scalar<T: Scalar>(inputs: &mut Inputs, attrs: &Attrs) -> Val<T> {
    let (low, high) = scalar_range::<T>(attrs);
    if T::KIND == Kind::Int {
        assert!(low > -16_777_216.0 && high < 16_777_216.0, "a {} parameter's range must lie within ±2^24", T::NAME);
    }
    let index = inputs.next;
    inputs.next += 1;
    with_circuit(|circuit| {
        let wire = circuit.source(index);
        match T::KIND {
            Kind::Bool => {
                circuit.assume_bit(wire);
                circuit.assume_grid(wire, 1.0);
            }
            Kind::Int => {
                circuit.assume_range(wire, low, high);
                circuit.assume_grid(wire, attrs.grid.unwrap_or(1.0).max(1.0));
            }
            Kind::Float => {
                circuit.assume_range(wire, low, high);
                if let Some(grid) = attrs.grid {
                    circuit.assume_grid(wire, grid);
                }
            }
        }
        Val::from_wire(wire)
    })
}

fn sample_scalar<T: Scalar>(rng: &mut Rng, attrs: &Attrs) -> T {
    let (low, high) = scalar_range::<T>(attrs);
    let pick = rng.next_u64() % 8;
    let value = match pick {
        0 => low,
        1 => high,
        2 if low <= 0.0 && high >= 0.0 => 0.0,
        _ => low + rng.unit() * (high - low),
    };
    let snapped = match (T::KIND, attrs.grid) {
        (Kind::Float, None) => value,
        (Kind::Float, Some(grid)) => ((value / grid).round() * grid).clamp(low, high),
        _ => value.round().clamp(low, high),
    };
    T::from_f64(snapped)
}

/// One of an enum's variants, uniformly, among those within a declared `#[range]`.
pub fn sample_variant<T: Scalar>(rng: &mut Rng, attrs: &Attrs, variants: &[T]) -> T {
    let (low, high) = scalar_range::<T>(attrs);
    let allowed: Vec<T> = variants.iter().copied().filter(|variant| (low..=high).contains(&variant.to_f64())).collect();
    assert!(!allowed.is_empty(), "no variant of {} lies in #[range({low}, {high})]", T::NAME);
    allowed[(rng.next_u64() % allowed.len() as u64) as usize]
}

macro_rules! scalar_types {
    ($($t:ident),*) => {$(
        impl HwType for $t {
            type Val = Val<$t>;
            const CELLS: usize = 1;
            fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Val<$t> {
                declare_scalar::<$t>(inputs, attrs)
            }
            fn encode(&self, out: &mut Vec<f32>) {
                out.push(self.to_f64() as f32);
            }
            fn decode(values: &mut dyn Iterator<Item = f32>) -> $t {
                <$t as Scalar>::from_f64(values.next().expect("too few output values") as f64)
            }
            fn outputs(value: &Val<$t>, out: &mut Vec<Wire>) {
                if value.is_wide() {
                    panic!(
                        "a {} result may leave ±2^24, which one f32 cell can't hold; return it narrowed, e.g. as bytes \
                         with `(x >> 8) as u8`, or masked within 24 bits",
                        stringify!($t)
                    );
                }
                out.push(value.wire());
            }
            fn sample(rng: &mut Rng, attrs: &Attrs) -> $t {
                sample_scalar::<$t>(rng, attrs)
            }
            fn zero() -> Val<$t> {
                Val::constant(<$t as Scalar>::from_f64(0.0))
            }
        }
    )*};
}
scalar_types!(bool, u8, u16, u32, u64, usize, i8, i16, i32, i64, f32);

impl<T: HwType, const N: usize> HwType for [T; N] {
    type Val = [T::Val; N];
    const CELLS: usize = N * T::CELLS;
    fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Self::Val {
        attrs.expect_cells(<Self as HwType>::CELLS);
        std::array::from_fn(|item| T::declare(inputs, &attrs.part(item * T::CELLS, T::CELLS)))
    }
    fn encode(&self, out: &mut Vec<f32>) {
        for item in self {
            item.encode(out);
        }
    }
    fn decode(values: &mut dyn Iterator<Item = f32>) -> Self {
        std::array::from_fn(|_| T::decode(values))
    }
    fn outputs(value: &Self::Val, out: &mut Vec<Wire>) {
        for item in value {
            T::outputs(item, out);
        }
    }
    fn sample(rng: &mut Rng, attrs: &Attrs) -> Self {
        attrs.expect_cells(<Self as HwType>::CELLS);
        std::array::from_fn(|item| T::sample(rng, &attrs.part(item * T::CELLS, T::CELLS)))
    }
    fn zero() -> Self::Val {
        std::array::from_fn(|_| T::zero())
    }
}

impl HwType for () {
    type Val = ();
    const CELLS: usize = 0;
    fn declare(_: &mut Inputs, _: &Attrs) {}
    fn encode(&self, _: &mut Vec<f32>) {}
    fn decode(_: &mut dyn Iterator<Item = f32>) {}
    fn outputs(_: &(), _: &mut Vec<Wire>) {}
    fn sample(_: &mut Rng, _: &Attrs) {}
    fn zero() {}
}

/// A function's parameter list, as a tuple, with one `Attrs` per parameter.
pub trait Params: Sized {
    const CELLS: usize;
    /// The source cells of argument `index`, in [`encode`](Self::encode) order.
    fn arg_cells(index: usize) -> std::ops::Range<usize>;
    fn encode(&self, out: &mut Vec<f32>);
    fn sample(rng: &mut Rng, attrs: &[Attrs]) -> Self;
}

impl Params for () {
    const CELLS: usize = 0;
    fn arg_cells(index: usize) -> std::ops::Range<usize> {
        panic!("a function without arguments has no argument {index}")
    }
    fn encode(&self, _: &mut Vec<f32>) {}
    fn sample(_: &mut Rng, _: &[Attrs]) {}
}

macro_rules! tuple_types {
    ($(($($name:ident $index:tt),+)),*) => {$(
        // The last component's offset is never read.
        #[allow(unused_assignments)]
        impl<$($name: HwType),+> HwType for ($($name,)+) {
            type Val = ($($name::Val,)+);
            const CELLS: usize = 0 $(+ $name::CELLS)+;
            fn declare(inputs: &mut Inputs, attrs: &Attrs) -> Self::Val {
                attrs.expect_cells(<Self as HwType>::CELLS);
                let mut offset = 0;
                ($({
                    let part = attrs.part(offset, $name::CELLS);
                    offset += $name::CELLS;
                    $name::declare(inputs, &part)
                },)+)
            }
            fn encode(&self, out: &mut Vec<f32>) {
                $(self.$index.encode(out);)+
            }
            fn decode(values: &mut dyn Iterator<Item = f32>) -> Self {
                ($($name::decode(values),)+)
            }
            fn outputs(value: &Self::Val, out: &mut Vec<Wire>) {
                $($name::outputs(&value.$index, out);)+
            }
            fn sample(rng: &mut Rng, attrs: &Attrs) -> Self {
                attrs.expect_cells(<Self as HwType>::CELLS);
                let mut offset = 0;
                ($({
                    let part = attrs.part(offset, $name::CELLS);
                    offset += $name::CELLS;
                    $name::sample(rng, &part)
                },)+)
            }
            fn zero() -> Self::Val {
                ($($name::zero(),)+)
            }
        }
    )*};
}
tuple_types!(
    (A0 0),
    (A0 0, A1 1),
    (A0 0, A1 1, A2 2),
    (A0 0, A1 1, A2 2, A3 3),
    (A0 0, A1 1, A2 2, A3 3, A4 4),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11)
);

macro_rules! tuple_params {
    ($(($($name:ident $index:tt),+)),*) => {$(
        impl<$($name: HwType),+> Params for ($($name,)+) {
            const CELLS: usize = 0 $(+ $name::CELLS)+;
            fn arg_cells(index: usize) -> std::ops::Range<usize> {
                let sizes = [$($name::CELLS),+];
                assert!(index < sizes.len(), "argument {index} is out of range: there are {}", sizes.len());
                let start: usize = sizes[..index].iter().sum();
                start..start + sizes[index]
            }
            fn encode(&self, out: &mut Vec<f32>) {
                $(self.$index.encode(out);)+
            }
            fn sample(rng: &mut Rng, attrs: &[Attrs]) -> Self {
                ($($name::sample(rng, &attrs[$index]),)+)
            }
        }
    )*};
}
tuple_params!(
    (A0 0),
    (A0 0, A1 1),
    (A0 0, A1 1, A2 2),
    (A0 0, A1 1, A2 2, A3 3),
    (A0 0, A1 1, A2 2, A3 3, A4 4),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11, A12 12),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11, A12 12, A13 13),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11, A12 12, A13 13, A14 14),
    (A0 0, A1 1, A2 2, A3 3, A4 4, A5 5, A6 6, A7 7, A8 8, A9 9, A10 10, A11 11, A12 12, A13 13, A14 14, A15 15)
);

/// Records encoded into source tables and checked against the declared ranges, by
/// [`Compiled::encode_records`], to run or differentiate many times without encoding them again.
#[derive(Clone, Debug)]
pub struct EncodedRecords {
    sources: Vec<f32>,
    count: usize,
}

impl EncodedRecords {
    pub fn len(&self) -> usize {
        self.count
    }

    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// The source tables, one after another.
    pub fn sources(&self) -> &[f32] {
        &self.sources
    }
}

fn threads_for(records: usize) -> usize {
    std::thread::available_parallelism().map_or(1, |count| count.get()).min(records.max(1))
}

/// A `#[handwired]` function's circuit, compiled for its outputs, with the
/// original function kept as the oracle.
pub struct Compiled<Args, R> {
    pub circuit: Circuit,
    pub program: TensorProgram,
    pub outputs: Vec<Wire>,
    oracle: fn(Args) -> R,
    attrs: &'static [Attrs],
}

/// Restores the previous build context, also when a build panics.
struct ContextGuard(Option<Option<Circuit>>);

impl Drop for ContextGuard {
    fn drop(&mut self) {
        if let Some(previous) = self.0.take() {
            CONTEXT.with(|context| *context.borrow_mut() = previous);
        }
    }
}

/// Why [`Compiled::verify`] failed.
#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub enum VerifyError {
    /// The circuit and the function disagree on these inputs, each written with `{:?}`.
    Mismatch { inputs: String, function: String, circuit: String },
    /// A sampled input is outside its declared range, e.g. from a custom sampler.
    OutOfRange { inputs: String, error: Error },
    /// On these inputs the function would run a `while` or `loop` more times than its `#[max_iterations(N)]`,
    /// where the circuit stops.
    IterationLimit { inputs: String, message: String },
}

impl std::fmt::Display for VerifyError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VerifyError::Mismatch { inputs, function, circuit } => {
                write!(formatter, "for inputs {inputs}, the function returns {function} and the circuit {circuit}")
            }
            VerifyError::OutOfRange { inputs, error } => write!(formatter, "inputs {inputs}: {error}"),
            VerifyError::IterationLimit { inputs, message } => write!(formatter, "inputs {inputs}: {message}"),
        }
    }
}

impl std::error::Error for VerifyError {}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Report {
    /// Inputs where the circuit matched the function.
    pub checked: usize,
    /// Inputs where the function panicked, e.g. on overflow, so any circuit value is allowed.
    pub skipped: usize,
}

const ITERATION_LIMIT: &str = "more times than its #[max_iterations(";

/// What a `#[handwired]` function's `while` or `loop` calls when it would run more times than its
/// `#[max_iterations(limit)]`: the circuit stops after `limit` passes, so the function panics, and `verify` reports
/// the inputs.
#[track_caller]
pub fn iteration_limit(limit: usize, at: &str) -> ! {
    panic!("{at}: this loop would run {ITERATION_LIMIT}{limit})]")
}

thread_local! {
    static QUIET_PANICS: Cell<bool> = const { Cell::new(false) };
}

/// Silence panic messages from oracle calls on this thread only.
fn install_quiet_hook() {
    static INSTALL: Once = Once::new();
    INSTALL.call_once(|| {
        let previous = panic::take_hook();
        panic::set_hook(Box::new(move |info| {
            if !QUIET_PANICS.with(Cell::get) {
                previous(info);
            }
        }));
    });
}

impl<Args: Params + Clone + Debug, R: HwType + PartialEq + Debug> Compiled<Args, R> {
    /// Build the circuit: `body` declares the parameters and returns the outputs.
    pub fn build(oracle: fn(Args) -> R, attrs: &'static [Attrs], body: impl FnOnce(&mut Inputs) -> R::Val) -> Self {
        let previous = CONTEXT.with(|context| context.replace(Some(Circuit::new(Args::CELLS))));
        let guard = ContextGuard(Some(previous));
        let output = body(&mut Inputs { next: 0 });
        let circuit = CONTEXT.with(|context| context.borrow_mut().take()).expect("the build context disappeared");
        drop(guard);
        let mut outputs = Vec::new();
        R::outputs(&output, &mut outputs);
        let program = circuit.compile(Some(&outputs));
        Compiled { circuit, program, outputs, oracle, attrs }
    }

    fn source(&self, args: &Args) -> Vec<f32> {
        let mut source = Vec::with_capacity(Args::CELLS);
        args.encode(&mut source);
        source
    }

    /// Run the compiled program.
    pub fn eval(&self, args: &Args) -> Result<R, Error> {
        let source = self.source(args);
        self.circuit.validate_source_ranges(&source)?;
        let mut executor = TensorExecutor::new(&self.program);
        Ok(R::decode(&mut executor.run(&source).iter().copied()))
    }

    /// Run the compiled program on many records at once, vectorized across records and split across all cores.
    /// To run or differentiate the same records again, encode them once with
    /// [`encode_records`](Self::encode_records) and use [`run_encoded`](Self::run_encoded).
    pub fn run_batch(&self, records: &[Args]) -> Result<Vec<R>, Error>
    where
        Args: Sync,
        R: Send,
    {
        Ok(self.run_encoded(&self.encode_records(records)?))
    }

    /// Encode `records` into source tables and check them against the declared ranges, on every core.
    pub fn encode_records(&self, records: &[Args]) -> Result<EncodedRecords, Error>
    where
        Args: Sync,
    {
        let cells = Args::CELLS;
        let mut sources = vec![0f32; records.len() * cells];
        let chunk = records.len().div_ceil(threads_for(records.len())).max(1);
        let results: Vec<Result<(), Error>> = std::thread::scope(|scope| {
            let handles: Vec<_> = records
                .chunks(chunk)
                .zip(sources.chunks_mut(chunk * cells.max(1)))
                .map(|(chunk_records, chunk_sources)| {
                    scope.spawn(move || {
                        let mut scratch = Vec::with_capacity(cells);
                        for (record, slot) in chunk_records.iter().zip(chunk_sources.chunks_mut(cells.max(1))) {
                            scratch.clear();
                            record.encode(&mut scratch);
                            self.circuit.validate_source_ranges(&scratch)?;
                            slot[..cells].copy_from_slice(&scratch);
                        }
                        Ok(())
                    })
                })
                .collect();
            handles.into_iter().map(|handle| handle.join().expect("an encoding thread panicked")).collect()
        });
        results.into_iter().collect::<Result<(), Error>>()?;
        Ok(EncodedRecords { sources, count: records.len() })
    }

    /// Write `value` as argument `arg` of every encoded record, checked against that argument's declared ranges, e.g.
    /// the parameters a tuning step changes while the rest of each record stays. `value` must have the argument's
    /// type; a different number of cells panics.
    pub fn broadcast<T: HwType>(&self, records: &mut EncodedRecords, arg: usize, value: &T) -> Result<(), Error> {
        let cells = Args::arg_cells(arg);
        let mut encoded = Vec::with_capacity(T::CELLS);
        value.encode(&mut encoded);
        assert_eq!(encoded.len(), cells.len(), "argument {arg} has {} cells, the value {}", cells.len(), encoded.len());
        for (offset, &cell_value) in encoded.iter().enumerate() {
            let wire = Wire::from_index(cells.start + offset);
            let range = self.circuit.range(wire);
            let as_f64 = cell_value as f64;
            if !(as_f64 >= range.min && as_f64 <= range.max) {
                return Err(Error::SourceOutOfRange { wire, value: as_f64, range });
            }
        }
        let width = Args::CELLS.max(1);
        let chunk = records.count.div_ceil(threads_for(records.count)).max(1);
        std::thread::scope(|scope| {
            for chunk_sources in records.sources.chunks_mut(chunk * width) {
                let (cells, encoded) = (cells.clone(), &encoded);
                scope.spawn(move || {
                    for source in chunk_sources.chunks_mut(width) {
                        source[cells.clone()].copy_from_slice(encoded);
                    }
                });
            }
        });
        Ok(())
    }

    /// Run the compiled program on records encoded by [`encode_records`](Self::encode_records).
    pub fn run_encoded(&self, records: &EncodedRecords) -> Vec<R>
    where
        R: Send,
    {
        let outputs = BatchExecutor::run_parallel(&self.program, &records.sources, None);
        let width = self.outputs.len().max(1);
        let chunk = records.count.div_ceil(threads_for(records.count)).max(1);
        std::thread::scope(|scope| {
            let handles: Vec<_> = outputs
                .chunks(chunk * width)
                .map(|chunk_outputs| {
                    scope.spawn(move || {
                        chunk_outputs
                            .chunks(width)
                            .map(|record| R::decode(&mut record.iter().copied()))
                            .collect::<Vec<R>>()
                    })
                })
                .collect();
            handles
                .into_iter()
                .flat_map(|handle| handle.join().expect("a decoding thread panicked"))
                .take(records.count)
                .collect()
        })
    }

    /// The gradient of output `output` (an index into the flattened outputs, e.g. 0 for a single `f32` result) with
    /// respect to every source cell, summed over `records` with per-record `weights` (all 1 when `None`), e.g.
    /// dLoss/dOutput for each record. Each record runs the exact reference evaluator and its backward pass; records
    /// are split across all cores. Cells are in [`Params::encode`] order, one per scalar.
    pub fn gradients(
        &self,
        records: &[Args],
        output: usize,
        weights: Option<&[f32]>,
        options: BackpropOptions,
    ) -> Result<Vec<f64>, Error>
    where
        Args: Sync,
    {
        self.gradients_encoded(&self.encode_records(records)?, output, weights, options)
    }

    /// [`gradients`](Self::gradients) of records encoded by [`encode_records`](Self::encode_records), e.g. the
    /// same training records at every tuning step.
    pub fn gradients_encoded(
        &self,
        records: &EncodedRecords,
        output: usize,
        weights: Option<&[f32]>,
        options: BackpropOptions,
    ) -> Result<Vec<f64>, Error> {
        assert!(output < self.outputs.len(), "output {output} is out of range: there are {}", self.outputs.len());
        if let Some(weights) = weights {
            assert_eq!(weights.len(), records.count, "one weight per record");
        }
        let cells = Args::CELLS;
        let output_wire = self.outputs[output];
        let chunk = records.count.div_ceil(threads_for(records.count)).max(1);
        let partials: Vec<Result<Vec<f64>, Error>> = std::thread::scope(|scope| {
            let handles: Vec<_> = records
                .sources
                .chunks(chunk * cells.max(1))
                .enumerate()
                .map(|(chunk_index, chunk_sources)| {
                    scope.spawn(move || {
                        let mut sums = vec![0f64; cells];
                        let mut seed = vec![0f32; self.circuit.width()];
                        for (offset, source) in chunk_sources.chunks(cells.max(1)).enumerate() {
                            let weight = weights.map_or(1.0, |weights| weights[chunk_index * chunk + offset]);
                            if weight == 0.0 {
                                continue;
                            }
                            let values = self.circuit.evaluate_unchecked(&source[..cells]);
                            seed[output_wire.index()] = weight;
                            let gradients = self.circuit.backpropagate_reference(&values, &seed, options)?;
                            for (sum, &gradient) in sums.iter_mut().zip(&gradients[..cells]) {
                                *sum += gradient as f64;
                            }
                        }
                        Ok(sums)
                    })
                })
                .collect();
            handles.into_iter().map(|handle| handle.join().expect("a gradient thread panicked")).collect()
        });
        let mut total = vec![0f64; cells];
        for partial in partials {
            for (sum, value) in total.iter_mut().zip(partial?) {
                *sum += value;
            }
        }
        Ok(total)
    }

    /// Run the reference evaluator, one operation at a time.
    pub fn eval_reference(&self, args: &Args) -> Result<R, Error> {
        let values = self.circuit.evaluate_reference(&self.source(args))?;
        Ok(R::decode(&mut self.outputs.iter().map(|wire| values[wire.index()])))
    }

    /// Run the original function.
    pub fn oracle(&self, args: Args) -> R {
        (self.oracle)(args)
    }

    /// Compare the compiled program with the original function on `cases`
    /// sampled inputs, including each parameter's range ends. Inputs where
    /// the function panics are skipped. Run it in a build with overflow
    /// checks, e.g. a debug build, so that overflow panics as the contract
    /// assumes instead of wrapping.
    pub fn verify(&self, cases: usize) -> Result<Report, VerifyError>
    where
        Args: Sync,
    {
        self.verify_seeded(cases, 0x5EED)
    }

    pub fn verify_seeded(&self, cases: usize, seed: u64) -> Result<Report, VerifyError>
    where
        Args: Sync,
    {
        self.verify_with(cases, seed, |rng| Args::sample(rng, self.attrs))
    }

    /// [`verify`](Self::verify) on inputs from `sample`, e.g. states reached by simulating a game, which uniform
    /// sampling of every parameter's range rarely hits.
    pub fn verify_with(
        &self,
        cases: usize,
        seed: u64,
        mut sample: impl FnMut(&mut Rng) -> Args,
    ) -> Result<Report, VerifyError>
    where
        Args: Sync,
    {
        if !cfg!(debug_assertions) {
            eprintln!(
                "warning: verify in a build without overflow checks: integer overflow wraps in the function instead \
                 of panicking, and the circuit isn't required to match there"
            );
        }
        install_quiet_hook();
        // Every input is drawn first, in order, so a seed gives the same inputs; then they're checked on every
        // core, and the first failure in sample order is the one reported.
        let mut rng = Rng::new(seed);
        let inputs: Vec<Args> = (0..cases).map(|_| sample(&mut rng)).collect();
        let chunk = inputs.len().div_ceil(threads_for(inputs.len())).max(1);
        let outcomes: Vec<Result<Report, VerifyError>> = std::thread::scope(|scope| {
            let handles: Vec<_> = inputs
                .chunks(chunk)
                .map(|chunk_inputs| scope.spawn(move || self.verify_inputs(chunk_inputs)))
                .collect();
            handles.into_iter().map(|handle| handle.join().expect("a verify thread panicked")).collect()
        });
        let mut report = Report { checked: 0, skipped: 0 };
        for outcome in outcomes {
            let part = outcome?;
            report.checked += part.checked;
            report.skipped += part.skipped;
        }
        Ok(report)
    }

    /// One thread's share of [`verify_with`](Self::verify_with), stopping at its first failure.
    fn verify_inputs(&self, inputs: &[Args]) -> Result<Report, VerifyError> {
        let mut executor = TensorExecutor::new(&self.program);
        let mut report = Report { checked: 0, skipped: 0 };
        for args in inputs {
            QUIET_PANICS.with(|quiet| quiet.set(true));
            let expected = panic::catch_unwind(AssertUnwindSafe(|| (self.oracle)(args.clone())));
            QUIET_PANICS.with(|quiet| quiet.set(false));
            let expected = match expected {
                Ok(expected) => expected,
                // A loop past its declared limit is where the circuit and the function part, not an allowed panic.
                Err(payload) => match payload.downcast_ref::<String>() {
                    Some(message) if message.contains(ITERATION_LIMIT) => {
                        return Err(VerifyError::IterationLimit {
                            inputs: format!("{args:?}"),
                            message: message.clone(),
                        });
                    }
                    _ => {
                        report.skipped += 1;
                        continue;
                    }
                },
            };
            let source = self.source(args);
            self.circuit
                .validate_source_ranges(&source)
                .map_err(|error| VerifyError::OutOfRange { inputs: format!("{args:?}"), error })?;
            let actual = R::decode(&mut executor.run(&source).iter().copied());
            if actual != expected {
                return Err(VerifyError::Mismatch {
                    inputs: format!("{args:?}"),
                    function: format!("{expected:?}"),
                    circuit: format!("{actual:?}"),
                });
            }
            report.checked += 1;
        }
        Ok(report)
    }
}