ML-KEM: Add AArch64 arithmetic backend

Context: The ML-KEM implementation in AWS-LC is imported from
mlkem-native. mlkem-native comes in a "C-only" version, but also
offers AArch64 and x86_64 backends for (a) arithmetic,
and (b) FIPS-202. Currently, only the "C-only" version is
imported into AWS-LC.

This commit adds a custom AArch64 backend to AWS-LC.
The backend is essentially the same as in mlkem-native, but its
assembly sources are taken from s2n-bignum and its headers are
written from scratch. The constant tables used in the backend
are copied from mlkem-native.

Compared to extending the mlkem-native->AWS-LC importer to include
mlkem-native's AArch64 backend, this approach sticks to s2n-bignum
as the sole source of verified assembly. It also provides greater
flexibility in maintaining and adjusting the backend, both the
assembly and the headers. For example, the assembly may be optimized
for Graviton cores in the future, or the dispatch in the metadata
files adjusted; the latter will mostly be relevant as we integrate
x86_64 assembly, for which we aim to use the same methodology.

To avoid a symbol clash with s2n-bignum, the mlkem-native namespace
is changed from `mlkem` to `mlkem_native`.

s2n-bignum is partially re-imported from the development branch
https://github.com/jargh/s2n-bignum-dev/tree/mlkem/, restricting to
the ML-KEM related files.

Signed-off-by: Hanno Becker <[email protected]>

Author: hanno-becker

crypto/fipsmodule/CMakeLists.txt

