Merge pull request #10981 from SparkiDev/mlkem_avx512

ML-KEM assembly: AVX512F and AVX512VBMI

.wolfssl_known_macro_extras

@@ -401,6 +401,8 @@ NO_ARDUINO_DEFAULT
 NO_ASM
 NO_ASN_OLD_TYPE_NAMES
 NO_AVX512_SUPPORT
+NO_AVX512_VBMI2_SUPPORT
+NO_AVX512_VBMI_SUPPORT
 NO_CAMELLIA_CBC
 NO_CERT
 NO_CERT_IN_TICKET
@@ -891,6 +893,7 @@ WOLFSSL_MANUALLY_SELECT_DEVICE_CONFIG
 WOLFSSL_MDK5
 WOLFSSL_MEM_FAIL_COUNT
 WOLFSSL_MICROCHIP_AESGCM
+WOLFSSL_MLKEM_ASM_TEST
 WOLFSSL_MLKEM_INVNTT_UNROLL
 WOLFSSL_MLKEM_NO_MALLOC
 WOLFSSL_MLKEM_NTT_UNROLL

wolfcrypt/src/cpuid.c

@@ -149,6 +149,12 @@
             if (cpuid_flag(7, 0, EBX, 29)) { new_cpuid_flags |= CPUID_SHA   ; }
             if (cpuid_flag(7, 0, ECX,  9)) { new_cpuid_flags |= CPUID_VAES  ; }
             if (cpuid_flag(7, 0, EBX, 16)) { new_cpuid_flags |= CPUID_AVX512; }
+            if (cpuid_flag(7, 0, ECX,  1)) {
+                new_cpuid_flags |= CPUID_AVX512_VBMI;
+            }
+            if (cpuid_flag(7, 0, ECX,  6)) {
+                new_cpuid_flags |= CPUID_AVX512_VBMI2;
+            }
             if (cpuid_is_intel())          { new_cpuid_flags |= CPUID_INTEL ; }
             if (cpuid_is_amd())            { new_cpuid_flags |= CPUID_AMD   ; }
             (void)wolfSSL_Atomic_Uint_CompareExchange

wolfcrypt/src/wc_mlkem.c

@@ -461,7 +461,7 @@ int wc_MlKemKey_Init(MlKemKey* key, int type, void* heap, int devId)
     #endif
     #endif
 
-        /* Zero out all data. */
+        /* Zero out the PRF object. */
         XMEMSET(&key->prf, 0, sizeof(key->prf));
 
         /* Initialize the hash algorithm object. */
@@ -1135,7 +1135,7 @@ int wc_MlKemKey_SharedSecretSize(MlKemKey* key, word32* len)
 
 #if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
     !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
-/* Encapsulate data and derive secret.
+/* Encrypt a message to cipher text with the encryption key.
  *
  * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE, m, r)
  * Uses the encryption key to encrypt a plaintext message using the randomness

wolfssl/wolfcrypt/cpuid.h

@@ -77,6 +77,8 @@ typedef word32 cpuid_flags_t;
     /* CPU vendor is AuthenticAMD.  Detected and exposed via IS_CPU_AMD() for
      * future vendor-specific dispatch; no current caller relies on it. */
     #define CPUID_AMD    0x2000
+    #define CPUID_AVX512_VBMI 0x4000  /* AVX-512 Vector Byte Manipulation */
+    #define CPUID_AVX512_VBMI2 0x8000 /* AVX-512 VBMI2 (vpcompressw etc.) */
 
     #define IS_INTEL_AVX1(f)    (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_AVX1)
     #define IS_INTEL_AVX2(f)    (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_AVX2)
@@ -90,6 +92,10 @@ typedef word32 cpuid_flags_t;
     #define IS_INTEL_SHA(f)     (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_SHA)
     #define IS_INTEL_VAES(f)    (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_VAES)
     #define IS_INTEL_AVX512(f)  (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_AVX512)
+    #define IS_INTEL_AVX512_VBMI(f) \
+        (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_AVX512_VBMI)
+    #define IS_INTEL_AVX512_VBMI2(f) \
+        (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_AVX512_VBMI2)
     #define IS_CPU_INTEL(f)     (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_INTEL)
     #define IS_CPU_AMD(f)       (WOLFSSL_ATOMIC_COERCE_UINT(f) & CPUID_AMD)
 

