aarch64: use target_feature=sha3 instructions if available

Author: ctz

admin/coverage

@@ -12,6 +12,12 @@ if [ `uname -p` == "x86_64" ] ; then
     env GRAVIOLA_CPU_DISABLE_sha=1 GRAVIOLA_CPU_DISABLE_bmi2=1 GRAVIOLA_CPU_DISABLE_avx512bw=1 GRAVIOLA_CPU_DISABLE_avx512f=1 GRAVIOLA_CPU_DISABLE_vaes=1 cargo test --locked
     env CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER="sde -dmr --" cargo test --locked
 fi
+case `uname -m` in
+    arm|arm64|aarch64)
+        env GRAVIOLA_CPU_DISABLE_sha3=1 GRAVIOLA_CPU_DISABLE_dit=1 cargo test --locked
+        ;;
+esac
+
 cargo run --example client https://jbp.io >/dev/null
 
 cargo llvm-cov report "$@"

admin/parse-asm/driver.py

@@ -228,7 +228,7 @@ def get_labels(self):
         return set(self.labels)
 
 
-RUST_FUNCTION_DECL = re.compile(r"^fn (?P<name>[a-zA-Z0-9_]+)\(")
+RUST_FUNCTION_DECL = re.compile(r"^(unsafe )?fn (?P<name>[a-zA-Z0-9_]+)\(")
 
 
 # A pass to collect (1) all macro definitions that reference labels, and (2) the
@@ -494,6 +494,7 @@ def emit_rust_function(
         return_value=None,
         assertions=[],
         allow_inline=True,
+        func_attrs="",
         hoist=None,
     ):
         self.macro_pass.emit_rust_function(name, rust_decl)
@@ -505,6 +506,7 @@ def emit_rust_function(
             return_value=return_value,
             assertions=assertions,
             allow_inline=allow_inline,
+            func_attrs=func_attrs,
             hoist=hoist,
         )
 
@@ -585,6 +587,7 @@ def emit_rust_function(
         return_value=None,
         assertions=[],
         allow_inline=True,
+        func_attrs="",
         hoist=None,
     ):
         self.expected_functions[name] = (
@@ -595,6 +598,7 @@ def emit_rust_function(
             allow_inline,
             hoist,
             assertions,
+            func_attrs,
         )
 
     def set_att_syntax(self, att_syntax):
@@ -794,10 +798,12 @@ def on_function(self, contexts, name):
                     allow_inline,
                     hoist,
                     assertions,
+                    func_attrs,
                 ) = defn
                 self.function_state.parameter_map = parameter_map
                 self.function_state.return_map = return_map
                 self.function_state.rust_decl = rust_decl
+                self.function_state.func_attrs = func_attrs
                 self.function_state.return_value = return_value
                 self.function_state.hoist = hoist
 
