From 6c58e74fca6cee64665647d1f2fb44576d5f9019 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B6rn=20Horstmann?= Date: Wed, 10 Apr 2024 22:22:18 +0200 Subject: [PATCH] Improve autovectorization of to_lowercase / to_uppercase functions Refactor the code in the `convert_while_ascii` helper function to make it more suitable for auto-vectorization and also process the full ascii prefix of the string. The generic case conversion logic will only be invoked starting from the first non-ascii character. The runtime on microbenchmarks with ascii-only inputs improves between 2x for short and 7x for long inputs on x86_64 and aarch64. The new implementation also encapsulates all unsafe inside the `convert_while_ascii` function. Fixes #123712 --- library/alloc/benches/str.rs | 2 + library/alloc/src/str.rs | 124 +++++++++++++++++++++-------------- library/alloc/tests/str.rs | 17 +++++ 3 files changed, 93 insertions(+), 50 deletions(-) diff --git a/library/alloc/benches/str.rs b/library/alloc/benches/str.rs index c148ab6b220a5..92a48e0e6b5a6 100644 --- a/library/alloc/benches/str.rs +++ b/library/alloc/benches/str.rs @@ -347,3 +347,5 @@ make_test!(rsplitn_space_char, s, s.rsplitn(10, ' ').count()); make_test!(split_space_str, s, s.split(" ").count()); make_test!(split_ad_str, s, s.split("ad").count()); + +make_test!(to_lowercase, s, s.to_lowercase()); diff --git a/library/alloc/src/str.rs b/library/alloc/src/str.rs index 3e23612d0c13c..3c0b581d958a4 100644 --- a/library/alloc/src/str.rs +++ b/library/alloc/src/str.rs @@ -10,6 +10,7 @@ use core::borrow::{Borrow, BorrowMut}; use core::iter::FusedIterator; use core::mem; +use core::mem::MaybeUninit; use core::ptr; use core::str::pattern::{DoubleEndedSearcher, Pattern, ReverseSearcher, Searcher}; use core::unicode::conversions; @@ -366,14 +367,9 @@ impl str { without modifying the original"] #[stable(feature = "unicode_case_mapping", since = "1.2.0")] pub fn to_lowercase(&self) -> String { - let out = convert_while_ascii(self.as_bytes(), u8::to_ascii_lowercase); + let (mut s, rest) = convert_while_ascii(self, u8::to_ascii_lowercase); - // Safety: we know this is a valid char boundary since - // out.len() is only progressed if ascii bytes are found - let rest = unsafe { self.get_unchecked(out.len()..) }; - - // Safety: We have written only valid ASCII to our vec - let mut s = unsafe { String::from_utf8_unchecked(out) }; + let prefix_len = s.len(); for (i, c) in rest.char_indices() { if c == 'Σ' { @@ -382,8 +378,7 @@ impl str { // in `SpecialCasing.txt`, // so hard-code it rather than have a generic "condition" mechanism. // See https://github.com/rust-lang/rust/issues/26035 - let out_len = self.len() - rest.len(); - let sigma_lowercase = map_uppercase_sigma(&self, i + out_len); + let sigma_lowercase = map_uppercase_sigma(self, prefix_len + i); s.push(sigma_lowercase); } else { match conversions::to_lower(c) { @@ -459,14 +454,7 @@ impl str { without modifying the original"] #[stable(feature = "unicode_case_mapping", since = "1.2.0")] pub fn to_uppercase(&self) -> String { - let out = convert_while_ascii(self.as_bytes(), u8::to_ascii_uppercase); - - // Safety: we know this is a valid char boundary since - // out.len() is only progressed if ascii bytes are found - let rest = unsafe { self.get_unchecked(out.len()..) }; - - // Safety: We have written only valid ASCII to our vec - let mut s = unsafe { String::from_utf8_unchecked(out) }; + let (mut s, rest) = convert_while_ascii(self, u8::to_ascii_uppercase); for c in rest.chars() { match conversions::to_upper(c) { @@ -615,50 +603,86 @@ pub unsafe fn from_boxed_utf8_unchecked(v: Box<[u8]>) -> Box { unsafe { Box::from_raw(Box::into_raw(v) as *mut str) } } -/// Converts the bytes while the bytes are still ascii. +/// Converts leading ascii bytes in `s` by calling the `convert` function. +/// /// For better average performance, this happens in chunks of `2*size_of::()`. -/// Returns a vec with the converted bytes. +/// +/// Returns a tuple of the converted prefix and the remainder starting from +/// the first non-ascii character. #[inline] #[cfg(not(test))] #[cfg(not(no_global_oom_handling))] -fn convert_while_ascii(b: &[u8], convert: fn(&u8) -> u8) -> Vec { - let mut out = Vec::with_capacity(b.len()); - +fn convert_while_ascii(s: &str, convert: fn(&u8) -> u8) -> (String, &str) { const USIZE_SIZE: usize = mem::size_of::(); const MAGIC_UNROLL: usize = 2; const N: usize = USIZE_SIZE * MAGIC_UNROLL; - const NONASCII_MASK: usize = usize::from_ne_bytes([0x80; USIZE_SIZE]); - let mut i = 0; - unsafe { - while i + N <= b.len() { - // Safety: we have checks the sizes `b` and `out` to know that our - let in_chunk = b.get_unchecked(i..i + N); - let out_chunk = out.spare_capacity_mut().get_unchecked_mut(i..i + N); - - let mut bits = 0; - for j in 0..MAGIC_UNROLL { - // read the bytes 1 usize at a time (unaligned since we haven't checked the alignment) - // safety: in_chunk is valid bytes in the range - bits |= in_chunk.as_ptr().cast::().add(j).read_unaligned(); - } - // if our chunks aren't ascii, then return only the prior bytes as init - if bits & NONASCII_MASK != 0 { - break; - } + let mut slice = s.as_bytes(); + let mut out = Vec::with_capacity(slice.len()); + let mut out_slice = out.spare_capacity_mut(); - // perform the case conversions on N bytes (gets heavily autovec'd) - for j in 0..N { - // safety: in_chunk and out_chunk is valid bytes in the range - let out = out_chunk.get_unchecked_mut(j); - out.write(convert(in_chunk.get_unchecked(j))); - } + let mut i = 0_usize; - // mark these bytes as initialised - i += N; + // process the input in chunks to enable auto-vectorization + let mut is_ascii = [false; N]; + while slice.len() >= N { + // Safety: out_slice was allocated with same lengths as input slice and gets updated with + // the same offsets + unsafe { + core::intrinsics::assume(slice.len() == out_slice.len()); } - out.set_len(i); + + let chunk = &slice[..N]; + + for j in 0..N { + is_ascii[j] = chunk[j] <= 127; + } + + // auto-vectorization for this check is a bit fragile, + // sum and comparing against the chunk size gives the best result, + // specifically a pmovmsk instruction on x86. + if is_ascii.iter().map(|x| *x as u8).sum::() as usize != N { + break; + } + + for j in 0..N { + out_slice[j] = MaybeUninit::new(convert(&chunk[j])); + } + + i += N; + slice = &slice[N..]; + out_slice = &mut out_slice[N..]; } - out + // handle the remainder as individual bytes + while !slice.is_empty() { + // Safety: out_slice was allocated with same lengths as input slice and gets updated with + // the same offsets + unsafe { + core::intrinsics::assume(slice.len() == out_slice.len()); + } + + let byte = slice[0]; + if byte > 127 { + break; + } + out_slice[0] = MaybeUninit::new(convert(&byte)); + i += 1; + slice = &slice[1..]; + out_slice = &mut out_slice[1..]; + } + + unsafe { + // SAFETY: i bytes have been initialized above + out.set_len(i); + + // SAFETY: We have written only valid ascii to the output vec + let ascii_string = String::from_utf8_unchecked(out); + + // SAFETY: we know this is a valid char boundary + // since we only skipped over leading ascii bytes + let rest = core::str::from_utf8_unchecked(slice); + + (ascii_string, rest) + } } diff --git a/library/alloc/tests/str.rs b/library/alloc/tests/str.rs index 0078f5eaa3d2b..4f26dab46d9d6 100644 --- a/library/alloc/tests/str.rs +++ b/library/alloc/tests/str.rs @@ -1826,6 +1826,19 @@ fn to_lowercase() { assert_eq!("Α'Σ".to_lowercase(), "α'ς"); assert_eq!("Α''Σ".to_lowercase(), "α''ς"); + assert_eq!("aΣ".to_lowercase(), "aς"); + assert_eq!("a'Σ".to_lowercase(), "a'ς"); + assert_eq!("a''Σ".to_lowercase(), "a''ς"); + + assert_eq!("ÄΣ".to_lowercase(), "äς"); + assert_eq!("ä'Σ".to_lowercase(), "ä'ς"); + assert_eq!("ä''Σ".to_lowercase(), "ä''ς"); + + // input lengths around the boundary of the chunk size used by the ascii prefix optimization + assert_eq!("abcdefghijklmnoΣ".to_lowercase(), "abcdefghijklmnoς"); + assert_eq!("abcdefghijklmnopΣ".to_lowercase(), "abcdefghijklmnopς"); + assert_eq!("abcdefghijklmnopqΣ".to_lowercase(), "abcdefghijklmnopqς"); + assert_eq!("ΑΣ Α".to_lowercase(), "ας α"); assert_eq!("Α'Σ Α".to_lowercase(), "α'ς α"); assert_eq!("Α''Σ Α".to_lowercase(), "α''ς α"); @@ -1840,6 +1853,10 @@ fn to_lowercase() { assert_eq!("Α 'Σ".to_lowercase(), "α 'σ"); assert_eq!("Α ''Σ".to_lowercase(), "α ''σ"); + assert_eq!("Ä Σ".to_lowercase(), "ä σ"); + assert_eq!("Ä 'Σ".to_lowercase(), "ä 'σ"); + assert_eq!("Ä ''Σ".to_lowercase(), "ä ''σ"); + assert_eq!("Σ".to_lowercase(), "σ"); assert_eq!("'Σ".to_lowercase(), "'σ"); assert_eq!("''Σ".to_lowercase(), "''σ");