chacha20/backends/
sse2.rs

1//! SSE2 backend.
2
3#![allow(unsafe_op_in_unsafe_fn, reason = "needs triage")]
4#![allow(clippy::cast_possible_truncation, reason = "needs triage")]
5#![allow(clippy::cast_possible_wrap, reason = "needs triage")]
6#![allow(clippy::cast_sign_loss, reason = "needs triage")]
7#![allow(clippy::undocumented_unsafe_blocks, reason = "TODO")]
8
9use crate::{Rounds, STATE_WORDS, Variant};
10
11#[cfg(feature = "rng")]
12use crate::ChaChaCore;
13
14#[cfg(feature = "cipher")]
15use crate::chacha::Block;
16#[cfg(feature = "cipher")]
17use cipher::{
18    BlockSizeUser, ParBlocksSizeUser, StreamCipherBackend, StreamCipherClosure,
19    consts::{U4, U64},
20};
21use core::marker::PhantomData;
22
23#[cfg(target_arch = "x86")]
24use core::arch::x86::*;
25#[cfg(target_arch = "x86_64")]
26use core::arch::x86_64::*;
27
28const PAR_BLOCKS: usize = 4;
29
30#[inline]
31#[target_feature(enable = "sse2")]
32#[cfg(feature = "cipher")]
33pub(crate) unsafe fn inner<R, F, V>(state: &mut [u32; STATE_WORDS], f: F)
34where
35    R: Rounds,
36    F: StreamCipherClosure<BlockSize = U64>,
37    V: Variant,
38{
39    let mut backend = Backend::<R, V>::new(state);
40    f.call(&mut backend);
41    backend.save_ctr(state);
42}
43
44#[inline]
45#[target_feature(enable = "sse2")]
46#[cfg(feature = "rng")]
47pub(crate) unsafe fn rng_inner<R, V>(core: &mut ChaChaCore<R, V>, buffer: &mut [u32; 64])
48where
49    R: Rounds,
50    V: Variant,
51{
52    let mut backend = Backend::<R, V>::new(&core.state);
53    backend.gen_ks_blocks(buffer);
54    backend.save_ctr(&mut core.state);
55}
56
57struct Backend<R: Rounds, V: Variant> {
58    v: [__m128i; 4],
59    _pd: PhantomData<(R, V)>,
60}
61
62impl<R: Rounds, V: Variant> Backend<R, V> {
63    unsafe fn new(state: &[u32; STATE_WORDS]) -> Self {
64        let state_ptr = state.as_ptr().cast::<__m128i>();
65        Self {
66            v: core::array::from_fn(|i| _mm_loadu_si128(state_ptr.add(i))),
67            _pd: PhantomData,
68        }
69    }
70
71    unsafe fn save_ctr(self, state: &mut [u32; STATE_WORDS]) {
72        let state_ptr = state.as_mut_ptr().cast::<__m128i>();
73        _mm_storeu_si128(state_ptr.add(3), self.v[3]);
74    }
75}
76
77#[cfg(feature = "cipher")]
78impl<R: Rounds, V: Variant> BlockSizeUser for Backend<R, V> {
79    type BlockSize = U64;
80}
81
82#[cfg(feature = "cipher")]
83impl<R: Rounds, V: Variant> ParBlocksSizeUser for Backend<R, V> {
84    type ParBlocksSize = U4;
85}
86
87#[cfg(feature = "cipher")]
88impl<R: Rounds, V: Variant> StreamCipherBackend for Backend<R, V> {
89    #[inline(always)]
90    fn gen_ks_block(&mut self, block: &mut Block) {
91        unsafe {
92            let res = rounds::<R, V>(&self.v);
93            self.v[3] = match size_of::<V::Counter>() {
94                4 => _mm_add_epi32(self.v[3], _mm_set_epi32(0, 0, 0, 1)),
95                8 => _mm_add_epi64(self.v[3], _mm_set_epi64x(0, 1)),
96                _ => unreachable!(),
97            };
98
99            let block_ptr = block.as_mut_ptr().cast::<__m128i>();
100            for i in 0..4 {
101                _mm_storeu_si128(block_ptr.add(i), res[0][i]);
102            }
103        }
104    }
105
106    #[inline(always)]
107    fn gen_par_ks_blocks(&mut self, blocks: &mut cipher::ParBlocks<Self>) {
108        unsafe {
109            let res = rounds::<R, V>(&self.v);
110            self.v[3] = match size_of::<V::Counter>() {
111                4 => _mm_add_epi32(self.v[3], _mm_set_epi32(0, 0, 0, PAR_BLOCKS as i32)),
112                8 => _mm_add_epi64(self.v[3], _mm_set_epi64x(0, PAR_BLOCKS as i64)),
113                _ => unreachable!(),
114            };
115
116            let blocks_ptr = blocks.as_mut_ptr().cast::<__m128i>();
117            for block in 0..PAR_BLOCKS {
118                for i in 0..4 {
119                    _mm_storeu_si128(blocks_ptr.add(i + block * PAR_BLOCKS), res[block][i]);
120                }
121            }
122        }
123    }
124}
125
126#[cfg(feature = "rng")]
127impl<R: Rounds, V: Variant> Backend<R, V> {
128    #[inline(always)]
129    fn gen_ks_blocks(&mut self, block: &mut [u32; 64]) {
130        const _: () = assert!(4 * PAR_BLOCKS * size_of::<__m128i>() == size_of::<[u32; 64]>());
131        unsafe {
132            let res = rounds::<R, V>(&self.v);
133            self.v[3] = _mm_add_epi64(self.v[3], _mm_set_epi64x(0, PAR_BLOCKS as i64));
134
135            let blocks_ptr = block.as_mut_ptr().cast::<__m128i>();
136            for block in 0..PAR_BLOCKS {
137                for i in 0..4 {
138                    _mm_storeu_si128(blocks_ptr.add(i + block * PAR_BLOCKS), res[block][i]);
139                }
140            }
141        }
142    }
143}
144
145#[inline]
146#[target_feature(enable = "sse2")]
147unsafe fn rounds<R: Rounds, V: Variant>(v: &[__m128i; 4]) -> [[__m128i; 4]; PAR_BLOCKS] {
148    let mut res = [*v; 4];
149    for block in 1..PAR_BLOCKS {
150        res[block][3] = match size_of::<V::Counter>() {
151            4 => _mm_add_epi32(res[block][3], _mm_set_epi32(0, 0, 0, block as i32)),
152            8 => _mm_add_epi64(res[block][3], _mm_set_epi64x(0, block as i64)),
153            _ => unreachable!(),
154        }
155    }
156
157    for _ in 0..R::COUNT {
158        double_quarter_round(&mut res);
159    }
160
161    for block in 0..PAR_BLOCKS {
162        for i in 0..3 {
163            res[block][i] = _mm_add_epi32(res[block][i], v[i]);
164        }
165        let ctr = match size_of::<V::Counter>() {
166            4 => _mm_add_epi32(v[3], _mm_set_epi32(0, 0, 0, block as i32)),
167            8 => _mm_add_epi64(v[3], _mm_set_epi64x(0, block as i64)),
168            _ => unreachable!(),
169        };
170        res[block][3] = _mm_add_epi32(res[block][3], ctr);
171    }
172
173    res
174}
175
176#[inline]
177#[target_feature(enable = "sse2")]
178unsafe fn double_quarter_round(v: &mut [[__m128i; 4]; PAR_BLOCKS]) {
179    add_xor_rot(v);
180    rows_to_cols(v);
181    add_xor_rot(v);
182    cols_to_rows(v);
183}
184
185/// The goal of this function is to transform the state words from:
186/// ```text
187/// [a0, a1, a2, a3]    [ 0,  1,  2,  3]
188/// [b0, b1, b2, b3] == [ 4,  5,  6,  7]
189/// [c0, c1, c2, c3]    [ 8,  9, 10, 11]
190/// [d0, d1, d2, d3]    [12, 13, 14, 15]
191/// ```
192///
193/// to:
194/// ```text
195/// [a0, a1, a2, a3]    [ 0,  1,  2,  3]
196/// [b1, b2, b3, b0] == [ 5,  6,  7,  4]
197/// [c2, c3, c0, c1]    [10, 11,  8,  9]
198/// [d3, d0, d1, d2]    [15, 12, 13, 14]
199/// ```
200///
201/// so that we can apply [`add_xor_rot`] to the resulting columns, and have it compute the
202/// "diagonal rounds" (as defined in RFC 7539) in parallel. In practice, this shuffle is
203/// non-optimal: the last state word to be altered in `add_xor_rot` is `b`, so the shuffle
204/// blocks on the result of `b` being calculated.
205///
206/// We can optimize this by observing that the four quarter rounds in `add_xor_rot` are
207/// data-independent: they only access a single column of the state, and thus the order of
208/// the columns does not matter. We therefore instead shuffle the other three state words,
209/// to obtain the following equivalent layout:
210/// ```text
211/// [a3, a0, a1, a2]    [ 3,  0,  1,  2]
212/// [b0, b1, b2, b3] == [ 4,  5,  6,  7]
213/// [c1, c2, c3, c0]    [ 9, 10, 11,  8]
214/// [d2, d3, d0, d1]    [14, 15, 12, 13]
215/// ```
216///
217/// See https://github.com/sneves/blake2-avx2/pull/4 for additional details. The earliest
218/// known occurrence of this optimization is in floodyberry's SSE4 ChaCha code from 2014:
219/// - https://github.com/floodyberry/chacha-opt/blob/0ab65cb99f5016633b652edebaf3691ceb4ff753/chacha_blocks_ssse3-64.S#L639-L643
220#[inline]
221#[target_feature(enable = "sse2")]
222unsafe fn rows_to_cols(blocks: &mut [[__m128i; 4]; PAR_BLOCKS]) {
223    for [a, _, c, d] in blocks.iter_mut() {
224        // c >>>= 32; d >>>= 64; a >>>= 96;
225        *c = _mm_shuffle_epi32(*c, 0b_00_11_10_01); // _MM_SHUFFLE(0, 3, 2, 1)
226        *d = _mm_shuffle_epi32(*d, 0b_01_00_11_10); // _MM_SHUFFLE(1, 0, 3, 2)
227        *a = _mm_shuffle_epi32(*a, 0b_10_01_00_11); // _MM_SHUFFLE(2, 1, 0, 3)
228    }
229}
230
231/// The goal of this function is to transform the state words from:
232/// ```text
233/// [a3, a0, a1, a2]    [ 3,  0,  1,  2]
234/// [b0, b1, b2, b3] == [ 4,  5,  6,  7]
235/// [c1, c2, c3, c0]    [ 9, 10, 11,  8]
236/// [d2, d3, d0, d1]    [14, 15, 12, 13]
237/// ```
238///
239/// to:
240/// ```text
241/// [a0, a1, a2, a3]    [ 0,  1,  2,  3]
242/// [b0, b1, b2, b3] == [ 4,  5,  6,  7]
243/// [c0, c1, c2, c3]    [ 8,  9, 10, 11]
244/// [d0, d1, d2, d3]    [12, 13, 14, 15]
245/// ```
246///
247/// reversing the transformation of [`rows_to_cols`].
248#[inline]
249#[target_feature(enable = "sse2")]
250unsafe fn cols_to_rows(blocks: &mut [[__m128i; 4]; PAR_BLOCKS]) {
251    for [a, _, c, d] in blocks.iter_mut() {
252        // c <<<= 32; d <<<= 64; a <<<= 96;
253        *c = _mm_shuffle_epi32(*c, 0b_10_01_00_11); // _MM_SHUFFLE(2, 1, 0, 3)
254        *d = _mm_shuffle_epi32(*d, 0b_01_00_11_10); // _MM_SHUFFLE(1, 0, 3, 2)
255        *a = _mm_shuffle_epi32(*a, 0b_00_11_10_01); // _MM_SHUFFLE(0, 3, 2, 1)
256    }
257}
258
259#[inline]
260#[target_feature(enable = "sse2")]
261unsafe fn add_xor_rot(blocks: &mut [[__m128i; 4]; PAR_BLOCKS]) {
262    for [a, b, c, d] in blocks.iter_mut() {
263        // a += b; d ^= a; d <<<= (16, 16, 16, 16);
264        *a = _mm_add_epi32(*a, *b);
265        *d = _mm_xor_si128(*d, *a);
266        *d = _mm_xor_si128(_mm_slli_epi32(*d, 16), _mm_srli_epi32(*d, 16));
267
268        // c += d; b ^= c; b <<<= (12, 12, 12, 12);
269        *c = _mm_add_epi32(*c, *d);
270        *b = _mm_xor_si128(*b, *c);
271        *b = _mm_xor_si128(_mm_slli_epi32(*b, 12), _mm_srli_epi32(*b, 20));
272
273        // a += b; d ^= a; d <<<= (8, 8, 8, 8);
274        *a = _mm_add_epi32(*a, *b);
275        *d = _mm_xor_si128(*d, *a);
276        *d = _mm_xor_si128(_mm_slli_epi32(*d, 8), _mm_srli_epi32(*d, 24));
277
278        // c += d; b ^= c; b <<<= (7, 7, 7, 7);
279        *c = _mm_add_epi32(*c, *d);
280        *b = _mm_xor_si128(*b, *c);
281        *b = _mm_xor_si128(_mm_slli_epi32(*b, 7), _mm_srli_epi32(*b, 25));
282    }
283}
chacha20/backends/sse2.rs

chacha20/backends/
sse2.rs
