Use function pointers for AVX-512 path

perf stat -d -d -d --delay 250 ./bench/run_func data/twitter.json.xz dumps 5000

    7349323334      instructions:u
    1503705855      branches:u

    7064976507      instructions:u
    1395888450      branches:u

Author: ijl

src/serialize/mod.rs

@@ -6,6 +6,6 @@ mod obtype;
 mod per_type;
 mod serializer;
 mod state;
-mod writer;
+pub mod writer;
 
 pub use serializer::serialize;

src/serialize/writer/json.rs

@@ -590,6 +590,22 @@ where
     }
 }
 
+#[cfg(all(feature = "unstable-simd", target_arch = "x86_64", feature = "avx512"))]
+type StrFormatter = unsafe fn(*mut u8, *const u8, usize) -> usize;
+
+#[cfg(all(feature = "unstable-simd", target_arch = "x86_64", feature = "avx512"))]
+static mut STR_FORMATTER_FN: StrFormatter =
+    crate::serialize::writer::str::format_escaped_str_impl_sse2_128;
+
+pub fn set_str_formatter_fn() {
+    unsafe {
+        #[cfg(all(feature = "unstable-simd", target_arch = "x86_64", feature = "avx512"))]
+        if std::is_x86_feature_detected!("avx512vl") {
+            STR_FORMATTER_FN = crate::serialize::writer::str::format_escaped_str_impl_512vl;
+        }
+    }
+}
+
 #[cfg(all(
     feature = "unstable-simd",
     target_arch = "x86_64",
@@ -622,21 +638,13 @@ where
     unsafe {
         reserve_str!(writer, value);
 
-        if std::is_x86_feature_detected!("avx512vl") {
-            let written = crate::serialize::writer::str::format_escaped_str_impl_512vl(
-                writer.as_mut_buffer_ptr(),
-                value.as_bytes().as_ptr(),
-                value.len(),
-            );
-            writer.set_written(written);
-        } else {
-            let written = crate::serialize::writer::str::format_escaped_str_impl_sse2_128(
-                writer.as_mut_buffer_ptr(),
-                value.as_bytes().as_ptr(),
-                value.len(),
-            );
-            writer.set_written(written);
-        }
+        let written = STR_FORMATTER_FN(
+            writer.as_mut_buffer_ptr(),
+            value.as_bytes().as_ptr(),
+            value.len(),
+        );
+
+        writer.set_written(written);
     }
 }
 
@@ -654,6 +662,7 @@ where
             value.as_bytes().as_ptr(),
             value.len(),
         );
+
         writer.set_written(written);
     }
 }

src/serialize/writer/mod.rs

@@ -6,4 +6,4 @@ mod json;
 mod str;
 
 pub use byteswriter::{BytesWriter, WriteExt};
-pub use json::{to_writer, to_writer_pretty};
+pub use json::{set_str_formatter_fn, to_writer, to_writer_pretty};

src/str/avx512.rs

@@ -93,10 +93,18 @@ pub fn unicode_from_str(buf: &str) -> *mut pyo3_ffi::PyObject {
         if unlikely!(buf.is_empty()) {
             return use_immortal!(crate::typeref::EMPTY_UNICODE);
         }
+        STR_CREATE_FN(buf)
+    }
+}
+
+pub type StrDeserializer = unsafe fn(&str) -> *mut pyo3_ffi::PyObject;
+
+static mut STR_CREATE_FN: StrDeserializer = super::scalar::str_impl_kind_scalar;
+
+pub fn set_str_create_fn() {
+    unsafe {
         if std::is_x86_feature_detected!("avx512vl") {
-            create_str_impl_avx512vl(buf)
-        } else {
-            super::scalar::unicode_from_str(buf)
+            STR_CREATE_FN = create_str_impl_avx512vl;
         }
     }
 }

src/str/mod.rs

@@ -7,9 +7,9 @@ mod pyunicode_new;
 mod scalar;
 
 #[cfg(not(feature = "avx512"))]
-pub use scalar::unicode_from_str;
+pub use scalar::{set_str_create_fn, unicode_from_str};
 
 #[cfg(feature = "avx512")]
-pub use avx512::unicode_from_str;
+pub use avx512::{set_str_create_fn, unicode_from_str};
 
 pub use ffi::*;

src/str/scalar.rs

@@ -3,10 +3,13 @@
 use crate::str::pyunicode_new::{
     pyunicode_ascii, pyunicode_fourbyte, pyunicode_onebyte, pyunicode_twobyte,
 };
-use crate::typeref::EMPTY_UNICODE;
 
-#[inline(always)]
-pub fn str_impl_kind_scalar(buf: &str, num_chars: usize) -> *mut pyo3_ffi::PyObject {
+#[inline(never)]
+pub fn str_impl_kind_scalar(buf: &str) -> *mut pyo3_ffi::PyObject {
+    let num_chars = bytecount::num_chars(buf.as_bytes());
+    if buf.len() == num_chars {
+        return pyunicode_ascii(buf.as_ptr(), num_chars);
+    }
     unsafe {
         let len = buf.len();
         assume!(len > 0);
@@ -33,15 +36,14 @@ pub fn str_impl_kind_scalar(buf: &str, num_chars: usize) -> *mut pyo3_ffi::PyObj
     }
 }
 
-#[inline(never)]
+#[cfg(not(feature = "avx512"))]
+#[inline(always)]
 pub fn unicode_from_str(buf: &str) -> *mut pyo3_ffi::PyObject {
     if unlikely!(buf.is_empty()) {
-        return use_immortal!(EMPTY_UNICODE);
-    }
-    let num_chars = bytecount::num_chars(buf.as_bytes());
-    if buf.len() == num_chars {
-        pyunicode_ascii(buf.as_ptr(), num_chars)
-    } else {
-        str_impl_kind_scalar(buf, num_chars)
+        return use_immortal!(crate::typeref::EMPTY_UNICODE);
     }
+    str_impl_kind_scalar(buf)
 }
+
+#[cfg(not(feature = "avx512"))]
+pub fn set_str_create_fn() {}

src/typeref.rs

@@ -140,6 +140,10 @@ fn _init_typerefs_impl() -> bool {
         assert!(crate::deserialize::KEY_MAP
             .set(crate::deserialize::KeyMap::default())
             .is_ok());
+
+        crate::serialize::writer::set_str_formatter_fn();
+        crate::str::set_str_create_fn();
+
         FRAGMENT_TYPE = orjson_fragmenttype_new();
         PyDateTime_IMPORT();
         NONE = Py_None();