Skip to main content

p3_dft/
butterflies.rs

1use core::mem::MaybeUninit;
2
3use itertools::izip;
4use p3_field::{Field, PackedField, PackedValue};
5
6/// A butterfly operation used in NTT to combine two values into a new pair.
7///
8/// This trait defines how to transform two elements (or vectors of elements)
9/// according to the structure of a butterfly gate.
10///
11/// In an NTT, butterflies are the core units that recursively combine values
12/// across layers. Each butterfly computes:
13/// ```text
14///   (a + b * twiddle, a - b * twiddle)   // DIT
15/// or
16///   (a + b, (a - b) * twiddle)           // DIF
17/// ```
18/// The transformation can be applied:
19/// - in-place (mutating input values)
20/// - to full rows of values (arrays of field elements)
21/// - out-of-place (writing results to separate destination buffers)
22///
23/// Different butterfly variants (DIT, DIF, or twiddle-free) define the exact formula.
24pub trait Butterfly<F: Field>: Copy + Send + Sync {
25    /// Applies the butterfly transformation to two packed field values.
26    ///
27    /// This method takes two inputs `x_1` and `x_2` and returns two outputs `(y_1, y_2)`
28    /// depending on the butterfly type.
29    /// ```text
30    /// Example (DIF):
31    ///   Input:  x_1 = a, x_2 = b
32    ///   Output: (a + b, (a - b) * twiddle)
33    /// ```
34    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF);
35
36    /// Applies the butterfly in-place to two packed values.
37    ///
38    /// Mutates both `x_1` and `x_2` directly, storing the result of `apply`.
39    #[inline]
40    fn apply_in_place<PF: PackedField<Scalar = F>>(&self, x_1: &mut PF, x_2: &mut PF) {
41        (*x_1, *x_2) = self.apply(*x_1, *x_2);
42    }
43
44    /// Applies the butterfly transformation to two rows of scalar field values.
45    ///
46    /// Each row is a slice of `F`. This function processes the rows in packed
47    /// chunks using SIMD where possible, and falls back to scalar operations
48    /// for the suffix (remaining elements).
49    ///
50    /// The transformation is done in-place.
51    #[inline]
52    fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
53        let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix_mut(row_1);
54        let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
55        debug_assert_eq!(shorts_1.len(), shorts_2.len());
56        debug_assert_eq!(suffix_1.len(), suffix_2.len());
57        for (x_1, x_2) in shorts_1.iter_mut().zip(shorts_2) {
58            self.apply_in_place(x_1, x_2);
59        }
60        for (x_1, x_2) in suffix_1.iter_mut().zip(suffix_2) {
61            self.apply_in_place(x_1, x_2);
62        }
63    }
64
65    /// Applies the butterfly out-of-place to two source rows.
66    ///
67    /// This version does not overwrite the source. Instead, it writes the
68    /// result of each butterfly to separate destination slices (which may
69    /// be uninitialized memory).
70    ///
71    /// This is useful when performing LDE's where the size of the output is larger than the size of the input.
72    ///
73    /// - `src_1`, `src_2`: input slices
74    /// - `dst_1`, `dst_2`: output slices to write to (must be MaybeUninit)
75    #[inline]
76    fn apply_to_rows_oop(
77        &self,
78        src_1: &[F],
79        dst_1: &mut [MaybeUninit<F>],
80        src_2: &[F],
81        dst_2: &mut [MaybeUninit<F>],
82    ) {
83        let (src_shorts_1, src_suffix_1) = F::Packing::pack_slice_with_suffix(src_1);
84        let (src_shorts_2, src_suffix_2) = F::Packing::pack_slice_with_suffix(src_2);
85        let (dst_shorts_1, dst_suffix_1) =
86            F::Packing::pack_maybe_uninit_slice_with_suffix_mut(dst_1);
87        let (dst_shorts_2, dst_suffix_2) =
88            F::Packing::pack_maybe_uninit_slice_with_suffix_mut(dst_2);
89        debug_assert_eq!(src_shorts_1.len(), src_shorts_2.len());
90        debug_assert_eq!(src_suffix_1.len(), src_suffix_2.len());
91        debug_assert_eq!(dst_shorts_1.len(), dst_shorts_2.len());
92        debug_assert_eq!(dst_suffix_1.len(), dst_suffix_2.len());
93        for (s_1, s_2, d_1, d_2) in izip!(src_shorts_1, src_shorts_2, dst_shorts_1, dst_shorts_2) {
94            let (res_1, res_2) = self.apply(*s_1, *s_2);
95            d_1.write(res_1);
96            d_2.write(res_2);
97        }
98        for (s_1, s_2, d_1, d_2) in izip!(src_suffix_1, src_suffix_2, dst_suffix_1, dst_suffix_2) {
99            let (res_1, res_2) = self.apply(*s_1, *s_2);
100            d_1.write(res_1);
101            d_2.write(res_2);
102        }
103    }
104}
105
106/// DIF (Decimation-In-Frequency) butterfly operation.
107///
108/// Used in the *output-ordering* variant of NTT.
109/// This butterfly computes:
110/// ```text
111///   output_1 = x1 + x2
112///   output_2 = (x1 - x2) * twiddle
113/// ```
114/// The twiddle factor is applied after subtraction.
115/// Suitable for DIF-style recursive transforms.
116#[derive(Copy, Clone)]
117#[repr(transparent)] // Allows safe transmutes from F to this.
118pub struct DifButterfly<F>(pub F);
119
120impl<F: Field> Butterfly<F> for DifButterfly<F> {
121    #[inline]
122    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
123        (x_1 + x_2, (x_1 - x_2) * self.0)
124    }
125
126    /// Override `apply_to_rows` to pre-broadcast the twiddle factor into a packed field
127    /// once before the inner loop, and manually unroll it to expose multiple independent
128    /// sub-then-mul chains to the compiler's scheduler, hiding the multiplication latency.
129    /// Mirrors the [`DitButterfly`] override.
130    #[inline]
131    fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
132        let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix_mut(row_1);
133        let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
134        debug_assert_eq!(shorts_1.len(), shorts_2.len());
135        debug_assert_eq!(suffix_1.len(), suffix_2.len());
136        let twiddle_packed = F::Packing::from(self.0);
137        let (c1, rem1) = shorts_1.as_chunks_mut::<4>();
138        let (c2, rem2) = shorts_2.as_chunks_mut::<4>();
139        for (p1, p2) in c1.iter_mut().zip(c2.iter_mut()) {
140            let a1 = p1[0];
141            let b1 = p1[1];
142            let c1_ = p1[2];
143            let d1 = p1[3];
144            let a2 = p2[0];
145            let b2 = p2[1];
146            let c2_ = p2[2];
147            let d2 = p2[3];
148            p1[0] = a1 + a2;
149            p1[1] = b1 + b2;
150            p1[2] = c1_ + c2_;
151            p1[3] = d1 + d2;
152            p2[0] = (a1 - a2) * twiddle_packed;
153            p2[1] = (b1 - b2) * twiddle_packed;
154            p2[2] = (c1_ - c2_) * twiddle_packed;
155            p2[3] = (d1 - d2) * twiddle_packed;
156        }
157        for (x_1, x_2) in rem1.iter_mut().zip(rem2.iter_mut()) {
158            let sum = *x_1 + *x_2;
159            *x_2 = (*x_1 - *x_2) * twiddle_packed;
160            *x_1 = sum;
161        }
162        for (x_1, x_2) in suffix_1.iter_mut().zip(suffix_2.iter_mut()) {
163            self.apply_in_place(x_1, x_2);
164        }
165    }
166}
167
168/// DIF (Decimation-In-Frequency) butterfly operation where `x_2` is guaranteed to be zero.
169///
170/// Useful in scenarios where the input has just been padded with zeros.
171///
172/// Used in the *output-ordering* variant of NTT.
173/// This butterfly computes:
174/// ```text
175///   output_1 = x1
176///   output_2 = x1 * twiddle
177/// ```
178#[derive(Copy, Clone)]
179#[repr(transparent)] // Allows safe transmutes from F to this.
180pub struct DifButterflyZeros<F>(pub F);
181
182impl<F: Field> Butterfly<F> for DifButterflyZeros<F> {
183    #[inline]
184    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
185        debug_assert!(x_2.as_slice().iter().all(|x| x.is_zero())); // Slightly convoluted but PF may not implement equality.
186        (x_1, x_1 * self.0)
187    }
188
189    #[inline]
190    fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
191        let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix(row_1);
192        let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
193        debug_assert_eq!(shorts_1.len(), shorts_2.len());
194        debug_assert_eq!(suffix_1.len(), suffix_2.len());
195        for (x_1, x_2) in shorts_1.iter().zip(shorts_2) {
196            debug_assert!(x_2.as_slice().iter().all(|x| x.is_zero())); // Slightly convoluted but PF may not implement equality.
197            *x_2 = *x_1 * self.0; // x_2 is guaranteed to be zero, so we just set it to x_1 * twiddle. 
198        }
199        for (x_1, x_2) in suffix_1.iter().zip(suffix_2) {
200            debug_assert!(x_2.is_zero());
201            *x_2 = *x_1 * self.0; // x_2 is guaranteed to be zero, so we just set it to x_1 * twiddle. 
202        }
203    }
204}
205
206/// DIT (Decimation-In-Time) butterfly operation.
207///
208/// Used in the *input-ordering* variant of NTT/FFT.
209/// This butterfly computes:
210/// ```text
211///   output_1 = x1 + x2 * twiddle
212///   output_2 = x1 - x2 * twiddle
213/// ```
214/// The twiddle factor is applied to x2 before combining.
215/// Suitable for DIT-style recursive transforms.
216#[derive(Copy, Clone)]
217#[repr(transparent)] // Allows safe transmutes from F to this.
218pub struct DitButterfly<F>(pub F);
219
220impl<F: Field> Butterfly<F> for DitButterfly<F> {
221    #[inline]
222    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
223        let x_2_twiddle = x_2 * self.0;
224        (x_1 + x_2_twiddle, x_1 - x_2_twiddle)
225    }
226
227    /// Override `apply_to_rows` to pre-broadcast the twiddle factor into a packed field
228    /// once before the inner loop, avoiding a scalar-to-vector broadcast on each packed
229    /// multiplication. For wide rows (e.g., 256 columns with AVX512 width=16, giving 16
230    /// packed iterations per row-pair), this eliminates 15 redundant broadcasts per call.
231    /// Manually unroll the inner packed loop to expose multiple independent mul chains
232    /// to the compiler's scheduler, hiding the ~12–15 cyc Montgomery mul latency.
233    #[inline]
234    fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
235        let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix_mut(row_1);
236        let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
237        debug_assert_eq!(shorts_1.len(), shorts_2.len());
238        debug_assert_eq!(suffix_1.len(), suffix_2.len());
239        let twiddle_packed = F::Packing::from(self.0);
240        let (c1, rem1) = shorts_1.as_chunks_mut::<4>();
241        let (c2, rem2) = shorts_2.as_chunks_mut::<4>();
242        for (p1, p2) in c1.iter_mut().zip(c2.iter_mut()) {
243            let a1 = p1[0];
244            let b1 = p1[1];
245            let c1_ = p1[2];
246            let d1 = p1[3];
247            let a2 = p2[0];
248            let b2 = p2[1];
249            let c2_ = p2[2];
250            let d2 = p2[3];
251            let a2t = a2 * twiddle_packed;
252            let b2t = b2 * twiddle_packed;
253            let c2t = c2_ * twiddle_packed;
254            let d2t = d2 * twiddle_packed;
255            p1[0] = a1 + a2t;
256            p2[0] = a1 - a2t;
257            p1[1] = b1 + b2t;
258            p2[1] = b1 - b2t;
259            p1[2] = c1_ + c2t;
260            p2[2] = c1_ - c2t;
261            p1[3] = d1 + d2t;
262            p2[3] = d1 - d2t;
263        }
264        for (x_1, x_2) in rem1.iter_mut().zip(rem2.iter_mut()) {
265            let x_2_twiddle = *x_2 * twiddle_packed;
266            let new_x1 = *x_1 + x_2_twiddle;
267            *x_2 = *x_1 - x_2_twiddle;
268            *x_1 = new_x1;
269        }
270        for (x_1, x_2) in suffix_1.iter_mut().zip(suffix_2.iter_mut()) {
271            self.apply_in_place(x_1, x_2);
272        }
273    }
274
275    /// Out-of-place variant with matching unroll factor.
276    #[inline]
277    fn apply_to_rows_oop(
278        &self,
279        src_1: &[F],
280        dst_1: &mut [MaybeUninit<F>],
281        src_2: &[F],
282        dst_2: &mut [MaybeUninit<F>],
283    ) {
284        let (src_shorts_1, src_suffix_1) = F::Packing::pack_slice_with_suffix(src_1);
285        let (src_shorts_2, src_suffix_2) = F::Packing::pack_slice_with_suffix(src_2);
286        let (dst_shorts_1, dst_suffix_1) =
287            F::Packing::pack_maybe_uninit_slice_with_suffix_mut(dst_1);
288        let (dst_shorts_2, dst_suffix_2) =
289            F::Packing::pack_maybe_uninit_slice_with_suffix_mut(dst_2);
290        debug_assert_eq!(src_shorts_1.len(), src_shorts_2.len());
291        debug_assert_eq!(src_suffix_1.len(), src_suffix_2.len());
292        debug_assert_eq!(dst_shorts_1.len(), dst_shorts_2.len());
293        debug_assert_eq!(dst_suffix_1.len(), dst_suffix_2.len());
294        let twiddle_packed = F::Packing::from(self.0);
295        let n = src_shorts_1.len();
296        let n4 = n - (n & 3);
297        let mut i = 0;
298        while i < n4 {
299            let a1 = src_shorts_1[i];
300            let b1 = src_shorts_1[i + 1];
301            let c1 = src_shorts_1[i + 2];
302            let d1 = src_shorts_1[i + 3];
303            let a2 = src_shorts_2[i];
304            let b2 = src_shorts_2[i + 1];
305            let c2 = src_shorts_2[i + 2];
306            let d2 = src_shorts_2[i + 3];
307            let a2t = a2 * twiddle_packed;
308            let b2t = b2 * twiddle_packed;
309            let c2t = c2 * twiddle_packed;
310            let d2t = d2 * twiddle_packed;
311            dst_shorts_1[i].write(a1 + a2t);
312            dst_shorts_2[i].write(a1 - a2t);
313            dst_shorts_1[i + 1].write(b1 + b2t);
314            dst_shorts_2[i + 1].write(b1 - b2t);
315            dst_shorts_1[i + 2].write(c1 + c2t);
316            dst_shorts_2[i + 2].write(c1 - c2t);
317            dst_shorts_1[i + 3].write(d1 + d2t);
318            dst_shorts_2[i + 3].write(d1 - d2t);
319            i += 4;
320        }
321        while i < n {
322            let s1 = src_shorts_1[i];
323            let s2 = src_shorts_2[i];
324            let x_2_twiddle = s2 * twiddle_packed;
325            dst_shorts_1[i].write(s1 + x_2_twiddle);
326            dst_shorts_2[i].write(s1 - x_2_twiddle);
327            i += 1;
328        }
329        for (s_1, s_2, d_1, d_2) in izip!(src_suffix_1, src_suffix_2, dst_suffix_1, dst_suffix_2) {
330            let (res_1, res_2) = self.apply(*s_1, *s_2);
331            d_1.write(res_1);
332            d_2.write(res_2);
333        }
334    }
335}
336
337/// DIT (Decimation-In-Time) butterfly operation with a post-multiplication scale factor.
338///
339/// This butterfly computes:
340/// ```text
341///   output_1 = (x1 + x2 * twiddle) * scale
342///   output_2 = (x1 - x2 * twiddle) * scale
343/// ```
344/// which is equivalent to:
345/// ```text
346///   output_1 = x1 * scale + x2 * (twiddle * scale)
347///   output_2 = x1 * scale - x2 * (twiddle * scale)
348/// ```
349///
350/// This is used to merge a uniform scaling step (e.g., 1/N normalization in inverse DFT)
351/// into a butterfly pass, avoiding a separate memory pass over the data.
352///
353/// The struct stores `scale` and `twiddle_times_scale = twiddle * scale` so that the
354/// `apply` method only needs 2 multiplications instead of 3.
355#[derive(Copy, Clone)]
356pub struct ScaledDitButterfly<F> {
357    pub twiddle: F,
358    pub scale: F,
359    /// Precomputed product `twiddle * scale` to reduce multiplications in the hot loop.
360    pub twiddle_times_scale: F,
361}
362
363impl<F: Field> ScaledDitButterfly<F> {
364    /// Construct a `ScaledDitButterfly`, precomputing `twiddle * scale`.
365    #[inline]
366    pub fn new(twiddle: F, scale: F) -> Self {
367        Self {
368            twiddle,
369            scale,
370            twiddle_times_scale: twiddle * scale,
371        }
372    }
373}
374
375impl<F: Field> Butterfly<F> for ScaledDitButterfly<F> {
376    #[inline]
377    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
378        // 2 multiplications instead of 3:
379        //   x1_s   = x1 * scale
380        //   x2_ts  = x2 * (twiddle * scale)   [precomputed]
381        //   out1   = x1_s + x2_ts
382        //   out2   = x1_s - x2_ts
383        let x_1_scale = x_1 * self.scale;
384        let x_2_twiddle_scale = x_2 * self.twiddle_times_scale;
385        (x_1_scale + x_2_twiddle_scale, x_1_scale - x_2_twiddle_scale)
386    }
387
388    /// Override `apply_to_rows` to pre-broadcast both `scale` and `twiddle_times_scale`
389    /// into packed fields once before the inner loop.
390    #[inline]
391    fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
392        let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix_mut(row_1);
393        let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
394        debug_assert_eq!(shorts_1.len(), shorts_2.len());
395        debug_assert_eq!(suffix_1.len(), suffix_2.len());
396        let scale_packed = F::Packing::from(self.scale);
397        let twiddle_times_scale_packed = F::Packing::from(self.twiddle_times_scale);
398        // ScaledDitButterfly has 2 muls per butterfly (scale + twiddle_scale), so unroll-4
399        // exposes 8 independent mul chains — better ILP than unroll-2's 4 chains.
400        let (c1, rem1) = shorts_1.as_chunks_mut::<4>();
401        let (c2, rem2) = shorts_2.as_chunks_mut::<4>();
402        for (p1, p2) in c1.iter_mut().zip(c2.iter_mut()) {
403            let a1 = p1[0];
404            let b1 = p1[1];
405            let c1_ = p1[2];
406            let d1 = p1[3];
407            let a2 = p2[0];
408            let b2 = p2[1];
409            let c2_ = p2[2];
410            let d2 = p2[3];
411            let a1s = a1 * scale_packed;
412            let b1s = b1 * scale_packed;
413            let c1s = c1_ * scale_packed;
414            let d1s = d1 * scale_packed;
415            let a2t = a2 * twiddle_times_scale_packed;
416            let b2t = b2 * twiddle_times_scale_packed;
417            let c2t = c2_ * twiddle_times_scale_packed;
418            let d2t = d2 * twiddle_times_scale_packed;
419            p1[0] = a1s + a2t;
420            p2[0] = a1s - a2t;
421            p1[1] = b1s + b2t;
422            p2[1] = b1s - b2t;
423            p1[2] = c1s + c2t;
424            p2[2] = c1s - c2t;
425            p1[3] = d1s + d2t;
426            p2[3] = d1s - d2t;
427        }
428        for (x_1, x_2) in rem1.iter_mut().zip(rem2.iter_mut()) {
429            let x_1_scale = *x_1 * scale_packed;
430            let x_2_twiddle_scale = *x_2 * twiddle_times_scale_packed;
431            *x_1 = x_1_scale + x_2_twiddle_scale;
432            *x_2 = x_1_scale - x_2_twiddle_scale;
433        }
434        for (x_1, x_2) in suffix_1.iter_mut().zip(suffix_2.iter_mut()) {
435            self.apply_in_place(x_1, x_2);
436        }
437    }
438}
439
440/// Butterfly with no twiddle factor (`twiddle = 1`).
441///
442/// This is used when no root-of-unity scaling is needed.
443/// It works for either DIT or DIF, and is often used at
444/// the final or base level of a transform tree.
445///
446/// This butterfly computes:
447/// ```text
448///   - output_1 = x1 + x2
449///   - output_2 = x1 - x2
450/// ```
451#[derive(Copy, Clone)]
452pub struct TwiddleFreeButterfly;
453
454impl<F: Field> Butterfly<F> for TwiddleFreeButterfly {
455    #[inline]
456    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
457        (x_1 + x_2, x_1 - x_2)
458    }
459}