halo2_base/gates/flex_gate/threads/
parallelize.rs

1use rayon::prelude::*;
2
3use crate::{utils::ScalarField, Context};
4
5use super::SinglePhaseCoreManager;
6
7/// Utility function to parallelize an operation involving [`Context`]s.
8pub fn parallelize_core<F, T, R, FR>(
9    builder: &mut SinglePhaseCoreManager<F>, // leaving `builder` for historical reasons, `pool` is a better name
10    input: Vec<T>,
11    f: FR,
12) -> Vec<R>
13where
14    F: ScalarField,
15    T: Send,
16    R: Send,
17    FR: Fn(&mut Context<F>, T) -> R + Send + Sync,
18{
19    // to prevent concurrency issues with context id, we generate all the ids first
20    let thread_count = builder.thread_count();
21    let mut ctxs =
22        (0..input.len()).map(|i| builder.new_context(thread_count + i)).collect::<Vec<_>>();
23    let outputs: Vec<_> =
24        input.into_par_iter().zip(ctxs.par_iter_mut()).map(|(input, ctx)| f(ctx, input)).collect();
25    // we collect the new threads to ensure they are a FIXED order, otherwise the circuit will not be deterministic
26    builder.threads.append(&mut ctxs);
27
28    outputs
29}