Escape str using SIMD

ci/azure-macos.yml

@@ -23,7 +23,7 @@ steps:
  PATH=$HOME/.cargo/bin:$PATH \
  MACOSX_DEPLOYMENT_TARGET=$(macosx_deployment_target) \
  PYO3_CROSS_LIB_DIR=$(python -c "import sysconfig;print(sysconfig.get_config_var('LIBDIR'))") \
- maturin build --release --strip --features=encoding_rs/simd-accel,no-panic,yyjson --interpreter $(interpreter) --target=universal2-apple-darwin
+ maturin build --release --strip --features=encoding_rs/simd-accel,no-panic,unstable-simd,yyjson --interpreter $(interpreter) --target=universal2-apple-darwin
  env:
  CC: "clang"
  CFLAGS: "-O2 -fstrict-aliasing -flto=full"

src/lib.rs

@@ -4,6 +4,7 @@
 #![cfg_attr(feature = "optimize", feature(optimize_attribute))]
 #![cfg_attr(feature = "strict_provenance", feature(strict_provenance))]
 #![cfg_attr(feature = "strict_provenance", warn(fuzzy_provenance_casts))]
+#![cfg_attr(feature = "unstable-simd", feature(portable_simd))]
 #![allow(unknown_lints)] // internal_features
 #![allow(internal_features)] // core_intrinsics
 #![allow(unused_unsafe)]

src/serialize/writer/json.rs

@@ -565,6 +565,85 @@ where
  }
 }
 
+#[cfg(all(
+ feature = "unstable-simd",
+ target_arch = "x86_64",
+ target_feature = "avx2"
+))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+where
+ W: ?Sized + io::Write + WriteExt,
+{
+ unsafe {
+ let num_reserved_bytes = value.len() * 8 + 32 + 3;
+ writer.reserve(num_reserved_bytes);
+
+ let written = crate::serialize::writer::simd::format_escaped_str_impl_256(
+ writer.as_mut_buffer_ptr(),
+ value.as_bytes().as_ptr(),
+ value.len(),
+ );
+ writer.set_written(written);
+ }
+ Ok(())
+}
+
+#[cfg(all(
+ feature = "unstable-simd",
+ target_arch = "x86_64",
+ not(target_feature = "avx2")
+))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+where
+ W: ?Sized + io::Write + WriteExt,
+{
+ unsafe {
+ let num_reserved_bytes = value.len() * 8 + 32 + 3;
+ writer.reserve(num_reserved_bytes);
+
+ if std::is_x86_feature_detected!("avx2") {
+ let written = crate::serialize::writer::simd::format_escaped_str_impl_256(
+ writer.as_mut_buffer_ptr(),
+ value.as_bytes().as_ptr(),
+ value.len(),
+ );
+ writer.set_written(written);
+ } else {
+ let written = crate::serialize::writer::simd::format_escaped_str_impl_128(
+ writer.as_mut_buffer_ptr(),
+ value.as_bytes().as_ptr(),
+ value.len(),
+ );
+ writer.set_written(written);
+ }
+ }
+ Ok(())
+}
+
+#[cfg(all(feature = "unstable-simd", not(target_arch = "x86_64")))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+where
+ W: ?Sized + io::Write + WriteExt,
+{
+ unsafe {
+ let num_reserved_bytes = value.len() * 8 + 32 + 3;
+ writer.reserve(num_reserved_bytes);
+
+ let written = crate::serialize::writer::simd::format_escaped_str_impl_128(
+ writer.as_mut_buffer_ptr(),
+ value.as_bytes().as_ptr(),
+ value.len(),
+ );
+
+ writer.set_written(written);
+ }
+ Ok(())
+}
+
+#[cfg(not(feature = "unstable-simd"))]
 #[inline(always)]
 fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
 where

src/serialize/writer/mod.rs

@@ -1,9 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod byteswriter;
+#[cfg(not(feature = "unstable-simd"))]
 mod escape;
 mod formatter;
 mod json;
+#[cfg(feature = "unstable-simd")]
+mod simd;
 
 pub use byteswriter::{BytesWriter, WriteExt};
 pub use json::{to_writer, to_writer_pretty};

src/serialize/writer/simd.rs

