Port ASCII optimization from simdjson: Check whole block at once (#18)

Port ASCII optimization from simdjson, checks whole block at once, or'ing the SIMD chunks and then testing them for the high bit.

Author: hkratz

src/implementation/macros.rs

@@ -37,23 +37,34 @@ macro_rules! static_cast_i8 {
     };
 }
 
+macro_rules! check_utf8 {
+    ($feat:expr, $t:ident) => {
+        #[target_feature(enable = $feat)]
+        #[inline]
+        unsafe fn check_utf8(&self, previous: &mut Utf8CheckingState<$t>) {
+            if likely!(self.is_ascii()) {
+                previous.error =
+                    Utf8CheckingState::<$t>::check_eof(previous.error, previous.incomplete)
+            } else {
+                self.check_block(previous);
+            }
+        }
+    };
+}
+
 /// check_bytes() strategy
 macro_rules! check_bytes {
     ($feat:expr, $t:ident) => {
         #[target_feature(enable = $feat)]
         #[inline]
         unsafe fn check_bytes(current: $t, previous: &mut Utf8CheckingState<$t>) {
-            if likely!(Self::is_ascii(current)) {
-                previous.error = Self::check_eof(previous.error, previous.incomplete)
-            } else {
-                let prev1 = Self::prev1(current, previous.prev);
-                let sc = Self::check_special_cases(current, prev1);
-                previous.error = Self::or(
-                    previous.error,
-                    Self::check_multibyte_lengths(current, previous.prev, sc),
-                );
-                previous.incomplete = Self::is_incomplete(current);
-            }
+            let prev1 = Self::prev1(current, previous.prev);
+            let sc = Self::check_special_cases(current, prev1);
+            previous.error = Self::or(
+                previous.error,
+                Self::check_multibyte_lengths(current, previous.prev, sc),
+            );
+            previous.incomplete = Self::is_incomplete(current);
             previous.prev = current
         }
     };

src/implementation/x86/avx2.rs

@@ -37,12 +37,6 @@ impl Utf8CheckingState<__m256i> {
         _mm256_or_si256(a, b)
     }
 
-    #[target_feature(enable = "avx2")]
-    #[inline]
-    unsafe fn is_ascii(input: __m256i) -> bool {
-        _mm256_movemask_epi8(input) == 0
-    }
-
     #[target_feature(enable = "avx2")]
     #[inline]
     unsafe fn check_eof(error: __m256i, incomplete: __m256i) -> __m256i {
@@ -272,8 +266,6 @@ impl Utf8CheckingState<__m256i> {
         _mm256_testz_si256(error, error) != 1
     }
 
-    #[target_feature(enable = "avx2")]
-    #[inline]
     check_bytes!("avx2", __m256i);
 }
 
@@ -302,11 +294,18 @@ impl SimdInput {
 
     #[target_feature(enable = "avx2")]
     #[inline]
-    unsafe fn check_utf8(&self, state: &mut Utf8CheckingState<__m256i>) {
+    unsafe fn check_block(&self, state: &mut Utf8CheckingState<__m256i>) {
         Utf8CheckingState::<__m256i>::check_bytes(self.v0, state);
         Utf8CheckingState::<__m256i>::check_bytes(self.v1, state);
     }
 
+    #[target_feature(enable = "avx2")]
+    #[inline]
+    unsafe fn is_ascii(&self) -> bool {
+        let res = _mm256_or_si256(self.v0, self.v1);
+        _mm256_movemask_epi8(res) == 0
+    }
+
     #[target_feature(enable = "avx2")]
     #[inline]
     unsafe fn check_eof(state: &mut Utf8CheckingState<__m256i>) {
@@ -318,6 +317,8 @@ impl SimdInput {
     unsafe fn check_utf8_errors(state: &Utf8CheckingState<__m256i>) -> bool {
         Utf8CheckingState::<__m256i>::has_error(state.error)
     }
+
+    check_utf8!("avx2", __m256i);
 }
 
 use crate::implementation::Temp2x64A32;

src/implementation/x86/sse42.rs

@@ -36,12 +36,6 @@ impl Utf8CheckingState<__m128i> {
         _mm_or_si128(a, b)
     }
 
-    #[target_feature(enable = "sse4.2")]
-    #[inline]
-    unsafe fn is_ascii(input: __m128i) -> bool {
-        _mm_movemask_epi8(input) == 0
-    }
-
     #[target_feature(enable = "sse4.2")]
     #[inline]
     unsafe fn check_eof(error: __m128i, incomplete: __m128i) -> __m128i {
@@ -205,8 +199,6 @@ impl Utf8CheckingState<__m128i> {
         _mm_testz_si128(error, error) != 1
     }
 
-    #[target_feature(enable = "sse4.2")]
-    #[inline]
     check_bytes!("sse4.2", __m128i);
 }
 
@@ -239,13 +231,22 @@ impl SimdInput {
 
     #[target_feature(enable = "sse4.2")]
     #[inline]
-    unsafe fn check_utf8(&self, state: &mut Utf8CheckingState<__m128i>) {
+    unsafe fn check_block(&self, state: &mut Utf8CheckingState<__m128i>) {
         Utf8CheckingState::<__m128i>::check_bytes(self.v0, state);
         Utf8CheckingState::<__m128i>::check_bytes(self.v1, state);
         Utf8CheckingState::<__m128i>::check_bytes(self.v2, state);
         Utf8CheckingState::<__m128i>::check_bytes(self.v3, state);
     }
 
+    #[target_feature(enable = "sse4.2")]
+    #[inline]
+    unsafe fn is_ascii(&self) -> bool {
+        let r1 = _mm_or_si128(self.v0, self.v1);
+        let r2 = _mm_or_si128(self.v2, self.v3);
+        let r = _mm_or_si128(r1, r2);
+        _mm_movemask_epi8(r) == 0
+    }
+
     #[target_feature(enable = "sse4.2")]
     #[inline]
     unsafe fn check_eof(state: &mut Utf8CheckingState<__m128i>) {
@@ -257,6 +258,8 @@ impl SimdInput {
     unsafe fn check_utf8_errors(state: &Utf8CheckingState<__m128i>) -> bool {
         Utf8CheckingState::<__m128i>::has_error(state.error)
     }
+
+    check_utf8!("sse4.2", __m128i);
 }
 
 use crate::implementation::Temp2x64A16;