ML-KEM: Import AArch64 backend from mlkem-native (#2498)

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 extends the mlkem-native->AWS-LC import to include the AArch64 arithmetic backend.

Author: hanno-becker

crypto/fipsmodule/CMakeLists.txt

@@ -295,6 +295,32 @@ if((((ARCH STREQUAL "x86_64") AND NOT MY_ASSEMBLER_IS_TOO_OLD_FOR_512AVX) OR
 
 endif()
 
+
+# mlkem-native assembly files can be compiled on Unix platforms for x86_64 and arm64 only.
+if((ARCH STREQUAL "aarch64") AND UNIX)
+
+  # Set the source directory for s2n-bignum assembly files
+  set(MLKEM_NATIVE_DIR "${AWSLC_SOURCE_DIR}/crypto/fipsmodule/ml_kem")
+
+  set(MLKEM_NATIVE_AARCH64_ASM_SOURCES
+
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/intt.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/ntt.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/poly_mulcache_compute_asm.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/poly_reduce_asm.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/poly_tobytes_asm.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/poly_tomont_asm.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/polyvec_basemul_acc_montgomery_cached_asm_k2.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/polyvec_basemul_acc_montgomery_cached_asm_k3.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/polyvec_basemul_acc_montgomery_cached_asm_k4.S
+    ${MLKEM_NATIVE_DIR}/mlkem/native/aarch64/src/rej_uniform_asm.S
+  )
+
+  list(APPEND BCM_ASM_SOURCES ${MLKEM_NATIVE_AARCH64_ASM_SOURCES})
+
+endif()
+
+
 if(FIPS_DELOCATE)
   if(FIPS_SHARED)
     message(FATAL_ERROR "Can't set both delocate and shared mode for FIPS build")

crypto/fipsmodule/ml_kem/META.yml

@@ -1,5 +1,5 @@
 name: mlkem-native
 source: pq-code-package/mlkem-native.git
-branch: v1.0.0
-commit: 048fc2a7a7b4ba0ad4c989c1ac82491aa94d5bfa
-imported-at: 2025-06-04T13:38:24+0100
+branch: main
+commit: a67a02ee3fa05713ae01572efce741caa63285e6
+imported-at: 2025-06-25T12:11:38+0100

crypto/fipsmodule/ml_kem/importer.sh

@@ -72,10 +72,18 @@ popd
 
 echo "Pull source code from remote repository..."
 
-# Copy mlkem-native source tree -- C-only, no FIPS-202
+# Copy mlkem-native source tree -- C source
 mkdir $SRC
 cp $TMP/mlkem/src/* $SRC
 
+# Copy AArch64 backend
+mkdir -p $SRC/native/aarch64/src
+# Backend API and specification assumed by mlkem-native frontend
+cp $TMP/mlkem/src/native/api.h $SRC/native
+# Copy AArch64 backend implementation
+cp $TMP/mlkem/src/native/aarch64/meta.h $SRC/native/aarch64
+cp $TMP/mlkem/src/native/aarch64/src/* $SRC/native/aarch64/src
+
 # We use the custom `mlkem_native_config.h`, so can remove the default one
 rm $SRC/config.h
 
@@ -84,11 +92,11 @@ cp $TMP/.clang-format $SRC
 
 # Copy and statically simplify BCM file
 # The static simplification is not necessary, but improves readability
-# by removing directives related to native backends that are irrelevant
-# for the C-only import.
+# by removing directives related to the FIPS-202 backend and the x86_64
+# arithmetic backend that are not yet imported.
 unifdef -DMLK_CONFIG_FIPS202_CUSTOM_HEADER                             \
-        -UMLK_CONFIG_USE_NATIVE_BACKEND_ARITH                          \
         -UMLK_CONFIG_USE_NATIVE_BACKEND_FIPS202                        \
+        -UMLK_SYS_X86_64                                               \
         $TMP/mlkem/mlkem_native.c                                      \
         > $SRC/mlkem_native_bcm.c
 
@@ -107,9 +115,30 @@ if [[ "$(uname)" == "Darwin" ]]; then
 else
   SED_I=(-i)
 fi
+
 echo "Fixup include paths"
 sed "${SED_I[@]}" 's/#include "src\/\([^"]*\)"/#include "\1"/' $SRC/mlkem_native_bcm.c
 
+echo "Fixup AArch64 assembly backend to use s2n-bignum macros"
+for file in $SRC/native/aarch64/src/*.S; do
+  echo "Processing $file"
+  tmp_file=$(mktemp)
+
+  # Flatten multiline preprocessor directives, then process with unifdef
+  sed -e ':a' -e 'N' -e '$!ba' -e 's/\\\n/ /g' "$file" | \
+  unifdef -DMLK_ARITH_BACKEND_AARCH64 -UMLK_CONFIG_MULTILEVEL_NO_SHARED -DMLK_CONFIG_MULTILEVEL_WITH_SHARED > "$tmp_file"
+  mv "$tmp_file" "$file"
+
+  # Replace common.h include and assembly macros
+  sed "${SED_I[@]}" 's/#include "\.\.\/\.\.\/\.\.\/common\.h"/#include "_internal_s2n_bignum.h"/' "$file"
+
+  func_name=$(grep -o '\.global MLK_ASM_NAMESPACE(\([^)]*\))' "$file" | sed 's/\.global MLK_ASM_NAMESPACE(\([^)]*\))/\1/')
+  if [ -n "$func_name" ]; then
+    sed "${SED_I[@]}" "s/\.global MLK_ASM_NAMESPACE($func_name)/        S2N_BN_SYM_VISIBILITY_DIRECTIVE(mlkem_$func_name)\n        S2N_BN_SYM_PRIVACY_DIRECTIVE(mlkem_$func_name)/" "$file"
+    sed "${SED_I[@]}" "s/MLK_ASM_FN_SYMBOL($func_name)/S2N_BN_SYMBOL(mlkem_$func_name):/" "$file"
+  fi
+done
+
 echo "Remove temporary artifacts ..."
 rm -rf $TMP
 

crypto/fipsmodule/ml_kem/mlkem/mlkem_native_bcm.c

@@ -5,6 +5,7 @@
 
 /*
  * WARNING: This file is auto-generated from scripts/autogen
+ *          in the mlkem-native repository.
  *          Do not modify it directly.
  */
 
@@ -68,6 +69,12 @@
 #include "verify.c"
 
 
+#if defined(MLK_CONFIG_USE_NATIVE_BACKEND_ARITH)
+#if defined(MLK_SYS_AARCH64)
+#include "native/aarch64/src/aarch64_zetas.c"
+#include "native/aarch64/src/rej_uniform_table.c"
+#endif
+#endif /* MLK_CONFIG_USE_NATIVE_BACKEND_ARITH */
 
 
 /* Macro #undef's
@@ -310,4 +317,47 @@
 #undef __loop__
 
 
+#if defined(MLK_CONFIG_USE_NATIVE_BACKEND_ARITH)
+/* mlkem/src/native/api.h */
+#undef MLK_INVNTT_BOUND
+#undef MLK_NATIVE_API_H
+#undef MLK_NTT_BOUND
+/* mlkem/src/native/meta.h */
+#undef MLK_NATIVE_META_H
+#if defined(MLK_SYS_AARCH64)
+/*
+ * Undefine macros from native code (Arith, AArch64)
+ */
+/* mlkem/src/native/aarch64/meta.h */
+#undef MLK_ARITH_BACKEND_AARCH64
+#undef MLK_NATIVE_AARCH64_META_H
+#undef MLK_USE_NATIVE_INTT
+#undef MLK_USE_NATIVE_NTT
+#undef MLK_USE_NATIVE_POLYVEC_BASEMUL_ACC_MONTGOMERY_CACHED
+#undef MLK_USE_NATIVE_POLY_MULCACHE_COMPUTE
+#undef MLK_USE_NATIVE_POLY_REDUCE
+#undef MLK_USE_NATIVE_POLY_TOBYTES
+#undef MLK_USE_NATIVE_POLY_TOMONT
+#undef MLK_USE_NATIVE_REJ_UNIFORM
+/* mlkem/src/native/aarch64/src/arith_native_aarch64.h */
+#undef MLK_NATIVE_AARCH64_SRC_ARITH_NATIVE_AARCH64_H
+#undef mlk_aarch64_invntt_zetas_layer12345
+#undef mlk_aarch64_invntt_zetas_layer67
+#undef mlk_aarch64_ntt_zetas_layer12345
+#undef mlk_aarch64_ntt_zetas_layer67
+#undef mlk_aarch64_zetas_mulcache_native
+#undef mlk_aarch64_zetas_mulcache_twisted_native
+#undef mlk_intt_asm
+#undef mlk_ntt_asm
+#undef mlk_poly_mulcache_compute_asm
+#undef mlk_poly_reduce_asm
+#undef mlk_poly_tobytes_asm
+#undef mlk_poly_tomont_asm
+#undef mlk_polyvec_basemul_acc_montgomery_cached_asm_k2
+#undef mlk_polyvec_basemul_acc_montgomery_cached_asm_k3
+#undef mlk_polyvec_basemul_acc_montgomery_cached_asm_k4
+#undef mlk_rej_uniform_asm
+#undef mlk_rej_uniform_table
+#endif /* MLK_SYS_AARCH64 */
+#endif /* MLK_CONFIG_USE_NATIVE_BACKEND_ARITH */
 #endif /* !MLK_CONFIG_MONOBUILD_KEEP_SHARED_HEADERS */

crypto/fipsmodule/ml_kem/mlkem/native/aarch64/meta.h

@@ -0,0 +1,101 @@
+/*
+ * Copyright (c) The mlkem-native project authors
+ * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+ */
+
+#ifndef MLK_NATIVE_AARCH64_META_H
+#define MLK_NATIVE_AARCH64_META_H
+
+/* Set of primitives that this backend replaces */
+#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
+
+/* Identifier for this backend so that source and assembly files
+ * in the build can be appropriately guarded. */
+#define MLK_ARITH_BACKEND_AARCH64
+
+
+#if !defined(__ASSEMBLER__)
+#include "src/arith_native_aarch64.h"
+
+static MLK_INLINE void mlk_ntt_native(int16_t data[MLKEM_N])
+{
+  mlk_ntt_asm(data, mlk_aarch64_ntt_zetas_layer12345,
+              mlk_aarch64_ntt_zetas_layer67);
+}
+
+static MLK_INLINE void mlk_intt_native(int16_t data[MLKEM_N])
+{
+  mlk_intt_asm(data, mlk_aarch64_invntt_zetas_layer12345,
+               mlk_aarch64_invntt_zetas_layer67);
+}
+
+static MLK_INLINE void mlk_poly_reduce_native(int16_t data[MLKEM_N])
+{
+  mlk_poly_reduce_asm(data);
+}
+
+static MLK_INLINE void mlk_poly_tomont_native(int16_t data[MLKEM_N])
+{
+  mlk_poly_tomont_asm(data);
+}
+
+static MLK_INLINE void mlk_poly_mulcache_compute_native(
+    int16_t x[MLKEM_N / 2], const int16_t y[MLKEM_N])
+{
+  mlk_poly_mulcache_compute_asm(x, y, mlk_aarch64_zetas_mulcache_native,
+                                mlk_aarch64_zetas_mulcache_twisted_native);
+}
+
+#if defined(MLK_CONFIG_MULTILEVEL_WITH_SHARED) || MLKEM_K == 2
+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)])
+{
+  mlk_polyvec_basemul_acc_montgomery_cached_asm_k2(r, a, b, b_cache);
+}
+#endif /* MLK_CONFIG_MULTILEVEL_WITH_SHARED || MLKEM_K == 2 */
+
+#if defined(MLK_CONFIG_MULTILEVEL_WITH_SHARED) || MLKEM_K == 3
+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)])
+{
+  mlk_polyvec_basemul_acc_montgomery_cached_asm_k3(r, a, b, b_cache);
+}
+#endif /* MLK_CONFIG_MULTILEVEL_WITH_SHARED || MLKEM_K == 3 */
+
+#if defined(MLK_CONFIG_MULTILEVEL_WITH_SHARED) || MLKEM_K == 4
+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)])
+{
+  mlk_polyvec_basemul_acc_montgomery_cached_asm_k4(r, a, b, b_cache);
+}
+#endif /* MLK_CONFIG_MULTILEVEL_WITH_SHARED || MLKEM_K == 4 */
+
+static MLK_INLINE void mlk_poly_tobytes_native(uint8_t r[MLKEM_POLYBYTES],
+                                               const int16_t a[MLKEM_N])
+{
+  mlk_poly_tobytes_asm(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)mlk_rej_uniform_asm(r, buf, buflen, mlk_rej_uniform_table);
+}
+#endif /* !__ASSEMBLER__ */
+
+#endif /* !MLK_NATIVE_AARCH64_META_H */