Import ML-DSA x86_64 NTT/INTT assembly optimizations

crypto/fipsmodule/CMakeLists.txt

@@ -391,6 +391,18 @@ if((ARCH STREQUAL "x86_64") AND UNIX AND NOT MY_ASSEMBLER_IS_TOO_OLD_FOR_AVX)
 
   set(S2N_BIGNUM_INCLUDE_DIR "${AWSLC_SOURCE_DIR}/third_party/s2n-bignum/s2n-bignum-imported/include")
 
+  # mldsa-native assembly files for x86_64 (NTT/INTT + nttunpack for custom order)
+  set(MLDSA_NATIVE_DIR "${AWSLC_SOURCE_DIR}/crypto/fipsmodule/ml_dsa")
+
+  set(MLDSA_NATIVE_X86_64_ASM_SOURCES
+
+    ${MLDSA_NATIVE_DIR}/mldsa/native/x86_64/src/intt.S
+    ${MLDSA_NATIVE_DIR}/mldsa/native/x86_64/src/ntt.S
+    ${MLDSA_NATIVE_DIR}/mldsa/native/x86_64/src/nttunpack.S
+  )
+
+  list(APPEND BCM_ASM_SOURCES ${MLDSA_NATIVE_X86_64_ASM_SOURCES})
+
 endif()
 
 

crypto/fipsmodule/ml_dsa/META.yml

@@ -1,5 +1,5 @@
 name: mldsa-native
 source: pq-code-package/mldsa-native.git
 branch: main
-commit: b61e84f0c73d4ed612ffcaea4282a9d682de3f46
-imported-at: 2026-01-16T13:12:01-0800
+commit: e3c632a2cf6c05a3525670db9c8ba0e3973056ac
+imported-at: 2026-02-12T21:45:42+0000

crypto/fipsmodule/ml_dsa/importer.sh

@@ -72,12 +72,26 @@ popd
 
 echo "Pull source code from remote repository..."
 
