String::to_lowercase does not get vectorized well contrary to code comments

[jhorstmann]

I'm looking into the performance of to_lowercase / to_uppercaseon mostly ascii strings, using a small microbenchmark added to library/alloc/benches/string.rs.

#[bench]
fn bench_to_lowercase(b: &mut Bencher) {
 let s = "Hello there, the quick brown fox jumped over the lazy dog! \
 Lorem ipsum dolor sit amet, consectetur. ";
 b.iter(|| s.to_lowercase())
}

Using linux perf tooling I see that the hot part of the code is the following large loop, which despite heavy use of sse2 instructions only seems to process 32 bytes per iteration.

 │ d0:┌─→mov r9,QWORD PTR [r14+r15*1] ▒
 │ │ movdqu xmm3,XMMWORD PTR [r14+r15*1] ▒
 0,12 │ │ pshufd xmm12,xmm3,0xee ▒
 2,56 │ │ movq rdx,xmm12 ▒
 │ │ mov rsi,rdx ▒
 │ │ or rsi,r9 ▒
 0,59 │ │ test rsi,rcx ▒
 │ │↓ jne 319 ▒
 2,10 │ │ mov rsi,r9 ▒
 │ │ mov rdi,r9 ▒
 │ │ mov r8,r9 ▒
 │ │ mov r10,r9 ▒
 2,23 │ │ shr r9d,0x8 ▒
 │ │ movd xmm12,r9d ▒
 │ │ shr r10,0x20 ◆
 0,12 │ │ pshufd xmm13,xmm3,0x44 ▒
 1,05 │ │ movdqa xmm14,xmm3 ▒
 │ │ psrlq xmm14,0x10 ▒
 │ │ psrlq xmm13,0x18 ▒
 0,12 │ │ movsd xmm13,xmm14 ▒
 1,75 │ │ movd xmm14,r10d ▒
 │ │ shr r8,0x28 ▒
 │ │ punpcklqdq xmm3,xmm12 ▒
 │ │ movd xmm12,r8d ▒
 1,52 │ │ andpd xmm13,xmm0 ▒
 │ │ pand xmm3,xmm0 ▒
 │ │ packuswb xmm3,xmm13 ▒
 0,35 │ │ pshufd xmm13,xmm14,0x50 ▒
 1,05 │ │ movdqa xmm14,XMMWORD PTR [rip+0x403c7] ▒
 │ │ pandn xmm14,xmm13 ▒
 │ │ psllq xmm12,0x28 ▒
 │ │ movdqa xmm13,XMMWORD PTR [rip+0x403c3] ▒
 2,94 │ │ pandn xmm13,xmm12 ▒
 │ │ shr rdi,0x30 ▒
 │ │ por xmm13,xmm14 ▒
 0,35 │ │ movd xmm12,edi ▒
 2,22 │ │ shr rsi,0x38 ▒
 │ │ packuswb xmm3,xmm1 ▒
 │ │ packuswb xmm3,xmm1 ▒
 0,47 │ │ por xmm13,xmm3 ▒
 4,21 │ │ psllq xmm12,0x30 ▒
 │ │ movdqa xmm3,xmm4 ▒
 │ │ pandn xmm3,xmm12 ▒
 0,35 │ │ movd xmm12,esi ▒
 2,47 │ │ mov esi,edx ▒
 │ │ shr esi,0x8 ▒
 │ │ pand xmm13,xmm4 ▒
 0,35 │ │ por xmm3,xmm13 ▒
 2,10 │ │ pand xmm3,xmm5 ▒
 │ │ psllq xmm12,0x38 ▒
 │ │ movdqa xmm13,xmm5 ▒
 │ │ pandn xmm13,xmm12 ▒
 2,34 │ │ por xmm13,xmm3 ▒
 │ │ movd xmm3,edx ▒
 │ │ pshufd xmm3,xmm3,0x44 ▒
 0,53 │ │ movdqa xmm12,xmm6 ▒
 2,47 │ │ pandn xmm12,xmm3 ▒
 │ │ movd xmm3,esi ▒
 │ │ mov esi,edx ▒
 0,23 │ │ shr esi,0x10 ▒
 2,64 │ │ pand xmm13,xmm6 ▒
 │ │ por xmm12,xmm13 ▒
 │ │ pslldq xmm3,0x9 ▒
 0,12 │ │ movdqa xmm13,xmm7 ▒
 2,45 │ │ pandn xmm13,xmm3 ▒
 │ │ movd xmm3,esi ▒
 │ │ mov esi,edx ▒
 0,51 │ │ shr esi,0x18 ▒
 2,60 │ │ pand xmm12,xmm7 ▒
 │ │ por xmm13,xmm12 ▒
 │ │ pslldq xmm3,0xa ▒
 │ │ movdqa xmm12,xmm8 ▒
 1,76 │ │ pandn xmm12,xmm3 ▒
 │ │ movd xmm3,esi ▒
 │ │ mov rsi,rdx ▒
 0,47 │ │ shr rsi,0x20 ▒
 2,34 │ │ pand xmm13,xmm8 ▒
 │ │ por xmm12,xmm13 ▒
 │ │ pslldq xmm3,0xb ▒
 0,23 │ │ movdqa xmm13,xmm9 ▒
 1,99 │ │ pandn xmm13,xmm3 ▒
 │ │ movd xmm3,esi ▒
 │ │ mov rsi,rdx ▒
 0,35 │ │ shr rsi,0x28 ▒
 2,97 │ │ pand xmm12,xmm9 ▒
 │ │ por xmm13,xmm12 ▒
 │ │ pshufd xmm3,xmm3,0x0 ▒
 0,12 │ │ movdqa xmm12,xmm10 ▒
 2,11 │ │ pandn xmm12,xmm3 ▒
 │ │ movd xmm3,esi ▒
 │ │ shr rdx,0x30 ▒
 │ │ pand xmm13,xmm10 ▒
 1,87 │ │ por xmm12,xmm13 ▒
 │ │ pand xmm12,xmm11 ▒
 │ │ pslldq xmm3,0xd ▒
 0,23 │ │ movdqa xmm13,xmm11 ▒
 2,23 │ │ pandn xmm13,xmm3 ▒
 │ │ por xmm13,xmm12 ▒
 │ │ pand xmm13,XMMWORD PTR [rip+0x40320] ▒
 0,12 │ │ movd xmm3,edx ▒
 2,80 │ │ pslldq xmm3,0xe ▒
 │ │ por xmm3,xmm13 ▒
 │ │ pand xmm3,XMMWORD PTR [rip+0x4031a] ▒
 0,23 │ │ movzx edx,BYTE PTR [r14+r15*1+0xf] ▒
 3,31 │ │ movd xmm12,edx ▒
 │ │ pslldq xmm12,0xf ▒
 │ │ por xmm12,xmm3 ▒
 0,12 │ │ movdqa xmm3,xmm12 ▒
 2,92 │ │ paddb xmm3,XMMWORD PTR [rip+0x31d97] # 1009a0 <anon.cf73386a2f5127d166baeac25be116f0.63.llvm.16014458289627072720+0x459> ▒
 │ │ movdqa xmm13,xmm3 ▒
 │ │ pminub xmm13,xmm15 ▒
 0,47 │ │ pcmpeqb xmm13,xmm3 ▒
 1,53 │ │ pand xmm13,xmm2 ▒
 0,36 │ │ por xmm13,xmm12 ▒
 │ │ movdqu XMMWORD PTR [rax+r15*1],xmm13 ▒
 0,23 │ │ lea rdx,[r15+0x10] ▒
 2,34 │ │ add r15,0x20 ▒
 0,12 │ │ cmp r15,rbx ▒
 │ │ mov r15,rdx ▒
 1,64 │ └──jbe d0 