wolfssl/wolfcrypt/sha3.h

@@ -401,6 +401,11 @@ WOLFSSL_LOCAL void BlockSha3(word64 *s);
         word32 len);
     WOLFSSL_LOCAL void sha3_blocksx8_out_avx512(word64* s, byte* out,
         word32 len);
+#ifndef NO_AVX512_SUPPORT
+    WOLFSSL_LOCAL void sha3_blocksx8_avx512(word64* s);
+    WOLFSSL_LOCAL void sha3_128_blocksx8_seed_avx512(word64* s, byte* seed);
+    WOLFSSL_LOCAL void sha3_256_blocksx8_seed_avx512(word64* s, byte* seed);
+#endif
 
     WOLFSSL_LOCAL void sha3_128_blocksx4_seed_avx2(word64* s, byte* seed);
     WOLFSSL_LOCAL void sha3_256_blocksx4_seed_avx2(word64* s, byte* seed);

wolfssl/wolfcrypt/wc_mlkem.h

@@ -387,7 +387,7 @@ struct MlKemKey {
     int  labelLen;
 #endif
 
-    /* A pseudo-random function object. */
+    /* A hash function object. */
     MLKEM_HASH_T hash;
     /* A pseudo-random function object. */
     MLKEM_PRF_T prf;
@@ -585,6 +585,26 @@ WOLFSSL_LOCAL
 int mlkem_check_reduced(const sword16* p, int k);
 
 #ifdef USE_INTEL_SPEEDUP
+/* AVX512 assembly for ML-KEM is built (and dispatched at runtime on capable
+ * CPUs) whenever the Intel speedups are enabled and AVX512 is not opted out.
+ * Matches the HAVE_INTEL_AVX512 guard around the generated assembly. */
+#ifndef NO_AVX512_SUPPORT
+    #define WOLFSSL_MLKEM_HAVE_INTEL_AVX512
+    /* AVX512VBMI (vpermb) functions are built and dispatched at runtime when
+     * CPUID reports VBMI. Opt out with NO_AVX512_VBMI_SUPPORT if the assembler
+     * cannot emit the vpermb-based *_vbmi routines. VBMI2 (vpcompressw) is a
+     * separate opt-out (NO_AVX512_VBMI2_SUPPORT), handled below. */
+    #ifndef NO_AVX512_VBMI_SUPPORT
+        #define WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+    #endif
+    /* AVX512VBMI2 (vpcompressw) is used only by the rejection samplers; every
+     * other AVX512 routine is plain AVX512F/BW or VBMI. Opt out with
+     * NO_AVX512_VBMI2_SUPPORT when the assembler cannot emit VBMI2 - the
+     * AVX512F/BW rej variants (vpcompressd) are then used instead. */
+    #ifndef NO_AVX512_VBMI2_SUPPORT
+        #define WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI2
+    #endif
+#endif
 WOLFSSL_LOCAL
 void mlkem_keygen_avx2(sword16* priv, sword16* pub, sword16* e,
     const sword16* a, int kp);
@@ -651,6 +671,131 @@ void mlkem_decompress_5_avx2(sword16* p, const byte* r);
 
 WOLFSSL_LOCAL
 int mlkem_cmp_avx2(const byte* a, const byte* b, int sz);
+
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512
+WOLFSSL_LOCAL
+void mlkem_keygen_avx512(sword16* priv, sword16* pub, sword16* e,
+    const sword16* a, int kp);
+WOLFSSL_LOCAL
+void mlkem_encapsulate_avx512(const sword16* pub, sword16* bp, sword16* v,
+    const sword16* at, sword16* sp, const sword16* ep, const sword16* epp,
+    const sword16* m, int kp);
+WOLFSSL_LOCAL
+void mlkem_decapsulate_avx512(const sword16* priv, sword16* mp, sword16* bp,
+    const sword16* v, int kp);
+WOLFSSL_LOCAL
+void mlkem_csubq_avx512(sword16* p);
+WOLFSSL_LOCAL
+int mlkem_cmp_avx512(const byte* a, const byte* b, int sz);
+WOLFSSL_LOCAL
+void mlkem_from_bytes_avx512(sword16* p, const byte* b);
+WOLFSSL_LOCAL
+void mlkem_to_bytes_avx512(byte* b, sword16* p);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+void mlkem_from_bytes_avx512_vbmi(sword16* p, const byte* b);
+WOLFSSL_LOCAL
+void mlkem_to_bytes_avx512_vbmi(byte* b, sword16* p);
+#endif
+WOLFSSL_LOCAL
+void mlkem_from_msg_avx512(sword16* p, const byte* msg);
+WOLFSSL_LOCAL
+void mlkem_to_msg_avx512(byte* msg, sword16* p);
+WOLFSSL_LOCAL
+void mlkem_compress_10_avx512(byte* r, const sword16* p, int n);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+void mlkem_compress_10_avx512_vbmi(byte* r, const sword16* p, int n);
+#endif
+WOLFSSL_LOCAL
+void mlkem_decompress_10_avx512(sword16* p, const byte* r, int n);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+void mlkem_decompress_10_avx512_vbmi(sword16* p, const byte* r, int n);
+#endif
+WOLFSSL_LOCAL
+void mlkem_compress_11_avx512(byte* r, const sword16* p, int n);
+WOLFSSL_LOCAL
+void mlkem_decompress_11_avx512(sword16* p, const byte* r, int n);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+void mlkem_decompress_11_avx512_vbmi(sword16* p, const byte* r, int n);
+#endif
+WOLFSSL_LOCAL
+void mlkem_compress_4_avx512(byte* r, const sword16* p);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+void mlkem_compress_4_avx512_vbmi(byte* r, const sword16* p);
+#endif
+WOLFSSL_LOCAL
+void mlkem_decompress_4_avx512(sword16* p, const byte* r);
+WOLFSSL_LOCAL
+void mlkem_compress_5_avx512(byte* r, const sword16* p);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+void mlkem_compress_5_avx512_vbmi(byte* r, const sword16* p);
+#endif
+WOLFSSL_LOCAL
+void mlkem_decompress_5_avx512(sword16* p, const byte* r);
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_n_avx512(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_avx512(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI2
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_n_avx512_vbmi2(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_avx512_vbmi2(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+#endif
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_n_avx512_vbmi(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_avx512_vbmi(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+#ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512_VBMI2
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_n_avx512_vbmi_vbmi2(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+WOLFSSL_LOCAL
+unsigned int mlkem_rej_uniform_avx512_vbmi_vbmi2(sword16* p, unsigned int len,
+    const byte* r, unsigned int rLen);
+#endif
+#endif
+WOLFSSL_LOCAL
+void mlkem_redistribute_8_rand_avx512(const word64* s, byte* r0, byte* r1,
+    byte* r2, byte* r3);
+WOLFSSL_LOCAL
+void mlkem_redistribute_16_rand_avx512(const word64* s, byte* r0, byte* r1,
+    byte* r2, byte* r3);
+WOLFSSL_LOCAL
+void mlkem_redistribute_17_rand_avx512(const word64* s, byte* r0, byte* r1,
+    byte* r2, byte* r3);
+WOLFSSL_LOCAL
+void mlkem_redistribute_21_rand_avx512(const word64* s, byte* r0, byte* r1,
+    byte* r2, byte* r3);
+WOLFSSL_LOCAL
+void mlkem_redistribute_8_rand_x8_avx512(const word64* s, byte* out,
+    word32 stride);
+WOLFSSL_LOCAL
+void mlkem_redistribute_16_rand_x8_avx512(const word64* s, byte* out,
+    word32 stride);
+WOLFSSL_LOCAL
+void mlkem_redistribute_17_rand_x8_avx512(const word64* s, byte* out,
+    word32 stride);
+WOLFSSL_LOCAL
+void mlkem_redistribute_21_rand_x8_avx512(const word64* s, byte* out,
+    word32 stride);
+WOLFSSL_LOCAL
+void mlkem_cbd_eta2_avx512(sword16* p, const byte* r);
+WOLFSSL_LOCAL
+void mlkem_cbd_eta3_avx512(sword16* p, const byte* r);
+#endif /* WOLFSSL_MLKEM_HAVE_INTEL_AVX512 */
 #elif defined(__aarch64__) && defined(WOLFSSL_ARMASM)
 WOLFSSL_LOCAL void mlkem_ntt(sword16* r);
 WOLFSSL_LOCAL void mlkem_invntt(sword16* r);