Bad codegen for comparing struct of two 16bit ints

[bjorn3]

View all comments

I tried this code:

#[derive(Clone, Copy, PartialEq)]
pub struct Foo {
    pub x: u16,
    pub y: u16,
}

#[no_mangle]
fn eq(a: &Foo, b: &Foo) -> bool {
    a == b
}

I expected to see this happen: a and b are loaded into a single register each and then the registers are compared against each other.

Instead, this happened: https://rust.godbolt.org/z/8zaxKG8vx

For -Copt-level=2:

eq:
        movzx   eax, word ptr [rdi]
        movzx   ecx, word ptr [rdi + 2]
        xor     ax, word ptr [rsi]
        movzx   edx, word ptr [rsi + 2]
        xor     edx, ecx
        or      dx, ax
        sete    al
        ret

For -Copt-level=3:

eq:
        movd    xmm0, dword ptr [rdi]
        movd    xmm1, dword ptr [rsi]
        pcmpeqw xmm1, xmm0
        pshuflw xmm0, xmm1, 80
        pshufd  xmm0, xmm0, 80
        movmskpd        eax, xmm0
        cmp     eax, 3
        sete    al
        ret

Meta

rustc --version --verbose:
Both 1.86 and nightly.

[bjorn3]

Using transmute::<_, u32>(read_unaligned(self)) and transmute::<_, u32>(read_unaligned(other)) inside of the PartialEq impl produces the expected codegen.

[folkertdev]

Trying this in C shows that it is also unable to optimize (the equivalent of) Foo::eq, but it is able to generate good code when memcmp is used:

https://godbolt.org/z/cqhM9hThK

#include <stdint.h>
#include <stdbool.h>
#include <string.h>

typedef struct {
    uint16_t x;
    uint16_t y;
} Foo;

bool eq(const Foo* a, const Foo* b) {
    return memcmp(a, b, sizeof(Foo)) == 0;
}

bool partial_eq(const Foo* a, const Foo* b) {
    return a->x == b->x && a->y == b->y;
}

eq:
        mov     eax, dword ptr [rdi]
        cmp     eax, dword ptr [rsi]
        sete    al
        ret

partial_eq:
        movzx   eax, word ptr [rdi]
        cmp     ax, word ptr [rsi]
        jne     .LBB1_1
        movzx   eax, word ptr [rdi + 2]
        cmp     ax, word ptr [rsi + 2]
        sete    al
        ret
.LBB1_1:
        xor     eax, eax
        ret

