add Rust FFI wrappers for AVX-512 and NEON

1 files changed, 65 insertions, 0 deletions

diff --git a/build.rs b/build.rs
new file mode 100644
index 0000000..f710508
--- /dev/null
+++ b/build.rs
@@ -0,0 +1,65 @@
+use std::env;
+
+fn defined(var: &str) -> bool {
+    env::var_os(var).is_some()
+}
+
+fn is_windows() -> bool {
+    let target = env::var("TARGET").unwrap();
+    let target_components: Vec<&str> = target.split("-").collect();
+    let target_os = target_components[2];
+    target_os == "windows"
+}
+
+fn new_build() -> cc::Build {
+    let mut build = cc::Build::new();
+    if !is_windows() {
+        build.flag("-std=c11");
+    }
+    build
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    if defined("CARGO_FEATURE_C_AVX512") {
+        let mut build = new_build();
+        build.file("src/c/blake3_avx512.c");
+        if is_windows() {
+            // Note that a lot of versions of MSVC don't support /arch:AVX512,
+            // and they'll discard it with a warning, hopefully leading to a
+            // build error.
+            build.flag("/arch:AVX512");
+        } else {
+            build.flag("-mavx512f");
+            build.flag("-mavx512vl");
+        }
+        build.compile("blake3_avx512");
+    }
+
+    if defined("CARGO_FEATURE_C_NEON") {
+        let mut build = new_build();
+        build.file("src/c/blake3_neon.c");
+        build.file("src/c/blake3_portable.c");
+        // ARMv7 platforms that support NEON generally need the following
+        // flags. AArch64 supports NEON by default and does not support -mpfu.
+        // build.flag("-mfpu=neon-vfpv4");
+        // build.flag("-mfloat-abi=hard");
+        build.compile("blake3_neon");
+    }
+
+    // The `cc` crate does not automatically emit rerun-if directives for the
+    // environment variables it supports, in particular for $CC. We expect to
+    // do a lot of benchmarking across different compilers, so we explicitly
+    // add the variables that we're likely to need.
+    println!("cargo:rerun-if-env-changed=CC");
+    println!("cargo:rerun-if-env-changed=CFLAGS");
+
+    // Ditto for source files, though these shouldn't change as often.
+    for file in std::fs::read_dir("src/c")? {
+        println!(
+            "cargo:rerun-if-changed={}",
+            file?.path().to_str().expect("utf-8")
+        );
+    }
+
+    Ok(())
+}