Comparison of Option<NonZeroU*> is not fully optimized

[kennytm]

Compare:

#![feature(nonzero)]
use std::num::NonZeroU32;

pub fn f(a: Option<NonZeroU32>, b: Option<NonZeroU32>) -> bool {
    a < b
}
pub fn g(a: u32, b: u32) -> bool {
    a < b
}

The functions f and g have equivalent effect, and should compile to the same code. However, the translated f is much more complex, involving needless comparison of the discriminant.

LLVM IR and ASM (in release mode)

; playground::f
; Function Attrs: norecurse nounwind readnone uwtable
define zeroext i1 @_ZN10playground1f17hd628eff49d2ff60dE(i32, i32) unnamed_addr #0 {
start:
  %2 = icmp ne i32 %0, 0
  %3 = icmp ne i32 %1, 0
  %4 = xor i1 %2, %3
  br i1 %4, label %bb7.i, label %bb6.i

bb6.i:                                            ; preds = %start
  %5 = icmp ult i32 %0, %1
  %spec.select = and i1 %2, %5
  ret i1 %spec.select

bb7.i:                                            ; preds = %start
  %6 = xor i1 %2, true
  %7 = and i1 %3, %6
  ret i1 %7
}

; playground::g
; Function Attrs: norecurse nounwind readnone uwtable
define zeroext i1 @_ZN10playground1g17hea1c394facd055aeE(i32 %a, i32 %b) unnamed_addr #0 {
start:
  %0 = icmp ult i32 %a, %b
  ret i1 %0
}

playground::f:
	test	edi, edi
	setne	al
	test	esi, esi
	setne	cl
	xor	cl, al
	je	.LBB0_1
	test	esi, esi
	setne	cl
	test	edi, edi
	sete	al
	and	al, cl
	ret

.LBB0_1:
	test	edi, edi
	setne	cl
	cmp	edi, esi
	setb	al
	and	al, cl
	ret

playground::g:
	cmp	edi, esi
	setb	al
	ret

Note that specializing PartialOrd etc for Option<NonZero*> is not a valid solution, since custom wrapper types of NonZero* will still have the same problem.

#[derive(PartialOrd, PartialEq)]
pub struct S(NonZeroU32);
pub fn f(a: Option<S>, b: Option<S>) -> bool {
    a < b
}

cc #49137.

[scottmcm]

Possibly-related: #49572

[oli-obk]

I'm assuming we need to change the layout here

rust/src/librustc_trans/mir/rvalue.rs

Line 435 in 5ebf748

layout: self.cx.layout_of(discr_ty)

In order to pass a layout that knows that the discriminant is not just an integer, but a subset of an integer.

Also maybe the discriminant intrinsic must be adjusted:

rust/src/librustc_trans/intrinsic.rs

Line 406 in 31ae4f9

