chore: add avx acceleration to poseidon hash function

algebraic/src/arch/x86_64/avx2_field_gl.rs

@@ -6,12 +6,12 @@
 //!
 use crate::ff::*;
 use crate::field_gl::{Fr, FrRepr as GoldilocksField};
+use crate::packed::PackedField;
 use core::arch::x86_64::*;
 use core::fmt;
 use core::fmt::{Debug, Formatter};
 use core::mem::transmute;
 use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-// use crate::packed::PackedField;
 
 /// AVX2 Goldilocks Field
 ///
@@ -24,8 +24,6 @@ use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAss
 #[repr(transparent)]
 pub struct Avx2GoldilocksField(pub [GoldilocksField; 4]);
 
-const WIDTH: usize = 4;
-
 impl Avx2GoldilocksField {
     #[inline]
     pub fn new(x: __m256i) -> Self {
@@ -35,30 +33,36 @@ impl Avx2GoldilocksField {
     pub fn get(&self) -> __m256i {
         unsafe { transmute(*self) }
     }
-    // }
-    // unsafe impl PackedField for Avx2GoldilocksField {
     #[inline]
-    pub fn from_slice(slice: &[GoldilocksField]) -> &Self {
-        assert_eq!(slice.len(), WIDTH);
+    pub fn square(&self) -> Avx2GoldilocksField {
+        Self::new(unsafe { square(self.get()) })
+    }
+}
+
+unsafe impl PackedField for Avx2GoldilocksField {
+    const WIDTH: usize = 4;
+    type Scalar = GoldilocksField;
+    const ZEROS: Self = Self([GoldilocksField([0]); 4]);
+    const ONES: Self = Self([GoldilocksField([1]); 4]);
+
+    #[inline]
+    fn from_slice(slice: &[GoldilocksField]) -> &Self {
+        assert_eq!(slice.len(), Self::WIDTH);
         unsafe { &*slice.as_ptr().cast() }
     }
     #[inline]
-    pub fn from_slice_mut(slice: &mut [GoldilocksField]) -> &mut Self {
-        assert_eq!(slice.len(), WIDTH);
+    fn from_slice_mut(slice: &mut [GoldilocksField]) -> &mut Self {
+        assert_eq!(slice.len(), Self::WIDTH);
         unsafe { &mut *slice.as_mut_ptr().cast() }
     }
     #[inline]
-    pub fn as_slice(&self) -> &[GoldilocksField] {
+    fn as_slice(&self) -> &[GoldilocksField] {
         &self.0[..]
     }
     #[inline]
-    pub fn as_slice_mut(&mut self) -> &mut [GoldilocksField] {
+    fn as_slice_mut(&mut self) -> &mut [GoldilocksField] {
         &mut self.0[..]
     }
-    #[inline]
-    pub fn square(&self) -> Avx2GoldilocksField {
-        Self::new(unsafe { square(self.get()) })
-    }
 
     #[inline]
     fn interleave(&self, other: Self, block_len: usize) -> (Self, Self) {
@@ -117,7 +121,7 @@ impl Debug for Avx2GoldilocksField {
 impl Default for Avx2GoldilocksField {
     #[inline]
     fn default() -> Self {
-        Self([GoldilocksField::from(0); 4])
+        Self::ZEROS
     }
 }
 
@@ -503,8 +507,8 @@ mod tests {
     use super::Avx2GoldilocksField;
     use crate::ff::*;
     use crate::field_gl::{Fr, FrRepr as GoldilocksField};
+    use crate::packed::PackedField;
     use std::time::Instant;
-    // use crate::packed::PackedField;
 
     fn test_vals_a() -> [GoldilocksField; 4] {
         [

algebraic/src/arch/x86_64/avx512_field_gl.rs

@@ -5,6 +5,7 @@
 //! RUSTFLAGS='-C target-feature=+avx512f,+avx512bw,+avx512cd,+avx512dq,+avx512vl' cargo build --release
 use crate::ff::*;
 use crate::field_gl::{Fr, FrRepr as GoldilocksField};
+use crate::packed::PackedField;
 use core::arch::x86_64::*;
 use core::fmt;
 use core::fmt::{Debug, Formatter};
@@ -34,27 +35,36 @@ impl Avx512GoldilocksField {
         unsafe { transmute(*self) }
     }
     #[inline]
-    pub fn from_slice(slice: &[GoldilocksField]) -> &Self {
+    pub fn square(&self) -> Avx512GoldilocksField {
+        Self::new(unsafe { square(self.get()) })
+    }
+}
+
+unsafe impl PackedField for Avx512GoldilocksField {
+    const WIDTH: usize = 8;
+
+    type Scalar = GoldilocksField;
+
+    const ZEROS: Self = Self([GoldilocksField([0]); 8]);
+    const ONES: Self = Self([GoldilocksField([1]); 8]);
+    #[inline]
+    fn from_slice(slice: &[GoldilocksField]) -> &Self {
         assert_eq!(slice.len(), WIDTH);
         unsafe { &*slice.as_ptr().cast() }
     }
     #[inline]
-    pub fn from_slice_mut(slice: &mut [GoldilocksField]) -> &mut Self {
+    fn from_slice_mut(slice: &mut [GoldilocksField]) -> &mut Self {
         assert_eq!(slice.len(), WIDTH);
         unsafe { &mut *slice.as_mut_ptr().cast() }
     }
     #[inline]
-    pub fn as_slice(&self) -> &[GoldilocksField] {
+    fn as_slice(&self) -> &[GoldilocksField] {
         &self.0[..]
     }
     #[inline]
-    pub fn as_slice_mut(&mut self) -> &mut [GoldilocksField] {
+    fn as_slice_mut(&mut self) -> &mut [GoldilocksField] {
         &mut self.0[..]
     }
-    #[inline]
-    pub fn square(&self) -> Avx512GoldilocksField {
-        Self::new(unsafe { square(self.get()) })
-    }
 
     #[inline]
     fn interleave(&self, other: Self, block_len: usize) -> (Self, Self) {
@@ -114,7 +124,7 @@ impl Debug for Avx512GoldilocksField {
 impl Default for Avx512GoldilocksField {
     #[inline]
     fn default() -> Self {
-        Self([GoldilocksField::from(0); 8])
+        Self::ZEROS
     }
 }
 
@@ -397,6 +407,7 @@ mod tests {
     use super::Avx512GoldilocksField;
     use crate::ff::*;
     use crate::field_gl::{Fr, FrRepr as GoldilocksField};
+    use crate::packed::PackedField;
     use std::time::Instant;
 
     fn test_vals_a() -> [GoldilocksField; 8] {

algebraic/src/lib.rs

@@ -1,6 +1,6 @@
 #![allow(clippy::unit_arg)]
-
-// #![feature(stdsimd)]
+#![feature(stdsimd)]
+#![feature(const_trait_impl)]
 
 #[macro_use]
 extern crate serde;
@@ -13,7 +13,7 @@ extern crate num_bigint;
 extern crate num_traits;
 extern crate rand;
 
-// pub mod arch;
+pub mod arch;
 
 pub mod circom_circuit;
 pub mod errors;
@@ -24,13 +24,15 @@ pub mod witness;
 
 pub mod utils;
 
+pub mod packable;
+pub mod packed;
+
 pub use bellman_ce::pairing::ff;
 pub use ff::*;
 pub use franklin_crypto::bellman as bellman_ce;
 
 #[cfg(test)]
 mod field_gl_test;
-// mod packed;
 
 #[cfg(target_arch = "wasm32")]
 extern crate wasm_bindgen;

algebraic/src/packable.rs

@@ -1,15 +1,18 @@
+use crate::ff::PrimeFieldRepr;
 use crate::packed::PackedField;
-use crate::types::Field;
 
 /// Points us to the default packing for a particular field. There may me multiple choices of
 /// PackedField for a particular Field (e.g. every Field is also a PackedField), but this is the
 /// recommended one. The recommended packing varies by target_arch and target_feature.
-pub trait Packable: Field {
+pub trait Packable: PrimeFieldRepr {
     type Packing: PackedField<Scalar = Self>;
 }
 
-impl<F: Field> Packable for F {
-    default type Packing = Self;
+impl<F> Packable for F
+where
+    F: PrimeFieldRepr + PackedField<Scalar = F>,
+{
+    type Packing = Self;
 }
 
 #[cfg(all(
@@ -23,8 +26,8 @@ impl<F: Field> Packable for F {
         target_feature = "avx512vl"
     ))
 ))]
-impl Packable for crate::goldilocks_field::GoldilocksField {
-    type Packing = crate::arch::x86_64::avx2_goldilocks_field::Avx2GoldilocksField;
+impl Packable for crate::field_gl::FrRepr {
+    type Packing = crate::arch::x86_64::avx2_field_gl::Avx2GoldilocksField;
 }
 
 #[cfg(all(
@@ -35,6 +38,6 @@ impl Packable for crate::goldilocks_field::GoldilocksField {
     target_feature = "avx512f",
     target_feature = "avx512vl"
 ))]
-impl Packable for crate::goldilocks_field::GoldilocksField {
-    type Packing = crate::arch::x86_64::avx512_goldilocks_field::Avx512GoldilocksField;
+impl Packable for crate::field_gl::FrRepr {
+    type Packing = crate::arch::x86_64::avx512_field_gl::Avx512GoldilocksField;
 }

algebraic/src/packed.rs

@@ -1,11 +1,9 @@
 use core::fmt::Debug;
-use core::iter::{Product, Sum};
+// use core::iter::{Product, Sum};
+use crate::ff::PrimeFieldRepr;
 use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
 use core::slice;
 
-// use crate::ops::Square;
-use crate::ff::*;
-
 /// # Safety
 /// - WIDTH is assumed to be a power of 2.
 /// - If P implements PackedField then P must be castable to/from [P::Scalar; P::WIDTH] without UB.
@@ -27,20 +25,20 @@ pub unsafe trait PackedField:
     + MulAssign<Self::Scalar>
     // + Square
     + Neg<Output = Self>
-    + Product
+    // + Product
     + Send
     + Sub<Self, Output = Self>
     + Sub<Self::Scalar, Output = Self>
     + SubAssign<Self>
     + SubAssign<Self::Scalar>
-    + Sum
+    // + Sum
     + Sync
 where
     Self::Scalar: Add<Self, Output = Self>,
     Self::Scalar: Mul<Self, Output = Self>,
     Self::Scalar: Sub<Self, Output = Self>,
 {
-    type Scalar: Field;
+    type Scalar: PrimeFieldRepr;
 
     const WIDTH: usize;
     const ZEROS: Self;
@@ -93,35 +91,4 @@ where
         unsafe { slice::from_raw_parts_mut(buf_ptr, n) }
     }
 
-    fn doubles(&self) -> Self {
-        *self * Self::Scalar::TWO
-    }
-}
-
-unsafe impl<F: Field> PackedField for F {
-    type Scalar = Self;
-
-    const WIDTH: usize = 1;
-    const ZEROS: Self = F::ZERO;
-    const ONES: Self = F::ONE;
-
-    fn from_slice(slice: &[Self::Scalar]) -> &Self {
-        &slice[0]
-    }
-    fn from_slice_mut(slice: &mut [Self::Scalar]) -> &mut Self {
-        &mut slice[0]
-    }
-    fn as_slice(&self) -> &[Self::Scalar] {
-        slice::from_ref(self)
-    }
-    fn as_slice_mut(&mut self) -> &mut [Self::Scalar] {
-        slice::from_mut(self)
-    }
-
-    fn interleave(&self, other: Self, block_len: usize) -> (Self, Self) {
-        match block_len {
-            1 => (*self, other),
-            _ => panic!("unsupported block length"),
-        }
-    }
 }

starky/src/arch/mod.rs

@@ -0,0 +1,2 @@
+#[cfg(target_arch = "x86_64")]
+pub mod x86_64;