-# Copy mldsa-native source tree -- C source only (no native backends for now)
+# Copy mldsa-native source tree -- C source
 mkdir $SRC
 cp $TMP/mldsa/src/* $SRC
 
+# Copy x86_64 backend (NTT/INTT only)
+mkdir -p $SRC/native/x86_64/src
+# Backend API and specification assumed by mldsa-native frontend
+cp $TMP/mldsa/src/native/api.h $SRC/native
+# Copy x86_64 backend implementation - NTT/INTT only
+cp $TMP/mldsa/src/native/x86_64/meta.h $SRC/native/x86_64
+# Copy only NTT/INTT assembly and supporting constants
+cp $TMP/mldsa/src/native/x86_64/src/ntt.S $SRC/native/x86_64/src/
+cp $TMP/mldsa/src/native/x86_64/src/intt.S $SRC/native/x86_64/src/
+cp $TMP/mldsa/src/native/x86_64/src/nttunpack.S $SRC/native/x86_64/src/
+cp $TMP/mldsa/src/native/x86_64/src/consts.c $SRC/native/x86_64/src/
+cp $TMP/mldsa/src/native/x86_64/src/consts.h $SRC/native/x86_64/src/
+cp $TMP/mldsa/src/native/x86_64/src/arith_native_x86_64.h $SRC/native/x86_64/src/
+
 # We use the custom `mldsa_native_config.h`, so can remove the default one
-rm $SRC/config.h
+rm -f $SRC/config.h
 
 # Copy formatting file
 cp $TMP/.clang-format $SRC
@@ -103,13 +117,145 @@ fi
 # of everything being inlined into that file.
 cp $TMP/mldsa/mldsa_native.h $SRC
 
-# Modify include paths to match position of mldsa_native_bcm.c
-# In mldsa-native, the include path is "mldsa/*", while here we
-# embed mldsa_native_bcm.c in the main source directory of mldsa-native,
-# hence the relative import path is just ".".
+# Remove non-NTT/INTT includes from the BCM file (files not imported)
+echo "Remove unused native backend includes from BCM file"
+# Remove all aarch64 includes (not imported)
+sed "${SED_I[@]}" '/native\/aarch64/d' $SRC/mldsa_native_bcm.c
+# Remove x86_64 AVX2 C files (not imported)
+sed "${SED_I[@]}" '/poly_caddq_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/poly_chknorm_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/poly_decompose_32_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/poly_decompose_88_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/poly_use_hint_32_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/poly_use_hint_88_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/polyz_unpack_17_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/polyz_unpack_19_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/rej_uniform_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/rej_uniform_eta2_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/rej_uniform_eta4_avx2.c/d' $SRC/mldsa_native_bcm.c
+sed "${SED_I[@]}" '/rej_uniform_table.c/d' $SRC/mldsa_native_bcm.c
+
+# Strip arith_native_x86_64.h to only NTT/INTT/nttunpack declarations
+echo "Strip arith_native_x86_64.h to only NTT/INTT/nttunpack declarations"
+# Create minimal header with only NTT/INTT/nttunpack
+cat > $SRC/native/x86_64/src/arith_native_x86_64.h << 'HEADER_EOF'
+/*
+ * Copyright (c) The mlkem-native project authors
+ * Copyright (c) The mldsa-native project authors
+ * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+ */
+
+#ifndef MLD_ARITH_NATIVE_X86_64_H
+#define MLD_ARITH_NATIVE_X86_64_H
+
+#include <stdint.h>
+#include "../../../params.h"
+#include "consts.h"
+
+/* NTT forward transformation */
+#define mld_ntt_avx2 MLD_NAMESPACE(ntt_avx2)
+void mld_ntt_avx2(int32_t r[MLDSA_N], const int32_t *qdata);
+
+/* Inverse NTT transformation */
+#define mld_invntt_avx2 MLD_NAMESPACE(invntt_avx2)
+void mld_invntt_avx2(int32_t r[MLDSA_N], const int32_t *qdata);
+
+/* NTT unpack - permute from bitreversed to custom order */
+#define mld_nttunpack_avx2 MLD_NAMESPACE(nttunpack_avx2)
+void mld_nttunpack_avx2(int32_t *r);
+
+#endif /* MLD_ARITH_NATIVE_X86_64_H */
+HEADER_EOF
+
+# Strip meta.h to only NTT/INTT with custom order - create minimal version
+echo "Strip meta.h to only NTT/INTT with custom order"
+cat > $SRC/native/x86_64/meta.h << 'META_EOF'
+/*
+* Copyright (c) The mlkem-native project authors
+ * Copyright (c) The mldsa-native project authors
+ * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
+ */
+
+#ifndef MLD_NATIVE_X86_64_META_H
+#define MLD_NATIVE_X86_64_META_H
+
+/* Identifier for this backend so that source and assembly files
+ * in the build can be appropriately guarded. */
+#define MLD_ARITH_BACKEND_X86_64_DEFAULT
+
+#define MLD_USE_NATIVE_NTT_CUSTOM_ORDER
+#define MLD_USE_NATIVE_NTT
+#define MLD_USE_NATIVE_INTT
+
+#if !defined(__ASSEMBLER__)
+#include <string.h>
+#include "../../common.h"
+#include "../api.h"
+#include "src/arith_native_x86_64.h"
+
+static MLD_INLINE void mld_poly_permute_bitrev_to_custom(int32_t data[MLDSA_N])
+{
+  if (mld_sys_check_capability(MLD_SYS_CAP_AVX2))
+  {
+    mld_nttunpack_avx2(data);
+  }
+}
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_ntt_native(int32_t data[MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_AVX2))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+
+  mld_ntt_avx2(data, mld_qdata);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+
+MLD_MUST_CHECK_RETURN_VALUE
+static MLD_INLINE int mld_intt_native(int32_t data[MLDSA_N])
+{
+  if (!mld_sys_check_capability(MLD_SYS_CAP_AVX2))
+  {
+    return MLD_NATIVE_FUNC_FALLBACK;
+  }
+  mld_invntt_avx2(data, mld_qdata);
+  return MLD_NATIVE_FUNC_SUCCESS;
+}
+
+#endif /* !__ASSEMBLER__ */
+
+#endif /* !MLD_NATIVE_X86_64_META_H */
+META_EOF
+
 echo "Fixup include paths"
 sed "${SED_I[@]}" 's/#include "src\/\([^"]*\)"/#include "\1"/' $SRC/mldsa_native_bcm.c
 
