diff --git a/library/core/src/slice/ascii.rs b/library/core/src/slice/ascii.rs index 9db07d8abbbef..0dc336125c8cd 100644 --- a/library/core/src/slice/ascii.rs +++ b/library/core/src/slice/ascii.rs @@ -62,7 +62,10 @@ impl [u8] { return false; } - #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] + #[cfg(any( + all(target_arch = "x86_64", target_feature = "sse2"), + all(target_arch = "aarch64", target_feature = "neon") + ))] { const CHUNK_SIZE: usize = 16; // The following function has two invariants: @@ -108,7 +111,10 @@ impl [u8] { /// /// The caller must guarantee that the slices are equal in length, and the /// slice lengths are greater than or equal to `N` bytes. - #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] + #[cfg(any( + all(target_arch = "x86_64", target_feature = "sse2"), + all(target_arch = "aarch64", target_feature = "neon") + ))] #[inline] const fn eq_ignore_ascii_case_chunks(&self, other: &[u8]) -> bool { // FIXME(const-hack): The while-loops that follow should be replaced by @@ -451,7 +457,8 @@ pub const fn is_ascii_simple(mut bytes: &[u8]) -> bool { /// (above) returns true, then we know the answer is false. #[cfg(not(any( all(target_arch = "x86_64", target_feature = "sse2"), - all(target_arch = "loongarch64", target_feature = "lsx") + all(target_arch = "loongarch64", target_feature = "lsx"), + all(target_arch = "aarch64", target_feature = "neon") )))] #[inline] #[rustc_allow_const_fn_unstable(const_eval_select)] // fallback impl has same behavior @@ -584,24 +591,81 @@ fn is_ascii_sse2(bytes: &[u8]) -> bool { rest.iter().all(|b| b.is_ascii()) } -/// ASCII test optimized to use the `pmovmskb` instruction on `x86-64`. -/// -/// Uses explicit SSE2 intrinsics to prevent LLVM from auto-vectorizing with -/// broken AVX-512 code that extracts mask bits one-by-one. -#[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] +/// Chunk size for NEON vectorized ASCII checking (4x 16-byte loads). +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const NEON_CHUNK_SIZE: usize = 64; + +/// Width of a single NEON vector, used to vectorize the tail left over by the +/// unrolled `NEON_CHUNK_SIZE` loop. +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const NEON_VECTOR_SIZE: usize = 16; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +#[inline] +fn is_ascii_neon(bytes: &[u8]) -> bool { + use crate::arch::aarch64::{vld1q_u8, vmaxvq_u8, vorrq_u8}; + + let (chunks, rest) = bytes.as_chunks::(); + + for chunk in chunks { + let ptr = chunk.as_ptr(); + // SAFETY: chunk is 64 bytes, and `vld1q_u8` has no alignment requirement. + let max = unsafe { + let a1 = vld1q_u8(ptr); + let a2 = vld1q_u8(ptr.add(16)); + let b1 = vld1q_u8(ptr.add(32)); + let b2 = vld1q_u8(ptr.add(48)); + // OR all chunks - if any byte has high bit set, combined will too. + let combined = vorrq_u8(vorrq_u8(a1, a2), vorrq_u8(b1, b2)); + // `vmaxvq_u8` is a horizontal reduction with a longer latency than + // `vorrq_u8`, so it runs once per 64 bytes rather than once per load. + vmaxvq_u8(combined) + }; + if max >= 128 { + return false; + } + } + + // The unrolled loop above leaves up to 63 bytes, so sweep those a vector at + // a time before falling back to a byte-at-a-time check. + let (vectors, rest) = rest.as_chunks::(); + + for vector in vectors { + // SAFETY: vector is 16 bytes, and `vld1q_u8` has no alignment requirement. + let max = unsafe { vmaxvq_u8(vld1q_u8(vector.as_ptr())) }; + if max >= 128 { + return false; + } + } + + // Handle remaining bytes + rest.iter().all(|b| b.is_ascii()) +} + +/// Uses explicit SIMD intrinsics to prevent LLVM from auto-vectorizing with +/// broken code (e.g., AVX-512 on x86-64 that extracts mask bits one-by-one). +#[cfg(any( + all(target_arch = "x86_64", target_feature = "sse2"), + all(target_arch = "aarch64", target_feature = "neon") +))] #[inline] #[rustc_allow_const_fn_unstable(const_eval_select)] const fn is_ascii(bytes: &[u8]) -> bool { const USIZE_SIZE: usize = size_of::(); const NONASCII_MASK: usize = usize::MAX / 255 * 0x80; + #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] + const SIMD_MIN_LEN: usize = SSE2_CHUNK_SIZE; + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + const SIMD_MIN_LEN: usize = NEON_CHUNK_SIZE; + const_eval_select!( @capture { bytes: &[u8] } -> bool: if const { is_ascii_simple(bytes) } else { // For small inputs, use usize-at-a-time processing to avoid SSE2 call overhead. - if bytes.len() < SSE2_CHUNK_SIZE { + if bytes.len() < SIMD_MIN_LEN { let chunks = bytes.chunks_exact(USIZE_SIZE); let remainder = chunks.remainder(); for chunk in chunks { @@ -613,7 +677,10 @@ const fn is_ascii(bytes: &[u8]) -> bool { return remainder.iter().all(|b| b.is_ascii()); } - is_ascii_sse2(bytes) + #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] + { is_ascii_sse2(bytes) } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { is_ascii_neon(bytes) } } ) } diff --git a/library/coretests/benches/str/eq_ignore_ascii_case.rs b/library/coretests/benches/str/eq_ignore_ascii_case.rs index 29129b933bc46..db75960664106 100644 --- a/library/coretests/benches/str/eq_ignore_ascii_case.rs +++ b/library/coretests/benches/str/eq_ignore_ascii_case.rs @@ -43,3 +43,23 @@ fn bench_large_str_eq(b: &mut Bencher) { let other = black_box(en::LARGE); b.iter(|| assert!(s.eq_ignore_ascii_case(other))) } + +#[bench] +fn bench_medium_str_early_mismatch(b: &mut Bencher) { + // Differs in the first byte, exercising the early-return path of the + // chunked comparison. + let other_owned = ["_", &en::MEDIUM[1..]].concat(); + let s = black_box(en::MEDIUM); + let other = black_box(&*other_owned); + b.iter(|| assert!(!s.eq_ignore_ascii_case(other))) +} + +#[bench] +fn bench_medium_str_tail_mismatch(b: &mut Bencher) { + // Differs in the last byte, exercising the `last_chunk` tail handling of + // the chunked comparison. + let other_owned = [&en::MEDIUM[..en::MEDIUM.len() - 1], "_"].concat(); + let s = black_box(en::MEDIUM); + let other = black_box(&*other_owned); + b.iter(|| assert!(!s.eq_ignore_ascii_case(other))) +}