Add is_ascii function optimized for x86-64 for [u8]

The new `is_ascii` function is optimized to use the
`pmovmskb` vector instruction which tests the high bit in a lane.
This corresponds to the same check of whether a byte is ASCII so
more bytes can be checked at a time for ASCII validity. This instruction
does not exist on other platforms so it is likely to regress performance
and is gated to all(target_arch = "x86_64", target_feature = "sse2").

Remove crate::mem import for functions which now exist in the prelude
Add codegen test

Author: okaneco

library/core/src/slice/ascii.rs

@@ -3,7 +3,7 @@
 use core::ascii::EscapeDefault;
 
 use crate::fmt::{self, Write};
-use crate::{ascii, iter, mem, ops};
+use crate::{ascii, iter, ops};
 
 #[cfg(not(test))]
 impl [u8] {
@@ -297,14 +297,6 @@ impl<'a> fmt::Debug for EscapeAscii<'a> {
     }
 }
 
-/// Returns `true` if any byte in the word `v` is nonascii (>= 128). Snarfed
-/// from `../str/mod.rs`, which does something similar for utf8 validation.
-#[inline]
-const fn contains_nonascii(v: usize) -> bool {
-    const NONASCII_MASK: usize = usize::repeat_u8(0x80);
-    (NONASCII_MASK & v) != 0
-}
-
 /// ASCII test *without* the chunk-at-a-time optimizations.
 ///
 /// This is carefully structured to produce nice small code -- it's smaller in
@@ -335,9 +327,17 @@ pub const fn is_ascii_simple(mut bytes: &[u8]) -> bool {
 ///
 /// If any of these loads produces something for which `contains_nonascii`
 /// (above) returns true, then we know the answer is false.
+#[cfg(not(all(target_arch = "x86_64", target_feature = "sse2")))]
 #[inline]
 const fn is_ascii(s: &[u8]) -> bool {
-    const USIZE_SIZE: usize = mem::size_of::<usize>();
+    /// Returns `true` if any byte in the word `v` is nonascii (>= 128). Snarfed
+    /// from `../str/mod.rs`, which does something similar for utf8 validation.
+    const fn contains_nonascii(v: usize) -> bool {
+        const NONASCII_MASK: usize = usize::repeat_u8(0x80);
+        (NONASCII_MASK & v) != 0
+    }
+
+    const USIZE_SIZE: usize = size_of::<usize>();
 
     let len = s.len();
     let align_offset = s.as_ptr().align_offset(USIZE_SIZE);
@@ -347,7 +347,7 @@ const fn is_ascii(s: &[u8]) -> bool {
     //
     // We also do this for architectures where `size_of::<usize>()` isn't
     // sufficient alignment for `usize`, because it's a weird edge case.
-    if len < USIZE_SIZE || len < align_offset || USIZE_SIZE < mem::align_of::<usize>() {
+    if len < USIZE_SIZE || len < align_offset || USIZE_SIZE < align_of::<usize>() {
         return is_ascii_simple(s);
     }
 
@@ -381,7 +381,7 @@ const fn is_ascii(s: &[u8]) -> bool {
     // have alignment information it should have given a `usize::MAX` for
     // `align_offset` earlier, sending things through the scalar path instead of
     // this one, so this check should pass if it's reachable.
-    debug_assert!(word_ptr.is_aligned_to(mem::align_of::<usize>()));
+    debug_assert!(word_ptr.is_aligned_to(align_of::<usize>()));
 
     // Read subsequent words until the last aligned word, excluding the last
     // aligned word by itself to be done in tail check later, to ensure that
@@ -420,3 +420,48 @@ const fn is_ascii(s: &[u8]) -> bool {
 
     !contains_nonascii(last_word)
 }
+
+/// ASCII test optimized to use the `pmovmskb` instruction available on `x86-64`
+/// platforms.
+///
+/// Other platforms are not likely to benefit from this code structure, so they
+/// use SWAR techniques to test for ASCII in `usize`-sized chunks.
+#[cfg(all(target_arch = "x86_64", target_feature = "sse2"))]
+#[inline]
+const fn is_ascii(bytes: &[u8]) -> bool {
+    // Process chunks of 32 bytes at a time in the fast path to enable
+    // auto-vectorization and use of `pmovmskb`. Two 128-bit vector registers
+    // can be OR'd together and then the resulting vector can be tested for
+    // non-ASCII bytes.
+    const CHUNK_SIZE: usize = 32;
+
+    let mut i = 0;
+
+    while i + CHUNK_SIZE <= bytes.len() {
+        let chunk_end = i + CHUNK_SIZE;
+
+        // Get LLVM to produce a `pmovmskb` instruction on x86-64 which
+        // creates a mask from the most significant bit of each byte.
+        // ASCII bytes are less than 128 (0x80), so their most significant
+        // bit is unset.
+        let mut count = 0;
+        while i < chunk_end {
+            count += bytes[i].is_ascii() as u8;
+            i += 1;
+        }
+
+        // All bytes should be <= 127 so count is equal to chunk size.
+        if count != CHUNK_SIZE as u8 {
+            return false;
+        }
+    }
+
+    // Process the remaining `bytes.len() % N` bytes.
+    let mut is_ascii = true;
+    while i < bytes.len() {
+        is_ascii &= bytes[i].is_ascii();
+        i += 1;
+    }
+
+    is_ascii
+}

tests/codegen/slice-is-ascii.rs

@@ -0,0 +1,16 @@
+//@ only-x86_64
+//@ compile-flags: -C opt-level=3
+#![crate_type = "lib"]
+
+/// Check that the fast-path of `is_ascii` uses a `pmovmskb` instruction.
+/// Platforms lacking an equivalent instruction use other techniques for
+/// optimizing `is_ascii`.
+// CHECK-LABEL: @is_ascii_autovectorized
+#[no_mangle]
+pub fn is_ascii_autovectorized(s: &[u8]) -> bool {
+    // CHECK: load <32 x i8>
+    // CHECK-NEXT: icmp slt <32 x i8>
+    // CHECK-NEXT: bitcast <32 x i1>
+    // CHECK-NEXT: icmp eq i32
+    s.is_ascii()
+}