+echo "Fixup x86_64 assembly backend to use s2n-bignum macros"
+for file in $SRC/native/x86_64/src/*.S; do
+  echo "Processing $file"
+  tmp_file=$(mktemp)
+
+  backend_define="MLD_ARITH_BACKEND_X86_64_DEFAULT"
+
+  # Flatten multiline preprocessor directives, then process with unifdef
+  sed -e ':a' -e 'N' -e '$!ba' -e 's/\\\n/ /g' "$file" | \
+    unifdef -D$backend_define -UMLD_CONFIG_MULTILEVEL_NO_SHARED -DMLD_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 MLD_ASM_NAMESPACE(\([^)]*\))' "$file" | sed 's/\.global MLD_ASM_NAMESPACE(\([^)]*\))/\1/')
+  if [ -n "$func_name" ]; then
+    sed "${SED_I[@]}" "s/\.global MLD_ASM_NAMESPACE($func_name)/        S2N_BN_SYM_VISIBILITY_DIRECTIVE(mldsa_$func_name)\n        S2N_BN_SYM_PRIVACY_DIRECTIVE(mldsa_$func_name)/" "$file"
+    sed "${SED_I[@]}" "s/MLD_ASM_FN_SYMBOL($func_name)/S2N_BN_SYMBOL(mldsa_$func_name):/" "$file"
+    # Remove MLD_ASM_FN_SIZE line (not present in s2n-bignum convention)
+    sed "${SED_I[@]}" "/MLD_ASM_FN_SIZE($func_name)/d" "$file"
+  fi
+done
+
 echo "Remove temporary artifacts ..."
 rm -rf $TMP
 
@@ -125,4 +271,4 @@ imported-at: $(date "+%Y-%m-%dT%H:%M:%S%z")
 EOF
 
 echo "Import complete!"
-echo "Imported mldsa-native commit: $GITHUB_COMMIT"
+echo "Imported mldsa-native commit: $GITHUB_COMMIT"

crypto/fipsmodule/ml_dsa/mldsa/.clang-format

@@ -26,4 +26,4 @@ Macros:
  - __contract__(x)={ void a; void b; void c; void d; void e; void f; } void abcdefghijklmnopqrstuvw()
  - __loop__(x)={} do
  # Make this artifically long to force line break
- - MLK_INTERNAL_API=void abcdefghijklmnopqrstuvwabcdefghijklmnopqrstuvwabcdefg();
+ - MLD_INTERNAL_API=void abcdefghijklmnopqrstuvwabcdefghijklmnopqrstuvwabcdefg();

crypto/fipsmodule/ml_dsa/mldsa/common.h

@@ -77,8 +77,24 @@
  */
 #if defined(MLD_SYS_X86_64)
 #define MLD_ASM_FN_SYMBOL(sym) MLD_ASM_NAMESPACE(sym) : MLD_CET_ENDBR
-#else
+#elif defined(MLD_SYS_ARMV81M_MVE)
+/* clang-format off */
+#define MLD_ASM_FN_SYMBOL(sym) \
+  .type MLD_ASM_NAMESPACE(sym), %function; \
+  MLD_ASM_NAMESPACE(sym) :
+/* clang-format on */
+#else /* !MLD_SYS_X86_64 && MLD_SYS_ARMV81M_MVE */
 #define MLD_ASM_FN_SYMBOL(sym) MLD_ASM_NAMESPACE(sym) :
+#endif /* !MLD_SYS_X86_64 && !MLD_SYS_ARMV81M_MVE */
+
+/*
+ * Output the size of an assembly function.
+ */
+#if defined(__ELF__)
+#define MLD_ASM_FN_SIZE(sym) \
+  .size MLD_ASM_NAMESPACE(sym), .- MLD_ASM_NAMESPACE(sym)
+#else
+#define MLD_ASM_FN_SIZE(sym)
 #endif
 
 /* We aim to simplify the user's life by supporting builds where
@@ -278,9 +294,9 @@
  * is resolved
  */
 #if defined(MLD_CONFIG_REDUCE_RAM)
-#define MLK_UNION_OR_STRUCT union
+#define MLD_UNION_OR_STRUCT union
 #else
-#define MLK_UNION_OR_STRUCT struct
+#define MLD_UNION_OR_STRUCT struct
 #endif
 
 /****************************** Error codes ***********************************/

crypto/fipsmodule/ml_dsa/mldsa/ct.h

