Merge pull request #31 from anton-rs/cl/aarch64-simd-keccak

✨ Add aarch64 SIMD keccak

.github/workflows/ci.yaml

@@ -20,7 +20,7 @@ jobs:
         run: sudo apt-get install xsltproc
       - uses: taiki-e/install-action@nextest
       - name: cargo test
-        run: cargo nextest run --workspace --all --all-features --locked
+        run: cargo nextest run --release --workspace --all --locked
   cargo-lint:
     runs-on: ubuntu-latest
     timeout-minutes: 20
@@ -37,9 +37,9 @@ jobs:
       - name: Install xsltproc
         run: sudo apt-get install xsltproc
       - name: cargo fmt
-        run: cargo +nightly fmt --all -- --check
+        run: cargo fmt --all -- --check
       - name: cargo clippy
-        run: cargo +nightly clippy --workspace --all --all-features --locked -- -D warnings
+        run: cargo clippy --workspace --all --locked -- -D warnings
   cargo-build:
     runs-on: ubuntu-latest
     timeout-minutes: 20
@@ -74,7 +74,7 @@ jobs:
       - name: doclint
         id: build
         continue-on-error: true
-        run: RUSTDOCFLAGS="-D warnings" cargo doc --all --no-deps --all-features --document-private-items
+        run: RUSTDOCFLAGS="-D warnings" cargo doc --all --no-deps --document-private-items
       - name: doctest
-        run: cargo test --doc --all --all-features --locked
+        run: cargo test --doc --all --locked
 

Cargo.toml

@@ -12,7 +12,6 @@ lto = true
 codegen-units = 1
 
 [profile.dev]
-opt-level = 1
 overflow-checks = false
 
 [profile.bench]

README.md

@@ -80,9 +80,9 @@ To get started, a few dependencies are required:
 
 ```sh
 # With `cargo-nextest`
-cargo nextest run --all --all-features
+cargo +nightly nextest run --release --all --all-features
 # Without `cargo-nextest`
-cargo t --all --all-features
+cargo +nightly t --release --all --all-features
 ```
 
 ### Linting and Formatting
@@ -94,7 +94,7 @@ cargo +nightly fmt --all -- && cargo +nightly clippy --all --all-features -- -D
 ### Running Benchmarks
 
 ```sh
-cargo bench --all --all-features
+cargo +nightly bench --all --all-features
 ```
 
 ## Documentation

crates/cannon/src/traces.rs