I don't see an easy way to improve the autovectorization of this code, but it should be relatively easy to explicitly vectorize it using portable_simd, and I would like to prepare such a PR if there are no objections. As far as I know, portable_simd is already in use inside core, for example by #103779.

[the8472]

The way the loop ors the bits for the is-ascii check into a single usize doesn't seem ideal for vectorization. keeping the lanes independent should yield better autovectorization results.

[jhorstmann]

@the8472 Thanks for the suggestion. I already started with a simd implementation, but now checked again if the autovectorization could be improved. With the following main loop

 while slice.len() >= N {
 let chunk = &slice[..N];

 let mut is_ascii = true;
 for j in 0..N {
 is_ascii &= chunk[j] <= 127;
 }
 if !is_ascii {
 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..];
 }

The assembly and performance is indeed better, but there is still some weird shuffling and shifting going on:

 0,08 │ 80:┌─→movdqu xmm3,XMMWORD PTR [rbx+rax*1] ▒
 0,07 │ │ pshufd xmm4,xmm3,0xee ▒
 0,37 │ │ por xmm4,xmm3 ▒
 6,22 │ │ pshufd xmm5,xmm4,0x55 ▒
 0,16 │ │ por xmm5,xmm4 ▒
 │ │ movdqa xmm4,xmm5 ▒
 2,02 │ │ psrld xmm4,0x10 ▒
 4,48 │ │ por xmm4,xmm5 ▒
 0,07 │ │ movdqa xmm5,xmm4 ▒
 0,08 │ │ psrlw xmm5,0x8 ▒
 0,60 │ │ por xmm5,xmm4 ▒
 6,58 │ │ movd ecx,xmm5 ▒
 │ │ test cl,cl ▒
 │ │↓ js ef ▒
 │ │ movdqa xmm4,xmm3 ▒
 1,80 │ │ paddb xmm4,xmm0 ▒
 4,70 │ │ movdqa xmm5,xmm4 ▒
 │ │ pminub xmm5,xmm1 ▒
 │ │ pcmpeqb xmm5,xmm4 ▒
 0,37 │ │ pand xmm5,xmm2 ▒
 8,50 │ │ por xmm5,xmm3 ▒
 │ │ movdqu XMMWORD PTR [r13+rax*1+0x0],xmm5 ▒
 │ │ add rax,0x10 ▒
 1,94 │ │ add r12,0xfffffffffffffff0 ▒
 1,05 │ ├──cmp r12,0xf ▒
 4,19 │ └──ja 80 

