p3_dft/
radix_2_dit.rs

1use alloc::sync::Arc;
2
3use p3_field::{Field, TwoAdicField};
4use p3_matrix::Matrix;
5use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixViewMut};
6use p3_matrix::util::reverse_matrix_index_bits;
7use p3_maybe_rayon::prelude::*;
8use p3_util::log2_strict_usize;
9
10use crate::butterflies::{Butterfly, DitButterfly, TwiddleFreeButterfly};
11use crate::twiddle_cache::TwiddleCache;
12use crate::util::coset_shift_cols;
13use crate::{Layout, TwoAdicSubgroupDft};
14
15/// Radix-2 Decimation-in-Time FFT over a two-adic subgroup.
16///
17/// This struct implements a fast Fourier transform (FFT) using the Radix-2
18/// Decimation-in-Time (DIT) algorithm over a two-adic multiplicative subgroup of a finite field.
19/// It is optimized for a batch setting where multiple FFT's are being computed simultaneously.
20///
21/// Internally, the implementation memoizes twiddle factors (powers of the root of unity)
22/// for reuse across multiple transforms. This avoids redundant computation
23/// when performing FFTs of the same size.
24#[derive(Default, Clone, Debug)]
25pub struct Radix2Dit<F: TwoAdicField> {
26    /// Memoized twiddle factors indexed by `log2(n)`, where `n` is the DFT length.
27    ///
28    /// This allows fast lookup and reuse of previously computed twiddle values
29    /// (powers of a two-adic generator), which are expensive to recompute.
30    twiddles: Arc<TwiddleCache<usize, [F]>>,
31}
32
33impl<F: TwoAdicField> Radix2Dit<F> {
34    /// Returns the twiddle factors for a DFT of size `2^log_h`.
35    /// If they haven't been computed yet, this function computes and caches them.
36    fn get_or_compute_twiddles(&self, log_h: usize) -> Arc<[F]> {
37        self.twiddles.get_or_compute(log_h, || {
38            let n = 1 << log_h;
39            let root = F::two_adic_generator(log_h);
40            Arc::from(root.powers().collect_n(n / 2))
41        })
42    }
43}
44
45impl<F: TwoAdicField> TwoAdicSubgroupDft<F> for Radix2Dit<F> {
46    type Evaluations = RowMajorMatrix<F>;
47
48    fn dft_batch(&self, mut mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
49        let h = mat.height();
50        let log_h = log2_strict_usize(h);
51
52        // Compute twiddle factors, or take memoized ones if already available.
53        let twiddles = self.get_or_compute_twiddles(log_h);
54
55        // DIT butterfly
56        reverse_matrix_index_bits(&mut mat);
57        for layer in 0..log_h {
58            dit_layer(&mut mat.as_view_mut(), layer, &twiddles);
59        }
60        mat
61    }
62
63    fn coset_lde_batch_with_transform<T>(
64        &self,
65        mat: RowMajorMatrix<F>,
66        added_bits: usize,
67        shift: F,
68        transform: T,
69    ) -> Self::Evaluations
70    where
71        T: FnOnce(&mut RowMajorMatrixViewMut<'_, F>, Layout),
72    {
73        let w = mat.width();
74        let h = mat.height();
75        let padded_h = h << added_bits;
76
77        // Recover polynomial coefficients.
78        let mut coeffs = self.idft_batch(mat);
79
80        // Let the caller inspect or modify the coefficient buffer.
81        transform(&mut coeffs.as_view_mut(), Layout::Natural);
82
83        // Apply the coset shift to only the live coefficient rows; the zero-padded
84        // tail does not need shifting because 0 * s^j = 0 for any j.
85        coset_shift_cols(&mut coeffs, shift);
86
87        // Grow the buffer to full padded capacity without touching existing elements.
88        // The tail must hold zeros; we zero-initialize only the new slots.
89        let live_len = w * h;
90        let total_len = w * padded_h;
91        coeffs.values.reserve_exact(total_len - live_len);
92        coeffs.values.resize(total_len, F::ZERO);
93        coeffs.width = w;
94
95        // Forward DFT over the full padded height.
96        self.dft_batch(coeffs)
97    }
98}
99
100/// Applies one layer of the Radix-2 DIT FFT butterfly network.
101///
102/// Splits the matrix into blocks of rows and performs in-place butterfly operations
103/// on each block. Uses a `TwiddleFreeButterfly` for the first pair and `DitButterfly`
104/// with precomputed twiddles for the rest.
105///
106/// # Arguments
107/// - `mat`: Mutable matrix view with height as a power of two.
108/// - `layer`: Index of the current FFT layer (starting at 0).
109/// - `twiddles`: Precomputed twiddle factors for this layer.
110fn dit_layer<F: Field>(mat: &mut RowMajorMatrixViewMut<'_, F>, layer: usize, twiddles: &[F]) {
111    // Get the number of rows in the matrix (must be a power of two)
112    let h = mat.height();
113    // Compute reversed layer index to access twiddle indices correctly
114    let log_h = log2_strict_usize(h);
115    let layer_rev = log_h - 1 - layer;
116
117    // Each butterfly operates on 2 rows; this is the number of rows in half a block
118    let half_block_size = 1 << layer;
119    // Each block contains 2^layer * 2 rows; full size of the butterfly block
120    let block_size = half_block_size * 2;
121
122    // Process the matrix in blocks of rows of size `block_size`
123    mat.par_row_chunks_exact_mut(block_size)
124        .for_each(|mut block_chunks| {
125            // Split each block vertically into top (hi) and bottom (lo) halves
126            let (mut hi_chunks, mut lo_chunks) = block_chunks.split_rows_mut(half_block_size);
127            // For each pair of rows (hi, lo), apply a butterfly
128            hi_chunks
129                .par_rows_mut()
130                .zip(lo_chunks.par_rows_mut())
131                .enumerate()
132                .for_each(|(ind, (hi_chunk, lo_chunk))| {
133                    if ind == 0 {
134                        // The first pair doesn't require a twiddle factor
135                        TwiddleFreeButterfly.apply_to_rows(hi_chunk, lo_chunk);
136                    } else {
137                        // Apply DIT butterfly using the twiddle factor at index `ind << layer_rev`
138                        DitButterfly(twiddles[ind << layer_rev]).apply_to_rows(hi_chunk, lo_chunk);
139                    }
140                });
141        });
142}
p3_dft/radix_2_dit.rs

p3_dft/
radix_2_dit.rs