@@ -1,3 +1,5 @@
+#![allow(unused_imports)]
+
 /// Performs a tracing debug if the `tracing` feature is enabled.
 #[macro_export]
 macro_rules! debug {

crates/mipsevm/Cargo.toml

@@ -17,6 +17,7 @@ rustc-hash = "1.1.0"
 serde = { version = "1.0.188", features = ["derive"] }
 serde_json = "1.0.107"
 xkcp-rs = { git = "https://github.com/DaniPopes/xkcp-rs", rev = "40447a5" }
+keccak256-aarch64-simd = { git = "https://github.com/clabby/keccak256-aarch64", rev = "5c4c8f8", optional = true }
 
 # misc
 anyhow = "1.0.75"
@@ -30,6 +31,7 @@ proptest = "1.2.0"
 
 [features]
 tracing = ["dep:tracing"]
+simd-keccak = ["dep:keccak256-aarch64-simd"]
 
 [[bench]]
 name = "memory"

crates/mipsevm/README.md

@@ -60,10 +60,18 @@ Supported 55 instructions:
 | `Logical`            | `xori`        | Bitwise XOR immediate.                       |
 
 To run:
-1. Load a program into a state, e.g. using `LoadELF`.
-2. Patch the program if necessary: e.g. using `PatchGo` for Go programs, `PatchStack` for empty initial stack, etc.
+1. Load a program into a state, e.g. using `patch::load_elf`.
+2. Patch the program if necessary: e.g. using `patch::patch_go` for Go programs, `patch::patch_stack` for empty initial stack, etc.
 4. Implement the `PreimageOracle` interface
-5. Instrument the emulator with the state, and pre-image oracle, using `NewInstrumentedState`
-6. Step through the instrumented state with `Step(proof)`,
+5. Instrument the emulator with the state, and pre-image oracle, using `InstrumentedState::new`
+6. Step through the instrumented state with `step(proof)`,
    where `proof==true` if witness data should be generated. Steps are faster with `proof==false`.
 7. Optionally repeat the step on-chain by calling `MIPS.sol` and `PreimageOracle.sol`, using the above witness data.
+
+## Features
+- `tracing`: Enables tracing within the VM. This is useful for debugging, but does not need to be enabled in production
+   environments for performance reasons, unless a store of logs is required.
+- `simd-keccak`: Exclusive to ARMv8-A processors. Uses the [`keccak256-aarch64-simd`](https://github.com/clabby/keccak256-aarch64/tree/master) crate
+  for performance-critical `keccak256` hashing, which provides a very significant speedup to merkleization. **Warning**:
+  This crate is *highly* experimental, and it is not suggested that this feature is enabled in production, unless you
+  understand the risks associated with enabling it.

crates/mipsevm/src/memory.rs

@@ -3,7 +3,7 @@
 use crate::{
     page::{self},
     types::SharedCachedPage,
-    utils::keccak_concat_fixed,
+    utils::keccak_concat_hashes,
     Address, Gindex, Page, PageIndex,
 };
 use anyhow::Result;
@@ -148,7 +148,7 @@ impl Memory {
 
         let left = self.merkleize_subtree(g_index << 1)?;
         let right = self.merkleize_subtree((g_index << 1) | 1)?;
-        let result = *keccak_concat_fixed(left, right);
+        let result = *keccak_concat_hashes(left, right);
 
         self.nodes.insert(g_index, Some(result));
 
@@ -507,9 +507,9 @@ mod test {
             (32..proof.len()).step_by(32).for_each(|i| {
                 let sib: [u8; 32] = proof[i..i + 32].try_into().unwrap();
                 if path & 1 != 0 {
-                    node = *keccak_concat_fixed(sib, node);
+                    node = *keccak_concat_hashes(sib, node);
                 } else {
-                    node = *keccak_concat_fixed(node, sib);
+                    node = *keccak_concat_hashes(node, sib);
                 }
                 path >>= 1;
             });
@@ -590,15 +590,15 @@ mod test {
                 .merkleize_subtree((1 << page::PAGE_KEY_SIZE) | 6)
                 .unwrap();
             let z = page::ZERO_HASHES[page::PAGE_ADDRESS_SIZE - 5];
-            let r1 = keccak_concat_fixed(
-                keccak_concat_fixed(
-                    keccak_concat_fixed(z.into(), z.into()).into(),
-                    keccak_concat_fixed(z.into(), p3.into()).into(),
+            let r1 = keccak_concat_hashes(
+                keccak_concat_hashes(
+                    keccak_concat_hashes(z.into(), z.into()).into(),
+                    keccak_concat_hashes(z.into(), p3.into()).into(),
                 )
                 .into(),
-                keccak_concat_fixed(
-                    keccak_concat_fixed(z.into(), p5.into()).into(),
-                    keccak_concat_fixed(p6.into(), z.into()).into(),
+                keccak_concat_hashes(
+                    keccak_concat_hashes(z.into(), p5.into()).into(),
+                    keccak_concat_hashes(p6.into(), z.into()).into(),
                 )
                 .into(),
             );

crates/mipsevm/src/page.rs

@@ -1,9 +1,12 @@
 //! This module contains the data structure for a [Page] within the MIPS emulator's [Memory].
 
-use crate::{utils::keccak256, utils::keccak_concat_fixed, Address, Gindex, Page};
+use crate::{utils::keccak_concat_hashes, Address, Gindex, Page};
 use anyhow::Result;
 use once_cell::sync::Lazy;
 
+#[cfg(not(feature = "simd-keccak"))]
+use crate::utils::keccak256;
+
 pub(crate) const PAGE_ADDRESS_SIZE: usize = 12;
 pub(crate) const PAGE_KEY_SIZE: usize = 32 - PAGE_ADDRESS_SIZE;
 pub(crate) const PAGE_SIZE: usize = 1 << PAGE_ADDRESS_SIZE;
@@ -16,7 +19,7 @@ pub(crate) const PAGE_KEY_MASK: usize = MAX_PAGE_COUNT - 1;
 pub(crate) static ZERO_HASHES: Lazy<[[u8; 32]; 256]> = Lazy::new(|| {
     let mut out = [[0u8; 32]; 256];
     for i in 1..256 {
-        out[i] = *keccak_concat_fixed(out[i - 1], out[i - 1])
+        out[i] = *keccak_concat_hashes(out[i - 1], out[i - 1])
     }
     out
 });
@@ -117,14 +120,25 @@ impl CachedPage {
         let hash = if g_index >= PAGE_SIZE_WORDS >> 1 {
             // This is a leaf node.
             let data_idx = (g_index - (PAGE_SIZE_WORDS >> 1)) << 6;
+            #[cfg(feature = "simd-keccak")]
+            {
+                let mut out = [0u8; 32];
+                keccak256_aarch64_simd::simd_keccak256_64b_single(
+                    &self.data[data_idx..data_idx + 64],
+                    &mut out,
+                );
+                out
+            }
+
+            #[cfg(not(feature = "simd-keccak"))]
             *keccak256(&self.data[data_idx..data_idx + 64])
         } else {
             // This is an internal node.
             let left_child = g_index << 1;
             let right_child = left_child + 1;
 
             // Ensure children are hashed.
-            *keccak_concat_fixed(
+            *keccak_concat_hashes(
                 self.merkleize_subtree(left_child as Gindex)?,
                 self.merkleize_subtree(right_child as Gindex)?,
             )
@@ -152,12 +166,12 @@ mod test {
         assert_eq!(node, expected_leaf, "Leaf nodes should not be hashed");
 
         let node = page.merkleize_subtree(g_index >> 1).unwrap();
-        let expected_parent = keccak_concat_fixed(ZERO_HASHES[0].into(), expected_leaf.into());
+        let expected_parent = keccak_concat_hashes(ZERO_HASHES[0].into(), expected_leaf.into());
         assert_eq!(node, expected_parent, "Parent should be correct");
 
         let node = page.merkleize_subtree(g_index >> 2).unwrap();
         let expected_grandparent =
-            keccak_concat_fixed(expected_parent.into(), ZERO_HASHES[1].into());
+            keccak_concat_hashes(expected_parent.into(), ZERO_HASHES[1].into());
         assert_eq!(node, expected_grandparent, "Grandparent should be correct");
 
         let pre = page.merkle_root().unwrap();

crates/mipsevm/src/traces.rs

@@ -1,3 +1,5 @@
+#![allow(unused_imports)]
+
 /// Performs a tracing debug if the `tracing` feature is enabled.
 #[macro_export]
 macro_rules! debug {

crates/mipsevm/src/utils.rs

@@ -15,12 +15,17 @@ where
     concatenated
 }
 
-/// Hash the concatenation of two fixed sized arrays.
+/// Hash the concatenation of two 32 byte digests.
 #[inline(always)]
-pub(crate) fn keccak_concat_fixed<const N: usize, const M: usize>(a: [u8; N], b: [u8; M]) -> B256
-where
-    [(); N + M]:,
-{
+pub(crate) fn keccak_concat_hashes(a: [u8; 32], b: [u8; 32]) -> B256 {
+    #[cfg(feature = "simd-keccak")]
+    {
+        let mut out = B256::ZERO;
+        keccak256_aarch64_simd::simd_keccak256_64b_single(&concat_fixed(a, b), out.as_mut());
+        out
+    }
+
+    #[cfg(not(feature = "simd-keccak"))]
     keccak256(concat_fixed(a, b).as_slice())
 }
 

crates/preimage/src/traces.rs

@@ -1,3 +1,5 @@
+#![allow(unused_imports)]
+
 /// Performs a tracing debug if the `tracing` feature is enabled.
 #[macro_export]
 macro_rules! debug {