Wasm32 SIMD implementation

This code is based on rust_sse2.rs of the same distribution, and is
subject to further improvements. Some comments are left intact even if
their applicability is questioned.

SIMD implementation is gated by `wasm32-simd` feature, portable version
is used otherwise.

Performance measurements with a primitive benchmark with ~16Kb of data:

| M1 native     | 11,610 ns |
| M1 WASM SIMD  | 13,355 ns |
| M1 WASM       | 22,037 ns |
| x64 native    |  6,713 ns |
| x64 WASM SIMD | 11,985 ns |
| x64 WASM      | 25,978 ns |

wasmtime v12.0.1 was used on both platforms.

Closes BLAKE3-team#187.

Cargo.toml

@@ -18,6 +18,11 @@ default = ["std"]
 # implementation uses C intrinsics and requires a C compiler.
 neon = []
 
+# The WASM-SIMD implementation does not participate in dynamic feature detection,
+# which is currently x86-only. If "wasm_simd" is on, WASM SIMD support is assumed.
+# Note that not all WASM implementation may support WASM-SIMD specification.
+wasm32_simd = []
+
 # This crate uses libstd for std::io trait implementations, and also for
 # runtime CPU feature detection. This feature is enabled by default. If you use
 # --no-default-features, the only way to use the SIMD implementations in this
@@ -94,7 +99,7 @@ zeroize_crate = { package = "zeroize", version = "1", default-features = false,
 
 [dev-dependencies]
 hex = "0.4.2"
-page_size = "0.5.0"
+page_size = "0.6.0"
 rand = "0.8.0"
 rand_chacha = "0.3.0"
 reference_impl = { path = "./reference_impl" }

build.rs

@@ -21,6 +21,10 @@ fn is_no_neon() -> bool {
     defined("CARGO_FEATURE_NO_NEON")
 }
 
+fn is_wasm32_simd() -> bool {
+    defined("CARGO_FEATURE_WASM32_SIMD")
+}
+
 fn is_ci() -> bool {
     defined("BLAKE3_CI")
 }
@@ -60,6 +64,10 @@ fn is_armv7() -> bool {
     target_components()[0] == "armv7"
 }
 
+fn is_wasm32() -> bool {
+    target_components()[0] == "wasm32"
+}
+
 // Windows targets may be using the MSVC toolchain or the GNU toolchain. The
 // right compiler flags to use depend on the toolchain. (And we don't want to
 // use flag_if_supported, because we don't want features to be silently
@@ -225,6 +233,13 @@ fn build_neon_c_intrinsics() {
     build.compile("blake3_neon");
 }
 
+fn build_wasm32_simd() {
+    assert!(is_wasm32());
+    // No C code to compile here. Set the cfg flags that enable the WASM SIMD.
+    // The regular Cargo build will compile it.
+    println!("cargo:rustc-cfg=blake3_wasm32_simd");
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     if is_pure() && is_neon() {
         panic!("It doesn't make sense to enable both \"pure\" and \"neon\".");
@@ -258,6 +273,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         build_neon_c_intrinsics();
     }
 
+    if is_wasm32() && is_wasm32_simd() {
+        build_wasm32_simd();
+    }
+
     // The `cc` crate doesn't 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

src/lib.rs

@@ -70,7 +70,6 @@
 #[cfg(feature = "zeroize")]
 extern crate zeroize_crate as zeroize; // Needed because `zeroize::Zeroize` assumes the crate is named `zeroize`.
 
-
 #[cfg(test)]
 mod test;
 
@@ -113,6 +112,10 @@ mod sse41;
 #[path = "ffi_sse41.rs"]
 mod sse41;
 
+#[cfg(blake3_wasm32_simd)]
+#[path = "wasm32_simd.rs"]
+mod wasm32_simd;
+
 #[cfg(feature = "traits-preview")]
 pub mod traits;
 

src/platform.rs

@@ -12,6 +12,8 @@ cfg_if::cfg_if! {
         }
     } else if #[cfg(blake3_neon)] {
         pub const MAX_SIMD_DEGREE: usize = 4;
+    } else if #[cfg(blake3_wasm32_simd)] {
+        pub const MAX_SIMD_DEGREE: usize = 8;
     } else {
         pub const MAX_SIMD_DEGREE: usize = 1;
     }
