openvm_stark_backend/utils/
batch_inverse.rs

1//! Copied from p3-field [src/batch_inverse.rs] to remove use of rayon
2
3use p3_field::{Field, FieldArray, PackedValue, PrimeCharacteristicRing};
4use tracing::instrument;
5
6/// Batch multiplicative inverses with Montgomery's trick
7/// This is Montgomery's trick. At a high level, we invert the product of the given field
8/// elements, then derive the individual inverses from that via multiplication.
9///
10/// The usual Montgomery trick involves calculating an array of cumulative products,
11/// resulting in a long dependency chain. To increase instruction-level parallelism, we
12/// compute WIDTH separate cumulative product arrays that only meet at the end.
13///
14/// # Panics
15/// This will panic if any of the inputs is zero.
16#[instrument(level = "debug", skip_all)]
17pub fn batch_multiplicative_inverse_serial<F: Field>(x: &[F]) -> Vec<F> {
18    let n = x.len();
19    let mut result = F::zero_vec(n);
20
21    batch_multiplicative_inverse_helper(x, &mut result);
22
23    result
24}
25
26/// Like `batch_multiplicative_inverse`, but writes the result to the given output buffer.
27fn batch_multiplicative_inverse_helper<F: Field>(x: &[F], result: &mut [F]) {
28    // Higher WIDTH increases instruction-level parallelism, but too high a value will cause us
29    // to run out of registers.
30    const WIDTH: usize = 4;
31
32    let n = x.len();
33    assert_eq!(result.len(), n);
34    if !n.is_multiple_of(WIDTH) {
35        // There isn't a very clean way to do this with FieldArray; for now just do it in serial.
36        // Another simple (though suboptimal) workaround would be to make two separate calls, one
37        // for the packed part and one for the remainder.
38        return batch_multiplicative_inverse_general(x, result, |x| x.inverse());
39    }
40
41    let x_packed = FieldArray::<F, 4>::pack_slice(x);
42    let result_packed = FieldArray::<F, 4>::pack_slice_mut(result);
43
44    let inv = |x_packed: FieldArray<F, 4>| {
45        let mut result = FieldArray::<F, 4>::default();
46        batch_multiplicative_inverse_general(&x_packed.0, &mut result.0, |x| x.inverse());
47        result
48    };
49    batch_multiplicative_inverse_general(x_packed, result_packed, inv);
50}
51
52/// A simple single-threaded implementation of Montgomery's trick. Since not all
53/// `PrimeCharacteristicRing`s support inversion, this takes a custom inversion function.
54pub(crate) fn batch_multiplicative_inverse_general<F, Inv>(x: &[F], result: &mut [F], inv: Inv)
55where
56    F: PrimeCharacteristicRing + Copy,
57    Inv: Fn(F) -> F,
58{
59    let n = x.len();
60    assert_eq!(result.len(), n);
61    if n == 0 {
62        return;
63    }
64
65    result[0] = F::ONE;
66    for i in 1..n {
67        result[i] = result[i - 1] * x[i - 1];
68    }
69
70    let product = result[n - 1] * x[n - 1];
71    let mut inv = inv(product);
72
73    for i in (0..n).rev() {
74        result[i] *= inv;
75        inv *= x[i];
76    }
77}
openvm_stark_backend/utils/batch_inverse.rs

openvm_stark_backend/utils/
batch_inverse.rs
