Add runtime dispatch (mld_ntt_native)

- Change mld_ntt_native() return type from void to int
- Add runtime capability checking with fallback support
- Implement dispatch logic in mld_poly_ntt() to try native first,
  fallback to C
- Add MLD_NATIVE_FUNC_SUCCESS/FALLBACK return codes
- Add mld_sys_check_capability() for system capability detection
- Add test configuration for AVX2, static ON/OFF, add to CI test.

Signed-off-by: willieyz <[email protected]>

Author: willieyz

.github/actions/config-variations/action.yml

@@ -123,3 +123,42 @@ runs:
         acvp: true
         opt: ${{ inputs.opt }}
         examples: false # Some examples use a custom config themselves
+    - name: "Custom native capability functions (static ON)"
+      if: ${{ inputs.tests == 'all' || contains(inputs.tests, 'native-cap-ON') }}
+      uses: ./.github/actions/multi-functest
+      with:
+        gh_token: ${{ inputs.gh_token }}
+        compile_mode: native
+        cflags: "-std=c11 -D_GNU_SOURCE -DMLD_CONFIG_FILE=\\\\\\\"../../test/custom_native_capability_config_1.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        ldflags: "-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        func: true
+        kat: true
+        acvp: true
+        opt: ${{ inputs.opt }}
+        examples: false # Some examples use a custom config themselves
+    - name: "Custom native capability functions (static OFF)"
+      if: ${{ inputs.tests == 'all' || contains(inputs.tests, 'native-cap-OFF') }}
+      uses: ./.github/actions/multi-functest
+      with:
+        gh_token: ${{ inputs.gh_token }}
+        compile_mode: native
+        cflags: "-std=c11 -D_GNU_SOURCE -DMLD_CONFIG_FILE=\\\\\\\"../../test/custom_native_capability_config_0.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        ldflags: "-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        func: true
+        kat: true
+        acvp: true
+        opt: ${{ inputs.opt }}
+        examples: false # Some examples use a custom config themselves
+    - name: "Custom native capability functions (CPUID AVX2 detection)"
+      if: ${{ (inputs.tests == 'all' || contains(inputs.tests, 'native-cap-CPUID_AVX2')) && runner.os == 'Linux' && runner.arch == 'X64' }}
+      uses: ./.github/actions/multi-functest
+      with:
+        gh_token: ${{ inputs.gh_token }}
+        compile_mode: native
+        cflags: "-std=c11 -mavx2 -mbmi2 -mpopcnt -D_GNU_SOURCE -DMLD_CONFIG_FILE=\\\\\\\"../../test/custom_native_capability_config_CPUID_AVX2.h\\\\\\\" -fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        ldflags: "-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all"
+        func: true
+        kat: true
+        acvp: true
+        opt: ${{ inputs.opt }}
+        examples: false # Some examples use a custom config themselves

BIBLIOGRAPHY.md

@@ -26,6 +26,9 @@ source code and documentation.
   - [test/break_pct_config.h](test/break_pct_config.h)
   - [test/custom_memcpy_config.h](test/custom_memcpy_config.h)
   - [test/custom_memset_config.h](test/custom_memset_config.h)
+  - [test/custom_native_capability_config_0.h](test/custom_native_capability_config_0.h)
+  - [test/custom_native_capability_config_1.h](test/custom_native_capability_config_1.h)
+  - [test/custom_native_capability_config_CPUID_AVX2.h](test/custom_native_capability_config_CPUID_AVX2.h)
   - [test/custom_randombytes_config.h](test/custom_randombytes_config.h)
   - [test/custom_stdlib_config.h](test/custom_stdlib_config.h)
   - [test/custom_zeroize_config.h](test/custom_zeroize_config.h)
@@ -68,6 +71,9 @@ source code and documentation.
   - [test/break_pct_config.h](test/break_pct_config.h)
   - [test/custom_memcpy_config.h](test/custom_memcpy_config.h)
   - [test/custom_memset_config.h](test/custom_memset_config.h)
