diff --git a/src/backend/vector/ifma/field.rs b/src/backend/vector/ifma/field.rs index 307f807..d344514 100644 --- a/src/backend/vector/ifma/field.rs +++ b/src/backend/vector/ifma/field.rs @@ -190,6 +190,120 @@ impl F51x4Reduced { blend_lanes(self.0[4], other.0[4], control), ]) } + + #[inline] + pub fn square(&self) -> F51x4Unreduced { + unsafe { + let x = &self.0; + + // Represent values with coeff. 2 + let mut z0_2 = u64x4::splat(0); + let mut z1_2 = u64x4::splat(0); + let mut z2_2 = u64x4::splat(0); + let mut z3_2 = u64x4::splat(0); + let mut z4_2 = u64x4::splat(0); + let mut z5_2 = u64x4::splat(0); + let mut z6_2 = u64x4::splat(0); + let mut z7_2 = u64x4::splat(0); + let mut z9_2 = u64x4::splat(0); + + // Represent values with coeff. 4 + let mut z2_4 = u64x4::splat(0); + let mut z3_4 = u64x4::splat(0); + let mut z4_4 = u64x4::splat(0); + let mut z5_4 = u64x4::splat(0); + let mut z6_4 = u64x4::splat(0); + let mut z7_4 = u64x4::splat(0); + let mut z8_4 = u64x4::splat(0); + + let mut z0_1 = u64x4::splat(0); + z0_1 = madd52lo(z0_1, x[0], x[0]); + + let mut z1_1 = u64x4::splat(0); + z1_2 = madd52lo(z1_2, x[0], x[1]); + z1_2 = madd52hi(z1_2, x[0], x[0]); + + z2_4 = madd52hi(z2_4, x[0], x[1]); + let mut z2_1 = z2_4 << 2; + z2_2 = madd52lo(z2_2, x[0], x[2]); + z2_1 = madd52lo(z2_1, x[1], x[1]); + + z3_4 = madd52hi(z3_4, x[0], x[2]); + let mut z3_1 = z3_4 << 2; + z3_2 = madd52lo(z3_2, x[1], x[2]); + z3_2 = madd52lo(z3_2, x[0], x[3]); + z3_2 = madd52hi(z3_2, x[1], x[1]); + + z4_4 = madd52hi(z4_4, x[1], x[2]); + z4_4 = madd52hi(z4_4, x[0], x[3]); + let mut z4_1 = z4_4 << 2; + z4_2 = madd52lo(z4_2, x[1], x[3]); + z4_2 = madd52lo(z4_2, x[0], x[4]); + z4_1 = madd52lo(z4_1, x[2], x[2]); + + z5_4 = madd52hi(z5_4, x[1], x[3]); + z5_4 = madd52hi(z5_4, x[0], x[4]); + let mut z5_1 = z5_4 << 2; + z5_2 = madd52lo(z5_2, x[2], x[3]); + z5_2 = madd52lo(z5_2, x[1], x[4]); + z5_2 = madd52hi(z5_2, x[2], x[2]); + + z6_4 = madd52hi(z6_4, x[2], x[3]); + z6_4 = madd52hi(z6_4, x[1], x[4]); + let mut z6_1 = z6_4 << 2; + z6_2 = madd52lo(z6_2, x[2], x[4]); + z6_1 = madd52lo(z6_1, x[3], x[3]); + + z7_4 = madd52hi(z7_4, x[2], x[4]); + let mut z7_1 = z7_4 << 2; + z7_2 = madd52lo(z7_2, x[3], x[4]); + z7_2 = madd52hi(z7_2, x[3], x[3]); + + z8_4 = madd52hi(z8_4, x[3], x[4]); + let mut z8_1 = z8_4 << 2; + z8_1 = madd52lo(z8_1, x[4], x[4]); + + let mut z9_1 = u64x4::splat(0); + z9_2 = madd52hi(z9_2, x[4], x[4]); + + z5_1 += z5_2 << 1; + z6_1 += z6_2 << 1; + z7_1 += z7_2 << 1; + z9_1 += z9_2 << 1; + + let mut t0 = u64x4::splat(0); + let mut t1 = u64x4::splat(0); + let r19 = u64x4::splat(19); + + t0 = madd52hi(t0, r19, z9_1); + t1 = madd52lo(t1, r19, z9_1 >> 52); + + z4_2 = madd52lo(z4_2, r19, z8_1 >> 52); + z3_2 = madd52lo(z3_2, r19, z7_1 >> 52); + z2_2 = madd52lo(z2_2, r19, z6_1 >> 52); + z1_2 = madd52lo(z1_2, r19, z5_1 >> 52); + + z0_2 = madd52lo(z0_2, r19, t0 + t1); + z1_2 = madd52hi(z1_2, r19, z5_1); + z2_2 = madd52hi(z2_2, r19, z6_1); + z3_2 = madd52hi(z3_2, r19, z7_1); + z4_2 = madd52hi(z4_2, r19, z8_1); + + z0_1 = madd52lo(z0_1, r19, z5_1); + z1_1 = madd52lo(z1_1, r19, z6_1); + z2_1 = madd52lo(z2_1, r19, z7_1); + z3_1 = madd52lo(z3_1, r19, z8_1); + z4_1 = madd52lo(z4_1, r19, z9_1); + + F51x4Unreduced([ + z0_1 + z0_2 + z0_2, + z1_1 + z1_2 + z1_2, + z2_1 + z2_2 + z2_2, + z3_1 + z3_2 + z3_2, + z4_1 + z4_2 + z4_2, + ]) + } + } } impl From for F51x4Unreduced { @@ -250,7 +364,6 @@ impl<'a> Mul<(u32, u32, u32, u32)> for &'a F51x4Reduced { scalars.2 as u64, scalars.3 as u64, ); - let mask = u64x4::splat((1 << 51) - 1); let r19 = u64x4::splat(19); let mut z0lo = u64x4::splat(0); @@ -546,6 +659,39 @@ mod test { } } + #[test] + fn square_matches_mul() { + // Invert a small field element to get a big one + let a = FieldElement51([2438, 24, 243, 0, 0]).invert(); + + let ax4: F51x4Reduced = F51x4Unreduced::new(&a, &a, &a, &a).into(); + let mut cx4 = &ax4 * &ax4; + let mut cx4_sq = ax4.square(); + + let splits = cx4.split(); + let splits_sq = cx4_sq.split(); + + for i in 0..4 { + assert_eq!(splits_sq[i], splits[i]); + } + } + + #[test] + fn iterated_square_matches_serial() { + // Invert a small field element to get a big one + let mut a = FieldElement51([2438, 24, 243, 0, 0]).invert(); + let mut ax4 = F51x4Unreduced::new(&a, &a, &a, &a); + for j in 0..1024 { + a = a.square(); + ax4 = F51x4Reduced::from(ax4).square(); + + let splits = ax4.split(); + for i in 0..4 { + assert_eq!(a, splits[i]); + } + } + } + #[test] fn iterated_u32_mul_matches_serial() { // Invert a small field element to get a big one