p3_dft/
butterflies.rs
1use core::mem::MaybeUninit;
2
3use itertools::izip;
4use p3_field::{Field, PackedField, PackedValue};
5
6pub trait Butterfly<F: Field>: Copy + Send + Sync {
7 fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF);
8
9 #[inline]
10 fn apply_in_place<PF: PackedField<Scalar = F>>(&self, x_1: &mut PF, x_2: &mut PF) {
11 (*x_1, *x_2) = self.apply(*x_1, *x_2);
12 }
13
14 #[inline]
15 fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
16 let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix_mut(row_1);
17 let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
18 debug_assert_eq!(shorts_1.len(), shorts_2.len());
19 debug_assert_eq!(suffix_1.len(), suffix_2.len());
20 for (x_1, x_2) in shorts_1.iter_mut().zip(shorts_2) {
21 self.apply_in_place(x_1, x_2);
22 }
23 for (x_1, x_2) in suffix_1.iter_mut().zip(suffix_2) {
24 self.apply_in_place(x_1, x_2);
25 }
26 }
27
28 #[inline]
30 fn apply_to_rows_oop(
31 &self,
32 src_1: &[F],
33 dst_1: &mut [MaybeUninit<F>],
34 src_2: &[F],
35 dst_2: &mut [MaybeUninit<F>],
36 ) {
37 let (src_shorts_1, src_suffix_1) = F::Packing::pack_slice_with_suffix(src_1);
38 let (src_shorts_2, src_suffix_2) = F::Packing::pack_slice_with_suffix(src_2);
39 let (dst_shorts_1, dst_suffix_1) =
40 F::Packing::pack_maybe_uninit_slice_with_suffix_mut(dst_1);
41 let (dst_shorts_2, dst_suffix_2) =
42 F::Packing::pack_maybe_uninit_slice_with_suffix_mut(dst_2);
43 debug_assert_eq!(src_shorts_1.len(), src_shorts_2.len());
44 debug_assert_eq!(src_suffix_1.len(), src_suffix_2.len());
45 debug_assert_eq!(dst_shorts_1.len(), dst_shorts_2.len());
46 debug_assert_eq!(dst_suffix_1.len(), dst_suffix_2.len());
47 for (s_1, s_2, d_1, d_2) in izip!(src_shorts_1, src_shorts_2, dst_shorts_1, dst_shorts_2) {
48 let (res_1, res_2) = self.apply::<F::Packing>(*s_1, *s_2);
49 d_1.write(res_1);
50 d_2.write(res_2);
51 }
52 for (s_1, s_2, d_1, d_2) in izip!(src_suffix_1, src_suffix_2, dst_suffix_1, dst_suffix_2) {
53 let (res_1, res_2) = self.apply::<F>(*s_1, *s_2);
54 d_1.write(res_1);
55 d_2.write(res_2);
56 }
57 }
58}
59
60#[derive(Copy, Clone)]
61pub struct DifButterfly<F>(pub F);
62impl<F: Field> Butterfly<F> for DifButterfly<F> {
63 #[inline]
64 fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
65 (x_1 + x_2, (x_1 - x_2) * self.0)
66 }
67}
68
69#[derive(Copy, Clone)]
70pub struct DitButterfly<F>(pub F);
71impl<F: Field> Butterfly<F> for DitButterfly<F> {
72 #[inline]
73 fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
74 let x_2_twiddle = x_2 * self.0;
75 (x_1 + x_2_twiddle, x_1 - x_2_twiddle)
76 }
77}
78
79#[derive(Copy, Clone)]
81pub struct TwiddleFreeButterfly;
82impl<F: Field> Butterfly<F> for TwiddleFreeButterfly {
83 #[inline]
84 fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
85 (x_1 + x_2, x_1 - x_2)
86 }
87}