1use core::ops::Mul;
4
5use p3_binary_field::poly_basis::HAS_HARDWARE_CLMUL;
6use p3_binary_field::{
7 BinaryField2, BinaryField4, BinaryField8, BinaryField16, BinaryField32, BinaryField64,
8 BinaryField128, Gf2, Ghash128, Poly64, TowerLevel,
9};
10use p3_field::{PackedValue, PrimeCharacteristicRing};
11
12pub trait ButterflyField: TowerLevel {
19 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self);
28}
29
30unsafe trait ByteCoordinates: TowerLevel {}
40
41unsafe impl ByteCoordinates for BinaryField8 {}
47unsafe impl ByteCoordinates for BinaryField16 {}
48unsafe impl ByteCoordinates for BinaryField32 {}
49unsafe impl ByteCoordinates for BinaryField64 {}
50unsafe impl ByteCoordinates for BinaryField128 {}
51
52#[derive(Copy, Clone, Debug, PartialEq, Eq)]
58enum TwiddleWidth {
59 Byte(u8),
61 Word(u16),
63 DoubleWord(u32),
65 Wide,
67}
68
69const fn classify(bits: u128) -> TwiddleWidth {
71 if bits <= u8::MAX as u128 {
72 TwiddleWidth::Byte(bits as u8)
73 } else if bits <= u16::MAX as u128 {
74 TwiddleWidth::Word(bits as u16)
75 } else if bits <= u32::MAX as u128 {
76 TwiddleWidth::DoubleWord(bits as u32)
77 } else {
78 TwiddleWidth::Wide
79 }
80}
81
82#[inline]
86pub(crate) fn packed_butterfly<F: TowerLevel, const INVERSE: bool>(
87 lo: &mut [F],
88 hi: &mut [F],
89 t: F,
90) {
91 let (lo, lo_tail) = F::Packing::pack_slice_with_suffix_mut(lo);
93 let (hi, hi_tail) = F::Packing::pack_slice_with_suffix_mut(hi);
94 let zero = t.is_zero();
95 butterfly_values::<_, INVERSE>(lo, hi, t.into(), zero);
96 butterfly_values::<_, INVERSE>(lo_tail, hi_tail, t, zero);
97}
98
99#[inline]
101fn butterfly_values<R: PrimeCharacteristicRing + Copy, const INVERSE: bool>(
102 lo: &mut [R],
103 hi: &mut [R],
104 t: R,
105 zero: bool,
106) {
107 if zero {
108 for (u, v) in lo.iter_mut().zip(hi) {
110 *v += *u;
111 }
112 } else if INVERSE {
113 for (u, v) in lo.iter_mut().zip(hi) {
115 *v += *u;
116 *u += t * *v;
117 }
118 } else {
119 for (u, v) in lo.iter_mut().zip(hi) {
121 *u += t * *v;
122 *v += *u;
123 }
124 }
125}
126
127#[inline]
133fn typed_butterfly<F, S, const INVERSE: bool>(lo: &mut [F], hi: &mut [F], t: S)
134where
135 F: Copy + PrimeCharacteristicRing + Mul<S, Output = F>,
136 S: Copy,
137{
138 if INVERSE {
139 for (u, v) in lo.iter_mut().zip(hi) {
141 *v += *u;
142 *u += *v * t;
143 }
144 } else {
145 for (u, v) in lo.iter_mut().zip(hi) {
147 *u += *v * t;
148 *v += *u;
149 }
150 }
151}
152
153#[inline]
157fn scale_by_byte<F, const INVERSE: bool>(lo: &mut [F], hi: &mut [F], t: F, width: TwiddleWidth)
158where
159 F: TowerLevel + Mul<BinaryField8, Output = F>,
160{
161 if let TwiddleWidth::Byte(s) = width {
162 typed_butterfly::<F, BinaryField8, INVERSE>(lo, hi, BinaryField8::from_repr(s));
163 } else {
164 packed_butterfly::<F, INVERSE>(lo, hi, t);
165 }
166}
167
168#[inline]
170fn scale_by_word<F, const INVERSE: bool>(lo: &mut [F], hi: &mut [F], t: F, width: TwiddleWidth)
171where
172 F: TowerLevel + Mul<BinaryField8, Output = F> + Mul<BinaryField16, Output = F>,
173{
174 if let TwiddleWidth::Word(s) = width {
175 typed_butterfly::<F, BinaryField16, INVERSE>(lo, hi, BinaryField16::from_repr(s));
176 } else {
177 scale_by_byte::<F, INVERSE>(lo, hi, t, width);
178 }
179}
180
181#[inline]
183fn scale_by_dword<F, const INVERSE: bool>(lo: &mut [F], hi: &mut [F], t: F, width: TwiddleWidth)
184where
185 F: TowerLevel
186 + Mul<BinaryField8, Output = F>
187 + Mul<BinaryField16, Output = F>
188 + Mul<BinaryField32, Output = F>,
189{
190 if let TwiddleWidth::DoubleWord(s) = width {
191 typed_butterfly::<F, BinaryField32, INVERSE>(lo, hi, BinaryField32::from_repr(s));
192 } else {
193 scale_by_word::<F, INVERSE>(lo, hi, t, width);
194 }
195}
196
197const WIDE_TYPED_PRODUCTS_PAY: bool = !HAS_HARDWARE_CLMUL;
212
213trait SubfieldScaled: ByteCoordinates {
220 fn scale_butterfly<const INVERSE: bool>(
222 lo: &mut [Self],
223 hi: &mut [Self],
224 t: Self,
225 width: TwiddleWidth,
226 );
227}
228
229impl SubfieldScaled for BinaryField8 {
230 #[inline]
232 fn scale_butterfly<const INVERSE: bool>(
233 lo: &mut [Self],
234 hi: &mut [Self],
235 t: Self,
236 width: TwiddleWidth,
237 ) {
238 scale_by_byte::<Self, INVERSE>(lo, hi, t, width);
239 }
240}
241
242impl SubfieldScaled for BinaryField16 {
243 #[inline]
245 fn scale_butterfly<const INVERSE: bool>(
246 lo: &mut [Self],
247 hi: &mut [Self],
248 t: Self,
249 width: TwiddleWidth,
250 ) {
251 scale_by_byte::<Self, INVERSE>(lo, hi, t, width);
252 }
253}
254
255impl SubfieldScaled for BinaryField32 {
256 #[inline]
258 fn scale_butterfly<const INVERSE: bool>(
259 lo: &mut [Self],
260 hi: &mut [Self],
261 t: Self,
262 width: TwiddleWidth,
263 ) {
264 scale_by_word::<Self, INVERSE>(lo, hi, t, width);
265 }
266}
267
268impl SubfieldScaled for BinaryField64 {
269 #[inline]
271 fn scale_butterfly<const INVERSE: bool>(
272 lo: &mut [Self],
273 hi: &mut [Self],
274 t: Self,
275 width: TwiddleWidth,
276 ) {
277 if WIDE_TYPED_PRODUCTS_PAY {
278 scale_by_dword::<Self, INVERSE>(lo, hi, t, width);
279 } else {
280 scale_by_byte::<Self, INVERSE>(lo, hi, t, width);
281 }
282 }
283}
284
285impl SubfieldScaled for BinaryField128 {
286 #[inline]
288 fn scale_butterfly<const INVERSE: bool>(
289 lo: &mut [Self],
290 hi: &mut [Self],
291 t: Self,
292 width: TwiddleWidth,
293 ) {
294 if WIDE_TYPED_PRODUCTS_PAY {
295 scale_by_dword::<Self, INVERSE>(lo, hi, t, width);
296 } else {
297 scale_by_byte::<Self, INVERSE>(lo, hi, t, width);
298 }
299 }
300}
301
302#[cfg(all(
306 target_arch = "x86_64",
307 target_feature = "gfni",
308 target_feature = "avx512f",
309 target_feature = "avx512bw"
310))]
311#[inline]
312fn subfield_prefix<F: ByteCoordinates, const INVERSE: bool>(
313 lo: &mut [F],
314 hi: &mut [F],
315 width: TwiddleWidth,
316) -> usize {
317 use core::arch::x86_64::__m512i;
318
319 use crate::affine;
320 use crate::lanes::ByteRegister;
321
322 const {
323 assert!(cfg!(target_endian = "little"));
325
326 assert!(size_of::<F>() == 1 << (F::LOG_BITS - 3));
328
329 assert!(64 % size_of::<F>() == 0);
331 }
332
333 let bytes = size_of_val(lo);
334
335 if bytes < <__m512i as ByteRegister>::BYTES {
338 return 0;
339 }
340
341 let (lo, hi) = unsafe {
347 (
348 core::slice::from_raw_parts_mut(lo.as_mut_ptr().cast::<u8>(), bytes),
349 core::slice::from_raw_parts_mut(hi.as_mut_ptr().cast::<u8>(), bytes),
350 )
351 };
352
353 let covered = match width {
357 TwiddleWidth::Byte(t) => {
358 affine::subfield_butterfly::<__m512i, 1, INVERSE>(lo, hi, &affine::byte_blocks(t))
359 }
360 TwiddleWidth::Word(t) if size_of::<F>() >= 2 => {
361 affine::subfield_butterfly::<__m512i, 2, INVERSE>(lo, hi, &affine::word_blocks(t))
362 }
363 TwiddleWidth::DoubleWord(t) if size_of::<F>() >= 4 => {
364 affine::subfield_butterfly::<__m512i, 4, INVERSE>(lo, hi, &affine::dword_blocks(t))
365 }
366 _ => 0,
367 };
368 covered / size_of::<F>()
369}
370
371#[cfg(all(
377 target_arch = "aarch64",
378 target_feature = "neon",
379 target_endian = "little"
380))]
381#[inline]
382fn subfield_prefix<F: ByteCoordinates, const INVERSE: bool>(
383 lo: &mut [F],
384 hi: &mut [F],
385 width: TwiddleWidth,
386) -> usize {
387 use core::arch::aarch64::uint8x16_t;
388
389 use crate::lanes::ByteRegister;
390 use crate::neon;
391
392 const {
396 assert!(size_of::<F>() == 1 << (F::LOG_BITS - 3));
398
399 assert!(<uint8x16_t as ByteRegister>::BYTES.is_multiple_of(size_of::<F>()));
401 }
402
403 let TwiddleWidth::Byte(t) = width else {
404 return 0;
405 };
406
407 let bytes = size_of_val(lo);
408 if bytes < neon::MIN_BYTES {
409 return 0;
410 }
411
412 let (lo, hi) = unsafe {
418 (
419 core::slice::from_raw_parts_mut(lo.as_mut_ptr().cast::<u8>(), bytes),
420 core::slice::from_raw_parts_mut(hi.as_mut_ptr().cast::<u8>(), bytes),
421 )
422 };
423
424 let t = BinaryField8::from_repr(t);
426 let columns: [u8; 8] =
427 core::array::from_fn(|b| (t * BinaryField8::from_repr(1 << b)).to_repr());
428
429 neon::byte_butterfly::<INVERSE>(lo, hi, &columns) / size_of::<F>()
430}
431
432#[cfg(not(any(
434 all(
435 target_arch = "x86_64",
436 target_feature = "gfni",
437 target_feature = "avx512f",
438 target_feature = "avx512bw"
439 ),
440 all(
441 target_arch = "aarch64",
442 target_feature = "neon",
443 target_endian = "little"
444 )
445)))]
446#[inline]
447const fn subfield_prefix<F: ByteCoordinates, const INVERSE: bool>(
448 _lo: &mut [F],
449 _hi: &mut [F],
450 _width: TwiddleWidth,
451) -> usize {
452 0
453}
454
455#[inline]
464fn coordinate_butterfly<F, const INVERSE: bool>(lo: &mut [F], hi: &mut [F], t: F)
465where
466 F: SubfieldScaled,
467 F::Repr: Into<u128>,
468{
469 assert_eq!(lo.len(), hi.len(), "butterfly lengths differ");
473
474 if t.is_zero() {
476 return packed_butterfly::<F, INVERSE>(lo, hi, t);
477 }
478
479 let width = classify(t.to_repr().into());
480 let covered = subfield_prefix::<F, INVERSE>(lo, hi, width);
481
482 F::scale_butterfly::<INVERSE>(&mut lo[covered..], &mut hi[covered..], t, width);
487}
488
489#[inline]
494fn plain_butterfly<F: TowerLevel, const INVERSE: bool>(lo: &mut [F], hi: &mut [F], t: F) {
495 assert_eq!(lo.len(), hi.len(), "butterfly lengths differ");
499 packed_butterfly::<F, INVERSE>(lo, hi, t);
500}
501
502impl ButterflyField for BinaryField8 {
503 #[inline]
504 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
505 coordinate_butterfly::<Self, INVERSE>(lo, hi, t);
506 }
507}
508
509impl ButterflyField for BinaryField16 {
510 #[inline]
511 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
512 coordinate_butterfly::<Self, INVERSE>(lo, hi, t);
513 }
514}
515
516impl ButterflyField for BinaryField32 {
517 #[inline]
518 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
519 coordinate_butterfly::<Self, INVERSE>(lo, hi, t);
520 }
521}
522
523impl ButterflyField for BinaryField64 {
524 #[inline]
525 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
526 coordinate_butterfly::<Self, INVERSE>(lo, hi, t);
527 }
528}
529
530impl ButterflyField for BinaryField128 {
531 #[inline]
532 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
533 coordinate_butterfly::<Self, INVERSE>(lo, hi, t);
534 }
535}
536
537impl ButterflyField for Gf2 {
538 #[inline]
539 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
540 plain_butterfly::<Self, INVERSE>(lo, hi, t);
541 }
542}
543
544impl ButterflyField for BinaryField2 {
545 #[inline]
546 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
547 plain_butterfly::<Self, INVERSE>(lo, hi, t);
548 }
549}
550
551impl ButterflyField for BinaryField4 {
552 #[inline]
553 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
554 plain_butterfly::<Self, INVERSE>(lo, hi, t);
555 }
556}
557
558impl ButterflyField for Ghash128 {
559 #[inline]
560 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
561 plain_butterfly::<Self, INVERSE>(lo, hi, t);
562 }
563}
564
565impl ButterflyField for Poly64 {
566 #[inline]
567 fn butterfly<const INVERSE: bool>(lo: &mut [Self], hi: &mut [Self], t: Self) {
568 plain_butterfly::<Self, INVERSE>(lo, hi, t);
569 }
570}
571
572#[cfg(test)]
573mod tests {
574 use alloc::vec::Vec;
575
576 use p3_binary_field::TowerLevel;
577 use p3_field::PrimeCharacteristicRing;
578 use proptest::prelude::*;
579
580 use super::{
581 BinaryField8, BinaryField16, BinaryField32, BinaryField64, BinaryField128, ButterflyField,
582 Ghash128, TwiddleWidth, classify,
583 };
584
585 const CORNERS: [u128; 13] = [
605 0,
606 1,
607 0x80,
608 0xff,
609 0x100,
610 0x8000,
611 0xffff,
612 0x1_0000,
613 0x8000_0000,
614 0xffff_ffff,
615 0x1_0000_0000,
616 1 << 127,
617 0x87,
618 ];
619
620 const LENGTHS: core::ops::RangeInclusive<usize> = 0..=70;
628
629 fn sample<F: TowerLevel>(len: usize, seed: u64) -> Vec<F> {
631 (0..len)
632 .map(|i| {
633 let bits = seed
635 .wrapping_mul(0x9e37_79b9_7f4a_7c15)
636 .wrapping_add(i as u64 + 1)
637 .wrapping_mul(0xbf58_476d_1ce4_e5b9);
638
639 F::from_le_byte_iter(bits.to_le_bytes().into_iter().cycle())
641 })
642 .collect()
643 }
644
645 fn reference<F: TowerLevel>(lo: &mut [F], hi: &mut [F], t: F, inverse: bool) {
647 for (u, v) in lo.iter_mut().zip(hi) {
648 if inverse {
649 *v += *u;
651 *u += t * *v;
652 } else {
653 *u += t * *v;
655 *v += *u;
656 }
657 }
658 }
659
660 fn agrees<F: ButterflyField>(len: usize, t: F) -> Result<(), TestCaseError> {
662 let (lo, hi) = (sample::<F>(len, 1), sample::<F>(len, 2));
664
665 let (mut want_lo, mut want_hi) = (lo.clone(), hi.clone());
667 reference(&mut want_lo, &mut want_hi, t, false);
668
669 let (mut got_lo, mut got_hi) = (lo.clone(), hi.clone());
670 F::butterfly::<false>(&mut got_lo, &mut got_hi, t);
671 prop_assert_eq!(&got_lo, &want_lo);
672 prop_assert_eq!(&got_hi, &want_hi);
673
674 F::butterfly::<true>(&mut got_lo, &mut got_hi, t);
676 prop_assert_eq!(&got_lo, &lo);
677 prop_assert_eq!(&got_hi, &hi);
678
679 let (mut want_lo, mut want_hi) = (lo.clone(), hi.clone());
681 reference(&mut want_lo, &mut want_hi, t, true);
682
683 let (mut got_lo, mut got_hi) = (lo, hi);
684 F::butterfly::<true>(&mut got_lo, &mut got_hi, t);
685 prop_assert_eq!(&got_lo, &want_lo);
686 prop_assert_eq!(&got_hi, &want_hi);
687 Ok(())
688 }
689
690 fn sweep<F: ButterflyField>(name: &str) {
692 for &bits in &CORNERS {
693 let t =
695 F::from_le_byte_iter(bits.to_le_bytes().into_iter().chain(core::iter::repeat(0)));
696 for len in LENGTHS {
697 agrees::<F>(len, t)
698 .unwrap_or_else(|e| panic!("{name}, twiddle {bits:#x}, len {len}: {e}"));
699 }
700 }
701 }
702
703 #[test]
704 fn every_level_agrees_with_the_element_loop_at_the_corners() {
705 sweep::<BinaryField8>("BinaryField8");
707 sweep::<BinaryField16>("BinaryField16");
708 sweep::<BinaryField32>("BinaryField32");
709 sweep::<BinaryField64>("BinaryField64");
710 sweep::<BinaryField128>("BinaryField128");
711 sweep::<Ghash128>("Ghash128");
712 }
713
714 #[test]
715 fn one_byte_twiddles_agree_around_every_register_boundary() {
716 const LENS: [usize; 14] = [0, 1, 15, 16, 17, 31, 32, 33, 255, 256, 257, 511, 512, 513];
718 for t in [0x01u8, 0x02, 0x03, 0x80, 0xa5, 0xff] {
719 for len in LENS {
720 agrees::<BinaryField8>(len, BinaryField8::from_repr(t)).unwrap();
721 agrees::<BinaryField16>(len, BinaryField16::from_repr(t as u16)).unwrap();
722 agrees::<BinaryField32>(len, BinaryField32::from_repr(t as u32)).unwrap();
723 agrees::<BinaryField64>(len, BinaryField64::from_repr(t as u64)).unwrap();
724 agrees::<BinaryField128>(len, BinaryField128::from_repr(t as u128)).unwrap();
725 }
726 }
727 }
728
729 #[test]
730 fn the_classification_follows_the_twiddle_magnitude() {
731 assert_eq!(classify(0), TwiddleWidth::Byte(0));
740 assert_eq!(classify(0xff), TwiddleWidth::Byte(0xff));
741 assert_eq!(classify(0x100), TwiddleWidth::Word(0x100));
742 assert_eq!(classify(0xffff), TwiddleWidth::Word(0xffff));
743 assert_eq!(classify(0x1_0000), TwiddleWidth::DoubleWord(0x1_0000));
744 assert_eq!(classify(0xffff_ffff), TwiddleWidth::DoubleWord(0xffff_ffff));
745 assert_eq!(classify(0x1_0000_0000), TwiddleWidth::Wide);
746 assert_eq!(classify(u128::MAX), TwiddleWidth::Wide);
747 }
748
749 #[test]
750 #[should_panic = "butterfly lengths differ"]
751 fn a_length_mismatch_is_refused() {
752 let mut lo = [BinaryField32::ONE; 4];
754 let mut hi = [BinaryField32::ONE; 3];
755 BinaryField32::butterfly::<false>(&mut lo, &mut hi, BinaryField32::ONE);
756 }
757
758 proptest! {
759 #![proptest_config(ProptestConfig::with_cases(64))]
760
761 #[test]
763 fn random_twiddles_agree_with_the_element_loop(
764 bits in any::<u128>(),
765 len in 0usize..70,
766 ) {
767 let bytes = bits.to_le_bytes();
768 agrees::<BinaryField32>(len, BinaryField32::from_le_byte_iter(bytes.into_iter()))?;
769 agrees::<BinaryField64>(len, BinaryField64::from_le_byte_iter(bytes.into_iter()))?;
770 agrees::<BinaryField128>(len, BinaryField128::from_le_byte_iter(bytes.into_iter()))?;
771 }
772
773 #[test]
775 fn byte_twiddles_agree_with_the_element_loop(t in any::<u8>(), len in 0usize..70) {
776 agrees::<BinaryField8>(len, BinaryField8::from_repr(t))?;
777 agrees::<BinaryField16>(len, BinaryField16::from_repr(t as u16))?;
778 agrees::<BinaryField32>(len, BinaryField32::from_repr(t as u32))?;
779 agrees::<BinaryField64>(len, BinaryField64::from_repr(t as u64))?;
780 agrees::<BinaryField128>(len, BinaryField128::from_repr(t as u128))?;
781 }
782
783 #[test]
785 fn byte_twiddles_agree_on_runs_past_the_register_threshold(
786 t in any::<u8>(),
787 len in 0usize..600,
788 ) {
789 agrees::<BinaryField8>(len, BinaryField8::from_repr(t))?;
790 agrees::<BinaryField16>(len, BinaryField16::from_repr(t as u16))?;
791 agrees::<BinaryField32>(len, BinaryField32::from_repr(t as u32))?;
792 agrees::<BinaryField64>(len, BinaryField64::from_repr(t as u64))?;
793 agrees::<BinaryField128>(len, BinaryField128::from_repr(t as u128))?;
794 }
795
796 #[test]
798 fn word_twiddles_agree_with_the_element_loop(t in 0x100u16.., len in 0usize..70) {
799 agrees::<BinaryField16>(len, BinaryField16::from_repr(t))?;
800 agrees::<BinaryField32>(len, BinaryField32::from_repr(t as u32))?;
801 agrees::<BinaryField64>(len, BinaryField64::from_repr(t as u64))?;
802 agrees::<BinaryField128>(len, BinaryField128::from_repr(t as u128))?;
803 }
804
805 #[test]
807 fn double_word_twiddles_agree_with_the_element_loop(
808 t in 0x1_0000u32..,
809 len in 0usize..70,
810 ) {
811 agrees::<BinaryField32>(len, BinaryField32::from_repr(t))?;
812 agrees::<BinaryField64>(len, BinaryField64::from_repr(t as u64))?;
813 agrees::<BinaryField128>(len, BinaryField128::from_repr(t as u128))?;
814 }
815 }
816}