@@ -117,33 +117,25 @@ pub fn chacha20_state(key: &[u8; 32], counter: u32, nonce: &[u8; 12]) -> [u32; 1
117117/// Fill `out` with successive ChaCha20 keystream blocks: block `b` uses block
118118/// counter `counter.wrapping_add(b)`, key and nonce fixed (RFC 8439 CTR mode).
119119///
120- /// This is the accelerated entry point the `encryption` AEAD hot path draws its
121- /// keystream from. It runs the **AVX-512 backend in 16-block strides** when
122- /// `avx512f` is detected at runtime, and the scalar [`chacha20_block`] reference
123- /// for the ragged tail and on every non-x86 / non-AVX-512 target. Every backend
124- /// is a drop-in for the scalar reference — see the module KAT and the
125- /// `chacha20_keystream_avx512_parity` test that pins byte-for-byte equality .
120+ /// Accelerated entry point for the `encryption` AEAD keystream. Compile-time
121+ /// tier dispatch ( the `ndarray::simd` model — no `is_x86_feature_detected!`):
122+ /// the AVX-512 16-block backend when built `+avx512f`, the wasm128 4-block
123+ /// backend when built `+simd128`, else the scalar [`chacha20_block`] reference,
124+ /// which also fills the ragged tail. Every backend is byte-parity-pinned to the
125+ /// scalar reference (`chacha20_keystream_dispatch_parity_scalar`) .
126126///
127- /// **Constant-time:** all backends are straight-line ARX (add / xor / rotate),
128- /// no secret-dependent control flow or memory indexing. Backend equivalence is
129- /// pinned by `chacha20_keystream_dispatch_parity_scalar` (byte-for-byte vs the
130- /// scalar reference across two AVX-512 strides + a scalar tail).
127+ /// Constant-time: straight-line ARX (add / xor / rotate), no secret-dependent
128+ /// control flow or memory indexing.
131129pub fn chacha20_keystream ( key : & [ u8 ; 32 ] , counter : u32 , nonce : & [ u8 ; 12 ] , out : & mut [ [ u8 ; 64 ] ] ) {
132130 let mut i = 0usize ;
133131
134- #[ cfg( target_arch = "x86_64" ) ]
132+ #[ cfg( all ( target_arch = "x86_64" , target_feature = "avx512f" ) ) ]
135133 {
136- if std:: is_x86_feature_detected!( "avx512f" ) {
137- while i + 16 <= out. len ( ) {
138- let state = chacha20_state ( key, counter. wrapping_add ( i as u32 ) , nonce) ;
139- // SAFETY: this arm is only reached after `is_x86_feature_detected!("avx512f")`
140- // returned true, so every AVX-512F intrinsic in `chacha20_block16_avx512`
141- // is supported on this CPU. The function reads only its `&state` argument
142- // and returns an owned array — no raw pointers escape, no aliasing.
143- let blocks = unsafe { chacha20_block16_avx512 ( & state) } ;
144- out[ i..i + 16 ] . copy_from_slice ( & blocks) ;
145- i += 16 ;
146- }
134+ while i + 16 <= out. len ( ) {
135+ let state = chacha20_state ( key, counter. wrapping_add ( i as u32 ) , nonce) ;
136+ let blocks = chacha20_block16_avx512 ( & state) ;
137+ out[ i..i + 16 ] . copy_from_slice ( & blocks) ;
138+ i += 16 ;
147139 }
148140 }
149141
@@ -167,82 +159,75 @@ pub fn chacha20_keystream(key: &[u8; 32], counter: u32, nonce: &[u8; 12], out: &
167159 }
168160}
169161
170- /// AVX-512 backend: compute 16 consecutive ChaCha20 keystream blocks in parallel
171- /// (block `l` uses `state[12].wrapping_add(l)` as its counter, all other words
172- /// shared). Word-sliced ("vertical") layout — lane `l` of working vector `w`
173- /// holds word `w` of block `l` — so every round is pure SIMD add/xor/rotate with
174- /// no cross-lane shuffles; the only transpose is the final little-endian
175- /// serialization.
176- ///
177- /// # Safety
178- /// Caller must ensure the CPU supports AVX-512F (guarded by
179- /// `is_x86_feature_detected!("avx512f")` at the single call site in
180- /// [`chacha20_keystream`]).
181- #[ cfg( target_arch = "x86_64" ) ]
182- #[ target_feature( enable = "avx512f" ) ]
183- unsafe fn chacha20_block16_avx512 ( state : & [ u32 ; 16 ] ) -> [ [ u8 ; 64 ] ; 16 ] {
162+ /// AVX-512 backend: 16 ChaCha20 blocks in parallel, word-sliced (lane `l` = block
163+ /// `l`), so each round is pure SIMD add/xor/rotate and only the final LE serialize
164+ /// transposes. cfg-gated to an `+avx512f` build → the intrinsics are statically
165+ /// available (no `#[target_feature]`, no runtime detect).
166+ #[ cfg( all( target_arch = "x86_64" , target_feature = "avx512f" ) ) ]
167+ fn chacha20_block16_avx512 ( state : & [ u32 ; 16 ] ) -> [ [ u8 ; 64 ] ; 16 ] {
184168 use core:: arch:: x86_64:: * ;
169+ // SAFETY: avx512f is statically enabled (cfg gate), so every AVX-512F
170+ // intrinsic below is available; all lane indices are fixed constants.
171+ unsafe {
172+ // One AVX-512 vertical quarter-round: same word indices as the scalar
173+ // `quarter_round`, applied across all 16 blocks at once. Pure ARX.
174+ macro_rules! qr512 {
175+ ( $v: expr, $a: expr, $b: expr, $c: expr, $d: expr) => { {
176+ $v[ $a] = _mm512_add_epi32( $v[ $a] , $v[ $b] ) ;
177+ $v[ $d] = _mm512_rol_epi32:: <16 >( _mm512_xor_si512( $v[ $d] , $v[ $a] ) ) ;
178+ $v[ $c] = _mm512_add_epi32( $v[ $c] , $v[ $d] ) ;
179+ $v[ $b] = _mm512_rol_epi32:: <12 >( _mm512_xor_si512( $v[ $b] , $v[ $c] ) ) ;
180+ $v[ $a] = _mm512_add_epi32( $v[ $a] , $v[ $b] ) ;
181+ $v[ $d] = _mm512_rol_epi32:: <8 >( _mm512_xor_si512( $v[ $d] , $v[ $a] ) ) ;
182+ $v[ $c] = _mm512_add_epi32( $v[ $c] , $v[ $d] ) ;
183+ $v[ $b] = _mm512_rol_epi32:: <7 >( _mm512_xor_si512( $v[ $b] , $v[ $c] ) ) ;
184+ } } ;
185+ }
185186
186- // One AVX-512 vertical quarter-round: same word indices as the scalar
187- // `quarter_round`, applied across all 16 blocks at once. Pure ARX.
188- macro_rules! qr512 {
189- ( $v: expr, $a: expr, $b: expr, $c: expr, $d: expr) => { {
190- $v[ $a] = _mm512_add_epi32( $v[ $a] , $v[ $b] ) ;
191- $v[ $d] = _mm512_rol_epi32:: <16 >( _mm512_xor_si512( $v[ $d] , $v[ $a] ) ) ;
192- $v[ $c] = _mm512_add_epi32( $v[ $c] , $v[ $d] ) ;
193- $v[ $b] = _mm512_rol_epi32:: <12 >( _mm512_xor_si512( $v[ $b] , $v[ $c] ) ) ;
194- $v[ $a] = _mm512_add_epi32( $v[ $a] , $v[ $b] ) ;
195- $v[ $d] = _mm512_rol_epi32:: <8 >( _mm512_xor_si512( $v[ $d] , $v[ $a] ) ) ;
196- $v[ $c] = _mm512_add_epi32( $v[ $c] , $v[ $d] ) ;
197- $v[ $b] = _mm512_rol_epi32:: <7 >( _mm512_xor_si512( $v[ $b] , $v[ $c] ) ) ;
198- } } ;
199- }
200-
201- // Build the 16 initial working vectors. Every word is broadcast identically
202- // across the 16 lanes EXCEPT the counter (word 12), which gets lane-index
203- // 0..=15 added (u32 wrapping) — the 16 consecutive block counters.
204- let lane_index = _mm512_setr_epi32 ( 0 , 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 , 13 , 14 , 15 ) ;
205- let mut orig = [ _mm512_setzero_si512 ( ) ; 16 ] ;
206- for ( w, o) in orig. iter_mut ( ) . enumerate ( ) {
207- * o = _mm512_set1_epi32 ( state[ w] as i32 ) ;
208- }
209- orig[ 12 ] = _mm512_add_epi32 ( _mm512_set1_epi32 ( state[ 12 ] as i32 ) , lane_index) ;
187+ // Build the 16 initial working vectors. Every word is broadcast identically
188+ // across the 16 lanes EXCEPT the counter (word 12), which gets lane-index
189+ // 0..=15 added (u32 wrapping) — the 16 consecutive block counters.
190+ let lane_index = _mm512_setr_epi32 ( 0 , 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 , 13 , 14 , 15 ) ;
191+ let mut orig = [ _mm512_setzero_si512 ( ) ; 16 ] ;
192+ for ( w, o) in orig. iter_mut ( ) . enumerate ( ) {
193+ * o = _mm512_set1_epi32 ( state[ w] as i32 ) ;
194+ }
195+ orig[ 12 ] = _mm512_add_epi32 ( _mm512_set1_epi32 ( state[ 12 ] as i32 ) , lane_index) ;
210196
211- let mut v = orig;
212- // 10 double-rounds = 20 rounds, identical schedule to the scalar reference.
213- for _ in 0 ..10 {
214- qr512 ! ( v, 0 , 4 , 8 , 12 ) ;
215- qr512 ! ( v, 1 , 5 , 9 , 13 ) ;
216- qr512 ! ( v, 2 , 6 , 10 , 14 ) ;
217- qr512 ! ( v, 3 , 7 , 11 , 15 ) ;
218- qr512 ! ( v, 0 , 5 , 10 , 15 ) ;
219- qr512 ! ( v, 1 , 6 , 11 , 12 ) ;
220- qr512 ! ( v, 2 , 7 , 8 , 13 ) ;
221- qr512 ! ( v, 3 , 4 , 9 , 14 ) ;
222- }
197+ let mut v = orig;
198+ // 10 double-rounds = 20 rounds, identical schedule to the scalar reference.
199+ for _ in 0 ..10 {
200+ qr512 ! ( v, 0 , 4 , 8 , 12 ) ;
201+ qr512 ! ( v, 1 , 5 , 9 , 13 ) ;
202+ qr512 ! ( v, 2 , 6 , 10 , 14 ) ;
203+ qr512 ! ( v, 3 , 7 , 11 , 15 ) ;
204+ qr512 ! ( v, 0 , 5 , 10 , 15 ) ;
205+ qr512 ! ( v, 1 , 6 , 11 , 12 ) ;
206+ qr512 ! ( v, 2 , 7 , 8 , 13 ) ;
207+ qr512 ! ( v, 3 , 4 , 9 , 14 ) ;
208+ }
223209
224- // Add the original input state back (RFC 8439 §2.3.1), then de-vectorize:
225- // `words[w][l]` = final word `w` of block `l`.
226- let mut words = [ [ 0u32 ; 16 ] ; 16 ] ;
227- for ( w, dst) in words. iter_mut ( ) . enumerate ( ) {
228- let summed = _mm512_add_epi32 ( v[ w] , orig[ w] ) ;
229- let mut tmp = [ 0i32 ; 16 ] ;
230- // SAFETY: `tmp` is a 16-element (64-byte) i32 array; `_mm512_storeu_si512`
231- // writes exactly one 512-bit (64-byte) vector, unaligned, in bounds.
232- _mm512_storeu_si512 ( tmp. as_mut_ptr ( ) . cast ( ) , summed) ;
233- for ( l, slot) in dst. iter_mut ( ) . enumerate ( ) {
234- * slot = tmp[ l] as u32 ;
210+ // Add the original input state back (RFC 8439 §2.3.1), then de-vectorize:
211+ // `words[w][l]` = final word `w` of block `l`.
212+ let mut words = [ [ 0u32 ; 16 ] ; 16 ] ;
213+ for ( w, dst) in words. iter_mut ( ) . enumerate ( ) {
214+ let summed = _mm512_add_epi32 ( v[ w] , orig[ w] ) ;
215+ let mut tmp = [ 0i32 ; 16 ] ;
216+ _mm512_storeu_si512 ( tmp. as_mut_ptr ( ) . cast ( ) , summed) ;
217+ for ( l, slot) in dst. iter_mut ( ) . enumerate ( ) {
218+ * slot = tmp[ l] as u32 ;
219+ }
235220 }
236- }
237221
238- // Serialize each block little-endian: block `l`, word `w` → bytes 4w..4w+4.
239- let mut out = [ [ 0u8 ; 64 ] ; 16 ] ;
240- for ( l, block) in out. iter_mut ( ) . enumerate ( ) {
241- for w in 0 ..16 {
242- block[ w * 4 ..w * 4 + 4 ] . copy_from_slice ( & words[ w] [ l] . to_le_bytes ( ) ) ;
222+ // Serialize each block little-endian: block `l`, word `w` → bytes 4w..4w+4.
223+ let mut out = [ [ 0u8 ; 64 ] ; 16 ] ;
224+ for ( l, block) in out. iter_mut ( ) . enumerate ( ) {
225+ for w in 0 ..16 {
226+ block[ w * 4 ..w * 4 + 4 ] . copy_from_slice ( & words[ w] [ l] . to_le_bytes ( ) ) ;
227+ }
243228 }
229+ out
244230 }
245- out
246231}
247232
248233/// wasm128 (browser) backend: compute 4 consecutive ChaCha20 keystream blocks in
0 commit comments