Skip to main content

openvm_cuda_backend/
utils.rs

1use std::mem::transmute;
2
3use itertools::Itertools;
4use openvm_stark_backend::utils::batch_multiplicative_inverse_serial;
5use p3_field::{BasedVectorSpace, ExtensionField, Field, PrimeField64};
6
7use crate::prelude::{D_EF, EF, F};
8
9// https://hackmd.io/@vbuterin/barycentric_evaluation#Special-case-roots-of-unity
10pub fn compute_barycentric_inv_lagrange_denoms<F: Field, EF: ExtensionField<F>>(
11    l_skip: usize,
12    omega_skip_pows: &[F],
13    z: EF,
14) -> Vec<EF> {
15    debug_assert_eq!(1 << l_skip, omega_skip_pows.len());
16    let denoms = omega_skip_pows
17        .iter()
18        .map(|&w_i| {
19            let denom = z - w_i;
20            if denom.is_zero() {
21                EF::ONE
22            } else {
23                denom
24            }
25        })
26        .collect_vec();
27    let mut inv_denoms = batch_multiplicative_inverse_serial(&denoms);
28    let zerofier = z.exp_power_of_2(l_skip) - F::ONE;
29    let denominator = F::from_usize(1 << l_skip);
30    let scale_factor = zerofier * denominator.inverse();
31    for v in &mut inv_denoms {
32        *v *= scale_factor;
33    }
34    inv_denoms
35}
36
37/// Reduce overflowing u64 to extension field elements. After reducing modulo `p`, the `u64` are in
38/// Montgomery form.
39#[inline]
40pub fn reduce_raw_u64_to_ef(accum: &[u64]) -> Vec<EF> {
41    debug_assert_eq!(accum.len() % D_EF, 0);
42    debug_assert_eq!(size_of::<F>(), size_of::<u32>());
43    accum
44        .chunks_exact(D_EF)
45        .map(|chunk| {
46            EF::from_basis_coefficients_fn(|i| {
47                let monty_raw = (chunk[i] % F::ORDER_U64) as u32;
48                // SAFETY:
49                // - BabyBear has same memory layout as u32
50                // - Internally stored in Montgomery form
51                unsafe { transmute::<u32, F>(monty_raw) }
52            })
53        })
54        .collect()
55}