diff --git a/src/backend/vector/ifma/field.rs b/src/backend/vector/ifma/field.rs index 12a72f4..be26a21 100644 --- a/src/backend/vector/ifma/field.rs +++ b/src/backend/vector/ifma/field.rs @@ -8,7 +8,7 @@ // - Henry de Valence use core::ops::{Add, Mul, Neg}; -use packed_simd::{i32x8, u32x8, u64x4, IntoBits}; +use packed_simd::{u64x4, FromBits, IntoBits}; use backend::serial::u64::field::FieldElement51; @@ -28,6 +28,38 @@ pub struct F51x4Unreduced(pub(crate) [u64x4; 5]); #[derive(Copy, Clone, Debug)] pub struct F51x4Reduced(pub(crate) [u64x4; 5]); +#[derive(Copy, Clone)] +pub enum Lanes { + AB, +} + +#[derive(Copy, Clone)] +pub enum Shuffle { + AAAA, +} + +#[inline(always)] +fn shuffle_lanes(x: u64x4, control: Shuffle) -> u64x4 { + unsafe { + use core::arch::x86_64::_mm256_permute4x64_epi64 as perm; + + match control { + Shuffle::AAAA => perm(x.into_bits(), 0b00_00_00_00).into_bits(), + } + } +} + +#[inline] +fn blend_lanes(x: u64x4, y: u64x4, control: Lanes) -> u64x4 { + unsafe { + use core::arch::x86_64::_mm256_blend_epi32 as blend; + + match control { + Lanes::AB => blend(x.into_bits(), y.into_bits(), 0b00_00_11_11).into_bits(), + } + } +} + impl F51x4Unreduced { pub fn new( x0: &FieldElement51, @@ -77,6 +109,52 @@ impl F51x4Unreduced { ]), ] } + + #[inline] + pub fn shuffle(&self, control: Shuffle) -> F51x4Unreduced { + F51x4Unreduced([ + shuffle_lanes(self.0[0], control), + shuffle_lanes(self.0[1], control), + shuffle_lanes(self.0[2], control), + shuffle_lanes(self.0[3], control), + shuffle_lanes(self.0[4], control), + ]) + } + + #[inline] + pub fn blend(&self, other: &F51x4Unreduced, control: Lanes) -> F51x4Unreduced { + F51x4Unreduced([ + blend_lanes(self.0[0], other.0[0], control), + blend_lanes(self.0[1], other.0[1], control), + blend_lanes(self.0[2], other.0[2], control), + blend_lanes(self.0[3], other.0[3], control), + blend_lanes(self.0[4], other.0[4], control), + ]) + } +} + +impl F51x4Reduced { + #[inline] + pub fn shuffle(&self, control: Shuffle) -> F51x4Reduced { + F51x4Reduced([ + shuffle_lanes(self.0[0], control), + shuffle_lanes(self.0[1], control), + shuffle_lanes(self.0[2], control), + shuffle_lanes(self.0[3], control), + shuffle_lanes(self.0[4], control), + ]) + } + + #[inline] + pub fn blend(&self, other: &F51x4Reduced, control: Lanes) -> F51x4Reduced { + F51x4Reduced([ + blend_lanes(self.0[0], other.0[0], control), + blend_lanes(self.0[1], other.0[1], control), + blend_lanes(self.0[2], other.0[2], control), + blend_lanes(self.0[3], other.0[3], control), + blend_lanes(self.0[4], other.0[4], control), + ]) + } } impl From for F51x4Unreduced { @@ -268,7 +346,6 @@ impl<'a, 'b> Mul<&'b F51x4Reduced> for &'a F51x4Reduced { let mut t1 = u64x4::splat(0); let r19 = u64x4::splat(19); - // Wave 6 t0 = madd52hi(t0, r19, z9); t1 = madd52lo(t1, r19, z9 >> 52); @@ -443,4 +520,41 @@ mod test { assert_eq!(c, splits[i]); } } + + #[test] + fn shuffle_AAAA() { + let x0 = FieldElement51::from_bytes(&[0x10; 32]); + let x1 = FieldElement51::from_bytes(&[0x11; 32]); + let x2 = FieldElement51::from_bytes(&[0x12; 32]); + let x3 = FieldElement51::from_bytes(&[0x13; 32]); + + let x = F51x4Unreduced::new(&x0, &x1, &x2, &x3); + + let y = x.shuffle(Shuffle::AAAA); + let splits = y.split(); + + assert_eq!(splits[0], x0); + assert_eq!(splits[1], x0); + assert_eq!(splits[2], x0); + assert_eq!(splits[3], x0); + } + + #[test] + fn blend_AB() { + let x0 = FieldElement51::from_bytes(&[0x10; 32]); + let x1 = FieldElement51::from_bytes(&[0x11; 32]); + let x2 = FieldElement51::from_bytes(&[0x12; 32]); + let x3 = FieldElement51::from_bytes(&[0x13; 32]); + + let x = F51x4Unreduced::new(&x0, &x1, &x2, &x3); + let z = F51x4Unreduced::new(&x3, &x2, &x1, &x0); + + let y = x.blend(&z, Lanes::AB); + let splits = y.split(); + + assert_eq!(splits[0], x3); + assert_eq!(splits[1], x2); + assert_eq!(splits[2], x2); + assert_eq!(splits[3], x3); + } }