The explicit simd version looks a bit better, mostly because the ascii check directly translates to a pmovmskb:

 const LANES: usize = 16;

 let simd_range_start = Simd::splat(range_start);
 let simd_range_end = Simd::splat(range_end);
 let simd_xor_value = Simd::splat(xor_value);

 while slice.len() >= LANES {
 let chunk = Simd::<u8, LANES>::from_slice(slice);
 let is_ascii = chunk.cast::<i8>().simd_ge(Simd::splat(0));

 if is_ascii.all() {
 let is_in_range = chunk.simd_ge(simd_range_start) & chunk.simd_le(simd_range_end);
 let converted = is_in_range.select(chunk ^ simd_xor_value, chunk);
 // SAFETY: output has enough capacity and we never read the uninitialized slice
 unsafe {
 let out_slice = core::slice::from_raw_parts_mut(out_ptr, LANES);
 converted.copy_to_slice(out_slice);
 out_ptr = out_ptr.add(LANES);
 }
 i += LANES;
 slice = &slice[LANES..];
 } else {
 break;
 }
 }

 │ 80:┌─→movdqu xmm3,XMMWORD PTR [r14+rcx*1] ▒
 0,12 │ │ pmovmskb edx,xmm3 ▒
 0,12 │ │ test edx,edx ▒
 │ │↓ jne d4 ▒
 5,82 │ │ movdqa xmm4,xmm3 ▒
 0,04 │ │ paddb xmm4,xmm0 ▒
 │ │ movdqa xmm5,xmm4 ▒
 0,08 │ │ pminub xmm5,xmm1 ▒
 10,05 │ │ pcmpeqb xmm5,xmm4 ▒
 │ │ movdqa xmm4,xmm5 ▒
 │ │ pandn xmm4,xmm3 ▒
 │ │ pxor xmm3,xmm2 ▒
 7,02 │ │ pand xmm3,xmm5 ▒
 0,04 │ │ por xmm3,xmm4 ▒
 │ │ movdqu XMMWORD PTR [rax+rcx*1],xmm3 ▒
 │ │ add rcx,0x10 ▒
 9,69 │ │ mov rdx,rsi ▒
 │ │ add rdx,0xfffffffffffffff0 ▒
 │ │ mov rsi,rdx ▒
 │ ├──cmp rdx,0xf ▒
 5,03 │ └──ja 80 

