From 5dc57da7449cf34e4c84a87f411dfa1f4a1d1b86 Mon Sep 17 00:00:00 2001 From: Marabii Date: Wed, 19 Aug 2026 21:33:03 +0200 Subject: [PATCH] Added AVX512BW support --- .github/workflows/big-endian.yml | 2 +- .github/workflows/quality.yaml | 2 +- .github/workflows/tests.yml | 4 +- Cargo.toml | 2 +- README.md | 4 +- src/impls/avx512bw/deser.rs | 177 ++++++++++++++++++ src/impls/avx512bw/mod.rs | 6 + src/impls/avx512bw/stage1.rs | 296 +++++++++++++++++++++++++++++++ src/impls/mod.rs | 3 + src/lib.rs | 35 +++- 10 files changed, 522 insertions(+), 9 deletions(-) create mode 100644 src/impls/avx512bw/deser.rs create mode 100644 src/impls/avx512bw/mod.rs create mode 100644 src/impls/avx512bw/stage1.rs diff --git a/.github/workflows/big-endian.yml b/.github/workflows/big-endian.yml index 25c1ab02..96208060 100644 --- a/.github/workflows/big-endian.yml +++ b/.github/workflows/big-endian.yml @@ -20,7 +20,7 @@ jobs: steps: - uses: actions/checkout@v3 - - uses: dtolnay/rust-toolchain@1.88 + - uses: dtolnay/rust-toolchain@1.89 - name: Install cross uses: taiki-e/install-action@v2 diff --git a/.github/workflows/quality.yaml b/.github/workflows/quality.yaml index e7068bc5..1c95302e 100644 --- a/.github/workflows/quality.yaml +++ b/.github/workflows/quality.yaml @@ -18,7 +18,7 @@ jobs: steps: - uses: actions/checkout@v3 - - uses: dtolnay/rust-toolchain@1.88 # do clippy chekcs with the minimum supported version + - uses: dtolnay/rust-toolchain@1.89 # do clippy chekcs with the minimum supported version with: components: rustfmt, clippy diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index e17b40d2..88f6f5db 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -24,7 +24,7 @@ jobs: steps: - uses: actions/checkout@v3 - - uses: dtolnay/rust-toolchain@1.88 + - uses: dtolnay/rust-toolchain@1.89 with: components: llvm-tools-preview @@ -120,7 +120,7 @@ jobs: steps: - uses: actions/checkout@v3 - - uses: dtolnay/rust-toolchain@1.88 + - uses: dtolnay/rust-toolchain@1.89 with: targets: wasm32-wasip1 diff --git a/Cargo.toml b/Cargo.toml index 671d0ce6..84c70c4a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ description = "High performance JSON parser based on a port of simdjson" repository = "https://github.com/simd-lite/simd-json" readme = "README.md" documentation = "https://docs.rs/simd-json" -rust-version = "1.88" +rust-version = "1.89" [dependencies] simdutf8 = { version = "0.1.4", features = ["public_imp", "aarch64_neon"] } diff --git a/README.md b/README.md index 8814ce1b..51752e32 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ However, in some design decisions—such as parsing to a DOM or a tape—ergonom performance. In other places Rust makes it harder to achieve the same level of performance. To take advantage of this library your system needs to support SIMD instructions. On `x86`, it will -select the best available supported instruction set (`avx2` or `sse4.2`) when the `runtime-detection` feature +select the best available supported instruction set (`avx512bw`, `avx2` or `sse4.2`) when the `runtime-detection` feature is enabled (default). On `aarch64` this library uses the `NEON` instruction set. On `wasm` this library uses the `simd128` instruction set when available. When no supported SIMD instructions are found, this library will use a fallback implementation, but this is significantly slower. @@ -75,7 +75,7 @@ for internal configuration and testing. ### `runtime-detection` (default) This feature allows selecting the optimal algorithm based on available features during runtime. It has no effect on -non-`x86` platforms. When neither `AVX2` nor `SSE4.2` is supported, it will fall back to a native Rust implementation. +non-`x86` platforms. When neither one of `AVX512BW` `AVX2` `SSE4.2` is supported, it will fall back to a native Rust implementation. Disabling this feature (with `default-features = false`) **and** setting `RUSTFLAGS="-C target-cpu=native` will result in better performance but the resulting binary will not be portable across `x86` processors. diff --git a/src/impls/avx512bw/deser.rs b/src/impls/avx512bw/deser.rs new file mode 100644 index 00000000..80dfd3ff --- /dev/null +++ b/src/impls/avx512bw/deser.rs @@ -0,0 +1,177 @@ +#[cfg(target_arch = "x86_64")] +use std::arch::x86_64 as arch; +use std::arch::x86_64::{ + __m512i, _mm512_cmpeq_epu8_mask, _mm512_loadu_si512, _mm512_set1_epi8, _mm512_storeu_si512, +}; + +use crate::{ + Deserializer, Result, SillyWrapper, + error::ErrorType, + safer_unchecked::GetSaferUnchecked, + stringparse::{ESCAPE_MAP, handle_unicode_codepoint}, +}; + +#[target_feature(enable = "avx512bw")] +#[allow( + clippy::if_not_else, + clippy::cast_possible_wrap, + clippy::too_many_lines +)] +#[cfg_attr(not(feature = "no-inline"), inline)] +pub(crate) unsafe fn parse_str<'invoke, 'de>( + input: SillyWrapper<'de>, + data: &'invoke [u8], + buffer: &'invoke mut [u8], + mut idx: usize, +) -> Result<&'de str> { + unsafe { + use ErrorType::{InvalidEscape, InvalidUnicodeCodepoint}; + + let input = input.input; + // Add 1 to skip the initial " + idx += 1; + //let mut read: usize = 0; + + // we include the terminal '"' so we know where to end + // This is safe since we check sub's length in the range access above and only + // create sub sliced form sub to `sub.len()`. + + let src: &[u8] = data.get_kinda_unchecked(idx..); + let mut src_i: usize = 0; + let mut len = src_i; + loop { + // _mm512_loadu_si512 does not require alignment + #[allow(clippy::cast_ptr_alignment)] + let v: __m512i = _mm512_loadu_si512(src.as_ptr().add(src_i).cast::<__m512i>()); + + // store to dest unconditionally - we can overwrite the bits we don't like + // later + let bs_bits: u64 = _mm512_cmpeq_epu8_mask(v, _mm512_set1_epi8(b'\\' as i8)); + + let quote_bits = _mm512_cmpeq_epu8_mask(v, _mm512_set1_epi8(b'"' as i8)); + + if (bs_bits.wrapping_sub(1) & quote_bits) != 0 { + // we encountered quotes first. Move dst to point to quotes and exit + // find out where the quote is... + let quote_dist: u32 = quote_bits.trailing_zeros(); + + /////////////////////// + // Above, check for overflow in case someone has a crazy string (>=4GB?) + // But only add the overflow check when the document itself exceeds 4GB + // Currently unneeded because we refuse to parse docs larger or equal to 4GB. + //////////////////////// + + // we advance the point, accounting for the fact that we have a NULl termination + + len += quote_dist as usize; + let v = + std::str::from_utf8_unchecked(std::slice::from_raw_parts(input.add(idx), len)); + return Ok(v); + + // we compare the pointers since we care if they are 'at the same spot' + // not if they are the same value + } + if (quote_bits.wrapping_sub(1) & bs_bits) == 0 { + // they are the same. Since they can't co-occur, it means we encountered + // neither. + src_i += 64; + len += 64; + } else { + // Move to the 'bad' character + let bs_dist: u32 = bs_bits.trailing_zeros(); + len += bs_dist as usize; + src_i += bs_dist as usize; + break; + } + } + + let mut dst_i: usize = 0; + + // To be more conform with upstream + loop { + // _mm512_loadu_si512 does not require alignment + #[allow(clippy::cast_ptr_alignment)] + let v: __m512i = _mm512_loadu_si512(src.as_ptr().add(src_i).cast::<__m512i>()); + + #[allow(clippy::cast_ptr_alignment)] + _mm512_storeu_si512(buffer.as_mut_ptr().add(dst_i).cast::<__m512i>(), v); + + // store to dest unconditionally - we can overwrite the bits we don't like + // later + let bs_bits: u64 = _mm512_cmpeq_epu8_mask(v, _mm512_set1_epi8(b'\\' as i8)); + + let quote_bits = _mm512_cmpeq_epu8_mask(v, _mm512_set1_epi8(b'"' as i8)); + if (bs_bits.wrapping_sub(1) & quote_bits) != 0 { + // we encountered quotes first. Move dst to point to quotes and exit + // find out where the quote is... + let quote_dist: u32 = quote_bits.trailing_zeros(); + + /////////////////////// + // Above, check for overflow in case someone has a crazy string (>=4GB?) + // But only add the overflow check when the document itself exceeds 4GB + // Currently unneeded because we refuse to parse docs larger or equal to 4GB. + //////////////////////// + + // we advance the point, accounting for the fact that we have a NULl termination + + dst_i += quote_dist as usize; + input + .add(idx + len) + .copy_from_nonoverlapping(buffer.as_ptr(), dst_i); + let v = std::str::from_utf8_unchecked(std::slice::from_raw_parts( + input.add(idx), + len + dst_i, + )); + return Ok(v); + + // we compare the pointers since we care if they are 'at the same spot' + // not if they are the same value + } + if (quote_bits.wrapping_sub(1) & bs_bits) != 0 { + // find out where the backspace is + let bs_dist: u32 = bs_bits.trailing_zeros(); + let escape_char: u8 = *src.get_kinda_unchecked(src_i + bs_dist as usize + 1); + // we encountered backslash first. Handle backslash + if escape_char == b'u' { + // move src/dst up to the start; they will be further adjusted + // within the unicode codepoint handling code. + src_i += bs_dist as usize; + dst_i += bs_dist as usize; + let (o, s) = handle_unicode_codepoint( + src.get_kinda_unchecked(src_i..), + buffer.get_kinda_unchecked_mut(dst_i..), + ) + .map_err(|_| Deserializer::error_c(src_i, 'u', InvalidUnicodeCodepoint))?; + + if o == 0 { + return Err(Deserializer::error_c(src_i, 'u', InvalidUnicodeCodepoint)); + } + // We moved o steps forward at the destination and 6 on the source + src_i += s; + dst_i += o; + } else { + // simple 1:1 conversion. Will eat bs_dist+2 characters in input and + // write bs_dist+1 characters to output + // note this may reach beyond the part of the buffer we've actually + // seen. I think this is ok + let escape_result: u8 = *ESCAPE_MAP.get_kinda_unchecked(escape_char as usize); + if escape_result == 0 { + return Err(Deserializer::error_c( + src_i, + escape_char as char, + InvalidEscape, + )); + } + *buffer.get_kinda_unchecked_mut(dst_i + bs_dist as usize) = escape_result; + src_i += bs_dist as usize + 2; + dst_i += bs_dist as usize + 1; + } + } else { + // they are the same. Since they can't co-occur, it means we encountered + // neither. + src_i += 64; + dst_i += 64; + } + } + } +} diff --git a/src/impls/avx512bw/mod.rs b/src/impls/avx512bw/mod.rs new file mode 100644 index 00000000..2e259cd2 --- /dev/null +++ b/src/impls/avx512bw/mod.rs @@ -0,0 +1,6 @@ +#![allow(unused_imports, dead_code)] +mod deser; +mod stage1; + +pub(crate) use deser::parse_str; +pub(crate) use stage1::SimdInput; diff --git a/src/impls/avx512bw/stage1.rs b/src/impls/avx512bw/stage1.rs new file mode 100644 index 00000000..15de659b --- /dev/null +++ b/src/impls/avx512bw/stage1.rs @@ -0,0 +1,296 @@ +#![allow(dead_code)] +use crate::{ + Stage1Parse, + macros::{static_cast_i32, static_cast_i64, static_cast_u32}, +}; + +#[cfg(target_arch = "x86_64")] +use std::arch::x86_64::{ + __m128i, __m512i, __mmask64, _mm_setr_epi8, _mm512_broadcast_i32x4, _mm512_cmpeq_epi8_mask, + _mm512_cmpeq_epu8_mask, _mm512_cmpge_epi8_mask, _mm512_cmpge_epu8_mask, _mm512_loadu_si512, + _mm512_set1_epi8, _mm512_srli_epi32, _mm512_storeu_si512, +}; +use std::arch::x86_64::{ + __m256i, _MM_CMPINT_NE, _mm_loadu_si128, _mm256_add_epi32, _mm256_set_epi32, + _mm256_storeu_si256, _mm512_add_epi32, _mm512_and_si512, _mm512_cmp_epi8_mask, + _mm512_cvtepu8_epi32, _mm512_maskz_compress_epi8, _mm512_set_epi32, _mm512_shuffle_epi8, +}; + +macro_rules! low_nibble_mask128 { + () => { + _mm_setr_epi8(16, 0, 0, 0, 0, 0, 0, 0, 0, 8, 12, 1, 2, 9, 0, 0) + }; +} + +macro_rules! high_nibble_mask128 { + () => { + _mm_setr_epi8(8, 0, 18, 4, 0, 1, 0, 1, 0, 0, 0, 3, 2, 1, 0, 0) + }; +} + +#[derive(Debug)] +pub(crate) struct SimdInput { + v0: __m512i, +} + +impl Stage1Parse for SimdInput { + type Utf8Validator = simdutf8::basic::imp::x86::avx2::ChunkedUtf8ValidatorImp; + type SimdRepresentation = __m512i; + #[cfg_attr(not(feature = "no-inline"), inline)] + // _mm512_loadu_si512 does not need alignment + #[allow(clippy::cast_ptr_alignment)] + #[target_feature(enable = "avx512bw")] + unsafe fn new(ptr: &[u8]) -> Self { + unsafe { + Self { + v0: _mm512_loadu_si512(ptr.as_ptr().cast::<__m512i>()), + } + } + } + + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_sign_loss)] + #[target_feature(enable = "avx512bw", enable = "pclmulqdq")] + #[cfg(target_arch = "x86_64")] + unsafe fn compute_quote_mask(quote_bits: u64) -> u64 { + unsafe { + use std::arch::x86_64::{_mm_clmulepi64_si128, _mm_set_epi64x, _mm_set1_epi8}; + + std::arch::x86_64::_mm_cvtsi128_si64(_mm_clmulepi64_si128( + _mm_set_epi64x(0, static_cast_i64!(quote_bits)), + _mm_set1_epi8(-1_i8 /* 0xFF */), + 0, + )) as u64 + } + } + + /// a straightforward comparison of a mask against input + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_possible_wrap, clippy::cast_sign_loss)] + #[target_feature(enable = "avx512bw")] + unsafe fn cmp_mask_against_input(&self, m: u8) -> u64 { + unsafe { + let mask: __m512i = Self::fill_s8(m as i8); + let res_0: __mmask64 = _mm512_cmpeq_epi8_mask(self.v0, mask); + res_0 as u64 + } + } + + // find all values less than or equal than the content of maxval (using unsigned arithmetic) + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_sign_loss)] + #[target_feature(enable = "avx512bw")] + unsafe fn unsigned_lteq_against_input(&self, maxval: __m512i) -> u64 { + unsafe { _mm512_cmpge_epu8_mask(maxval, self.v0) } + } + + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_sign_loss)] + #[target_feature(enable = "avx512bw")] + unsafe fn find_whitespace_and_structurals(&self, whitespace: &mut u64, structurals: &mut u64) { + unsafe { + // do a 'shufti' to detect structural JSON characters + // they are + // * `{` 0x7b + // * `}` 0x7d + // * `:` 0x3a + // * `[` 0x5b + // * `]` 0x5d + // * `,` 0x2c + // these go into the first 3 buckets of the comparison (1/2/4) + + // we are also interested in the four whitespace characters: + // * space 0x20 + // * linefeed 0x0a + // * horizontal tab 0x09 + // * carriage return 0x0d + // these go into the next 2 buckets of the comparison (8/16) + + let low_nibble_mask: __m512i = _mm512_broadcast_i32x4(low_nibble_mask128!()); + let high_nibble_mask: __m512i = _mm512_broadcast_i32x4(high_nibble_mask128!()); + + let structural_shufti_mask: __m512i = Self::fill_s8(0x7); + let whitespace_shufti_mask: __m512i = Self::fill_s8(0x18); + + let msb_mak = Self::fill_s8(0x7f); + let filtered_lower_4 = _mm512_and_si512(self.v0, msb_mak); + let result_lower_4 = _mm512_shuffle_epi8(low_nibble_mask, filtered_lower_4); + let filtered_high_4 = _mm512_and_si512(_mm512_srli_epi32::<4>(self.v0), msb_mak); + let result_high_4 = _mm512_shuffle_epi8(high_nibble_mask, filtered_high_4); + + let result = _mm512_and_si512(result_high_4, result_lower_4); + + *structurals = _mm512_cmp_epi8_mask( + _mm512_and_si512(result, structural_shufti_mask), + Self::fill_s8(0), + _MM_CMPINT_NE, // Not Equal to zero means it's a structural + ); + + *whitespace = _mm512_cmp_epi8_mask( + _mm512_and_si512(result, whitespace_shufti_mask), + Self::fill_s8(0), + _MM_CMPINT_NE, // Not Equal to zero means it's whitespace + ); + } + } + + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_possible_wrap, clippy::cast_ptr_alignment)] + #[target_feature(enable = "avx512bw")] + unsafe fn flatten_bits(base: &mut Vec, idx: u32, bits: u64) { + unsafe { + if std::is_x86_feature_detected!("avx512vbmi2") { + FlattenBits::flatten_bits_fast(base, idx, bits); + } else { + FlattenBits::flatten_bits_slow(base, idx, bits); + } + } + } + + #[cfg_attr(not(feature = "no-inline"), inline)] + #[target_feature(enable = "avx512bw")] + unsafe fn fill_s8(n: i8) -> __m512i { + unsafe { _mm512_set1_epi8(n) } + } +} + +struct FlattenBits; + +impl FlattenBits { + // flatten out values in 'bits' assuming that they are are to have values of idx + // plus their position in the bitvector, and store these indexes at + // base_ptr[base] incrementing base as we go + // will potentially store extra values beyond end of valid bits, so base_ptr + // needs to be large enough to handle this + //TODO: usize was u32 here does this matter? + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_possible_wrap, clippy::cast_ptr_alignment)] + #[target_feature(enable = "avx512bw")] + unsafe fn flatten_bits_slow(base: &mut Vec, idx: u32, mut bits: u64) { + unsafe { + let cnt: usize = bits.count_ones() as usize; + let mut l: usize = base.len(); + let idx_minus_64 = idx.wrapping_sub(64); + let idx_64_v = _mm512_set_epi32( + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + ); + + // We're doing some trickery here. + // We reserve 64 extra entries, because we've at most 64 bit to set + // then we truncate the base to the next base (that we calculated above) + // We later indiscriminatory write over the len we set but that's OK + // since we ensure we reserve the needed space + base.reserve(64); + let final_len = l + cnt; + + while bits != 0 { + let v0 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v1 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v2 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v3 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v4 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v5 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v6 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v7 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v8 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v9 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v10 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v11 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v12 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v13 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v14 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + let v15 = bits.trailing_zeros() as i32; + bits &= bits.wrapping_sub(1); + + let v: __m512i = _mm512_set_epi32( + v15, v14, v13, v12, v11, v10, v9, v8, v7, v6, v5, v4, v3, v2, v1, v0, + ); + let v: __m512i = _mm512_add_epi32(idx_64_v, v); + _mm512_storeu_si512(base.as_mut_ptr().add(l).cast::<__m512i>(), v); + l += 16; + } + // We have written all the data + base.set_len(final_len); + } + } + + #[cfg_attr(not(feature = "no-inline"), inline)] + #[allow(clippy::cast_possible_wrap, clippy::cast_ptr_alignment)] + #[target_feature(enable = "avx512vbmi2")] + unsafe fn flatten_bits_fast(base: &mut Vec, idx: u32, bits: u64) { + const INDICES: [u8; 64] = [ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, + 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, + 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, + ]; + + unsafe { + let cnt: usize = bits.count_ones() as usize; + let l: usize = base.len(); + let idx_minus_64 = idx.wrapping_sub(64); + let idx_64_v = _mm512_set_epi32( + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + static_cast_i32!(idx_minus_64), + ); + base.reserve(64); + let base_indices = _mm512_loadu_si512(INDICES.as_ptr().cast()); + let result_vec_u8 = _mm512_maskz_compress_epi8(bits, base_indices); + + let mut tmp = [0u8; 64]; + _mm512_storeu_si512(tmp.as_mut_ptr().cast(), result_vec_u8); + + for i in 0..4 { + let chunk = _mm_loadu_si128(tmp.as_ptr().add(i * 16).cast()); + let widened = _mm512_cvtepu8_epi32(chunk); + let v: __m512i = _mm512_add_epi32(idx_64_v, widened); + _mm512_storeu_si512(base.as_mut_ptr().add(l + i * 16).cast(), v); + } + + base.set_len(cnt + l); + } + } +} diff --git a/src/impls/mod.rs b/src/impls/mod.rs index 8f4b08e4..137497c0 100644 --- a/src/impls/mod.rs +++ b/src/impls/mod.rs @@ -10,6 +10,9 @@ pub(crate) mod native; /// rust native implementation pub(crate) mod portable; +#[cfg(target_arch = "x86_64")] +pub(crate) mod avx512bw; + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] pub(crate) mod avx2; diff --git a/src/lib.rs b/src/lib.rs index c6f50246..d5bfd294 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -56,7 +56,7 @@ mod impls; pub mod cow; /// The maximum padding size required by any SIMD implementation -pub(crate) const SIMDJSON_PADDING: usize = 32; // take upper limit mem::size_of::<__m256i>() +pub(crate) const SIMDJSON_PADDING: usize = 64; // take upper limit mem::size_of::<__m512i>() /// It's 64 for all (Is this correct?) pub(crate) const SIMDINPUT_LENGTH: usize = 64; @@ -366,6 +366,8 @@ pub enum Implementation { SSE42, /// AVX2 implementation AVX2, + /// AVX512BW implementation + AVX512BW, /// ARM NEON implementation NEON, /// WEBASM SIMD128 implementation @@ -379,6 +381,7 @@ impl std::fmt::Display for Implementation { Implementation::StdSimd => write!(f, "std::simd"), Implementation::SSE42 => write!(f, "SSE42"), Implementation::AVX2 => write!(f, "AVX2"), + Implementation::AVX512BW => write!(f, "AVX512BW"), Implementation::NEON => write!(f, "NEON"), Implementation::SIMD128 => write!(f, "SIMD128"), } @@ -393,7 +396,9 @@ impl Deserializer<'_> { ))] #[must_use] pub fn algorithm() -> Implementation { - if std::is_x86_feature_detected!("avx2") { + if std::is_x86_feature_detected!("avx512bw") { + Implementation::AVX512BW + } else if std::is_x86_feature_detected!("avx2") { Implementation::AVX2 } else if std::is_x86_feature_detected!("sse4.2") { Implementation::SSE42 @@ -476,6 +481,11 @@ impl<'de> Deserializer<'de> { any(target_arch = "x86_64", target_arch = "x86"), ))] pub(crate) fn parse_str_fn() -> ParseStrFn { + #[cfg(target_arch = "x86_64")] + if std::is_x86_feature_detected!("avx512bw") { + return impls::avx512bw::parse_str; + } + if std::is_x86_feature_detected!("avx2") { impls::avx2::parse_str } else if std::is_x86_feature_detected!("sse4.2") { @@ -645,6 +655,20 @@ impl Deserializer<'_> { // The wrappers below carry the ISA's `target_feature` so that LLVM can inline // the `#[target_feature]`-annotated SIMD primitives into the stage-1 loop; // without them every primitive stays an outlined call per 64-byte block. + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx512bw", enable = "pclmulqdq")] + unsafe fn find_structural_bits_avx512bw( + input: &[u8], + structural_indexes: &mut Vec, + ) -> core::result::Result<(), error::ErrorType> { + unsafe { + Deserializer::_find_structural_bits::( + input, + structural_indexes, + ) + } + } + #[target_feature(enable = "avx2", enable = "pclmulqdq")] unsafe fn find_structural_bits_avx2( input: &[u8], @@ -673,6 +697,13 @@ impl Deserializer<'_> { #[cfg_attr(not(feature = "no-inline"), inline)] fn get_fastest_available_implementation() -> FindStructuralBitsFn { + #[cfg(target_arch = "x86_64")] + if std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("pclmulqdq") + { + return find_structural_bits_avx512bw; + } + if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("pclmulqdq") {