"discriminant_value" => {

Maybe we should add compiler-internal support for subsetting integers? So TypeVariants::Int also has a range field that we can set?

cc @eddyb
cc @rust-lang/wg-codegen

[hanna-kruppe]

I don't see how layout factors in, Option<NonZeroU32> is already represented as a single 32 bit integer with 0 as None. The generic PartialOrd is just not fully optimized for this monomorphization. That's an optimization LLVM could just do, it just currently doesn't.

[hanna-kruppe]

To clarify, the Alive link is just to confirm the overall transformation is correct within LLVM IR. Obviously this very specific pattern should not be added verbatim to instcombine (it wouldn't even help since I replaced the branching with select to make Alive accept it). But since this optimization applies only to particular monomorphizations and we can't reasonably special case this in rustc, we should really report this as a LLVM bug, so that they can look into which combination of passes could be improved to catch this.

[oli-obk]

Oh... I assumed that we're missing some assume annotations that would allow this optimization from happening by itself.

[eddyb]

What's happening in <Option as PartialOrd>::le, even? Is this the output of auto-derive? What's the actual code that optimizes to the shown LLVM IR? I'm wondering if this is a derive inefficiency.

[kennytm]

as of 1.57-nightly, the ASM generated for f is becoming even worse.

; playground::f
; Function Attrs: norecurse nounwind nonlazybind readnone uwtable willreturn
define zeroext i1 @_ZN10playground1f17h98eba4f4ca80f23cE(i32 %0, i32 %1) unnamed_addr #0 {
start:
  %2 = icmp ne i32 %0, 0
  %3 = icmp ne i32 %1, 0
  %4 = xor i1 %2, %3
  br i1 %4, label %bb2.i.i, label %bb1.i.i

bb2.i.i:                                          ; preds = %start
  %5 = xor i1 %2, true
  %_3.i.i.i.i = and i1 %3, %5
  %.0.i.i.i.i = select i1 %_3.i.i.i.i, i8 -1, i8 1
  br label %_ZN4core3cmp10PartialOrd2lt17h869f0d8a39125a29E.exit

bb1.i.i:                                          ; preds = %start
  %.not.i.i = icmp eq i32 %0, 0
  %.not5.i.i = icmp eq i32 %1, 0
  %or.cond.i.i = or i1 %.not.i.i, %.not5.i.i
  br i1 %or.cond.i.i, label %_ZN4core3cmp10PartialOrd2lt17h869f0d8a39125a29E.exit, label %bb5.i.i

bb5.i.i:                                          ; preds = %bb1.i.i
  %_3.i.i.i.i.i = icmp ult i32 %0, %1
  %_6.i.i.i.i.i = icmp ne i32 %0, %1
  %spec.select.i.i.i.i.i = zext i1 %_6.i.i.i.i.i to i8
  %.0.i.i.i.i.i = select i1 %_3.i.i.i.i.i, i8 -1, i8 %spec.select.i.i.i.i.i
  br label %_ZN4core3cmp10PartialOrd2lt17h869f0d8a39125a29E.exit

_ZN4core3cmp10PartialOrd2lt17h869f0d8a39125a29E.exit: ; preds = %bb2.i.i, %bb1.i.i, %bb5.i.i
  %.2.i.i = phi i8 [ %.0.i.i.i.i, %bb2.i.i ], [ 0, %bb1.i.i ], [ %.0.i.i.i.i.i, %bb5.i.i ]
  %cond.i = icmp eq i8 %.2.i.i, -1
  ret i1 %cond.i
}

example::f:
        test    edi, edi
        setne   al
        test    esi, esi
        setne   cl
        cmp     al, cl
        je      .LBB0_2
        test    edi, edi
        setne   al
        add     al, al
        add     al, -1
        test    esi, esi
        movzx   ecx, al
        mov     eax, 1
        cmovne  eax, ecx
.LBB0_5:
        cmp     al, -1
        sete    al
        ret
.LBB0_2:
        xor     eax, eax
        test    edi, edi
        je      .LBB0_5
        test    esi, esi
        je      .LBB0_5
        xor     ecx, ecx
        cmp     edi, esi
        setne   cl
        mov     eax, 255
        cmovae  eax, ecx
        cmp     al, -1
        sete    al
        ret

the regression happens on 1.52.0 which upgraded to LLVM 12.

---

the GCC backend produces less instructions, but still far from ideal:

example::f:
        xor     eax, eax
        test    esi, esi
        je      .L6
        test    edi, edi
        jne     .L12
        mov     eax, 1
        ret
.L12:
        cmp     edi, esi
        setb    al
.L6:
        ret

example::g:
        cmp     edi, esi
        setb    al
        ret

[saethlin]

Most of the above comments are about PartialOrd, but this also afflicts PartialEq. Including on the latest nightly with NewPM.

pub fn h(a: Option<NonZeroU32>, b: Option<NonZeroU32>) -> bool {
    a == b
}

example::h:
        test    edi, edi
        setne   al
        test    esi, esi
        setne   cl
        xor     cl, al
        je      .LBB0_2
        xor     eax, eax
        ret
.LBB0_2:
        test    esi, esi
        sete    al
        test    edi, edi
        sete    cl
        or      cl, al
        cmp     edi, esi
        sete    al
        or      al, cl
        ret

IMO the fact that LLVM can't do the expected optimization to just compare as two u32s is more egregious.