p3_binary_dft/
poly.rs

1//! The Lin–Chung–Han transform carried out in the polynomial basis of `GF(2^128)`.
2
3use alloc::vec::Vec;
4
5use p3_binary_field::{BinaryField128, TowerLevel, poly_basis};
6use p3_matrix::Matrix;
7use p3_matrix::dense::RowMajorMatrix;
8use p3_maybe_rayon::prelude::*;
9use p3_util::log2_strict_usize;
10
11use crate::domain::{domain_point, subspace_polynomial};
12use crate::lch::{BUTTERFLY_GRAIN, LchNtt};
13use crate::traits::AdditiveNtt;
14
15/// [`LchNtt`] over `BinaryField128`, with the data held in the polynomial basis throughout.
16///
17/// A tower-basis product converts both operands into the polynomial basis and the result back,
18/// sixteen dependent table lookups apiece, which is most of what a butterfly costs. Converting
19/// the whole matrix once on the way in and once on the way out pays `2` conversions per element
20/// instead of `3ℓ/2`, and every twiddle multiply in between is a bare carryless multiply and a
21/// reduction. Additions are `XOR` in both bases, so they are unaffected.
22///
23/// Without a carryless-multiply instruction that product is a bit-serial loop and slower than
24/// the tower arithmetic it replaces, so on such a target the transform runs in the tower basis
25/// instead. The choice is a constant and only one arm survives compilation.
26#[derive(Clone, Debug, Default)]
27pub struct PolyBasisNtt {
28    tower: LchNtt<BinaryField128>,
29}
30
31/// The twiddle of block `blk` of stage `j`, in the polynomial basis.
32#[inline]
33fn twiddle(base: BinaryField128, blk: usize) -> u128 {
34    poly_basis::from_tower(base + domain_point::<BinaryField128>(blk << 1))
35}
36
37impl AdditiveNtt<BinaryField128> for PolyBasisNtt {
38    fn shifted_ntt_batch(
39        &self,
40        mut mat: RowMajorMatrix<BinaryField128>,
41        shift: BinaryField128,
42    ) -> RowMajorMatrix<BinaryField128> {
43        if !poly_basis::HAS_HARDWARE_CLMUL {
44            return self.tower.shifted_ntt_batch(mat, shift);
45        }
46
47        let width = mat.width();
48        let log_n = log2_strict_usize(mat.height());
49        // `BinaryField128` is `#[repr(transparent)]` over `u128`, so this reuses `mat.values`'s
50        // allocation instead of allocating a second codeword buffer: `to_repr` is the identity
51        // bit pattern, and the specialised `Vec` collect below reinterprets the buffer in place.
52        let mut values: Vec<u128> = core::mem::take(&mut mat.values)
53            .into_iter()
54            .map(BinaryField128::to_repr)
55            .collect();
56        values
57            .par_iter_mut()
58            .for_each(|v| *v = poly_basis::from_tower(BinaryField128::from_repr(*v)));
59
60        for j in (0..log_n).rev() {
61            let half = (1 << j) * width;
62            let base = subspace_polynomial::<BinaryField128>(j, shift);
63            values
64                .par_chunks_mut(half << 1)
65                .enumerate()
66                .for_each(|(blk, block)| {
67                    let t = twiddle(base, blk);
68                    let zero = t == 0;
69                    let (lo, hi) = block.split_at_mut(half);
70                    let butterfly = |lo: &mut [u128], hi: &mut [u128]| {
71                        if zero {
72                            for (u, v) in lo.iter_mut().zip(hi) {
73                                *v ^= *u;
74                            }
75                        } else {
76                            for (u, v) in lo.iter_mut().zip(hi) {
77                                *u ^= poly_basis::mul(t, *v);
78                                *v ^= *u;
79                            }
80                        }
81                    };
82                    if half <= BUTTERFLY_GRAIN {
83                        butterfly(lo, hi);
84                    } else {
85                        lo.par_chunks_mut(BUTTERFLY_GRAIN)
86                            .zip(hi.par_chunks_mut(BUTTERFLY_GRAIN))
87                            .for_each(|(lo, hi)| butterfly(lo, hi));
88                    }
89                });
90        }
91
92        values
93            .par_iter_mut()
94            .for_each(|v| *v = poly_basis::to_tower(*v).to_repr());
95        mat.values = values.into_iter().map(BinaryField128::from_repr).collect();
96        mat
97    }
98
99    fn shifted_intt_batch(
100        &self,
101        mut mat: RowMajorMatrix<BinaryField128>,
102        shift: BinaryField128,
103    ) -> RowMajorMatrix<BinaryField128> {
104        if !poly_basis::HAS_HARDWARE_CLMUL {
105            return self.tower.shifted_intt_batch(mat, shift);
106        }
107
108        let width = mat.width();
109        let log_n = log2_strict_usize(mat.height());
110        // See `shifted_ntt_batch`: reuses `mat.values`'s allocation instead of a second buffer.
111        let mut values: Vec<u128> = core::mem::take(&mut mat.values)
112            .into_iter()
113            .map(BinaryField128::to_repr)
114            .collect();
115        values
116            .par_iter_mut()
117            .for_each(|v| *v = poly_basis::from_tower(BinaryField128::from_repr(*v)));
118
119        for j in 0..log_n {
120            let half = (1 << j) * width;
121            let base = subspace_polynomial::<BinaryField128>(j, shift);
122            values
123                .par_chunks_mut(half << 1)
124                .enumerate()
125                .for_each(|(blk, block)| {
126                    let t = twiddle(base, blk);
127                    let zero = t == 0;
128                    let (lo, hi) = block.split_at_mut(half);
129                    let butterfly = |lo: &mut [u128], hi: &mut [u128]| {
130                        if zero {
131                            for (u, v) in lo.iter_mut().zip(hi) {
132                                *v ^= *u;
133                            }
134                        } else {
135                            for (u, v) in lo.iter_mut().zip(hi) {
136                                *v ^= *u;
137                                *u ^= poly_basis::mul(t, *v);
138                            }
139                        }
140                    };
141                    if half <= BUTTERFLY_GRAIN {
142                        butterfly(lo, hi);
143                    } else {
144                        lo.par_chunks_mut(BUTTERFLY_GRAIN)
145                            .zip(hi.par_chunks_mut(BUTTERFLY_GRAIN))
146                            .for_each(|(lo, hi)| butterfly(lo, hi));
147                    }
148                });
149        }
150
151        values
152            .par_iter_mut()
153            .for_each(|v| *v = poly_basis::to_tower(*v).to_repr());
154        mat.values = values.into_iter().map(BinaryField128::from_repr).collect();
155        mat
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use p3_binary_field::{BinaryField128, TowerLevel};
162    use p3_matrix::dense::RowMajorMatrix;
163    use proptest::prelude::*;
164
165    use super::PolyBasisNtt;
166    use crate::naive::NaiveAdditiveNtt;
167    use crate::traits::AdditiveNtt;
168
169    /// Builds a matrix whose entries are distinct functions of the seed and the position.
170    fn matrix(log_n: usize, width: usize, seed: u64) -> RowMajorMatrix<BinaryField128> {
171        RowMajorMatrix::new(
172            (0..(width << log_n))
173                .map(|i| {
174                    let bits = seed
175                        .wrapping_mul(0x9e37_79b9_7f4a_7c15)
176                        .wrapping_add(i as u64);
177                    BinaryField128::from_le_byte_iter(bits.to_le_bytes().into_iter().cycle())
178                })
179                .collect(),
180            width,
181        )
182    }
183
184    proptest! {
185        #![proptest_config(ProptestConfig::with_cases(32))]
186
187        /// The polynomial-basis transform is the same map as the reference oracle.
188        #[test]
189        fn poly_basis_matches_naive(
190            log_n in 0usize..=8,
191            width in 1usize..=5,
192            seed in any::<u64>(),
193            shift in any::<u64>(),
194        ) {
195            let coeffs = matrix(log_n, width, seed);
196            let shift = BinaryField128::from_le_byte_iter(
197                shift.to_le_bytes().into_iter().cycle(),
198            );
199
200            let fast = PolyBasisNtt::default().shifted_ntt_batch(coeffs.clone(), shift);
201            let slow = NaiveAdditiveNtt::<BinaryField128>::default()
202                .shifted_ntt_batch(coeffs, shift);
203            prop_assert_eq!(fast, slow);
204        }
205
206        #[test]
207        fn poly_basis_round_trips(
208            log_n in 0usize..=8,
209            width in 1usize..=3,
210            seed in any::<u64>(),
211            shift in any::<u64>(),
212        ) {
213            let coeffs = matrix(log_n, width, seed);
214            let shift = BinaryField128::from_le_byte_iter(
215                shift.to_le_bytes().into_iter().cycle(),
216            );
217            let ntt = PolyBasisNtt::default();
218            let evals = ntt.shifted_ntt_batch(coeffs.clone(), shift);
219            prop_assert_eq!(ntt.shifted_intt_batch(evals, shift), coeffs);
220        }
221
222        /// `PolyBasisNtt`'s low-degree extension agrees with the oracle's, on a coset too, and
223        /// the input rows reappear as the prefix: the correspondence Phase 3 folds along.
224        #[test]
225        fn poly_basis_lde_matches_naive(
226            log_n in 0usize..=6,
227            added in 0usize..=3,
228            width in 1usize..=3,
229            seed in any::<u64>(),
230            shift in any::<u64>(),
231        ) {
232            let coeffs = matrix(log_n, width, seed);
233            let shift = BinaryField128::from_le_byte_iter(
234                shift.to_le_bytes().into_iter().cycle(),
235            );
236
237            let lde = PolyBasisNtt::default().shifted_lde_batch(coeffs.clone(), added, shift);
238            let naive = NaiveAdditiveNtt::<BinaryField128>::default()
239                .shifted_lde_batch(coeffs.clone(), added, shift);
240            prop_assert_eq!(&lde, &naive);
241            prop_assert_eq!(&lde.values[..coeffs.values.len()], &coeffs.values[..]);
242        }
243    }
244
245    /// A height that is not a power of two has no well-defined `ℓ`. `ℓ` exceeding the bit width
246    /// of `BinaryField128` is covered generically at [`LchNtt`](crate::LchNtt), where the level
247    /// is a type parameter and the panic is cheap to reach; at a fixed `BinaryField128` it is
248    /// only reachable through a matrix of `2^129` rows, which is not a test worth writing.
249    #[test]
250    #[should_panic]
251    fn shifted_ntt_batch_rejects_a_non_power_of_two_height() {
252        let coeffs = RowMajorMatrix::new(matrix(0, 1, 0).values.repeat(3), 1);
253        let _ = PolyBasisNtt::default().ntt_batch(coeffs);
254    }
255}
p3_binary_dft/poly.rs

p3_binary_dft/
poly.rs
