stwo_gpu/core/poly/
utils.rs

1use std_shims::Vec;
2
3use super::line::LineDomain;
4use crate::core::circle::CirclePoint;
5use crate::core::fields::m31::BaseField;
6use crate::core::fields::{ExtensionOf, Field};
7
8/// Folds values recursively in `O(n)` by a hierarchical application of folding factors.
9///
10/// i.e. folding `n = 8` values with `folding_factors = [x, y, z]`:
11///
12/// ```text
13///               n2=n1+x*n2
14///           /               \
15///     n1=n3+y*n4          n2=n5+y*n6
16///      /      \            /      \
17/// n3=a+z*b  n4=c+z*d  n5=e+z*f  n6=g+z*h
18///   /  \      /  \      /  \      /  \
19///  a    b    c    d    e    f    g    h
20/// ```
21///
22/// # Panics
23///
24/// Panics if the number of values is not a power of two or if an incorrect number of of folding
25/// factors is provided.
26// TODO(Andrew): Can be made to run >10x faster by unrolling lower layers of recursion
27pub fn fold<F: Field, E: ExtensionOf<F>>(values: &[F], folding_factors: &[E]) -> E {
28    let n = values.len();
29    assert_eq!(n, 1 << folding_factors.len());
30    if n == 1 {
31        return values[0].into();
32    }
33    let (lhs_values, rhs_values) = values.split_at(n / 2);
34    let (folding_factor, folding_factors) = folding_factors.split_first().unwrap();
35    let lhs_val = fold(lhs_values, folding_factors);
36    let rhs_val = fold(rhs_values, folding_factors);
37    lhs_val + rhs_val * *folding_factor
38}
39
40/// Computes the folding alphas for evaluation by folding.
41/// The folding alphas are the basis for the `len` dimensional FFT-basis:
42/// y, x, pi(x), pi^2(x), ..., pi^{len-2}(x).
43/// Returns the folding alphas in reverse order.
44#[allow(clippy::uninit_vec)]
45pub fn get_folding_alphas<F: Field + ExtensionOf<BaseField>>(
46    point: CirclePoint<F>,
47    len: usize,
48) -> Vec<F> {
49    if len == 0 {
50        return Vec::new();
51    }
52    // Manually set the length of the vector so we can directly assign the elements instead of push
53    // and then reverse.
54    let mut alphas = Vec::with_capacity(len);
55    unsafe {
56        alphas.set_len(len);
57    }
58
59    // Fill the vector with the correct values.
60    alphas[len - 1] = point.y;
61    if len > 1 {
62        let mut x = point.x;
63        for i in (0..len - 1).rev() {
64            alphas[i] = x;
65            x = CirclePoint::double_x(x);
66        }
67    }
68
69    alphas
70}
71
72/// Repeats each value sequentially `duplicity` many times.
73///
74/// # Examples
75///
76/// ```rust
77/// # use stwo::core::poly::utils::repeat_value;
78/// assert_eq!(repeat_value(&[1, 2, 3], 2), vec![1, 1, 2, 2, 3, 3]);
79/// ```
80pub fn repeat_value<T: Copy>(values: &[T], duplicity: usize) -> Vec<T> {
81    let n = values.len();
82    let mut res: Vec<T> = Vec::with_capacity(n * duplicity);
83
84    // Fill each chunk with its corresponding value.
85    for &v in values {
86        for _ in 0..duplicity {
87            res.push(v)
88        }
89    }
90
91    res
92}
93
94/// Computes the line twiddles for a [`CircleDomain`] or a [`LineDomain`] from the precomputed
95/// twiddles tree.
96///
97/// [`CircleDomain`]: super::circle::CircleDomain
98pub fn domain_line_twiddles_from_tree<T>(
99    domain: impl Into<LineDomain>,
100    twiddle_buffer: &[T],
101) -> Vec<&[T]> {
102    let domain = domain.into();
103    assert!(
104        domain.coset().size() <= twiddle_buffer.len(),
105        "Not enough twiddles!"
106    );
107    (0..domain.coset().log_size())
108        .map(|i| {
109            let len = 1 << i;
110            &twiddle_buffer[twiddle_buffer.len() - len * 2..twiddle_buffer.len() - len]
111        })
112        .rev()
113        .collect()
114}
115
116#[cfg(test)]
117mod tests {
118    use std_shims::vec;
119
120    use super::repeat_value;
121    use crate::core::poly::circle::CanonicCoset;
122    use crate::core::poly::line::LineDomain;
123    use crate::core::poly::utils::domain_line_twiddles_from_tree;
124
125    #[test]
126    fn repeat_value_0_times_works() {
127        assert!(repeat_value(&[1, 2, 3], 0).is_empty());
128    }
129
130    #[test]
131    fn repeat_value_2_times_works() {
132        assert_eq!(repeat_value(&[1, 2, 3], 2), vec![1, 1, 2, 2, 3, 3]);
133    }
134
135    #[test]
136    fn repeat_value_3_times_works() {
137        assert_eq!(repeat_value(&[1, 2], 3), vec![1, 1, 1, 2, 2, 2]);
138    }
139
140    #[test]
141    fn test_domain_line_twiddles_works() {
142        let domain: LineDomain = CanonicCoset::new(4).circle_domain().into();
143        let twiddles = domain_line_twiddles_from_tree(domain, &[0, 1, 2, 3, 4, 5, 6, 7]);
144        assert_eq!(twiddles.len(), 3);
145        assert_eq!(twiddles[0], &[0, 1, 2, 3]);
146        assert_eq!(twiddles[1], &[4, 5]);
147        assert_eq!(twiddles[2], &[6]);
148    }
149
150    #[test]
151    #[should_panic]
152    fn test_domain_line_twiddles_fails() {
153        let domain: LineDomain = CanonicCoset::new(5).circle_domain().into();
154        domain_line_twiddles_from_tree(domain, &[0, 1, 2, 3, 4, 5, 6, 7]);
155    }
156}
stwo_gpu/core/poly/utils.rs

stwo_gpu/core/poly/
utils.rs