@@ -291,6 +291,20 @@ if((((ARCH STREQUAL "x86_64") AND NOT MY_ASSEMBLER_IS_TOO_OLD_FOR_512AVX) OR
                 ${S2N_BIGNUM_DIR}/generic/bignum_copy_row_from_table_16.S
                 ${S2N_BIGNUM_DIR}/generic/bignum_copy_row_from_table_32.S
     )
+
+    # ML-KEM core arithmetic
+    list(APPEND BCM_ASM_SOURCES
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_basemul_k2.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_basemul_k3.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_basemul_k4.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_intt.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_mulcache_compute.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_ntt.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_poly_reduce.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_poly_tobytes.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_poly_tomont.S
+                ${S2N_BIGNUM_DIR}/mlkem/mlkem_rej_uniform_VARIABLE_TIME.S)
+
   endif()
 
   if(BORINGSSL_PREFIX)

crypto/fipsmodule/ml_kem/aarch64/README.md

@@ -0,0 +1 @@
+This directory contains an AArch64 arithmetic backend for mlkem-native. The core assembly routines are imported from [s2n-bignum](https://github.com/awslabs/s2n-bignum/).

crypto/fipsmodule/ml_kem/aarch64/constants.c

@@ -0,0 +1,81 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR ISC
+
+#ifndef ML_KEM_AARCH64_BACKEND_H
+#define ML_KEM_AARCH64_BACKEND_H
+
+#include "../mlkem/common.h"
+
+#define MLK_USE_NATIVE_NTT
+#define MLK_USE_NATIVE_INTT
+#define MLK_USE_NATIVE_POLY_REDUCE
+#define MLK_USE_NATIVE_POLY_TOMONT
+#define MLK_USE_NATIVE_POLY_MULCACHE_COMPUTE
+#define MLK_USE_NATIVE_POLYVEC_BASEMUL_ACC_MONTGOMERY_CACHED
+#define MLK_USE_NATIVE_POLY_TOBYTES
+#define MLK_USE_NATIVE_REJ_UNIFORM
+
+extern const int16_t mlk_aarch64_ntt_zetas_layer12345[];
+extern const int16_t mlk_aarch64_ntt_zetas_layer67[];
+extern const int16_t mlk_aarch64_invntt_zetas_layer12345[];
+extern const int16_t mlk_aarch64_invntt_zetas_layer67[];
+extern const uint8_t mlk_rej_uniform_table[];
+extern const int16_t mlk_aarch64_zetas_mulcache_native[];
+extern const int16_t mlk_aarch64_zetas_mulcache_twisted_native[];
+
+#include "s2n-bignum.h"
+
+static MLK_INLINE void mlk_ntt_native(int16_t data[MLKEM_N]) {
+    mlkem_ntt(data, mlk_aarch64_ntt_zetas_layer12345, mlk_aarch64_ntt_zetas_layer67);
+}
+
+static MLK_INLINE void mlk_intt_native(int16_t data[MLKEM_N]) {
+    mlkem_intt(data, mlk_aarch64_invntt_zetas_layer12345, mlk_aarch64_invntt_zetas_layer67);
+}
+
+static MLK_INLINE void mlk_poly_reduce_native(int16_t data[MLKEM_N]) {
+    mlkem_poly_reduce(data);
+}
+
+static MLK_INLINE void mlk_poly_tomont_native(int16_t data[MLKEM_N]) {
+    mlkem_poly_tomont(data);
+}
+
+static MLK_INLINE void mlk_poly_mulcache_compute_native(int16_t x[MLKEM_N / 2], const int16_t y[MLKEM_N]) {
+    mlkem_mulcache_compute(x, y, mlk_aarch64_zetas_mulcache_native,
+	mlk_aarch64_zetas_mulcache_twisted_native);
+}
+
+static MLK_INLINE void mlk_polyvec_basemul_acc_montgomery_cached_k2_native(
+    int16_t r[MLKEM_N], const int16_t a[2 * MLKEM_N],
+    const int16_t b[2 * MLKEM_N], const int16_t b_cache[2 * (MLKEM_N / 2)]) {
+  mlkem_basemul_k2(r, a, b, b_cache);
+}
+
+static MLK_INLINE void mlk_polyvec_basemul_acc_montgomery_cached_k3_native(
+    int16_t r[MLKEM_N], const int16_t a[3 * MLKEM_N],
+    const int16_t b[3 * MLKEM_N], const int16_t b_cache[3 * (MLKEM_N / 2)]) {
+  mlkem_basemul_k3(r, a, b, b_cache);
+}
+
+static MLK_INLINE void mlk_polyvec_basemul_acc_montgomery_cached_k4_native(
+    int16_t r[MLKEM_N], const int16_t a[4 * MLKEM_N],
+    const int16_t b[4 * MLKEM_N], const int16_t b_cache[4 * (MLKEM_N / 2)]) {
+  mlkem_basemul_k4(r, a, b, b_cache);
+}
+
+static MLK_INLINE void mlk_poly_tobytes_native(uint8_t r[MLKEM_POLYBYTES],
+					       const int16_t a[MLKEM_N]) {
+  mlkem_poly_tobytes(r, a);
+}
+
+static MLK_INLINE int mlk_rej_uniform_native(int16_t *r, unsigned len,
+					     const uint8_t *buf,
+					     unsigned buflen) {
+  if (len != MLKEM_N || buflen % 24 != 0) {
+      return -1;
+  }
+  return (int) mlkem_rej_uniform_VARIABLE_TIME(r, buf, buflen, mlk_rej_uniform_table);
+}
+
+#endif /* ML_KEM_AARCH64_BACKEND_H */

crypto/fipsmodule/ml_kem/aarch64/meta.h

@@ -26,6 +26,11 @@
 
 #include "./ml_kem.h"
 
+// AArch64 backend
+#if defined(OPENSSL_AARCH64) && !defined(OPENSSL_NO_ASM)
+#include "aarch64/constants.c"
+#endif
+
 typedef struct {
   uint8_t *buffer;
   size_t *length;
@@ -92,7 +97,7 @@ int ml_kem_512_keypair_deterministic_no_self_test(uint8_t *public_key  /* OUT */
   if (!check_buffer(pkey) || !check_buffer(skey)) {
     return 1;
   }
-  const int res = mlkem512_keypair_derand(pkey.buffer, skey.buffer, seed);
+  const int res = mlkem_native512_keypair_derand(pkey.buffer, skey.buffer, seed);
 #if defined(AWSLC_FIPS)
   /* PCT failure is the only failure condition for key generation. */
   if (res != 0) {
@@ -110,7 +115,7 @@ int ml_kem_512_keypair(uint8_t *public_key /* OUT */,
                        size_t *secret_len  /* IN_OUT */) {
   output_buffer pkey = {public_key, public_len, MLKEM512_PUBLIC_KEY_BYTES};
   output_buffer skey = {secret_key, secret_len, MLKEM512_SECRET_KEY_BYTES};
-  return ml_kem_common_keypair(mlkem512_keypair, pkey, skey);
+  return ml_kem_common_keypair(mlkem_native512_keypair, pkey, skey);
 }
 
 int ml_kem_512_encapsulate_deterministic(uint8_t *ciphertext       /* OUT */,
@@ -131,7 +136,7 @@ int ml_kem_512_encapsulate_deterministic_no_self_test(uint8_t *ciphertext
                                                       const uint8_t *seed       /* IN */) {
   output_buffer ctext = {ciphertext, ciphertext_len, MLKEM512_CIPHERTEXT_BYTES};
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM512_SHARED_SECRET_LEN};
-  return ml_kem_common_encapsulate_deterministic(mlkem512_enc_derand, ctext, ss, public_key, seed);
+  return ml_kem_common_encapsulate_deterministic(mlkem_native512_enc_derand, ctext, ss, public_key, seed);
 }
 
 int ml_kem_512_encapsulate(uint8_t *ciphertext       /* OUT */,
@@ -141,7 +146,7 @@ int ml_kem_512_encapsulate(uint8_t *ciphertext       /* OUT */,
                            const uint8_t *public_key /* IN */) {
   output_buffer ctext = {ciphertext, ciphertext_len, MLKEM512_CIPHERTEXT_BYTES};
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM512_SHARED_SECRET_LEN};
-  return ml_kem_common_encapsulate(mlkem512_enc, ctext, ss, public_key);
+  return ml_kem_common_encapsulate(mlkem_native512_enc, ctext, ss, public_key);
 }
 
 int ml_kem_512_decapsulate(uint8_t *shared_secret    /* OUT */,
@@ -157,7 +162,7 @@ int ml_kem_512_decapsulate_no_self_test(uint8_t *shared_secret    /* OUT */,
                                         const uint8_t *ciphertext /* IN  */,
                                         const uint8_t *secret_key /* IN  */) {
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM512_SHARED_SECRET_LEN};
-  return ml_kem_common_decapsulate(mlkem512_dec, ss, ciphertext, secret_key);
+  return ml_kem_common_decapsulate(mlkem_native512_dec, ss, ciphertext, secret_key);
 }
 
 
@@ -181,7 +186,7 @@ int ml_kem_768_keypair_deterministic_no_self_test(uint8_t *public_key /* OUT */,
   if (!check_buffer(pkey) || !check_buffer(skey)) {
     return 1;
   }
-  const int res = mlkem768_keypair_derand(pkey.buffer, skey.buffer, seed);
+  const int res = mlkem_native768_keypair_derand(pkey.buffer, skey.buffer, seed);
 #if defined(AWSLC_FIPS)
   /* PCT failure is the only failure condition for key generation. */
   if (res != 0) {
@@ -199,7 +204,7 @@ int ml_kem_768_keypair(uint8_t *public_key /* OUT */,
                        size_t *secret_len  /* IN_OUT */) {
   output_buffer pkey = {public_key, public_len, MLKEM768_PUBLIC_KEY_BYTES};
   output_buffer skey = {secret_key, secret_len, MLKEM768_SECRET_KEY_BYTES};
-  return ml_kem_common_keypair(mlkem768_keypair, pkey, skey);
+  return ml_kem_common_keypair(mlkem_native768_keypair, pkey, skey);
 }
 
 int ml_kem_768_encapsulate_deterministic(uint8_t *ciphertext       /* OUT */,
@@ -220,7 +225,7 @@ int ml_kem_768_encapsulate_deterministic_no_self_test(uint8_t *ciphertext
                                                       const uint8_t *seed       /* IN */) {
   output_buffer ctext = {ciphertext, ciphertext_len, MLKEM768_CIPHERTEXT_BYTES};
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM768_SHARED_SECRET_LEN};
-  return ml_kem_common_encapsulate_deterministic(mlkem768_enc_derand, ctext, ss, public_key, seed);
+  return ml_kem_common_encapsulate_deterministic(mlkem_native768_enc_derand, ctext, ss, public_key, seed);
 }
 
 int ml_kem_768_encapsulate(uint8_t *ciphertext       /* OUT */,
@@ -230,7 +235,7 @@ int ml_kem_768_encapsulate(uint8_t *ciphertext       /* OUT */,
                            const uint8_t *public_key /* IN */) {
   output_buffer ctext = {ciphertext, ciphertext_len, MLKEM768_CIPHERTEXT_BYTES};
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM768_SHARED_SECRET_LEN};
-  return ml_kem_common_encapsulate(mlkem768_enc, ctext, ss, public_key);
+  return ml_kem_common_encapsulate(mlkem_native768_enc, ctext, ss, public_key);
 }
 
 int ml_kem_768_decapsulate(uint8_t *shared_secret    /* OUT */,
@@ -246,7 +251,7 @@ int ml_kem_768_decapsulate_no_self_test(uint8_t *shared_secret    /* OUT */,
                                         const uint8_t *ciphertext /* IN */,
                                         const uint8_t *secret_key /* IN */) {
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM768_SHARED_SECRET_LEN};
-  return ml_kem_common_decapsulate(mlkem768_dec, ss, ciphertext, secret_key);
+  return ml_kem_common_decapsulate(mlkem_native768_dec, ss, ciphertext, secret_key);
 }
 
 int ml_kem_1024_keypair_deterministic(uint8_t *public_key /* OUT */,
@@ -268,7 +273,7 @@ int ml_kem_1024_keypair_deterministic_no_self_test(uint8_t *public_key /* OUT */
   if (!check_buffer(pkey) || !check_buffer(skey)) {
     return 1;
   }
-  const int res = mlkem1024_keypair_derand(pkey.buffer, skey.buffer, seed);
+  const int res = mlkem_native1024_keypair_derand(pkey.buffer, skey.buffer, seed);
 #if defined(AWSLC_FIPS)
   /* PCT failure is the only failure condition for key generation. */
   if (res != 0) {
@@ -286,7 +291,7 @@ int ml_kem_1024_keypair(uint8_t *public_key /* OUT */,
                         size_t *secret_len  /* IN_OUT */) {
   output_buffer pkey = {public_key, public_len, MLKEM1024_PUBLIC_KEY_BYTES};
   output_buffer skey = {secret_key, secret_len, MLKEM1024_SECRET_KEY_BYTES};
-  return ml_kem_common_keypair(mlkem1024_keypair, pkey, skey);
+  return ml_kem_common_keypair(mlkem_native1024_keypair, pkey, skey);
 }
 
 int ml_kem_1024_encapsulate_deterministic(uint8_t *ciphertext       /* OUT */,
@@ -307,7 +312,7 @@ int ml_kem_1024_encapsulate_deterministic_no_self_test(uint8_t *ciphertext
                                                        const uint8_t *seed       /* IN */) {
   output_buffer ctext = {ciphertext, ciphertext_len, MLKEM1024_CIPHERTEXT_BYTES};
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM1024_SHARED_SECRET_LEN};
-  return ml_kem_common_encapsulate_deterministic(mlkem1024_enc_derand, ctext, ss, public_key, seed);
+  return ml_kem_common_encapsulate_deterministic(mlkem_native1024_enc_derand, ctext, ss, public_key, seed);
 }
 
 int ml_kem_1024_encapsulate(uint8_t *ciphertext       /* OUT */,
@@ -317,7 +322,7 @@ int ml_kem_1024_encapsulate(uint8_t *ciphertext       /* OUT */,
                             const uint8_t *public_key /* IN */) {
   output_buffer ctext = {ciphertext, ciphertext_len, MLKEM1024_CIPHERTEXT_BYTES};
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM1024_SHARED_SECRET_LEN};
-  return ml_kem_common_encapsulate(mlkem1024_enc, ctext, ss, public_key);
+  return ml_kem_common_encapsulate(mlkem_native1024_enc, ctext, ss, public_key);
 }
 
 int ml_kem_1024_decapsulate(uint8_t *shared_secret    /* OUT */,
@@ -333,7 +338,7 @@ int ml_kem_1024_decapsulate_no_self_test(uint8_t *shared_secret    /* OUT */,
                                          const uint8_t *ciphertext /* IN */,
                                          const uint8_t *secret_key /* IN */) {
   output_buffer ss = {shared_secret, shared_secret_len, MLKEM1024_SHARED_SECRET_LEN};
-  return ml_kem_common_decapsulate(mlkem1024_dec, ss, ciphertext, secret_key);
+  return ml_kem_common_decapsulate(mlkem_native1024_dec, ss, ciphertext, secret_key);
 }
 
 int ml_kem_common_keypair(int (*keypair)(uint8_t * public_key, uint8_t *secret_key),

crypto/fipsmodule/ml_kem/ml_kem.c

@@ -9,7 +9,7 @@
 // Namespacing: All symbols are of the form mlkem*. Level-specific
 // symbols are further prefixed with their security level, e.g.
 // mlkem512*, mlkem768*, mlkem1024*.
-#define MLK_CONFIG_NAMESPACE_PREFIX mlkem
+#define MLK_CONFIG_NAMESPACE_PREFIX mlkem_native
 
 // Replace mlkem-native's FIPS 202 headers with glue code to
 // AWS-LC's own FIPS 202 implementation.
@@ -68,4 +68,9 @@ static MLK_INLINE void mlk_randombytes(void *ptr, size_t len) {
 #define MLK_CONFIG_NO_ASM
 #endif
 
+#if defined(OPENSSL_AARCH64) && !defined(OPENSSL_NO_ASM)
+#define MLK_CONFIG_USE_NATIVE_BACKEND_ARITH
+#define MLK_CONFIG_ARITH_BACKEND_FILE "../aarch64/meta.h"
+#endif
+
 #endif // MLkEM_NATIVE_CONFIG_H

crypto/fipsmodule/ml_kem/mlkem_native_config.h

@@ -1,5 +1,5 @@
 name: s2n-bignum-imported
-source: awslabs/s2n-bignum.git
-commit: 54e1fa5756d6b13961c2f61d90f75426aa25d373
-target: main
-imported-at: 2025-04-28T17:22:07+0000
+source: jargh/s2n-bignum-dev.git
+commit: ae84a59689cb50ad9b9c6e25cd34037d5b1fb2b4
+target: mlkem
+imported-at: 2025-06-23T13:38:02+0000

third_party/s2n-bignum/META.yml

@@ -0,0 +1,39 @@
+#############################################################################
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0
+#############################################################################
+
+# If actually on an ARM8 machine, just use the GNU assembler (as). Otherwise
+# use a cross-assembling version so that the code can still be assembled
+# and the proofs checked against the object files (though you won't be able
+# to run code without additional emulation infrastructure). The aarch64
+# cross-assembling version can be installed manually by something like:
+#
+#  sudo apt-get install binutils-aarch64-linux-gnu
+
+UNAME_RESULT=$(shell uname -p)
+
+ifeq ($(UNAME_RESULT),aarch64)
+GAS=as
+else
+GAS=aarch64-linux-gnu-as
+endif
+
+# List of object files
+
+OBJ = mlkem_basemul_k2.o \
+      mlkem_basemul_k3.o \
+      mlkem_basemul_k4.o \
+      mlkem_intt.o \
+      mlkem_mulcache_compute.o \
+      mlkem_ntt.o \
+      mlkem_poly_reduce.o \
+      mlkem_poly_tobytes.o \
+      mlkem_poly_tomont.o \
+      mlkem_rej_uniform_VARIABLE_TIME.o
+
+%.o : %.S ; $(CC) -E -I../../include $< | $(GAS) -o $@ -
+
+default: $(OBJ);
+
+clean:; rm -f *.o *.correct unopt/*.o