+  - [test/custom_native_capability_config_0.h](test/custom_native_capability_config_0.h)
+  - [test/custom_native_capability_config_1.h](test/custom_native_capability_config_1.h)
+  - [test/custom_native_capability_config_CPUID_AVX2.h](test/custom_native_capability_config_CPUID_AVX2.h)
   - [test/custom_randombytes_config.h](test/custom_randombytes_config.h)
   - [test/custom_stdlib_config.h](test/custom_stdlib_config.h)
   - [test/custom_zeroize_config.h](test/custom_zeroize_config.h)

dev/aarch64_clean/meta.h

@@ -32,12 +32,14 @@
 
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/arith_native_aarch64.h"
 
-static MLD_INLINE void mld_ntt_native(int32_t data[MLDSA_N])
+static MLD_INLINE int mld_ntt_native(int32_t data[MLDSA_N])
 {
   mld_ntt_asm(data, mld_aarch64_ntt_zetas_layer123456,
               mld_aarch64_ntt_zetas_layer78);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 
 static MLD_INLINE void mld_intt_native(int32_t data[MLDSA_N])

dev/x86_64/meta.h

@@ -33,16 +33,23 @@
 #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])
 {
   mld_nttunpack_avx2((__m256i *)(data));
 }
 
-static MLD_INLINE void mld_ntt_native(int32_t data[MLDSA_N])
+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((__m256i *)data, mld_qdata.vec);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 static MLD_INLINE void mld_intt_native(int32_t data[MLDSA_N])
 {

mldsa/mldsa_native.c

@@ -392,8 +392,10 @@
 #undef MLD_RESTRICT
 #undef MLD_SYS_AARCH64
 #undef MLD_SYS_AARCH64_EB
+#undef MLD_SYS_APPLE
 #undef MLD_SYS_BIG_ENDIAN
 #undef MLD_SYS_H
+#undef MLD_SYS_LINUX
 #undef MLD_SYS_LITTLE_ENDIAN
 #undef MLD_SYS_PPC64LE
 #undef MLD_SYS_RISCV64
@@ -508,6 +510,8 @@
 #if defined(MLD_CONFIG_USE_NATIVE_BACKEND_ARITH)
 /* mldsa/src/native/api.h */
 #undef MLD_NATIVE_API_H
+#undef MLD_NATIVE_FUNC_FALLBACK
+#undef MLD_NATIVE_FUNC_SUCCESS
 /* mldsa/src/native/meta.h */
 #undef MLD_NATIVE_META_H
 #if defined(MLD_SYS_AARCH64)

mldsa/src/config.h

@@ -298,6 +298,36 @@
    #endif
 */
 
+
+/******************************************************************************
+ * Name:        MLD_CONFIG_CUSTOM_CAPABILITY_FUNC
+ *
+ * Description: mldsa-native backends may rely on specific hardware features.
+ *              Those backends will only be included in an mldsa-native build
+ *              if support for the respective features is enabled at
+ *              compile-time. However, when building for a heteroneous set
+ *              of CPUs to run the resulting binary/library on, feature
+ *              detection at _runtime_ is needed to decided whether a backend
+ *              can be used or not.
+ *
+ *              Set this option and define `mld_sys_check_capability` if you
+ *              want to use a custom method to dispatch between implementations.
+ *
+ *              If this option is not set, mldsa-native uses compile-time
+ *              feature detection only to decide which backend to use.
+ *
+ *              If you compile mldsa-native on a system with different
+ *              capabilities than the system that the resulting binary/library
+ *              will be run on, you must use this option.
+ *
+ *****************************************************************************/
+/* #define MLD_CONFIG_CUSTOM_CAPABILITY_FUNC
+   static MLD_INLINE int mld_sys_check_capability(mld_sys_cap cap)
+   {
+       ... your implementation ...
+   }
+*/
+
 /******************************************************************************
  * Name:        MLD_CONFIG_KEYGEN_PCT
  *

mldsa/src/native/aarch64/meta.h

@@ -32,12 +32,14 @@
 
 
 #if !defined(__ASSEMBLER__)
+#include "../api.h"
 #include "src/arith_native_aarch64.h"
 
-static MLD_INLINE void mld_ntt_native(int32_t data[MLDSA_N])
+static MLD_INLINE int mld_ntt_native(int32_t data[MLDSA_N])
 {
   mld_ntt_asm(data, mld_aarch64_ntt_zetas_layer123456,
               mld_aarch64_ntt_zetas_layer78);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 
 static MLD_INLINE void mld_intt_native(int32_t data[MLDSA_N])

mldsa/src/native/api.h

@@ -22,6 +22,14 @@
 #include "../cbmc.h"
 #include "../common.h"
 
+/* Backends must return MLD_NATIVE_FUNC_SUCCESS upon success. */
+#define MLD_NATIVE_FUNC_SUCCESS (0)
+/* Backends may return MLD_NATIVE_FUNC_FALLBACK to signal to the frontend that
+ * the target/parameters are unsupported; typically, this would be because of
+ * dependencies on CPU features not detected on the host CPU. In this case,
+ * the frontend falls back to the default C implementation. */
+#define MLD_NATIVE_FUNC_FALLBACK (-1)
+
 /*
  * This is the C<->native interface allowing for the drop-in of
  * native code for performance critical arithmetic components of ML-DSA.
@@ -52,7 +60,7 @@
  *
  * Arguments:   - int32_t p[MLDSA_N]: pointer to in/output polynomial
  **************************************************/
-static MLD_INLINE void mld_ntt_native(int32_t p[MLDSA_N]);
+static MLD_INLINE int mld_ntt_native(int32_t p[MLDSA_N]);
 #endif /* MLD_USE_NATIVE_NTT */
 
 

mldsa/src/native/x86_64/meta.h

@@ -33,16 +33,23 @@
 #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])
 {
   mld_nttunpack_avx2((__m256i *)(data));
 }
 
-static MLD_INLINE void mld_ntt_native(int32_t data[MLDSA_N])
+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((__m256i *)data, mld_qdata.vec);
+  return MLD_NATIVE_FUNC_SUCCESS;
 }
 static MLD_INLINE void mld_intt_native(int32_t data[MLDSA_N])
 {

mldsa/src/poly.c

@@ -142,23 +142,24 @@ void mld_poly_shiftl(mld_poly *a)
   mld_assert_bound(a->coeffs, MLDSA_N, 0, MLDSA_Q);
 }
 
-#if !defined(MLD_USE_NATIVE_NTT)
 MLD_INTERNAL_API
 void mld_poly_ntt(mld_poly *a)
 {
   mld_assert_abs_bound(a->coeffs, MLDSA_N, MLDSA_Q);
+#if defined(MLD_USE_NATIVE_NTT)
+  {
+    int ret;
+    ret = mld_ntt_native(a->coeffs);
+    if (ret == MLD_NATIVE_FUNC_SUCCESS)
+    {
+      mld_assert_abs_bound(a->coeffs, MLDSA_N, MLD_NTT_BOUND);
+      return;
+    }
+  }
+#endif /* MLD_USE_NATIVE_NTT */
   mld_ntt(a->coeffs);
   mld_assert_abs_bound(a->coeffs, MLDSA_N, MLD_NTT_BOUND);
 }
-#else  /* !MLD_USE_NATIVE_NTT */
-MLD_INTERNAL_API
-void mld_poly_ntt(mld_poly *p)
-{
-  mld_assert_abs_bound(p->coeffs, MLDSA_N, MLDSA_Q);
-  mld_ntt_native(p->coeffs);
-  mld_assert_abs_bound(p->coeffs, MLDSA_N, MLD_NTT_BOUND);
-}
-#endif /* MLD_USE_NATIVE_NTT */
 
 #if !defined(MLD_USE_NATIVE_INTT)
 MLD_INTERNAL_API