Implement compress for NEON

BLAKE3-team · Jan 7, 2024 · 6f2d724 · 6f2d724
1 parent 4d32708
commit 6f2d724
Show file tree

Hide file tree

Showing 7 changed files with 397 additions and 45 deletions.
diff --git a/benches/bench.rs b/benches/bench.rs
@@ -85,6 +85,14 @@ fn bench_single_compression_avx512(b: &mut Bencher) {
     }
 }
 
+#[bench]
+#[cfg(feature = "neon")]
+fn bench_single_compression_neon(b: &mut Bencher) {
+    if let Some(platform) = Platform::neon() {
+        bench_single_compression_fn(b, platform);
+    }
+}
+
 fn bench_many_chunks_fn(b: &mut Bencher, platform: Platform) {
     let degree = platform.simd_degree();
     let mut inputs = Vec::new();

diff --git a/c/blake3_c_rust_bindings/src/lib.rs b/c/blake3_c_rust_bindings/src/lib.rs
@@ -289,6 +289,21 @@ pub mod ffi {
     pub mod neon {
         extern "C" {
             // NEON low level functions
+            pub fn blake3_compress_xof_neon(
+                cv: *const u32,
+                block: *const u8,
+                block_len: u8,
+                counter: u64,
+                flags: u8,
+                out: *mut u8,
+            );
+            pub fn blake3_compress_in_place_neon(
+                cv: *mut u32,
+                block: *const u8,
+                block_len: u8,
+                counter: u64,
+                flags: u8,
+            );
             pub fn blake3_hash_many_neon(
                 inputs: *const *const u8,
                 num_inputs: usize,

diff --git a/c/blake3_dispatch.c b/c/blake3_dispatch.c
@@ -188,6 +188,12 @@ void blake3_compress_in_place(uint32_t cv[8],
   }
 #endif
 #endif
+
+#if BLAKE3_USE_NEON == 1
+  blake3_compress_in_place_neon(cv, block, block_len, counter, flags);
+  return;
+#endif
+
   blake3_compress_in_place_portable(cv, block, block_len, counter, flags);
 }
 
@@ -217,6 +223,12 @@ void blake3_compress_xof(const uint32_t cv[8],
   }
 #endif
 #endif
+
+#if BLAKE3_USE_NEON == 1
+  blake3_compress_xof_neon(cv, block, block_len, counter, flags, out);
+  return;
+#endif
+
   blake3_compress_xof_portable(cv, block, block_len, counter, flags, out);
 }
 

diff --git a/c/blake3_impl.h b/c/blake3_impl.h
@@ -274,6 +274,16 @@ void blake3_hash_many_avx512(const uint8_t *const *inputs, size_t num_inputs,
 #endif
 
 #if BLAKE3_USE_NEON == 1
+void blake3_compress_in_place_neon(uint32_t cv[8],
+                                     const uint8_t block[BLAKE3_BLOCK_LEN],
+                                     uint8_t block_len, uint64_t counter,
+                                     uint8_t flags);
+
+void blake3_compress_xof_neon(const uint32_t cv[8],
+                                const uint8_t block[BLAKE3_BLOCK_LEN],
+                                uint8_t block_len, uint64_t counter,
+                                uint8_t flags, uint8_t out[64]);
+
 void blake3_hash_many_neon(const uint8_t *const *inputs, size_t num_inputs,
                            size_t blocks, const uint32_t key[8],
                            uint64_t counter, bool increment_counter,