tbb support in blake3_c_rust_bindings

5 files changed, 156 insertions, 0 deletions

diff --git a/c/blake3_c_rust_bindings/Cargo.toml b/c/blake3_c_rust_bindings/Cargo.toml
index e70f16b..4c3a252 100644
--- a/c/blake3_c_rust_bindings/Cargo.toml
+++ b/c/blake3_c_rust_bindings/Cargo.toml
@@ -16,6 +16,8 @@ prefer_intrinsics = []
 # Activate NEON bindings. We don't currently do any CPU feature detection for
 # this. If this Cargo feature is on, the NEON gets used.
 neon = []
+# Enable TBB-based multithreading.
+tbb = []
 
 [dev-dependencies]
 arrayref = "0.3.5"
diff --git a/c/blake3_c_rust_bindings/benches/bench.rs b/c/blake3_c_rust_bindings/benches/bench.rs
index 5011347..b6c5704 100644
--- a/c/blake3_c_rust_bindings/benches/bench.rs
+++ b/c/blake3_c_rust_bindings/benches/bench.rs
@@ -371,6 +371,90 @@ fn bench_incremental_1024_kib(b: &mut Bencher) {
     bench_incremental(b, 1024 * KIB);
 }
 
+#[cfg(feature = "tbb")]
+fn bench_tbb(b: &mut Bencher, len: usize) {
+    let mut input = RandomInput::new(b, len);
+    b.iter(|| {
+        let mut hasher = blake3_c_rust_bindings::Hasher::new();
+        hasher.update_tbb(input.get());
+        let mut out = [0; 32];
+        hasher.finalize(&mut out);
+        out
+    });
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0001_block(b: &mut Bencher) {
+    bench_tbb(b, BLOCK_LEN);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0001_kib(b: &mut Bencher) {
+    bench_tbb(b, 1 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0002_kib(b: &mut Bencher) {
+    bench_tbb(b, 2 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0004_kib(b: &mut Bencher) {
+    bench_tbb(b, 4 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0008_kib(b: &mut Bencher) {
+    bench_tbb(b, 8 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0016_kib(b: &mut Bencher) {
+    bench_tbb(b, 16 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0032_kib(b: &mut Bencher) {
+    bench_tbb(b, 32 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0064_kib(b: &mut Bencher) {
+    bench_tbb(b, 64 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0128_kib(b: &mut Bencher) {
+    bench_tbb(b, 128 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0256_kib(b: &mut Bencher) {
+    bench_tbb(b, 256 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_0512_kib(b: &mut Bencher) {
+    bench_tbb(b, 512 * KIB);
+}
+
+#[bench]
+#[cfg(feature = "tbb")]
+fn bench_tbb_1024_kib(b: &mut Bencher) {
+    bench_tbb(b, 1024 * KIB);
+}
+
 // This checks that update() splits up its input in increasing powers of 2, so
 // that it can recover a high degree of parallelism when the number of bytes
 // hashed so far is uneven. The performance of this benchmark should be
diff --git a/c/blake3_c_rust_bindings/build.rs b/c/blake3_c_rust_bindings/build.rs
index c5dc927..e176dac 100644
--- a/c/blake3_c_rust_bindings/build.rs
+++ b/c/blake3_c_rust_bindings/build.rs
@@ -52,6 +52,21 @@ fn new_build() -> cc::Build {
     build
 }
 
+fn new_cpp_build() -> cc::Build {
+    let mut build = cc::Build::new();
+    build.cpp(true);
+    if is_windows_msvc() {
+        build.flag("/std:c++20");
+        build.flag("/EHs-c-");
+        build.flag("/GR-");
+    } else {
+        build.flag("-std=c++20");
+        build.flag("-fno-exceptions");
+        build.flag("-fno-rtti");
+    }
+    build
+}
+
 fn c_dir_path(filename: &str) -> String {
     // The `cross` tool doesn't support reading files in parent directories. As a hacky workaround
     // in `cross_test.sh`, we move the c/ directory around and set BLAKE3_C_DIR_OVERRIDE. Regular
@@ -68,8 +83,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     base_build.file(c_dir_path("blake3.c"));
     base_build.file(c_dir_path("blake3_dispatch.c"));
     base_build.file(c_dir_path("blake3_portable.c"));
+    if cfg!(feature = "tbb") {
+        base_build.define("BLAKE3_USE_TBB", "1");
+    }
     base_build.compile("blake3_base");
 
+    if cfg!(feature = "tbb") {
+        let mut tbb_build = new_cpp_build();
+        tbb_build.define("BLAKE3_USE_TBB", "1");
+        tbb_build.file(c_dir_path("blake3_tbb.cpp"));
+        tbb_build.compile("blake3_tbb");
+        println!("cargo::rustc-link-lib=tbb");
+    }
+
     if is_x86_64() && !defined("CARGO_FEATURE_PREFER_INTRINSICS") {
         // On 64-bit, use the assembly implementations, unless the
         // "prefer_intrinsics" feature is enabled.
diff --git a/c/blake3_c_rust_bindings/src/lib.rs b/c/blake3_c_rust_bindings/src/lib.rs
index ce7185e..c2b3989 100644
--- a/c/blake3_c_rust_bindings/src/lib.rs
+++ b/c/blake3_c_rust_bindings/src/lib.rs
@@ -82,6 +82,17 @@ impl Hasher {
         }
     }
 
+    #[cfg(feature = "tbb")]
+    pub fn update_tbb(&mut self, input: &[u8]) {
+        unsafe {
+            ffi::blake3_hasher_update_tbb(
+                &mut self.0,
+                input.as_ptr() as *const c_void,
+                input.len(),
+            );
+        }
+    }
+
     pub fn finalize(&self, output: &mut [u8]) {
         unsafe {
             ffi::blake3_hasher_finalize(&self.0, output.as_mut_ptr(), output.len());
@@ -140,6 +151,12 @@ pub mod ffi {
             input: *const ::std::os::raw::c_void,
             input_len: usize,
         );
+        #[cfg(feature = "tbb")]
+        pub fn blake3_hasher_update_tbb(
+            self_: *mut blake3_hasher,
+            input: *const ::std::os::raw::c_void,
+            input_len: usize,
+        );
         pub fn blake3_hasher_finalize(self_: *const blake3_hasher, out: *mut u8, out_len: usize);
         pub fn blake3_hasher_finalize_seek(
             self_: *const blake3_hasher,
diff --git a/c/blake3_c_rust_bindings/src/test.rs b/c/blake3_c_rust_bindings/src/test.rs
index 38cbd5f..6ff61ec 100644
--- a/c/blake3_c_rust_bindings/src/test.rs
+++ b/c/blake3_c_rust_bindings/src/test.rs
@@ -479,6 +479,15 @@ fn test_compare_reference_impl() {
             let mut test_out = [0; OUT];
             test_hasher.finalize(&mut test_out);
             assert_eq!(test_out[..], expected_out[..]);
+
+            #[cfg(feature = "tbb")]
+            {
+                let mut tbb_hasher = crate::Hasher::new();
+                tbb_hasher.update_tbb(input);
+                let mut tbb_out = [0; OUT];
+                tbb_hasher.finalize(&mut tbb_out);
+                assert_eq!(tbb_out[..], expected_out[..]);
+            }
         }
 
         // keyed
@@ -493,6 +502,15 @@ fn test_compare_reference_impl() {
             let mut test_out = [0; OUT];
             test_hasher.finalize(&mut test_out);
             assert_eq!(test_out[..], expected_out[..]);
+
+            #[cfg(feature = "tbb")]
+            {
+                let mut tbb_hasher = crate::Hasher::new_keyed(&TEST_KEY);
+                tbb_hasher.update_tbb(input);
+                let mut tbb_out = [0; OUT];
+                tbb_hasher.finalize(&mut tbb_out);
+                assert_eq!(tbb_out[..], expected_out[..]);
+            }
         }
 
         // derive_key
@@ -516,6 +534,15 @@ fn test_compare_reference_impl() {
             let mut test_out_raw = [0; OUT];
             test_hasher_raw.finalize(&mut test_out_raw);
             assert_eq!(test_out_raw[..], expected_out[..]);
+
+            #[cfg(feature = "tbb")]
+            {
+                let mut tbb_hasher = crate::Hasher::new_derive_key(context);
+                tbb_hasher.update_tbb(input);
+                let mut tbb_out = [0; OUT];
+                tbb_hasher.finalize(&mut tbb_out);
+                assert_eq!(tbb_out[..], expected_out[..]);
+            }
         }
     }
 }