Bad codegen for _mm_packus_epi16

[hsivonen]

Consider https://rust.godbolt.org/z/dY8schnG4

Input:

#![feature(portable_simd)]
use std::simd::ToBytes;
use core::simd::u8x16;
use core::simd::u16x8;
use core::arch::x86_64::_mm_packus_epi16;

#[inline(always)]
pub fn simd_pack_sse2(a: u16x8, b: u16x8) -> u8x16 {
    unsafe {
        _mm_packus_epi16(a.into(), b.into()).into()
    }
}

#[inline(always)]
pub fn simd_pack_portable(a: u16x8, b: u16x8) -> u8x16 {
    let first: u8x16 = a.to_le_bytes();
    let second: u8x16 = b.to_le_bytes();
    let (ret, _) = first.deinterleave(second);
    ret
}

#[inline(always)]
fn split_u16_stride(stride: &[u16; 16]) -> (&[u16; 8], &[u16; 8]) {
    let (chunks, _) = stride.as_chunks::<8>();
    (&chunks[0], &chunks[1])
}

#[unsafe(no_mangle)]
pub fn pack_portable(src: &[u16; 16], dst: &mut [u8; 16]) {
    let (a, b) = split_u16_stride(src);
    let packed = simd_pack_portable((*a).into(), (*b).into());
    *dst = packed.to_array();
}

#[unsafe(no_mangle)]
pub fn pack_sse2(src: &[u16; 16], dst: &mut [u8; 16]) {
    let (a, b) = split_u16_stride(src);
    let packed = simd_pack_sse2((*a).into(), (*b).into());
    *dst = packed.to_array();
}

Output:

.LCPI0_0:
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
pack_portable:
        movdqu  xmm0, xmmword ptr [rdi]
        movdqu  xmm1, xmmword ptr [rdi + 16]
        movdqa  xmm2, xmmword ptr [rip + .LCPI0_0]
        pand    xmm1, xmm2
        pand    xmm0, xmm2
        packuswb        xmm0, xmm1
        movdqu  xmmword ptr [rsi], xmm0
        ret

.LCPI1_0:
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
        .short  255
pack_sse2:
        movdqu  xmm0, xmmword ptr [rdi]
        movdqu  xmm1, xmmword ptr [rdi + 16]
        pxor    xmm2, xmm2
        pmaxsw  xmm0, xmm2
        movdqa  xmm3, xmmword ptr [rip + .LCPI1_0]
        pminsw  xmm0, xmm3
        pmaxsw  xmm1, xmm2
        pminsw  xmm1, xmm3
        packuswb        xmm0, xmm1
        movdqu  xmmword ptr [rsi], xmm0
        ret

Expected pack_sse2 to compile like this instead:

pack_sse2:
        movdqu  xmm0, xmmword ptr [rdi]
        movdqu  xmm1, xmmword ptr [rdi + 16]
        packuswb        xmm0, xmm1
        movdqu  xmmword ptr [rsi], xmm0
        ret

By code inspection, #2033 is the likely regressor. It seems problematic that it landed without tests to ensure that the optimizer actually recognizes the pattern.

[hsivonen]

CC @okaneco @folkertdev

[RalfJung]

It seems problematic that it landed without tests to ensure that the optimizer actually recognizes the pattern.

Every single intrinsic has tests to ensure that the optimizer recognizes the pattern, that's what the assert_instr attributes everywhere do. However, we can't test whether LLVM recognizes the pattern in every possible context, so issues like this can still sometimes occur.

Please be a bit more cautious with your judgments before researching or asking what testing is actually happening.

[hsivonen]

It seems problematic that it landed without tests to ensure that the optimizer actually recognizes the pattern.

Every single intrinsic has tests to ensure that the optimizer recognizes the pattern, that's what the assert_instr attributes everywhere do. However, we can't test whether LLVM recognizes the pattern in every possible context, so issues like this can still sometimes occur.

Please be a bit more cautious with your judgments before researching or asking what testing is actually happening.

I'm sorry. I failed to locate such a test. Looking at things more closely now, do I understand correctly that the assert_instr proc macro requires the output to have a packuswb instruction but doesn't check that there aren't suspiciously many other instructions?

This does not appear to be a matter of the context of my earlier example causing a pattern recognition failure, since the extra instructions occur even with a direct wrapper: https://rust.godbolt.org/z/qjhEn7aao (Rust 1.96 seems to have been the first version with the regression.)

Rust 1.95.0 was OK: https://rust.godbolt.org/z/vxv56ecj1

As for a hopefully constructive suggestion for how to improve testing going forward, it might be useful to dump the instruction counts for the snippets generated by the tests and check that those counts don't increase over time. Also, when the expectation is for a pattern this elaborate to optimize away, it might be worthwhile to have substrings that must not occur in the assembly. In this case, the source has min, max, and shuffle operations that are meant to optimize to something different, so it might be worthwhile to have a way to test that the generated assembly does not contain "min", "max", or "shuf" to rule out instructions like pmaxsw, pminsw, and pshufb.