That suggests that codegen for Eq could be improved by using memcmp in some cases. (fields must be copy, no padding , probably other conditions that I'm missing).

[hanna-kruppe]

At the time derives are expanded, field types aren’t known at all. The derive just gets tokens and it might not even see the same tokens as later stages of the compiler. So the derive can’t do the necessary checks (notably including: is bitwise equality even correct for all field types). In theory the built-in derives could expand to some magic intrinsic that’s lowered much later when types are known, but that’s a big hammer. I don’t see any fundamental reason why LLVM shouldn’t be able to do this optimization (at least for Rust) and that would help non-derived code as well. But I haven’t checked Alive2.

[zachs18]

Note that a manual (a.x == b.x) && (a.y == b.y) compiles to the same suboptimal codegen (which makes sense, since that's what the derive expands to), but (a.x == b.x) & (a.y == b.y) (non-short-circuiting &) compiles to 32-bit load and compare.

Looks like LLVM in all three cases does 2 vectorized loads of 2 x i16, and then a compare of the vectors, but how it ANDs the resulting 2 x i1 is different: OP and && seems to do like if cmp[0] { cmp[1] } else { false }, while & does something else.

(Note the OP version's LLVM IR is different from the && version, but IIUC only in var names and debug info; they get deduplicated anyway if both are present)

// common prefix
  %0 = load <2 x i16>, ptr %a, align 2
  %1 = load <2 x i16>, ptr %b, align 2
  %2 = icmp eq <2 x i16> %0, %1

// OP and && verison
  %3 = extractelement <2 x i1> %2, i64 0
  %4 = extractelement <2 x i1> %2, i64 1
  %_0.sroa.0.0 = select i1 %3, i1 %4, i1 false
  ret i1 %_0.sroa.0.0

// & verison
  %shift = shufflevector <2 x i1> %2, <2 x i1> poison, <2 x i32> <i32 1, i32 poison>
  %3 = and <2 x i1> %2, %shift
  %_0 = extractelement <2 x i1> %3, i64 0
  ret i1 %_0

godbolt link with asm and llvm ir

(Maybe LLVM is assuming that the second fields could be poison in the && case, so avoiding using their comparison if the first comparison is false?)

[bjorn3]

but (a.x == b.x) & (a.y == b.y) (non-short-circuiting &) compiles to 32-bit load and compare.

Unfortunately not in the code I minimized this from that does roughly

#[unsafe(no_mangle)]
pub fn foo(a: &Foo) -> bool {
    foo_eq_3(a, &Foo { x: 1, y: 1 })
}

That still results in two separate loads even with non-short-circuiting &:

https://rust.godbolt.org/z/1e47McvW1

foo:
        movzx   eax, word ptr [rdi]
        xor     eax, 1
        movzx   ecx, word ptr [rdi + 2]
        xor     ecx, 1
        or      cx, ax
        sete    al
        ret

[jrmuizel]

memorysafety/rav1d#1400 (comment) suggests that adding !undef to the load would fix this.

[hanna-kruppe]

I assume they meant !noundef. As noted in a later comment rustc already adds that metadata when it emits LLVM IR, but SimplifyCFG drops them when turning the two branches on icmps into a select.

That's not simply a missed optimization opportunity. While the load of the second field can't load poison/undef, that property is control-dependent. For example, when matching on Option<u32> behind a reference, loading the u32 field in the Some case is also !noundef, but that part of memory will generally be undef/poison on the None branch. Thus, hoisting !noundef loads generally has to drop the !noundef metadata.

Solving this seems hard: I don't think LLVM has a way to express "loading through this pointer always reads initialized bytes" (like how dereferenceable says "loading through this pointer is always okay but may yield poison/undef"). I'm also no sure if we could apply such an attribute to Rust references -- it may break common unsafe patterns that form references to uninitialized memory. The reference conservatively calls that UB for now, but even there it's noted that this is controversial and may end up being non-UB. Miri doesn't enforce any such rule, I believe. So Rust may be unable to justify the desired optimization for the same reason why C compilers can't do it.

[mmastrac]

This might be a tangent, but it feels related and gives us some very useful building blocks:

If we had something like Plain (previously discussed in the forum) named AllBitPatternsValid<Repr=...> that was derived automatically for structs, and a const stdlib utility function std::mem::all_bit_patterns_valid<T> and std::mem::bitwise_compare<T>, could PartialEq generate code like so?

Fuller example at: https://play.rust-lang.org/?version=stable&mode=release&edition=2024&gist=bc3313d6c653a1c9844c49eaf735528a

unsafe trait AllBitPatternsValid {
    type Repr: AllBitPatternsValid;
}

pub const fn bitwise_compare<T: AllBitPatternsValid>(a: &T, b: &T) -> bool { ... }
pub const fn all_bit_patterns_valid<T: Sized>() -> bool { ... }

impl PartialEq for MyStruct {
    #[inline(always)]
    fn eq(self: &MyStruct, other: &MyStruct) -> bool {
        if all_bit_patterns_valid::<Self>() {
            bitwise_compare(self, other)
        } else {
            self.x == other.x && self.y == other.y
        }
    }
}

This approach wouldn’t directly address the !noundef metadata issue discussed earlier, but it could enable the compiler to perform safe, low-level comparisons for types where all bit patterns are valid, potentially improving performance.

Comparing the assembly when all_bit_patterns_valid returns false vs when all_bit_patterns_valid returns true, we can see that the compiler happily generates the wider comparisons:

-	movq	(%rsp), %rax
-	movq	8(%rsp), %rdi
-	movzwl	(%rax), %ecx
-	cmpw	(%rdi), %cx
-	jne	.LBB4_6
-	movzwl	2(%rax), %eax
-	cmpw	2(%rdi), %ax
-	jne	.LBB4_6
+	movq	8(%rsp), %rax
+	movq	16(%rsp), %rdi
+	movl	(%rax), %ebx
+	movl	(%rdi), %ebp

It's not a full solution. Obviously it would certainly require some nuance around types like Box that are technically AllBitPatternsValid but require pointer chasing to compute equality -- perhaps something like Plain or Pod might be a better approach, but it does seem like these building blocks could be very useful.