@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright 2023-2024 liuq19, ijl
+// adapted from sonic-rs' src/util/string.rs
+
+use std::simd::cmp::{SimdPartialEq, SimdPartialOrd};
+
+macro_rules! impl_escape_unchecked {
+ ($src:expr, $dst:expr, $nb:expr, $omask:expr, $cn:expr) => {
+ $nb -= $cn;
+ $dst = $dst.add($cn);
+ $src = $src.add($cn);
+ let mut mask = $omask << $cn;
+ loop {
+ $nb -= 1;
+ mask = mask << 1;
+ let replacement = if *($src) == b'"' {
+ (*b"\\\"\0\0\0\0\0\0", 2)
+ } else if *($src) == b'\\' {
+ (*b"\\\\\0\0\0\0\0\0", 2)
+ } else {
+ match *($src) {
+ 0 => (*b"\\u0000\0\0", 6),
+ 1 => (*b"\\u0001\0\0", 6),
+ 2 => (*b"\\u0002\0\0", 6),
+ 3 => (*b"\\u0003\0\0", 6),
+ 4 => (*b"\\u0004\0\0", 6),
+ 5 => (*b"\\u0005\0\0", 6),
+ 6 => (*b"\\u0006\0\0", 6),
+ 7 => (*b"\\u0007\0\0", 6),
+ 8 => (*b"\\b\0\0\0\0\0\0", 2),
+ 9 => (*b"\\t\0\0\0\0\0\0", 2),
+ 10 => (*b"\\n\0\0\0\0\0\0", 2),
+ 11 => (*b"\\u000b\0\0", 6),
+ 12 => (*b"\\f\0\0\0\0\0\0", 2),
+ 13 => (*b"\\r\0\0\0\0\0\0", 2),
+ 14 => (*b"\\u000e\0\0", 6),
+ 15 => (*b"\\u000f\0\0", 6),
+ 16 => (*b"\\u0010\0\0", 6),
+ 17 => (*b"\\u0011\0\0", 6),
+ 18 => (*b"\\u0012\0\0", 6),
+ 19 => (*b"\\u0013\0\0", 6),
+ 20 => (*b"\\u0014\0\0", 6),
+ 21 => (*b"\\u0015\0\0", 6),
+ 22 => (*b"\\u0016\0\0", 6),
+ 23 => (*b"\\u0017\0\0", 6),
+ 24 => (*b"\\u0018\0\0", 6),
+ 25 => (*b"\\u0019\0\0", 6),
+ 26 => (*b"\\u001a\0\0", 6),
+ 27 => (*b"\\u001b\0\0", 6),
+ 28 => (*b"\\u001c\0\0", 6),
+ 29 => (*b"\\u001d\0\0", 6),
+ 30 => (*b"\\u001e\0\0", 6),
+ 31 => (*b"\\u001f\0\0", 6),
+ _ => unreachable!(),
+ }
+ };
+ std::ptr::copy_nonoverlapping(replacement.0.as_ptr(), $dst, 8);
+ $dst = $dst.add(replacement.1 as usize);
+ $src = $src.add(1);
+ if likely!(mask & (1 << (STRIDE - 1)) != 1) {
+ break;
+ }
+ }
+ };
+}
+macro_rules! impl_format_simd {
+ ($odptr:expr, $value_ptr:expr, $value_len:expr) => {
+ let mut dptr = $odptr;
+ let dstart = $odptr;
+ let mut sptr = $value_ptr;
+ let mut nb: usize = $value_len;
+
+ let blash = StrVector::from_array([b'\\'; STRIDE]);
+ let quote = StrVector::from_array([b'"'; STRIDE]);
+ let x20 = StrVector::from_array([32; STRIDE]);
+
+ unsafe {
+ *dptr = b'"';
+ dptr = dptr.add(1);
+
+ while nb >= STRIDE {
+ let v = StrVector::from_slice(std::slice::from_raw_parts(sptr, STRIDE));
+ v.copy_to_slice(std::slice::from_raw_parts_mut(dptr, STRIDE));
+ let mask =
+ (v.simd_eq(blash) | v.simd_eq(quote) | v.simd_lt(x20)).to_bitmask() as u32;
+
+ if likely!(mask == 0) {
+ nb -= STRIDE;
+ dptr = dptr.add(STRIDE);
+ sptr = sptr.add(STRIDE);
+ } else {
+ let cn = mask.trailing_zeros() as usize;
+ impl_escape_unchecked!(sptr, dptr, nb, mask, cn);
+ }
+ }
+
+ while nb > 0 {
+ let v = StrVector::from_slice(std::slice::from_raw_parts(sptr, STRIDE));
+ v.copy_to_slice(std::slice::from_raw_parts_mut(dptr, STRIDE));
+ let mask = (v.simd_eq(blash) | v.simd_eq(quote) | v.simd_lt(x20)).to_bitmask()
+ as u32
+ & (STRIDE_SATURATION >> (STRIDE - nb));
+ if likely!(mask == 0) {
+ dptr = dptr.add(nb);
+ break;
+ } else {
+ let cn = mask.trailing_zeros() as usize;
+ impl_escape_unchecked!(sptr, dptr, nb, mask, cn);
+ }
+ }
+
+ *dptr = b'"';
+ dptr = dptr.add(1);
+ }
+
+ return dptr as usize - dstart as usize;
+ };
+}
+
+#[inline(never)]
+#[cfg(not(target_feature = "avx2"))]
+#[cfg_attr(target_arch = "x86_64", cold)]
+pub unsafe fn format_escaped_str_impl_128(
+ odptr: *mut u8,
+ value_ptr: *const u8,
+ value_len: usize,
+) -> usize {
+ const STRIDE: usize = 16;
+ const STRIDE_SATURATION: u32 = u16::MAX as u32;
+ type StrVector = std::simd::u8x16;
+
+ impl_format_simd!(odptr, value_ptr, value_len);
+}
+
+#[cfg(target_arch = "x86_64")]
+#[inline(never)]
+#[cfg_attr(not(target_feature = "avx2"), target_feature(enable = "avx2"))]
+#[cfg_attr(not(target_feature = "avx2"), target_feature(enable = "bmi2"))]
+pub unsafe fn format_escaped_str_impl_256(
+ odptr: *mut u8,
+ value_ptr: *const u8,
+ value_len: usize,
+) -> usize {
+ const STRIDE: usize = 32;
+ const STRIDE_SATURATION: u32 = u32::MAX;
+ type StrVector = std::simd::u8x32;
+
+ impl_format_simd!(odptr, value_ptr, value_len);
+}

src/util.rs

@@ -70,6 +70,22 @@ macro_rules! unlikely {
  };
 }
 
+#[allow(unused_macros)]
+#[cfg(feature = "intrinsics")]
+macro_rules! likely {
+ ($exp:expr) => {
+ std::intrinsics::likely($exp)
+ };
+}
+
+#[allow(unused_macros)]
+#[cfg(not(feature = "intrinsics"))]
+macro_rules! likely {
+ ($exp:expr) => {
+ $exp
+ };
+}
+
 macro_rules! nonnull {
  ($exp:expr) => {
  unsafe { std::ptr::NonNull::new_unchecked($exp) }