diff --git a/src/field/packed_avx2/packed_prime_field.rs b/src/field/packed_avx2/packed_prime_field.rs index b892da4a..ed87f347 100644 --- a/src/field/packed_avx2/packed_prime_field.rs +++ b/src/field/packed_avx2/packed_prime_field.rs @@ -43,6 +43,13 @@ impl PackedPrimeField { 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::(self.get(), rhs)) + } } impl Add for PackedPrimeField { @@ -286,6 +293,14 @@ unsafe fn canonicalize_s(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(x: __m256i, y: __m256i) -> __m256i { + let y_s = shift(y); + let res_s = add_no_canonicalize_64_64s_s::(x, y_s); + shift(res_s) +} + #[inline] unsafe fn add(x: __m256i, y: __m256i) -> __m256i { let y_s = shift(y); diff --git a/src/hash/poseidon.rs b/src/hash/poseidon.rs index 35d9c2ed..d5cb13a0 100644 --- a/src/hash/poseidon.rs +++ b/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] { diff --git a/src/hash/poseidon_avx2.rs b/src/hash/poseidon_avx2.rs index c5e197b6..643317f8 100644 --- a/src/hash/poseidon_avx2.rs +++ b/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( + 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]); + } +}