p3_monty_31/dft/
forward.rs

1#![allow(clippy::use_self)]
2
3//! Discrete Fourier Transform, in-place, decimation-in-frequency
4//!
5//! Straightforward recursive algorithm, "unrolled" up to size 256.
6//!
7//! Inspired by Bernstein's djbfft: https://cr.yp.to/djbfft.html
8
9extern crate alloc;
10
11use alloc::vec::Vec;
12
13use itertools::izip;
14use p3_field::{Field, PackedFieldPow2, PackedValue, PrimeCharacteristicRing, TwoAdicField};
15use p3_util::log2_strict_usize;
16
17use crate::utils::monty_reduce;
18use crate::{FieldParameters, MontyField31, TwoAdicData};
19
20impl<MP: FieldParameters + TwoAdicData> MontyField31<MP> {
21    /// Given a field element `gen` of order n where `n = 2^lg_n`,
22    /// return a vector of vectors `table` where table[i] is the
23    /// vector of twiddle factors for an fft of length n/2^i. The
24    /// values g_i^k for k >= i/2 are skipped as these are just the
25    /// negatives of the other roots (using g_i^{i/2} = -1).  The
26    /// value gen^0 = 1 is included to aid consistency between the
27    /// packed and non-packed variants.
28    pub fn roots_of_unity_table(n: usize) -> Vec<Vec<Self>> {
29        let lg_n = log2_strict_usize(n);
30        let generator = Self::two_adic_generator(lg_n);
31        let half_n = 1 << (lg_n - 1);
32        // nth_roots = [1, g, g^2, g^3, ..., g^{n/2 - 1}]
33        let nth_roots = generator.powers().collect_n(half_n);
34
35        (0..(lg_n - 1))
36            .map(|i| nth_roots.iter().step_by(1 << i).copied().collect())
37            .rev()
38            .collect()
39    }
40
41    pub fn get_missing_twiddles(req_lg_n: usize, cur_lg_n: usize) -> Vec<Vec<Self>> {
42        // Get the main generator for the largest required FFT size.
43        let main_generator = Self::two_adic_generator(req_lg_n);
44
45        (cur_lg_n..req_lg_n)
46            .map(|level| {
47                // For a given 'level', we're generating twiddles for a DIF pass
48                // where the number of butterflies is m = 2^level.
49                let count = 1 << level;
50
51                // The generator for this smaller FFT size is a power of the main generator.
52                //
53                // The exponent is 2^(req_lg_n - (level + 1)).
54                let sub_generator_exp = 1 << (req_lg_n - level - 1);
55                let sub_generator = main_generator.exp_u64(sub_generator_exp as u64);
56
57                // Now, we can collect the 'count' powers of this specific sub-generator.
58                sub_generator.powers().collect_n(count)
59            })
60            .collect()
61    }
62}
63
64#[inline(always)]
65fn forward_butterfly<T: PrimeCharacteristicRing + Copy>(x: T, y: T, roots: T) -> (T, T) {
66    let t = x - y;
67    (x + y, t * roots)
68}
69
70#[inline(always)]
71fn forward_butterfly_interleaved<const HALF_RADIX: usize, T: PackedFieldPow2>(
72    x: T,
73    y: T,
74    roots: T,
75) -> (T, T) {
76    let (x, y) = x.interleave(y, HALF_RADIX);
77    let (x, y) = forward_butterfly(x, y, roots);
78    x.interleave(y, HALF_RADIX)
79}
80
81#[inline]
82fn forward_pass_packed<T: PackedFieldPow2>(input: &mut [T], roots: &[T::Scalar]) {
83    let packed_roots = T::pack_slice(roots);
84    let n = input.len();
85    let (xs, ys) = unsafe { input.split_at_mut_unchecked(n / 2) };
86
87    izip!(xs, ys, packed_roots)
88        .for_each(|(x, y, &roots)| (*x, *y) = forward_butterfly(*x, *y, roots));
89}
90
91#[inline]
92fn forward_iterative_layer_1<T: PackedFieldPow2>(input: &mut [T], roots: &[T::Scalar]) {
93    let packed_roots = T::pack_slice(roots);
94    let n = input.len();
95    let (top_half, bottom_half) = unsafe { input.split_at_mut_unchecked(n / 2) };
96    let (xs, ys) = unsafe { top_half.split_at_mut_unchecked(n / 4) };
97    let (zs, ws) = unsafe { bottom_half.split_at_mut_unchecked(n / 4) };
98
99    izip!(xs, ys, zs, ws, packed_roots).for_each(|(x, y, z, w, &root)| {
100        (*x, *y) = forward_butterfly(*x, *y, root);
101        (*z, *w) = forward_butterfly(*z, *w, root);
102    });
103}
104
105#[inline]
106fn forward_iterative_packed<const HALF_RADIX: usize, T: PackedFieldPow2>(
107    input: &mut [T],
108    roots: &[T::Scalar],
109) {
110    // roots[0] == 1
111    // roots <-- [1, roots[1], ..., roots[HALF_RADIX-1], 1, roots[1], ...]
112    let roots = T::from_fn(|i| roots[i % HALF_RADIX]);
113
114    input.chunks_exact_mut(2).for_each(|pair| {
115        let (x, y) = forward_butterfly_interleaved::<HALF_RADIX, _>(pair[0], pair[1], roots);
116        pair[0] = x;
117        pair[1] = y;
118    });
119}
120
121#[inline]
122fn forward_iterative_packed_radix_2<T: PackedFieldPow2>(input: &mut [T]) {
123    input.chunks_exact_mut(2).for_each(|pair| {
124        let x = pair[0];
125        let y = pair[1];
126        let (mut x, y) = x.interleave(y, 1);
127        let t = x - y; // roots[0] == 1
128        x += y;
129        let (x, y) = x.interleave(t, 1);
130        pair[0] = x;
131        pair[1] = y;
132    });
133}
134
135impl<MP: FieldParameters + TwoAdicData> MontyField31<MP> {
136    #[inline]
137    fn forward_iterative_layer(
138        packed_input: &mut [<Self as Field>::Packing],
139        roots: &[Self],
140        m: usize,
141    ) {
142        debug_assert_eq!(roots.len(), m);
143        let packed_roots = <Self as Field>::Packing::pack_slice(roots);
144
145        // lg_m >= 4, so m = 2^lg_m >= 2^4, hence packing_width divides m
146        let packed_m = m / <Self as Field>::Packing::WIDTH;
147        packed_input
148            .chunks_exact_mut(2 * packed_m)
149            .for_each(|layer_chunk| {
150                let (xs, ys) = unsafe { layer_chunk.split_at_mut_unchecked(packed_m) };
151
152                izip!(xs, ys, packed_roots)
153                    .for_each(|(x, y, &root)| (*x, *y) = forward_butterfly(*x, *y, root));
154            });
155    }
156
157    #[inline]
158    fn forward_iterative_packed_radix_16(input: &mut [<Self as Field>::Packing]) {
159        // Rather surprisingly, a version similar where the separate
160        // loops in each call to forward_iterative_packed() are
161        // combined into one, was not only not faster, but was
162        // actually a bit slower.
163
164        // Radix 16
165        if <Self as Field>::Packing::WIDTH >= 16 {
166            forward_iterative_packed::<8, _>(input, MP::ROOTS_16.as_ref());
167        } else {
168            Self::forward_iterative_layer(input, MP::ROOTS_16.as_ref(), 8);
169        }
170
171        // Radix 8
172        if <Self as Field>::Packing::WIDTH >= 8 {
173            forward_iterative_packed::<4, _>(input, MP::ROOTS_8.as_ref());
174        } else {
175            Self::forward_iterative_layer(input, MP::ROOTS_8.as_ref(), 4);
176        }
177
178        // Radix 4
179        let roots4 = [MP::ROOTS_8.as_ref()[0], MP::ROOTS_8.as_ref()[2]];
180        if <Self as Field>::Packing::WIDTH >= 4 {
181            forward_iterative_packed::<2, _>(input, &roots4);
182        } else {
183            Self::forward_iterative_layer(input, &roots4, 2);
184        }
185
186        // Radix 2
187        forward_iterative_packed_radix_2(input);
188    }
189
190    /// Breadth-first DIF FFT for smallish vectors (must be >= 64)
191    #[inline]
192    fn forward_iterative(packed_input: &mut [<Self as Field>::Packing], root_table: &[Vec<Self>]) {
193        assert!(packed_input.len() >= 2);
194        let packing_width = <Self as Field>::Packing::WIDTH;
195        let n = packed_input.len() * packing_width;
196        let lg_n = log2_strict_usize(n);
197        debug_assert_eq!(root_table.len(), lg_n - 1);
198
199        // Stop loop early to do radix 16 separately. This value is determined by the largest
200        // packing width we will encounter, which is 16 at the moment for AVX512. Specifically
201        // it is log_2(max{possible packing widths}) = lg(16) = 4.
202        const LAST_LOOP_LAYER: usize = 4;
203
204        // How many layers have we specialised before the main loop
205        const NUM_SPECIALISATIONS: usize = 2;
206
207        // Needed to avoid overlap of the 2 specialisations at the start
208        // with the radix-16 specialisation at the end of the loop
209        assert!(lg_n >= LAST_LOOP_LAYER + NUM_SPECIALISATIONS);
210
211        // Specialise the first NUM_SPECIALISATIONS iterations; improves performance a little.
212        forward_pass_packed(packed_input, &root_table[lg_n - 2]); // lg_m == lg_n - 1, s == 0
213        forward_iterative_layer_1(packed_input, &root_table[lg_n - 3]); // lg_m == lg_n - 2, s == 1
214
215        // loop from lg_n-2 down to 4.
216        for lg_m in (LAST_LOOP_LAYER..(lg_n - NUM_SPECIALISATIONS)).rev() {
217            let m = 1 << lg_m;
218
219            let roots = &root_table[lg_m - 1];
220            debug_assert_eq!(roots.len(), m);
221
222            Self::forward_iterative_layer(packed_input, roots, m);
223        }
224
225        // Last 4 layers
226        Self::forward_iterative_packed_radix_16(packed_input);
227    }
228
229    #[inline(always)]
230    fn forward_butterfly(x: Self, y: Self, w: Self) -> (Self, Self) {
231        let t = MP::PRIME + x.value - y.value;
232        (
233            x + y,
234            Self::new_monty(monty_reduce::<MP>(t as u64 * w.value as u64)),
235        )
236    }
237
238    #[inline]
239    fn forward_pass(input: &mut [Self], roots: &[Self]) {
240        let half_n = input.len() / 2;
241        assert_eq!(roots.len(), half_n);
242
243        // Safe because 0 <= half_n < a.len()
244        let (xs, ys) = unsafe { input.split_at_mut_unchecked(half_n) };
245
246        let s = xs[0] + ys[0];
247        let t = xs[0] - ys[0];
248        xs[0] = s;
249        ys[0] = t;
250
251        izip!(&mut xs[1..], &mut ys[1..], &roots[1..]).for_each(|(x, y, &root)| {
252            (*x, *y) = Self::forward_butterfly(*x, *y, root);
253        });
254    }
255
256    #[inline(always)]
257    fn forward_2(a: &mut [Self]) {
258        assert_eq!(a.len(), 2);
259
260        let s = a[0] + a[1];
261        let t = a[0] - a[1];
262        a[0] = s;
263        a[1] = t;
264    }
265
266    #[inline(always)]
267    fn forward_4(a: &mut [Self]) {
268        assert_eq!(a.len(), 4);
269
270        // Expanding the calculation of t3 saves one instruction
271        let t1 = MP::PRIME + a[1].value - a[3].value;
272        let t3 = Self::new_monty(monty_reduce::<MP>(
273            t1 as u64 * MP::ROOTS_8.as_ref()[2].value as u64,
274        ));
275        let t5 = a[1] + a[3];
276        let t4 = a[0] + a[2];
277        let t2 = a[0] - a[2];
278
279        // Return in bit-reversed order
280        a[0] = t4 + t5;
281        a[1] = t4 - t5;
282        a[2] = t2 + t3;
283        a[3] = t2 - t3;
284    }
285
286    #[inline(always)]
287    fn forward_8(a: &mut [Self]) {
288        assert_eq!(a.len(), 8);
289
290        Self::forward_pass(a, MP::ROOTS_8.as_ref());
291
292        // Safe because a.len() == 8
293        let (a0, a1) = unsafe { a.split_at_mut_unchecked(a.len() / 2) };
294        Self::forward_4(a0);
295        Self::forward_4(a1);
296    }
297
298    #[inline(always)]
299    fn forward_16(a: &mut [Self]) {
300        assert_eq!(a.len(), 16);
301
302        Self::forward_pass(a, MP::ROOTS_16.as_ref());
303
304        // Safe because a.len() == 16
305        let (a0, a1) = unsafe { a.split_at_mut_unchecked(a.len() / 2) };
306        Self::forward_8(a0);
307        Self::forward_8(a1);
308    }
309
310    #[inline(always)]
311    fn forward_32(a: &mut [Self], root_table: &[Vec<Self>]) {
312        assert_eq!(a.len(), 32);
313
314        Self::forward_pass(a, &root_table[root_table.len() - 1]);
315
316        // Safe because a.len() == 32
317        let (a0, a1) = unsafe { a.split_at_mut_unchecked(a.len() / 2) };
318        Self::forward_16(a0);
319        Self::forward_16(a1);
320    }
321
322    /// Assumes `input.len() >= 64`.
323    #[inline]
324    fn forward_fft_recur(input: &mut [<Self as Field>::Packing], root_table: &[Vec<Self>]) {
325        const ITERATIVE_FFT_THRESHOLD: usize = 1024;
326
327        let n = input.len() * <Self as Field>::Packing::WIDTH;
328        if n <= ITERATIVE_FFT_THRESHOLD {
329            Self::forward_iterative(input, root_table);
330        } else {
331            assert_eq!(n, 1 << (root_table.len() + 1));
332            forward_pass_packed(input, &root_table[root_table.len() - 1]);
333
334            // Safe because input.len() > ITERATIVE_FFT_THRESHOLD
335            let (a0, a1) = unsafe { input.split_at_mut_unchecked(input.len() / 2) };
336
337            Self::forward_fft_recur(a0, &root_table[..root_table.len() - 1]);
338            Self::forward_fft_recur(a1, &root_table[..root_table.len() - 1]);
339        }
340    }
341
342    #[inline]
343    pub fn forward_fft(input: &mut [Self], root_table: &[Vec<Self>]) {
344        let n = input.len();
345        if n == 1 {
346            return;
347        }
348        assert_eq!(n, 1 << (root_table.len() + 1));
349        match n {
350            32 => Self::forward_32(input, root_table),
351            16 => Self::forward_16(input),
352            8 => Self::forward_8(input),
353            4 => Self::forward_4(input),
354            2 => Self::forward_2(input),
355            _ => {
356                let packed_input = <Self as Field>::Packing::pack_slice_mut(input);
357                Self::forward_fft_recur(packed_input, root_table);
358            }
359        }
360    }
361}
p3_monty_31/dft/forward.rs

p3_monty_31/dft/
forward.rs
