AArch64: Add MLK_SYS_AARCH64_FAST_SHA3 for fast SHA3 CPUs

Introduce MLK_SYS_AARCH64_FAST_SHA3 to signal AArch64 CPUs that implement
SHA3 instructions on all Neon execution units. When set, the FIPS202 backend
auto-selection picks the pure-Neon v84a Keccak-f1600 and Keccak-f1600x2
implementations, matching the existing Apple-CPU path. Set the flag for the
new Graviton5 (c9g) benchmark entry.

- Resolves #1783

Signed-off-by: Matthias J. Kannwischer <matthias@zerorisc.com>

Author: mkannwischer

.github/workflows/bench.yml

@@ -130,7 +130,7 @@ jobs:
           - name: Graviton5
             ec2_instance_type: c9g.medium
             ec2_ami: ubuntu-latest (aarch64)
-            archflags: -march=armv9-a+sha3
+            archflags: -march=armv9-a+sha3 -DMLK_SYS_AARCH64_FAST_SHA3
             cflags: "-flto -DMLK_FORCE_AARCH64"
             ldflags: "-flto"
             perf: PERF

dev/fips202/aarch64/auto.h

@@ -24,13 +24,15 @@
 /*
  * Keccak-f1600
  *
- * - On Arm-based Apple CPUs, we pick a pure Neon implementation.
+ * - On Arm-based Apple CPUs, or if MLK_SYS_AARCH64_FAST_SHA3 is set,
+ *   we pick a pure Neon implementation.
  * - Otherwise, unless MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER is set,
  *   we use lazy-rotation scalar assembly from @[HYBRID].
  * - Otherwise, if MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER is set, we
  *   fall back to the standard C implementation.
  */
-#if defined(__ARM_FEATURE_SHA3) && defined(__APPLE__)
+#if defined(__ARM_FEATURE_SHA3) && \
+    (defined(__APPLE__) || defined(MLK_SYS_AARCH64_FAST_SHA3))
 #include "x1_v84a.h"
 #elif !defined(MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER)
 #include "x1_scalar.h"
@@ -42,20 +44,21 @@
  * The optimal implementation is highly CPU-specific; see @[HYBRID].
  *
  * For now, if v8.4-A is not implemented, we fall back to Keccak-f1600.
- * If v8.4-A is implemented and we are on an Apple CPU, we use a plain
- * Neon-based implementation.
- * If v8.4-A is implemented and we are not on an Apple CPU, we use a
- * scalar/Neon/Neon hybrid.
- * The reason for this distinction is that Apple CPUs appear to implement
- * the SHA3 instructions on all SIMD units, while Arm CPUs prior to Cortex-X4
- * don't, and ordinary Neon instructions are still needed.
+ * If v8.4-A is implemented and we are on an Apple CPU or
+ * MLK_SYS_AARCH64_FAST_SHA3 is set, we use a plain Neon-based
+ * implementation.
+ * Otherwise, if v8.4-A is implemented, we use a scalar/Neon/Neon hybrid.
+ * The reason for this distinction is that Apple CPUs (and CPUs flagged with
+ * MLK_SYS_AARCH64_FAST_SHA3) implement the SHA3 instructions on all SIMD
+ * units, while Arm CPUs prior to Cortex-X4 don't, and ordinary Neon
+ * instructions are still needed.
  */
 #if defined(__ARM_FEATURE_SHA3)
 /*
- * For Apple-M cores, we use a plain implementation leveraging SHA3
- * instructions only.
+ * For Apple-M cores (and CPUs flagged with MLK_SYS_AARCH64_FAST_SHA3), we
+ * use a plain implementation leveraging SHA3 instructions only.
  */
-#if defined(__APPLE__)
+#if defined(__APPLE__) || defined(MLK_SYS_AARCH64_FAST_SHA3)
 #include "x2_v84a.h"
 #else
 #include "x4_v8a_v84a_scalar.h"

mlkem/src/fips202/native/aarch64/auto.h

@@ -24,13 +24,15 @@
 /*
  * Keccak-f1600
  *
- * - On Arm-based Apple CPUs, we pick a pure Neon implementation.
+ * - On Arm-based Apple CPUs, or if MLK_SYS_AARCH64_FAST_SHA3 is set,
+ *   we pick a pure Neon implementation.
  * - Otherwise, unless MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER is set,
  *   we use lazy-rotation scalar assembly from @[HYBRID].
  * - Otherwise, if MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER is set, we
  *   fall back to the standard C implementation.
  */
-#if defined(__ARM_FEATURE_SHA3) && defined(__APPLE__)
+#if defined(__ARM_FEATURE_SHA3) && \
+    (defined(__APPLE__) || defined(MLK_SYS_AARCH64_FAST_SHA3))
 #include "x1_v84a.h"
 #elif !defined(MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER)
 #include "x1_scalar.h"
@@ -42,20 +44,21 @@
  * The optimal implementation is highly CPU-specific; see @[HYBRID].
  *
  * For now, if v8.4-A is not implemented, we fall back to Keccak-f1600.
- * If v8.4-A is implemented and we are on an Apple CPU, we use a plain
- * Neon-based implementation.
- * If v8.4-A is implemented and we are not on an Apple CPU, we use a
- * scalar/Neon/Neon hybrid.
- * The reason for this distinction is that Apple CPUs appear to implement
- * the SHA3 instructions on all SIMD units, while Arm CPUs prior to Cortex-X4
- * don't, and ordinary Neon instructions are still needed.
+ * If v8.4-A is implemented and we are on an Apple CPU or
+ * MLK_SYS_AARCH64_FAST_SHA3 is set, we use a plain Neon-based
+ * implementation.
+ * Otherwise, if v8.4-A is implemented, we use a scalar/Neon/Neon hybrid.
+ * The reason for this distinction is that Apple CPUs (and CPUs flagged with
+ * MLK_SYS_AARCH64_FAST_SHA3) implement the SHA3 instructions on all SIMD
+ * units, while Arm CPUs prior to Cortex-X4 don't, and ordinary Neon
+ * instructions are still needed.
  */
 #if defined(__ARM_FEATURE_SHA3)
 /*
- * For Apple-M cores, we use a plain implementation leveraging SHA3
- * instructions only.
+ * For Apple-M cores (and CPUs flagged with MLK_SYS_AARCH64_FAST_SHA3), we
+ * use a plain implementation leveraging SHA3 instructions only.
  */
-#if defined(__APPLE__)
+#if defined(__APPLE__) || defined(MLK_SYS_AARCH64_FAST_SHA3)
 #include "x2_v84a.h"
 #else
 #include "x4_v8a_v84a_scalar.h"

scripts/autogen

@@ -2882,6 +2882,7 @@ def get_config_options():
         "MLK_FORCE_RISCV64",
         "MLK_FORCE_RISCV32",
         "MLK_SYS_AARCH64_SLOW_BARREL_SHIFTER",
+        "MLK_SYS_AARCH64_FAST_SHA3",
         "MLKEM_DEBUG",  # TODO: Rename?
         "MLK_BREAK_PCT",  # Use in PCT breakage test]
         "MLK_CHECK_APIS",