ML-DSA: AVX2 target feature

libcrux-ml-dsa/src/hash_functions.rs

@@ -299,12 +299,14 @@ pub(crate) mod simd256 {
 
     impl shake128::XofX4 for Shake128x4 {
         /// Init the state and absorb 4 blocks in parallel.
+        #[inline(always)]
         fn init_absorb(input0: &[u8], input1: &[u8], input2: &[u8], input3: &[u8]) -> Self {
             let mut state = x4::incremental::init();
             x4::incremental::shake128_absorb_final(&mut state, &input0, &input1, &input2, &input3);
             Self { state }
         }
 
+        #[inline(always)]
         fn squeeze_first_five_blocks(
             &mut self,
             out0: &mut [u8; shake128::FIVE_BLOCKS_SIZE],
@@ -321,6 +323,7 @@ pub(crate) mod simd256 {
             );
         }
 
+        #[inline(always)]
         fn squeeze_next_block(
             &mut self,
         ) -> (
@@ -353,23 +356,27 @@ pub(crate) mod simd256 {
         state: portable::KeccakState,
     }
     impl shake256::Xof for Shake256 {
+        #[inline(always)]
         fn shake256<const OUTPUT_LENGTH: usize>(input: &[u8], out: &mut [u8; OUTPUT_LENGTH]) {
             portable::shake256(out, input);
         }
 
+        #[inline(always)]
         fn init_absorb(input: &[u8]) -> Self {
             let mut state = portable::incremental::shake256_init();
             portable::incremental::shake256_absorb_final(&mut state, input);
 
             Self { state }
         }
 
+        #[inline(always)]
         fn squeeze_first_block(&mut self) -> [u8; shake256::BLOCK_SIZE] {
             let mut out = [0u8; shake256::BLOCK_SIZE];
             portable::incremental::shake256_squeeze_first_block(&mut self.state, &mut out);
             out
         }
 
+        #[inline(always)]
         fn squeeze_next_block(&mut self) -> [u8; shake256::BLOCK_SIZE] {
             let mut out = [0u8; shake256::BLOCK_SIZE];
             portable::incremental::shake256_squeeze_next_block(&mut self.state, &mut out);
@@ -383,12 +390,14 @@ pub(crate) mod simd256 {
     }
 
     impl shake256::XofX4 for Shake256x4 {
+        #[inline(always)]
         fn init_absorb(input0: &[u8], input1: &[u8], input2: &[u8], input3: &[u8]) -> Self {
             let mut state = x4::incremental::init();
             x4::incremental::shake256_absorb_final(&mut state, &input0, &input1, &input2, &input3);
             Self { state }
         }
 
+        #[inline(always)]
         fn squeeze_first_block(
             &mut self,
         ) -> (
@@ -412,6 +421,7 @@ pub(crate) mod simd256 {
             (out0, out1, out2, out3)
         }
 
+        #[inline(always)]
         fn squeeze_next_block(
             &mut self,
         ) -> (
@@ -435,6 +445,7 @@ pub(crate) mod simd256 {
             (out0, out1, out2, out3)
         }
 
+        #[inline(always)]
         fn shake256<const OUT_LEN: usize>(
             input0: &[u8],
             input1: &[u8],

libcrux-ml-dsa/src/lib.rs

@@ -1,4 +1,5 @@
 #![no_std]
+#![deny(unsafe_code)]
 
 mod arithmetic;
 mod constants;

libcrux-ml-dsa/src/ml_dsa_generic.rs

@@ -36,6 +36,7 @@ pub(crate) struct Signature<
 }
 
 /// Generate a key pair.
+#[inline(always)]
 pub(crate) fn generate_key_pair<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,
@@ -114,6 +115,7 @@ pub enum SigningError {
 }
 
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn sign_pre_hashed<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,
@@ -174,6 +176,7 @@ pub(crate) fn sign_pre_hashed<
 }
 
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn sign<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,
@@ -231,6 +234,7 @@ pub(crate) fn sign<
 /// If no `domain_separation_context` is supplied, it is assumed that
 /// `message` already contains the domain separation.
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn sign_internal<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,
@@ -437,6 +441,7 @@ pub(crate) fn sign_internal<
 /// for details on the domain separation for regular ML-DSA. Line
 /// 23 of Algorithm 4 (and line 18 of Algorithm 5,resp.) describe domain separation for the HashMl-DSA
 /// variant.
+#[inline(always)]
 fn derive_message_representative(
     verification_key_hash: [u8; 64],
     domain_separation_context: Option<DomainSeparationContext>,
@@ -463,6 +468,7 @@ fn derive_message_representative(
 /// If no `domain_separation_context` is supplied, it is assumed that
 /// `message` already contains the domain separation.
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn verify_internal<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,
@@ -563,6 +569,7 @@ pub(crate) fn verify_internal<
 }
 
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn verify<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,
@@ -612,6 +619,7 @@ pub(crate) fn verify<
 }
 
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn verify_pre_hashed<
     SIMDUnit: Operations,
     Shake128X4: shake128::XofX4,

libcrux-ml-dsa/src/ml_dsa_generic/instantiations.rs

@@ -305,12 +305,7 @@ instantiate! {portable,
 
 // AVX2 generic implementation.
 #[cfg(feature = "simd256")]
-instantiate! {avx2,
-    crate::simd::avx2::AVX2SIMDUnit,
-    crate::hash_functions::simd256::Shake128x4,
-    crate::hash_functions::simd256::Shake256,
-    crate::hash_functions::simd256::Shake256x4
-}
+pub mod avx2;
 
 // NEON generic implementation.
 #[cfg(feature = "simd128")]