str short, 128-bit SIMD, AVX512VL split

.github/workflows/debug.yaml

@@ -8,9 +8,10 @@ jobs:
       fail-fast: false
       matrix:
         profile: [
-          { rust: "1.72", features: "" },
-          { rust: "1.72", features: "--features=yyjson" },
-          { rust: "nightly-2024-04-15", features: "--features=yyjson,unstable-simd" },
+          { rust: "1.72", features: "", env "" },
+          { rust: "1.72", features: "--features=yyjson", env: "" },
+          { rust: "nightly-2024-04-15", features: "--features=yyjson,unstable-simd", env: "ORJSON_DISABLE_AVX512=1" },
+          { rust: "nightly-2024-04-15", features: "--features=yyjson,unstable-simd", env: "" },
         ]
         python: [
           { version: '3.13' },
@@ -37,7 +38,7 @@ jobs:
 
     - name: build
       run: |
-        PATH="$HOME/.cargo/bin:$PATH" maturin build --release \
+        PATH="$HOME/.cargo/bin:$PATH" "${{ matrix.profile.env }}" maturin build --release \
           --out=dist \
           --profile=dev \
           --interpreter python${{ matrix.python.version }} \

build.rs

@@ -32,6 +32,11 @@ fn main() {
     if env::var("ORJSON_DISABLE_SIMD").is_err() {
         if let Some(true) = version_check::supports_feature("portable_simd") {
             println!("cargo:rustc-cfg=feature=\"unstable-simd\"");
+
+            #[cfg(target_arch = "x86_64")]
+            if env::var("ORJSON_DISABLE_AVX512").is_err() {
+                println!("cargo:rustc-cfg=feature=\"avx512\"");
+            }
         }
     }
 

src/lib.rs

@@ -1,5 +1,7 @@
 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+#![cfg_attr(feature = "avx512", feature(avx512_target_feature))]
+#![cfg_attr(feature = "avx512", feature(stdarch_x86_avx512))]
 #![cfg_attr(feature = "intrinsics", feature(core_intrinsics))]
 #![cfg_attr(feature = "optimize", feature(optimize_attribute))]
 #![cfg_attr(feature = "strict_provenance", feature(strict_provenance))]

src/serialize/writer/json.rs

@@ -2,6 +2,7 @@
 // This is an adaptation of `src/value/ser.rs` from serde-json.
 
 use crate::serialize::writer::formatter::{CompactFormatter, Formatter, PrettyFormatter};
+use crate::serialize::writer::str::*;
 use crate::serialize::writer::WriteExt;
 use serde::ser::{self, Impossible, Serialize};
 use serde_json::error::{Error, Result};
@@ -148,9 +149,10 @@ where
         unreachable!();
     }
 
-    #[inline]
+    #[inline(always)]
     fn serialize_str(self, value: &str) -> Result<()> {
-        format_escaped_str(&mut self.writer, value).map_err(Error::io)
+        format_escaped_str(&mut self.writer, value);
+        Ok(())
     }
 
     fn serialize_bytes(self, value: &[u8]) -> Result<()> {
@@ -396,7 +398,7 @@ where
     type SerializeStruct = Impossible<(), Error>;
     type SerializeStructVariant = Impossible<(), Error>;
 
-    #[inline]
+    #[inline(always)]
     fn serialize_str(self, value: &str) -> Result<()> {
         self.ser.serialize_str(value)
     }
@@ -553,46 +555,77 @@ where
     }
 }
 
-#[cfg(feature = "unstable-simd")]
+macro_rules! reserve_str {
+    ($writer:expr, $value:expr) => {
+        $writer.reserve($value.len() * 8 + 32);
+    };
+}
+
+#[cfg(all(
+    feature = "unstable-simd",
+    any(not(target_arch = "x86_64"), not(feature = "avx512"))
+))]
 #[inline(always)]
-fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+fn format_escaped_str<W>(writer: &mut W, value: &str)
 where
     W: ?Sized + io::Write + WriteExt,
 {
     unsafe {
-        let num_reserved_bytes = value.len() * 8 + 32;
-        writer.reserve(num_reserved_bytes);
+        reserve_str!(writer, value);
 
-        let written = crate::serialize::writer::escape::format_escaped_str_impl_128(
+        let written = 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", target_arch = "x86_64", feature = "avx512"))]
+#[inline(always)]
+fn format_escaped_str<W>(writer: &mut W, value: &str)
+where
+    W: ?Sized + io::Write + WriteExt,
+{
+    unsafe {
+        reserve_str!(writer, value);
+
+        if std::is_x86_feature_detected!("avx512vl") {
+            let written = format_escaped_str_impl_512vl(
+                writer.as_mut_buffer_ptr(),
+                value.as_bytes().as_ptr(),
+                value.len(),
+            );
+            writer.set_written(written);
+        } else {
+            let written = format_escaped_str_impl_128(
+                writer.as_mut_buffer_ptr(),
+                value.as_bytes().as_ptr(),
+                value.len(),
+            );
+            writer.set_written(written);
+        };
+    }
 }
 
 #[cfg(not(feature = "unstable-simd"))]
 #[inline(always)]
-fn format_escaped_str<W>(writer: &mut W, value: &str) -> io::Result<()>
+fn format_escaped_str<W>(writer: &mut W, value: &str)
 where
     W: ?Sized + io::Write + WriteExt,
 {
     unsafe {
-        let num_reserved_bytes = value.len() * 8 + 32;
-        writer.reserve(num_reserved_bytes);
+        reserve_str!(writer, value);
 
-        let written = crate::serialize::writer::escape::format_escaped_str_scalar(
+        let written = format_escaped_str_scalar(
             writer.as_mut_buffer_ptr(),
             value.as_bytes().as_ptr(),
             value.len(),
         );
-
         writer.set_written(written);
     }
-    Ok(())
 }
 
 #[inline]

src/serialize/writer/mod.rs

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod byteswriter;
-mod escape;
 mod formatter;
 mod json;
+mod str;
 
 pub use byteswriter::{BytesWriter, WriteExt};
 pub use json::{to_writer, to_writer_pretty};