1use core::arch::x86_64::*;
  2
  3// The following AVX2 implementation has been copied from plonky2:
  4// https://github.com/0xPolygonZero/plonky2/blob/main/plonky2/src/hash/arch/x86_64/poseidon_goldilocks_avx2_bmi2.rs
  5
  6// Preliminary notes:
  7// 1. AVX does not support addition with carry but 128-bit (2-word) addition can be easily emulated.
  8//    The method recognizes that for a + b overflowed iff (a + b) < a: i. res_lo = a_lo + b_lo ii.
  9//    carry_mask = res_lo < a_lo iii. res_hi = a_hi + b_hi - carry_mask
 10//
 11//    Notice that carry_mask is subtracted, not added.
 12//
 13//    This is because AVX comparison instructions return -1 (all bits 1) for true and 0 for false.
 14//
 15// 2. AVX does not have unsigned 64-bit comparisons. Those can be emulated with signed comparisons
 16//    by recognizing that a <u b iff a + (1 << 63) <s b + (1 << 63), where the addition wraps around
 17//    and the comparisons are unsigned and signed respectively. The shift function adds/subtracts 1
 18//    << 63 to enable this trick. Addition with carry example: i. a_lo_s = shift(a_lo) ii. res_lo_s
 19//    = a_lo_s + b_lo iii. carry_mask = res_lo_s <s a_lo_s iv. res_lo = shift(res_lo_s) v. res_hi =
 20//    a_hi + b_hi - carry_mask
 21//
 22//    The suffix _s denotes a value that has been shifted by 1 << 63. The result of addition
 23//    is shifted if exactly one of the operands is shifted, as is the case on
 24//    line ii. Line iii. performs a signed comparison res_lo_s <s a_lo_s on shifted values to
 25//    emulate unsigned comparison res_lo <u a_lo on unshifted values. Finally, line iv. reverses the
 26//    shift so the result can be returned.
 27//
 28//    When performing a chain of calculations, we can often
 29//    save instructions by letting the shift propagate through and only undoing it when necessary.
 30//    For example, to compute the addition of three two-word (128-bit) numbers we can do:
 31//    i. a_lo_s = shift(a_lo)
 32//    ii. tmp_lo_s = a_lo_s + b_lo
 33//    iii. tmp_carry_mask = tmp_lo_s <s a_lo_s
 34//    iv. tmp_hi = a_hi + b_hi - tmp_carry_mask
 35//    v. res_lo_s = tmp_lo_s + c_lo vi. res_carry_mask = res_lo_s <s tmp_lo_s
 36//    vii. res_lo = shift(res_lo_s)
 37//    viii. res_hi = tmp_hi + c_hi - res_carry_mask
 38//
 39//    Notice that the above 3-value addition still only requires two calls to shift, just like our
 40//    2-value addition.
 41
 42#[inline(always)]
 43pub fn branch_hint() {
 44    // NOTE: These are the currently supported assembly architectures. See the
 45    // [nightly reference](https://doc.rust-lang.org/nightly/reference/inline-assembly.html) for
 46    // the most up-to-date list.
 47    #[cfg(any(
 48        target_arch = "aarch64",
 49        target_arch = "arm",
 50        target_arch = "riscv32",
 51        target_arch = "riscv64",
 52        target_arch = "x86",
 53        target_arch = "x86_64",
 54    ))]
 55    unsafe {
 56        core::arch::asm!("", options(nomem, nostack, preserves_flags));
 57    }
 58}
 59
 60macro_rules! map3 {
 61    ($f:ident:: < $l:literal > , $v:ident) => {
 62        ($f::<$l>($v.0), $f::<$l>($v.1), $f::<$l>($v.2))
 63    };
 64    ($f:ident:: < $l:literal > , $v1:ident, $v2:ident) => {
 65        ($f::<$l>($v1.0, $v2.0), $f::<$l>($v1.1, $v2.1), $f::<$l>($v1.2, $v2.2))
 66    };
 67    ($f:ident, $v:ident) => {
 68        ($f($v.0), $f($v.1), $f($v.2))
 69    };
 70    ($f:ident, $v0:ident, $v1:ident) => {
 71        ($f($v0.0, $v1.0), $f($v0.1, $v1.1), $f($v0.2, $v1.2))
 72    };
 73    ($f:ident,rep $v0:ident, $v1:ident) => {
 74        ($f($v0, $v1.0), $f($v0, $v1.1), $f($v0, $v1.2))
 75    };
 76
 77    ($f:ident, $v0:ident,rep $v1:ident) => {
 78        ($f($v0.0, $v1), $f($v0.1, $v1), $f($v0.2, $v1))
 79    };
 80}
 81
 82#[inline(always)]
 83unsafe fn square3(
 84    x: (__m256i, __m256i, __m256i),
 85) -> ((__m256i, __m256i, __m256i), (__m256i, __m256i, __m256i)) {
 86    let x_hi = {
 87        // Move high bits to low position. The high bits of x_hi are ignored. Swizzle is faster than
 88        // bitshift. This instruction only has a floating-point flavor, so we cast to/from float.
 89        // This is safe and free.
 90        let x_ps = map3!(_mm256_castsi256_ps, x);
 91        let x_hi_ps = map3!(_mm256_movehdup_ps, x_ps);
 92        map3!(_mm256_castps_si256, x_hi_ps)
 93    };
 94
 95    // All pairwise multiplications.
 96    let mul_ll = map3!(_mm256_mul_epu32, x, x);
 97    let mul_lh = map3!(_mm256_mul_epu32, x, x_hi);
 98    let mul_hh = map3!(_mm256_mul_epu32, x_hi, x_hi);
 99
100    // Bignum addition, but mul_lh is shifted by 33 bits (not 32).
101    let mul_ll_hi = map3!(_mm256_srli_epi64::<33>, mul_ll);
102    let t0 = map3!(_mm256_add_epi64, mul_lh, mul_ll_hi);
103    let t0_hi = map3!(_mm256_srli_epi64::<31>, t0);
104    let res_hi = map3!(_mm256_add_epi64, mul_hh, t0_hi);
105
106    // Form low result by adding the mul_ll and the low 31 bits of mul_lh (shifted to the high
107    // position).
108    let mul_lh_lo = map3!(_mm256_slli_epi64::<33>, mul_lh);
109    let res_lo = map3!(_mm256_add_epi64, mul_ll, mul_lh_lo);
110
111    (res_lo, res_hi)
112}
113
114#[inline(always)]
115unsafe fn mul3(
116    x: (__m256i, __m256i, __m256i),
117    y: (__m256i, __m256i, __m256i),
118) -> ((__m256i, __m256i, __m256i), (__m256i, __m256i, __m256i)) {
119    let epsilon = _mm256_set1_epi64x(0xffffffff);
120    let x_hi = {
121        // Move high bits to low position. The high bits of x_hi are ignored. Swizzle is faster than
122        // bitshift. This instruction only has a floating-point flavor, so we cast to/from float.
123        // This is safe and free.
124        let x_ps = map3!(_mm256_castsi256_ps, x);
125        let x_hi_ps = map3!(_mm256_movehdup_ps, x_ps);
126        map3!(_mm256_castps_si256, x_hi_ps)
127    };
128    let y_hi = {
129        let y_ps = map3!(_mm256_castsi256_ps, y);
130        let y_hi_ps = map3!(_mm256_movehdup_ps, y_ps);
131        map3!(_mm256_castps_si256, y_hi_ps)
132    };
133
134    // All four pairwise multiplications
135    let mul_ll = map3!(_mm256_mul_epu32, x, y);
136    let mul_lh = map3!(_mm256_mul_epu32, x, y_hi);
137    let mul_hl = map3!(_mm256_mul_epu32, x_hi, y);
138    let mul_hh = map3!(_mm256_mul_epu32, x_hi, y_hi);
139
140    // Bignum addition
141    // Extract high 32 bits of mul_ll and add to mul_hl. This cannot overflow.
142    let mul_ll_hi = map3!(_mm256_srli_epi64::<32>, mul_ll);
143    let t0 = map3!(_mm256_add_epi64, mul_hl, mul_ll_hi);
144    // Extract low 32 bits of t0 and add to mul_lh. Again, this cannot overflow.
145    // Also, extract high 32 bits of t0 and add to mul_hh.
146    let t0_lo = map3!(_mm256_and_si256, t0, rep epsilon);
147    let t0_hi = map3!(_mm256_srli_epi64::<32>, t0);
148    let t1 = map3!(_mm256_add_epi64, mul_lh, t0_lo);
149    let t2 = map3!(_mm256_add_epi64, mul_hh, t0_hi);
150    // Lastly, extract the high 32 bits of t1 and add to t2.
151    let t1_hi = map3!(_mm256_srli_epi64::<32>, t1);
152    let res_hi = map3!(_mm256_add_epi64, t2, t1_hi);
153
154    // Form res_lo by combining the low half of mul_ll with the low half of t1 (shifted into high
155    // position).
156    let t1_lo = {
157        let t1_ps = map3!(_mm256_castsi256_ps, t1);
158        let t1_lo_ps = map3!(_mm256_moveldup_ps, t1_ps);
159        map3!(_mm256_castps_si256, t1_lo_ps)
160    };
161    let res_lo = map3!(_mm256_blend_epi32::<0xaa>, mul_ll, t1_lo);
162
163    (res_lo, res_hi)
164}
165
166/// Addition, where the second operand is `0 <= y < 0xffffffff00000001`.
167#[inline(always)]
168unsafe fn add_small(
169    x_s: (__m256i, __m256i, __m256i),
170    y: (__m256i, __m256i, __m256i),
171) -> (__m256i, __m256i, __m256i) {
172    let res_wrapped_s = map3!(_mm256_add_epi64, x_s, y);
173    let mask = map3!(_mm256_cmpgt_epi32, x_s, res_wrapped_s);
174    let wrapback_amt = map3!(_mm256_srli_epi64::<32>, mask); // EPSILON if overflowed else 0.
175    let res_s = map3!(_mm256_add_epi64, res_wrapped_s, wrapback_amt);
176    res_s
177}
178
179#[inline(always)]
180unsafe fn maybe_adj_sub(res_wrapped_s: __m256i, mask: __m256i) -> __m256i {
181    // The subtraction is very unlikely to overflow so we're best off branching.
182    // The even u32s in `mask` are meaningless, so we want to ignore them. `_mm256_testz_pd`
183    // branches depending on the sign bit of double-precision (64-bit) floats. Bit cast `mask` to
184    // floating-point (this is free).
185    let mask_pd = _mm256_castsi256_pd(mask);
186    // `_mm256_testz_pd(mask_pd, mask_pd) == 1` iff all sign bits are 0, meaning that underflow
187    // did not occur for any of the vector elements.
188    if _mm256_testz_pd(mask_pd, mask_pd) == 1 {
189        res_wrapped_s
190    } else {
191        branch_hint();
192        // Highly unlikely: underflow did occur. Find adjustment per element and apply it.
193        let adj_amount = _mm256_srli_epi64::<32>(mask); // EPSILON if underflow.
194        _mm256_sub_epi64(res_wrapped_s, adj_amount)
195    }
196}
197
198/// Addition, where the second operand is much smaller than `0xffffffff00000001`.
199#[inline(always)]
200unsafe fn sub_tiny(
201    x_s: (__m256i, __m256i, __m256i),
202    y: (__m256i, __m256i, __m256i),
203) -> (__m256i, __m256i, __m256i) {
204    let res_wrapped_s = map3!(_mm256_sub_epi64, x_s, y);
205    let mask = map3!(_mm256_cmpgt_epi32, res_wrapped_s, x_s);
206    let res_s = map3!(maybe_adj_sub, res_wrapped_s, mask);
207    res_s
208}
209
210#[inline(always)]
211unsafe fn reduce3(
212    (lo0, hi0): ((__m256i, __m256i, __m256i), (__m256i, __m256i, __m256i)),
213) -> (__m256i, __m256i, __m256i) {
214    let sign_bit = _mm256_set1_epi64x(i64::MIN);
215    let epsilon = _mm256_set1_epi64x(0xffffffff);
216    let lo0_s = map3!(_mm256_xor_si256, lo0, rep sign_bit);
217    let hi_hi0 = map3!(_mm256_srli_epi64::<32>, hi0);
218    let lo1_s = sub_tiny(lo0_s, hi_hi0);
219    let t1 = map3!(_mm256_mul_epu32, hi0, rep epsilon);
220    let lo2_s = add_small(lo1_s, t1);
221    let lo2 = map3!(_mm256_xor_si256, lo2_s, rep sign_bit);
222    lo2
223}
224
225#[inline(always)]
226unsafe fn mul_reduce(
227    a: (__m256i, __m256i, __m256i),
228    b: (__m256i, __m256i, __m256i),
229) -> (__m256i, __m256i, __m256i) {
230    reduce3(mul3(a, b))
231}
232
233#[inline(always)]
234unsafe fn square_reduce(state: (__m256i, __m256i, __m256i)) -> (__m256i, __m256i, __m256i) {
235    reduce3(square3(state))
236}
237
238#[inline(always)]
239unsafe fn exp_acc(
240    high: (__m256i, __m256i, __m256i),
241    low: (__m256i, __m256i, __m256i),
242    exp: usize,
243) -> (__m256i, __m256i, __m256i) {
244    let mut result = high;
245    for _ in 0..exp {
246        result = square_reduce(result);
247    }
248    mul_reduce(result, low)
249}
250
251#[inline(always)]
252unsafe fn do_apply_sbox(state: (__m256i, __m256i, __m256i)) -> (__m256i, __m256i, __m256i) {
253    let state2 = square_reduce(state);
254    let state4_unreduced = square3(state2);
255    let state3_unreduced = mul3(state2, state);
256    let state4 = reduce3(state4_unreduced);
257    let state3 = reduce3(state3_unreduced);
258    let state7_unreduced = mul3(state3, state4);
259    let state7 = reduce3(state7_unreduced);
260    state7
261}
262
263#[inline(always)]
264unsafe fn do_apply_inv_sbox(state: (__m256i, __m256i, __m256i)) -> (__m256i, __m256i, __m256i) {
265    // compute base^10540996611094048183 using 72 multiplications per array element
266    // 10540996611094048183 = b1001001001001001001001001001000110110110110110110110110110110111
267
268    // compute base^10
269    let t1 = square_reduce(state);
270
271    // compute base^100
272    let t2 = square_reduce(t1);
273
274    // compute base^100100
275    let t3 = exp_acc(t2, t2, 3);
276
277    // compute base^100100100100
278    let t4 = exp_acc(t3, t3, 6);
279
280    // compute base^100100100100100100100100
281    let t5 = exp_acc(t4, t4, 12);
282
283    // compute base^100100100100100100100100100100
284    let t6 = exp_acc(t5, t3, 6);
285
286    // compute base^1001001001001001001001001001000100100100100100100100100100100
287    let t7 = exp_acc(t6, t6, 31);
288
289    // compute base^1001001001001001001001001001000110110110110110110110110110110111
290    let a = square_reduce(square_reduce(mul_reduce(square_reduce(t7), t6)));
291    let b = mul_reduce(t1, mul_reduce(t2, state));
292    mul_reduce(a, b)
293}
294
295#[inline(always)]
296unsafe fn avx2_load(state: &[u64; 12]) -> (__m256i, __m256i, __m256i) {
297    (
298        _mm256_loadu_si256((&state[0..4]).as_ptr().cast::<__m256i>()),
299        _mm256_loadu_si256((&state[4..8]).as_ptr().cast::<__m256i>()),
300        _mm256_loadu_si256((&state[8..12]).as_ptr().cast::<__m256i>()),
301    )
302}
303
304#[inline(always)]
305unsafe fn avx2_store(buf: &mut [u64; 12], state: (__m256i, __m256i, __m256i)) {
306    _mm256_storeu_si256((&mut buf[0..4]).as_mut_ptr().cast::<__m256i>(), state.0);
307    _mm256_storeu_si256((&mut buf[4..8]).as_mut_ptr().cast::<__m256i>(), state.1);
308    _mm256_storeu_si256((&mut buf[8..12]).as_mut_ptr().cast::<__m256i>(), state.2);
309}
310
311#[inline(always)]
312pub unsafe fn apply_sbox(buffer: &mut [u64; 12]) {
313    let mut state = avx2_load(&buffer);
314    state = do_apply_sbox(state);
315    avx2_store(buffer, state);
316}
317
318#[inline(always)]
319pub unsafe fn apply_inv_sbox(buffer: &mut [u64; 12]) {
320    let mut state = avx2_load(&buffer);
321    state = do_apply_inv_sbox(state);
322    avx2_store(buffer, state);
323}