1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105
// Copyright 2024 Ulvetanna Inc.
use crate::Error;
use binius_field::PackedField;
use binius_utils::bail;
use rayon::prelude::*;
use std::cmp::max;
/// Tensor Product expansion of values with partial eq indicator evaluated at extra_query_coordinates
///
/// Let $n$ be log_n_values, $p$, $k$ be the lengths of `packed_values` and `extra_query_coordinates`.
/// Requires
/// * $n \geq k$
/// * p = max(1, 2^{n+k} / P::WIDTH)
/// Let $v$ be a vector corresponding to the first $2^n$ scalar values of `values`.
/// Let $r = (r_0, \ldots, r_{k-1})$ be the vector of `extra_query_coordinates`.
///
/// # Formal Definition
/// `values` is updated to contain the result of:
/// $v \otimes (1 - r_0, r_0) \otimes \ldots \otimes (1 - r_{k-1}, r_{k-1})$
/// which is now a vector of length $2^{n+k}$. If 2^{n+k} < P::WIDTH, then
/// the result is packed into a single element of `values` where only the first
/// 2^{n+k} elements have meaning.
///
/// # Interpretation
/// Let $f$ be an $n$ variate multilinear polynomial that has evaluations over
/// the $n$ dimensional hypercube corresponding to $v$.
/// Then `values` is updated to contain the evaluations of $g$ over the $n+k$-dimensional
/// hypercube where
/// * $g(x_0, \ldots, x_{n+k-1}) = f(x_0, \ldots, x_{n-1}) * eq(x_n, \ldots, x_{n+k-1}, r)$
pub fn tensor_prod_eq_ind<P: PackedField>(
log_n_values: usize,
packed_values: &mut [P],
extra_query_coordinates: &[P::Scalar],
) -> Result<(), Error> {
let new_n_vars = log_n_values + extra_query_coordinates.len();
if packed_values.len() != max(1, (1 << new_n_vars) / P::WIDTH) {
bail!(Error::InvalidPackedValuesLength);
}
for (i, r_i) in extra_query_coordinates.iter().enumerate() {
let prev_length = 1 << (log_n_values + i);
if prev_length < P::WIDTH {
let q = &mut packed_values[0];
for h in 0..prev_length {
let x = q.get(h);
let prod = x * r_i;
q.set(h, x - prod);
q.set(prev_length | h, prod);
}
} else {
let prev_packed_length = prev_length / P::WIDTH;
let packed_r_i = P::broadcast(*r_i);
let (xs, ys) = packed_values.split_at_mut(prev_packed_length);
assert!(xs.len() <= ys.len());
// These magic numbers were chosen experimentally to have a reasonable performance
// for the calls with small number of elements.
xs.par_iter_mut()
.zip(ys.par_iter_mut())
.with_min_len(64)
.for_each(|(x, y): (&mut P, &mut P)| {
// x = x * (1 - packed_r_i) = x - x * packed_r_i
// y = x * packed_r_i
// Notice that we can reuse the multiplication: (x * packed_r_i)
let prod = (*x) * packed_r_i;
*x -= prod;
*y = prod;
});
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use binius_field::{
packed::{iter_packed_slice, set_packed_slice},
Field, PackedBinaryField4x32b,
};
use itertools::Itertools;
#[test]
fn test_tensor_prod_eq_ind() {
type P = PackedBinaryField4x32b;
type F = <P as PackedField>::Scalar;
let v0 = F::new(1);
let v1 = F::new(2);
let query = vec![v0, v1];
let mut result = vec![P::default(); 1 << (query.len() - P::LOG_WIDTH)];
set_packed_slice(&mut result, 0, F::ONE);
tensor_prod_eq_ind(0, &mut result, &query).unwrap();
let result = iter_packed_slice(&result).collect_vec();
assert_eq!(
result,
vec![
(F::ONE - v0) * (F::ONE - v1),
v0 * (F::ONE - v1),
(F::ONE - v0) * v1,
v0 * v1
]
);
}
}