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:

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.