@@ -32,6 +34,8 @@ cfg_if::cfg_if! {
         }
     } else if #[cfg(blake3_neon)] {
         pub const MAX_SIMD_DEGREE_OR_2: usize = 4;
+    } else if #[cfg(blake3_wasm32_simd)] {
+        pub const MAX_SIMD_DEGREE_OR_2: usize = 4;
     } else {
         pub const MAX_SIMD_DEGREE_OR_2: usize = 2;
     }
@@ -51,6 +55,9 @@ pub enum Platform {
     AVX512,
     #[cfg(blake3_neon)]
     NEON,
+    #[cfg(blake3_wasm32_simd)]
+    #[allow(non_camel_case_types)]
+    WASM32_SIMD,
 }
 
 impl Platform {
@@ -80,6 +87,10 @@ impl Platform {
         {
             return Platform::NEON;
         }
+        #[cfg(blake3_wasm32_simd)]
+        {
+            return Platform::WASM32_SIMD;
+        }
         Platform::Portable
     }
 
@@ -97,6 +108,9 @@ impl Platform {
             Platform::AVX512 => 16,
             #[cfg(blake3_neon)]
             Platform::NEON => 4,
+            #[cfg(blake3_wasm32_simd)]
+            // TODO is it 8 or 4??? SSE4 has 4...
+            Platform::WASM32_SIMD => 4,
         };
         debug_assert!(degree <= MAX_SIMD_DEGREE);
         degree
@@ -131,6 +145,11 @@ impl Platform {
             // No NEON compress_in_place() implementation yet.
             #[cfg(blake3_neon)]
             Platform::NEON => portable::compress_in_place(cv, block, block_len, counter, flags),
+            // Safe because is compiled for wasm32
+            #[cfg(blake3_wasm32_simd)]
+            Platform::WASM32_SIMD => unsafe {
+                crate::wasm32_simd::compress_in_place(cv, block, block_len, counter, flags)
+            },
         }
     }
 
@@ -163,6 +182,11 @@ impl Platform {
             // No NEON compress_xof() implementation yet.
             #[cfg(blake3_neon)]
             Platform::NEON => portable::compress_xof(cv, block, block_len, counter, flags),
+            #[cfg(blake3_wasm32_simd)]
+            // TODO Safe because compiled for wasm32
+            Platform::WASM32_SIMD => unsafe {
+                crate::wasm32_simd::compress_xof(cv, block, block_len, counter, flags)
+            },
         }
     }
 
@@ -269,6 +293,20 @@ impl Platform {
                     out,
                 )
             },
+            // Assumed to be safe if the "wasm32_simd" feature is on.
+            #[cfg(blake3_wasm32_simd)]
+            Platform::WASM32_SIMD => unsafe {
+                crate::wasm32_simd::hash_many(
+                    inputs,
+                    key,
+                    counter,
+                    increment_counter,
+                    flags,
+                    flags_start,
+                    flags_end,
+                    out,
+                )
+            },
         }
     }
 
@@ -320,6 +358,12 @@ impl Platform {
         // Assumed to be safe if the "neon" feature is on.
         Some(Self::NEON)
     }
+
+    #[cfg(blake3_wasm32_simd)]
+    pub fn wasm32_simd() -> Option<Self> {
+        // Assumed to be safe if the "neon" feature is on.
+        Some(Self::WASM32_SIMD)
+    }
 }
 
 // Note that AVX-512 is divided into multiple featuresets, and we use two of