But even with such testing, relying on a pattern this large seems brittle based on observing how much simpler patterns can optimize nicely in isolation but not when put together and considering that clang doesn't appear to rely on optimizer pattern matching but on __builtin_ia32_packuswb128.

[hsivonen]

This operation is performance-important for string creation in SpiderMonkey and DOM text node creation in Gecko. Measuring the performance change of that encoding_rs::mem operation with cargo bench in isolation of Gecko from Rust 1.95.0 to Rust 1.96.1 on Skylake, I see Rust 1.96.1 at about 0.70x the perf of 1.95.0 both with SSE2 codegen only and with AVX2 codegen (expected to produce the same code shape but with v-prefixed instruction variants). Rust-default opt_level=3 vs. opt_level=2, what Firefox uses, does not show a difference.

Looking at the discussion on #2033, it seems to me that the motivation of the change was to be able to mark the operation as const. Would it be possible to roll the change back when the intrinsic is not being const evaluated while leaving the min, max, shuffle formulation for const evaluation?

[hsivonen]

If I remove the ISA-specfic specialization and use the simd_pack_portable formulation on 1.96.1, then the perf is at 0.91x of 1.95.0 (with _mm_packus_epi16) with AVX2 codegen.

[RalfJung]

FWIW none of the stdarch intrinsics are actually stably const at the moment and from the discussion in rust-lang/rust#149648, it's not clear that there is a good case for making some subset of the intrinsics const (it's hard to pick a good subset).

In my eyes the motivation is that using the simd_* intrinsics instead of target-specific LLVM intrinsics makes it a lot easier for this code to be compiled with the cranelift and GCC backends and to be supported in Miri.

[hsivonen]

It seems to me that the Miri case could be addressed by cfg retaining the min, max, shuffle formulation for Miri.

Before seeing const mentioned in the discussion, I guessed GCC and Cranelift could be part of it, but since I don't develop either, I don't know whether maintaining a pattern match for this large a pattern is better than something like __builtin_ia32_packuswb128 that clang uses.

Relying on the optimizer to recognize such a large pattern seems to pose the risk that some earlier optimization consumes part of the pattern away and then the big pattern no longer matches.

(Apart from encoding_rs, in Firefox ppv-lite86 or jxl_simd use this. A GitHub search indicates that this intrinsic is used in various image/video processing crates. It seems applicable at least to yuv conversion and possibly other pixel format conversions.)

[folkertdev]

Some general context: we prefer the portable formulation for a variety of reasons: Miri, the other backends, and general readability. We've had good luck fixing mixed optimizations in LLVM in the past, and once a pattern makes it into the LLVM test suite it generally won't regress. (or rather, regressions are bugs and can be fixed in LLVM point releases).

The testing being kind of meh here is a long-standing issue. Ideally we'd lock in the specific LLVM IR (like rustc and LLVM themselves do) but that requires a fair amount of infrastructure that nobody has gotten around to.

---

It seems to me that the Miri case could be addressed by cfg retaining the min, max, shuffle formulation for Miri.

That is kind of risky because then Miri might have subtly different behavior.

Relying on the optimizer to recognize such a large pattern seems to pose the risk that some earlier optimization consumes part of the pattern away and then the big pattern no longer matches.

On the LLVM IR level the pattern is reasonably clear, and easy to match on. Just, somebody has to write that code. Yes it's a pile of heuristics in practice to determine the right order but given actual usage examples they can just be locked in and pretty much guaranteed to work.

In this particular case, inverting the shuffle and the clamp gives the expected assembly, see https://rust.godbolt.org/z/9vxP3vbnK.

[okaneco]

This regression was reported here rust-lang/rust#159464. It doesn't look like anyone reported upstream to LLVM with the LLVM IR, but the IR for swapping the shuffle and clamp looks simpler as it only has one set of min/max intrinsics.

I can try to put up a PR for swapping the concatenation and clamp if that's desired. It was noted that at the time, it broke the avx512 intrinsics tests though rust-lang/rust#159464 (comment)

[RalfJung]

Ouch, so we knew about this for more than a month then.

We should probably revert #2033 then as a short-term fix. The swap of shuffle and clamp was mentioned in that issue but hasn't been done yet so I assume it is non-trivial. It's too late for 1.98 but we can backport that revert to the beta branch that will become 1.99.

[folkertdev]

Yeah, a revert is probably the best call for now. I believe for _mm512_packs_epi16 and friends the shuffle-then-clamp approach only has upsides, but we can land that separately.

[RalfJung]

The testing being kind of meh here is a long-standing issue. Ideally we'd lock in the specific LLVM IR (like rustc and LLVM themselves do) but that requires a fair amount of infrastructure that nobody has gotten around to.

For complicated patterns such as these, maybe we should just have codegen-llvm tests in rustc? It's not ideal but it's better than nothing.

[okaneco]

Being able to write CHECK-NOT tests for instances of unwanted instructions would be nice to catch these regressions.

This unmerged PR was another case of trying to use generic Rust SIMD intrinsics to replace LLVM intrinsics. It passed stdarch tests but would have been a regression due to a conditional branch which needed to be manually verified by inspecting assembly output.