Remove common adds when icmp eqing two selects [InstCombine]

[scottmcm]

Take the following IR, extracted from part of how Rust generates == on enums:

define noundef i64 @get_discr(i8 noundef %x) unnamed_addr {
start:
  %_4 = icmp ule i8 %x, 1
  %small = select i1 %_4, i8 3, i8 %x
  %_5 = sub i8 %small, 2
  %_0 = zext i8 %_5 to i64
  ret i64 %_0
}

define noundef zeroext i1 @discr_eq(i8 noundef %a, i8 noundef %b) unnamed_addr {
start:
  %_3 = call noundef i64 @get_discr(i8 noundef %a)
  %_4 = call noundef i64 @get_discr(i8 noundef %b)
  %_0 = icmp eq i64 %_3, %_4
  ret i1 %_0
}

Today, when optimizing get_discr InstCombine will change the https://llvm.godbolt.org/z/6Eq87h6h7

(x <= 1 ? 3 : x) - 2

to

(x <= 1 ? 1 : x - 2)

which is reasonable on its own.

The problem, though, is that that makes things worse later in discr_eq.

It's unable to undo that transformation after the inlining, and thus ends up as https://llvm.godbolt.org/z/31MdEhfcv

define noundef zeroext i1 @discr_eq(i8 noundef %a, i8 noundef %b) unnamed_addr #0 {
  %0 = add i8 %a, -2
  %_4.inv.i = icmp ugt i8 %a, 1
  %_5.i = select i1 %_4.inv.i, i8 %0, i8 1
  %1 = add i8 %b, -2
  %_4.inv.i1 = icmp ugt i8 %b, 1
  %_5.i2 = select i1 %_4.inv.i1, i8 %1, i8 1
  %_0 = icmp eq i8 %_5.i, %_5.i2
  ret i1 %_0
}

But there's no reason to do those adds any more -- they could be pulled outside the selects where they'd cancel out.

Phrased as C, it's doing

(x <= 1 ? 1 : x - 2) == (y <= 1 ? 1 : y - 2)

when instead it could be doing

(x <= 1 ? 3 : x) == (y <= 1 ? 3 : y)

So InstCombine should be able to detect this case, and rewrite it back to

define noundef zeroext i1 @discr_eq(i8 noundef %a, i8 noundef %b) unnamed_addr #0 {
  %_4.inv.i = icmp ugt i8 %a, 1
  %_5.i = select i1 %_4.inv.i, i8 %a, i8 3
  %_4.inv.i1 = icmp ugt i8 %b, 1
  %_5.i2 = select i1 %_4.inv.i1, i8 %b, i8 3
  %_0 = icmp eq i8 %_5.i, %_5.i2
  ret i1 %_0
}

Alive2 proof of correctness for that: https://alive2.llvm.org/ce/z/AWCmDs

[scottmcm]

Hmm, I tried fixing this on the rust side (for how we emit enum discriminants, rust-lang/rust#139729), but it looks like the "move inside the select" part runs before the "it's an add on both sides of eq" part: https://llvm.godbolt.org/z/7q3bsKMP9

So that might be a good thing to fix too?

[dtcxzyw]

it looks like the "move inside the select" part runs before the "it's an add on both sides of eq" part

Yeah. It is the canonical form.
See also

llvm-project/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

Lines 863 to 864 in 3ed8363

	if (Instruction *NV = foldBinOpIntoSelectOrPhi(Add))
	return NV;

.