AVX2 vectorization of Poseidon constant layer (#243)

* AVX2 vectorization of Poseidon constant layer

* Silence unused import warning + minor style

* Daniel comment

src/field/packed_avx2/packed_prime_field.rs

@@ -43,6 +43,13 @@ impl<F: ReducibleAVX2> PackedPrimeField<F> {
         let ptr = (&self.0).as_ptr().cast::<__m256i>();
         unsafe { _mm256_loadu_si256(ptr) }
     }
+
+    /// Addition that assumes x + y < 2^64 + F::ORDER. May return incorrect results if this
+    /// condition is not met, hence it is marked unsafe.
+    #[inline]
+    pub unsafe fn add_canonical_u64(&self, rhs: __m256i) -> Self {
+        Self::new(add_canonical_u64::<F>(self.get(), rhs))
+    }
 }
 
 impl<F: ReducibleAVX2> Add<Self> for PackedPrimeField<F> {
@@ -286,6 +293,14 @@ unsafe fn canonicalize_s<F: PrimeField>(x_s: __m256i) -> __m256i {
     _mm256_add_epi64(x_s, wrapback_amt)
 }
 
+/// Addition that assumes x + y < 2^64 + F::ORDER.
+#[inline]
+unsafe fn add_canonical_u64<F: PrimeField>(x: __m256i, y: __m256i) -> __m256i {
+    let y_s = shift(y);
+    let res_s = add_no_canonicalize_64_64s_s::<F>(x, y_s);
+    shift(res_s)
+}
+
 #[inline]
 unsafe fn add<F: PrimeField>(x: __m256i, y: __m256i) -> __m256i {
     let y_s = shift(y);

src/hash/poseidon.rs

@@ -1,6 +1,9 @@
 //! Implementation of the Poseidon hash function, as described in
 //! https://eprint.iacr.org/2019/458.pdf
 
+#[cfg(target_feature = "avx2")]
+use std::convert::TryInto;
+
 use unroll::unroll_for_loops;
 
 use crate::field::crandall_field::CrandallField;
@@ -26,6 +29,9 @@ const MAX_WIDTH: usize = 12; // we only have width 8 and 12, and 12 is bigger. :
 // random numbers.
 #[rustfmt::skip]
 const ALL_ROUND_CONSTANTS: [u64; MAX_WIDTH * N_ROUNDS]  = [
+    // WARNING: These must be in 0..CrandallField::ORDER (i.e. canonical form). If this condition is
+    // not met, some platform-specific implementation of constant_layer may return incorrect
+    // results.
     0xb585f767417ee042, 0x7746a55f77c10331, 0xb2fb0d321d356f7a, 0x0f6760a486f1621f,
     0xe10d6666b36abcdf, 0x8cae14cb455cc50b, 0xd438539cf2cee334, 0xef781c7d4c1fd8b4,
     0xcdc4a23a0aca4b1f, 0x277fa208d07b52e3, 0xe17653a300493d38, 0xc54302f27c287dc1,
@@ -494,6 +500,14 @@ impl Poseidon<8> for CrandallField {
          0x800fc4e2c9f585d8, 0xda6cfb436cf6973e, 0x3fc702a71c42c8df, ],
     ];
 
+    #[cfg(target_feature="avx2")]
+    #[inline(always)]
+    fn constant_layer(state: &mut [Self; 8], round_ctr: usize) {
+        // This assumes that every element of ALL_ROUND_CONSTANTS is in 0..CrandallField::ORDER.
+        unsafe { crate::hash::poseidon_avx2::crandall_poseidon_const_avx2::<2>(state,
+            ALL_ROUND_CONSTANTS[8 * round_ctr..8 * round_ctr + 8].try_into().unwrap()); }
+    }
+
     #[cfg(target_feature="avx2")]
     #[inline(always)]
     fn mds_layer(state_: &[CrandallField; 8]) -> [CrandallField; 8] {
@@ -711,6 +725,14 @@ impl Poseidon<12> for CrandallField {
          0xe7ad8152c5d50bed, 0xb5d5efb12203ef9a, 0x8a041eb885fb24f5, ],
     ];
 
+    #[cfg(target_feature="avx2")]
+    #[inline(always)]
+    fn constant_layer(state: &mut [Self; 12], round_ctr: usize) {
+        // This assumes that every element of ALL_ROUND_CONSTANTS is in 0..CrandallField::ORDER.
+        unsafe { crate::hash::poseidon_avx2::crandall_poseidon_const_avx2::<3>(state,
+            ALL_ROUND_CONSTANTS[12 * round_ctr..12 * round_ctr + 12].try_into().unwrap()); }
+    }
+
     #[cfg(target_feature="avx2")]
     #[inline(always)]
     fn mds_layer(state_: &[CrandallField; 12]) -> [CrandallField; 12] {

src/hash/poseidon_avx2.rs

@@ -2,6 +2,8 @@ use core::arch::x86_64::*;
 
 use crate::field::crandall_field::CrandallField;
 use crate::field::field_types::PrimeField;
+use crate::field::packed_avx2::PackedCrandallAVX2;
+use crate::field::packed_field::PackedField;
 
 const EPSILON: u64 = 0u64.wrapping_sub(CrandallField::ORDER);
 const SIGN_BIT: u64 = 1 << 63;
@@ -202,3 +204,19 @@ unsafe fn add_no_canonicalize_64_64s(x: __m256i, y_s: __m256i) -> __m256i {
     let res = _mm256_add_epi64(res_wrapped, wrapback_amt);
     res
 }
+
+/// Poseidon constant layer for Crandall. Assumes that every element in round_constants is in
+/// 0..CrandallField::ORDER; when this is not true it may return garbage. It's marked unsafe for
+/// this reason.
+#[inline(always)]
+pub unsafe fn crandall_poseidon_const_avx2<const PACKED_WIDTH: usize>(
+    state: &mut [CrandallField; 4 * PACKED_WIDTH],
+    round_constants: [u64; 4 * PACKED_WIDTH],
+) {
+    let packed_state = PackedCrandallAVX2::pack_slice_mut(state);
+    let packed_round_constants =
+        std::slice::from_raw_parts((&round_constants).as_ptr().cast::<__m256i>(), PACKED_WIDTH);
+    for i in 0..PACKED_WIDTH {
+        packed_state[i] = packed_state[i].add_canonical_u64(packed_round_constants[i]);
+    }
+}