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