ML-KEM: Import AArch64 backend from mlkem-native

[hanno-becker]

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license and the ISC license.

---

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.

Summary: This commit extends the mlkem-native->AWS-LC import to include the AArch64 arithmetic backend.

Details:

• crypto/fipsmodule/ml_kem/importer.sh now imports
  the arithmetic backend API header native/api.h as well
  as the native backend native/aarch64/*.

• The backend is imported as-is, with one exception:
  importer.sh converts the preprocessor directives used by
  mlkem-native into the ones used by s2n-bignum. This is to
  piggy-back on the adjustments made to the delocator to work
  with s2n-bignum assembly; otherwise, similar adjustments would
  likely be needed for mlkem-native assembly files.

• All imported functions are formally verified for functional
  correctness using HOL-Light. The proofs run as part of
  mlkem-native's CI. The HOL-Light specs are manually translated
  into CBMC specs in the header accompanying the ASM, and
  all higher level CBMC proofs conducted against those specs.
  Again, those are part of the mlkem-native CI.

• A backend header crypto/fipsmodule/ml_kem/mlkem_native_backend.h
  is added, activating the AArch64 arithmetic backend on Linux and
  MacOS AArch64 system, except if the NO_ASM directive is set
  (same as for s2n-bignum).

  Once the x86_64 arithmetic backend is ready for integration,
  it will be added to mlkem_native_backend.h as well.

• The backend header is registered in the configuration file
  crypto/fipsmodule/ml_kem/mlkem_native_config.h.

• The importer.sh is re-run.

---

Performance

Platform      | Algorithm   | Operation | New (ops/sec) | Old (ops/sec) | Improvement
--------------|-------------|-----------|---------------|---------------|-------------
Graviton 4    | ML-KEM-512  | keygen    |    29,145.9   |    19,666.0   |   +48.2%
              |             | encaps    |    91,043.3   |    64,240.2   |   +41.7%
              |             | decaps    |    75,242.9   |    52,486.8   |   +43.4%
              | ML-KEM-768  | keygen    |    18,190.1   |    12,240.0   |   +48.6%
              |             | encaps    |    57,600.4   |    39,490.6   |   +45.8%
              |             | decaps    |    48,334.1   |    33,298.1   |   +45.2%
              | ML-KEM-1024 | keygen    |    12,094.7   |     8,112.9   |   +49.1%
              |             | encaps    |    37,795.3   |    27,194.0   |   +39.0%
              |             | decaps    |    32,194.3   |    23,314.0   |   +38.1%

Graviton 3    | ML-KEM-512  | keygen    |    24,518.9   |    16,239.5   |   +51.0%
              |             | encaps    |    76,458.9   |    53,104.8   |   +44.0%
              |             | decaps    |    63,591.6   |    43,364.4   |   +46.6%
              | ML-KEM-768  | keygen    |    15,241.8   |    10,127.7   |   +50.5%
              |             | encaps    |    48,145.3   |    34,219.0   |   +40.7%
              |             | decaps    |    40,504.1   |    28,606.5   |   +41.6%
              | ML-KEM-1024 | keygen    |    10,179.0   |     6,695.3   |   +52.0%
              |             | encaps    |    31,879.8   |    23,317.0   |   +36.7%
              |             | decaps    |    27,168.6   |    19,866.5   |   +36.8%

Graviton 2    | ML-KEM-512  | keygen    |    16,570.3   |    10,901.8   |   +52.0%
              |             | encaps    |    51,617.2   |    35,299.9   |   +46.2%
              |             | decaps    |    42,396.2   |    28,533.3   |   +48.6%
              | ML-KEM-768  | keygen    |    10,348.5   |     6,848.2   |   +51.1%
              |             | encaps    |    32,580.5   |    22,890.7   |   +42.3%
              |             | decaps    |    27,176.9   |    18,999.9   |   +43.0%
              | ML-KEM-1024 | keygen    |     6,992.2   |     4,618.4   |   +51.4%
              |             | encaps    |    21,909.4   |    15,703.5   |   +39.5%
              |             | decaps    |    18,565.0   |    13,320.0   |   +39.4%

Example for header patching during import:

Before (mlkem-native):

#include "../../../common.h"
#if defined(MLK_ARITH_BACKEND_AARCH64) &&  !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)

.text
.balign 4
.global MLK_ASM_NAMESPACE(ntt_asm)
MLK_ASM_FN_SYMBOL(ntt_asm)

After (imported):

#include "_internal_s2n_bignum.h"

.text
.balign 4
        S2N_BN_SYM_VISIBILITY_DIRECTIVE(mlkem_ntt_asm)
        S2N_BN_SYM_PRIVACY_DIRECTIVE(mlkem_ntt_asm)
S2N_BN_SYMBOL(mlkem_ntt_asm):

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.73%. Comparing base (a435f20) to head (7d888dc).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2498   +/-   ##
=======================================
  Coverage   78.72%   78.73%           
=======================================
  Files         645      645           
  Lines      110641   110641           
  Branches    15648    15648           
=======================================
+ Hits        87099    87109   +10     
+ Misses      22843    22832   -11     
- Partials      699      700    +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[andrewhop]

Where does mlkem-native do the assembly dispatching? This raises an interesting gap we might have in our CI: building the arm implementation once and testing on all CPUs like we do for x86 with the Intel SDE.

[andrewhop]

Can we update aws-lc's code to not need to modify the mlkem-native code? I would like to keep this as simple as git clone and copying some directories.

[hanno-becker]

The delocator seems to have issues with nested relative #include's, so trying to have it process the ASM as it is in mlkem-native seems to be a challenge, at best. I'm not convinced that patching the delocator and build is better than patching the assembly headers to be compatible with the existing tooling.

What we're currently doing is replacing

#include "../../../common.h"
#if defined(MLK_ARITH_BACKEND_AARCH64) &&  !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED)

.text
.balign 4
.global MLK_ASM_NAMESPACE(ntt_asm)
MLK_ASM_FN_SYMBOL(ntt_asm)

by

#include "_internal_s2n_bignum.h"

.text
.balign 4
        S2N_BN_SYM_VISIBILITY_DIRECTIVE(mlkem_ntt_asm)
        S2N_BN_SYM_PRIVACY_DIRECTIVE(mlkem_ntt_asm)
S2N_BN_SYMBOL(mlkem_ntt_asm):

This doesn't look too bad to me. Also, there is still room for simplification, since in mlkem-native we could use a different macro structure (e.g. wrapping .global MLK_ASM_NAMESPACE(...) as MLK_ASM_GLOBAL(...) which would further simplify the search-and-replace transformation to the s2n-bignum directives.

[dkostic]

Is the new backend code supposed to work on all aarch64 processors? Or are there specific versions/capabilities we need to be present for the code to work?

[hanno-becker]

@dkostic The backend is supposed to work with all AArch64 processors. Since Neon is mandatory with AArch64 and we do not use any extensions, no runtime checks are required (at least, as far as I know).

[torben-hansen]

Could be simplified to

list(APPEND BCM_ASM_SOURCES

Not important.

[hanno-becker]

@dkostic @torben-hansen Do you want me to rebase or will you take care of that during merge?