Do you have a preference here between autovectorization and explicit simd?

[the8472]

If the SIMD impl results in reasonable code across architectures, including some without SIMD then that should be fine.
I think it would be the first time where we use portable simd unconditionally, so it could use some additional scrutiny.

If it's only meant to target x86 with SSE2 then that'd mean having multiple paths and in that case we'd still have to tweak the generic path anyway. At that point we might as well just optimize the latter further.

The assembly and performance is indeed better, but there is still some weird shuffling and shifting going on:

You can probably get rid of those too by keeping multiple bools in a small array. That way the optimizer will more easily see that it can shove them into independent simd lanes. Increasing the unroll count might help too to fit better into simd registers.

We do get pretty decent autovectorization in other places in the standard library by massaging things into a state that basically looks like arrays-instead-of-SIMD-types. E.g.

rust/library/core/src/str/iter.rs

Lines 61 to 73 in 40cf1f9

	while remainder > CHUNK_SIZE
	&& let Some(chunk) = chunks.next()
	{
	bytes_skipped += CHUNK_SIZE;
	let mut start_bytes = [false; CHUNK_SIZE];
	for i in 0..CHUNK_SIZE {
	start_bytes[i] = !super::validations::utf8_is_cont_byte(chunk[i]);
	}
	remainder -= start_bytes.into_iter().map(|i| i as u8).sum::<u8>() as usize;
	}

[jhorstmann]

You can probably get rid of those too by keeping multiple bools in a small array. That way the optimizer will more easily see that it can shove them into independent simd lanes.

It seems llvm is too "smart" for this trick, on its own that generates the exact same code. The sum trick does work, although it seems a bit fragile. The following creates a pmovmskb instruction just like the simd version, but a slightly different alternative with a non_ascii array and comparing against 0 does not.

 let mut is_ascii = [false; N];

 for j in 0..N {
 is_ascii[j] = chunk[j] <= 127;
 }

 if is_ascii.into_iter().map(|x| x as u8).sum::<u8>() as usize != N {
 break;
 }

If the SIMD impl results in reasonable code across architectures, including some without SIMD then that should be fine.

It only needs to be better than the autovectorized version on those platforms ;)

If it's only meant to target x86 with SSE2 then that'd mean having multiple paths and in that case we'd still have to tweak the generic path anyway. At that point we might as well just optimize the latter further.

My initial idea was to gate the simd code on either SSE2 or Neon (possibly also Altivec and RiscV). I'd also add a scalar loop for the non-multiple-of-N remainder, so all-ascii strings are fully handled by specialized code. Currently this remainder goes through the generic char_indices and conversions::to_lower code path. On other platforms, only the scalar loop would be used, since I assume the autovectorization would also generate suboptimal code.

But I agree that if similar code quality can be achieved with autovectorization, that would be preferable. I'll open a PR after a little bit more polishing the code.