@@ -83,44 +83,54 @@ extern volatile uint64_t mld_ct_opt_blocker_u64;
  * Its validity relies on the assumption that the global opt-blocker
  * constant mld_ct_opt_blocker_u64 is not modified.
  */
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint64_t mld_ct_get_optblocker_u64(void)
 __contract__(ensures(return_value == 0)) { return mld_ct_opt_blocker_u64; }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE int64_t mld_ct_get_optblocker_i64(void)
 __contract__(ensures(return_value == 0)) { return (int64_t)mld_ct_get_optblocker_u64(); }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint32_t mld_ct_get_optblocker_u32(void)
 __contract__(ensures(return_value == 0)) { return (uint32_t)mld_ct_get_optblocker_u64(); }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint8_t mld_ct_get_optblocker_u8(void)
 __contract__(ensures(return_value == 0)) { return (uint8_t)mld_ct_get_optblocker_u64(); }
 
 /* Opt-blocker based implementation of value barriers */
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE int64_t mld_value_barrier_i64(int64_t b)
 __contract__(ensures(return_value == b)) { return (b ^ mld_ct_get_optblocker_i64()); }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint32_t mld_value_barrier_u32(uint32_t b)
 __contract__(ensures(return_value == b)) { return (b ^ mld_ct_get_optblocker_u32()); }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint8_t mld_value_barrier_u8(uint8_t b)
 __contract__(ensures(return_value == b)) { return (b ^ mld_ct_get_optblocker_u8()); }
 
 
 #else  /* !MLD_USE_ASM_VALUE_BARRIER */
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE int64_t mld_value_barrier_i64(int64_t b)
 __contract__(ensures(return_value == b))
 {
   __asm__ volatile("" : "+r"(b));
   return b;
 }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint32_t mld_value_barrier_u32(uint32_t b)
 __contract__(ensures(return_value == b))
 {
   __asm__ volatile("" : "+r"(b));
   return b;
 }
 
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint8_t mld_value_barrier_u8(uint8_t b)
 __contract__(ensures(return_value == b))
 {
@@ -147,6 +157,7 @@ __contract__(ensures(return_value == b))
  *              - x >= 2^31: returns x - 2^31
  *
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_ALWAYS_INLINE int32_t mld_cast_uint32_to_int32(uint32_t x)
 {
   /*
@@ -172,6 +183,7 @@ static MLD_ALWAYS_INLINE int32_t mld_cast_uint32_to_int32(uint32_t x)
  * Returns:     For int64_t x, the unique y in uint32_t
  *              so that x == y mod 2^32.
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_ALWAYS_INLINE uint32_t mld_cast_int64_to_uint32(int64_t x)
 {
   return (uint32_t)(x & (int64_t)UINT32_MAX);
@@ -185,6 +197,7 @@ static MLD_ALWAYS_INLINE uint32_t mld_cast_int64_to_uint32(int64_t x)
  * Returns:     For int32_t x, the unique y in uint32_t
  *              so that x == y mod 2^32.
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_ALWAYS_INLINE uint32_t mld_cast_int32_to_uint32(int32_t x)
 {
   return mld_cast_int64_to_uint32((int64_t)x);
@@ -203,6 +216,7 @@ static MLD_ALWAYS_INLINE uint32_t mld_cast_int32_to_uint32(int32_t x)
  *
  *
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE int32_t mld_ct_sel_int32(int32_t a, int32_t b, uint32_t cond)
 __contract__(
   requires(cond == 0x0 || cond == 0xFFFFFFFF)
@@ -223,6 +237,7 @@ __contract__(
  * Arguments:   uint32_t x: Value to be converted into a mask
  *
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint32_t mld_ct_cmask_nonzero_u32(uint32_t x)
 __contract__(ensures(return_value == ((x == 0) ? 0 : 0xFFFFFFFF)))
 {
@@ -239,6 +254,7 @@ __contract__(ensures(return_value == ((x == 0) ? 0 : 0xFFFFFFFF)))
  * Arguments:   uint8_t x: Value to be converted into a mask
  *
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint8_t mld_ct_cmask_nonzero_u8(uint8_t x)
 __contract__(ensures(return_value == ((x == 0) ? 0 : 0xFF)))
 {
@@ -254,6 +270,7 @@ __contract__(ensures(return_value == ((x == 0) ? 0 : 0xFF)))
  * Arguments:   int32_t x: Value to be converted into a mask
  *
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint32_t mld_ct_cmask_neg_i32(int32_t x)
 __contract__(
   ensures(return_value == ((x < 0) ? 0xFFFFFFFF : 0))
@@ -272,6 +289,7 @@ __contract__(
  * Arguments:   int32_t x: Input value
  *
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE int32_t mld_ct_abs_i32(int32_t x)
 __contract__(
   requires(x >= -INT32_MAX)
@@ -294,6 +312,7 @@ __contract__(
  *
  * Returns 0 if the byte arrays are equal, 0xFF otherwise.
  **************************************************/
+MLD_MUST_CHECK_RETURN_VALUE
 static MLD_INLINE uint8_t mld_ct_memcmp(const uint8_t *a, const uint8_t *b,
                                         const size_t len)
 __contract__(