Skip to content

Commit

Permalink
Wasm32 SIMD implementation
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
monoid committed Sep 11, 2023
1 parent f22d66b commit 909c6fc
Show file tree
Hide file tree
Showing 5 changed files with 859 additions and 1 deletion.
5 changes: 5 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 19 additions & 0 deletions build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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\".");
Expand Down Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;

Expand Down
44 changes: 44 additions & 0 deletions src/platform.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand All @@ -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 {
Expand Down Expand Up @@ -80,6 +87,10 @@ impl Platform {
{
return Platform::NEON;
}
#[cfg(blake3_wasm32_simd)]
{
return Platform::WASM32_SIMD;
}
Platform::Portable
}

Expand All @@ -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
Expand Down Expand Up @@ -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)
},
}
}

Expand Down Expand Up @@ -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)
},
}
}

Expand Down Expand Up @@ -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,
)
},
}
}

Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 909c6fc

Please sign in to comment.