@@ -819,11 +825,15 @@ def on_function(self, contexts, name):
                     print("#[inline(never)]", file=self.output)
 
                 print(
-                    """pub(crate) %s {
+                    """%spub(crate) %s {
                     %s
                     // SAFETY: inline assembly. see [crate::low::inline_assembly_safety] for safety info.
                     unsafe { core::arch::asm!("""
-                    % (self.function_state.rust_decl, locals),
+                    % (
+                        self.function_state.func_attrs,
+                        self.function_state.rust_decl,
+                        locals,
+                    ),
                     file=self.output,
                 )
 
@@ -907,6 +917,7 @@ def on_asm(self, contexts, opcode, operands):
 
         self.visit_operands(operands)
 
+        operands = operands.replace("{", "{{").replace("}", "}}")
         operands = self.expand_rust_macros_in_asm(operands)
 
         parts = ['"    %-15s "' % opcode]

admin/parse-asm/driver.sha3.py

@@ -35,3 +35,19 @@
             rust_decl="fn sha3_keccak_f1600(a: &mut [u64; 25], rc: &[u64; 24])",
         )
         parse_file(input, d)
+
+    with (
+        open("../../thirdparty/s2n-bignum/arm/sha3/sha3_keccak_f1600_alt.S") as input,
+        open("../../graviola/src/low/aarch64/sha3_keccak_f1600_alt.rs", "w") as output,
+    ):
+        d = RustDriver(output, Architecture_aarch64)
+        d.emit_rust_function(
+            "sha3_keccak_f1600_alt",
+            parameter_map=[
+                ("inout", "a.as_mut_ptr() => _"),
+                ("inout", "rc.as_ptr() => _"),
+            ],
+            func_attrs='#[target_feature(enable = "sha3")]\n',
+            rust_decl="unsafe fn sha3_keccak_f1600(a: &mut [u64; 25], rc: &[u64; 24])",
+        )
+        parse_file(input, d)

admin/parse-asm/parse.py

@@ -11,7 +11,7 @@
 
 MACRO = re.compile(r"^(?P<name>[a-z0-9_]+)\((?P<args>[a-z0-9_,\[\]\+\* \#]*)\);?$")
 ASM = re.compile(
-    r"^(?P<opcode>[a-z][a-z0-9\.]*)\s?(?P<operands>[A-Za-z0-9_,\s\(\)\[\]\+\*\-~\t#\.!%$@]*) ?;? ?(//(?P<comment>[A-Za-z0-9 =\/@#\*\+\(\)^\.\<\>\-_:,\!\?\|])*)?$"
+    r"^(?P<opcode>[a-z][a-z0-9\.]*)\s?(?P<operands>[A-Za-z0-9_,\s\(\)\[\]{}\+\*\-~\t#\.!%$@]*) ?;? ?(//(?P<comment>[A-Za-z0-9 =\/@#\*\+\(\)^\.\<\>\-_:,\!\?\|])*)?$"
 )
 CONST = re.compile(r"\s?(?P<type>\.(quad|long))\s+(?P<value>((0x[0-9a-fA-F]+),?)+)")
 LABEL = re.compile(r"^(?P<name>(\.L)?[a-zA-Z0-9_]+):$")
@@ -249,7 +249,7 @@ def tokenise_gen(s):
         label = re.compile(r"^\.?[a-zA-Z][a-zA-Z0-9_]+")
         comment = re.compile(r"^/\*.*?\*/")
         number = re.compile(r"^[\$#]?(-?0x[0-9a-fA-F]+|-?[0-9]+)")
-        operator = re.compile(r'^["\(\)\[\]\+\*/\-,;:#\.!]')
+        operator = re.compile(r'^["\(\)\[\]{}\+\*/\-,;:#\.!][{}]?')
         whitespace = re.compile(r"^\s+")
 
         while s:

graviola/src/low/aarch64/cpu.rs

@@ -1,8 +1,6 @@
 // Written for Graviola by Joe Birr-Pixton, 2024.
 // SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0
 
-use std::arch::is_aarch64_feature_detected;
-
 pub(crate) fn enter_cpu_state() -> u32 {
     dit::maybe_enable()
 }
@@ -157,28 +155,77 @@ pub(in crate::low) unsafe fn ct_compare_bytes(a: *const u8, b: *const u8, len: u
     }
 }
 
+/// This macro interdicts is_aarch64_feature_detected to allow testability.
+macro_rules! have_cpu_feature {
+    ("neon") => {
+        crate::low::aarch64::cpu::test_toggle(
+            "neon",
+            std::arch::is_aarch64_feature_detected!("neon"),
+        )
+    };
+    ("aes") => {
+        crate::low::aarch64::cpu::test_toggle("aes", std::arch::is_aarch64_feature_detected!("aes"))
+    };
+    ("pmull") => {
+        crate::low::aarch64::cpu::test_toggle(
+            "pmull",
+            std::arch::is_aarch64_feature_detected!("pmull"),
+        )
+    };
+    ("sha2") => {
+        crate::low::aarch64::cpu::test_toggle(
+            "sha2",
+            std::arch::is_aarch64_feature_detected!("sha2"),
+        )
+    };
+    ("sha3") => {
+        crate::low::aarch64::cpu::test_toggle(
+            "sha3",
+            std::arch::is_aarch64_feature_detected!("sha3"),
+        )
+    };
+    ("dit") => {
+        crate::low::aarch64::cpu::test_toggle("dit", std::arch::is_aarch64_feature_detected!("dit"))
+    };
+}
+
+/// Token type reflecting the check for the SHA3 CPU feature
+///
+/// A value of this type is proof that the CPU dynamic feature check has happened.
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct HaveSha3(());
+
+impl HaveSha3 {
+    pub(crate) fn check() -> Option<Self> {
+        match have_cpu_feature!("sha3") {
+            true => Some(Self(())),
+            false => None,
+        }
+    }
+}
+
 pub(crate) fn verify_cpu_features() {
     assert!(
-        is_aarch64_feature_detected!("neon"),
+        have_cpu_feature!("neon"),
         "graviola requires neon CPU support"
     );
     assert!(
-        is_aarch64_feature_detected!("aes"),
+        have_cpu_feature!("aes"),
         "graviola requires aes CPU support"
     );
     assert!(
-        is_aarch64_feature_detected!("pmull"),
+        have_cpu_feature!("pmull"),
         "graviola requires pmull CPU support"
     );
     assert!(
-        is_aarch64_feature_detected!("sha2"),
+        have_cpu_feature!("sha2"),
         "graviola requires sha2 CPU support"
     );
 }
 
 mod dit {
     pub(super) fn maybe_enable() -> u32 {
-        if super::is_aarch64_feature_detected!("dit") {
+        if have_cpu_feature!("dit") {
             // SAFETY: in this branch, we verified `dit` cpu feature is supported
             match unsafe { read() } {
                 0 => {
@@ -229,6 +276,21 @@ mod dit {
     }
 }
 
+#[cfg(not(debug_assertions))]
+pub(crate) fn test_toggle(_id: &str, detected: bool) -> bool {
+    detected
+}
+
+#[cfg(debug_assertions)]
+pub(crate) fn test_toggle(id: &str, detected: bool) -> bool {
+    if std::env::var(format!("GRAVIOLA_CPU_DISABLE_{id}")).is_ok() {
+        println!("DEBUG: denying cpuid {id:?}");
+        false
+    } else {
+        detected
+    }
+}
+
 /// Read-only prefetch hint.
 pub(in crate::low) fn prefetch_ro<T>(ptr: *const T) {
     // SAFETY: inline assembly

graviola/src/low/aarch64/mod.rs

@@ -68,3 +68,5 @@ pub(crate) mod p384_montjadd;
 pub(crate) mod p384_montjdouble;
 pub(crate) mod sha256;
 pub(crate) mod sha3_keccak_f1600;
+pub(crate) mod sha3_keccak_f1600_alt;
+pub(crate) mod sha3_keccak_f1600_mux;