[InstCombine] Improve eq/ne by parts to handle `ult/ugt` equality pattern.

[goldstein.w.n]

Summary

(icmp eq/ne (lshr x, C), (lshr y, C) gets optimized to `(icmp
ult/uge (xor x, y), (1 << C)`. This can cause the current equal by
parts detection to miss the high-bits as it may get optimized to the
new pattern.

This commit adds support for detecting / combining the ult/ugt
pattern.

[goldstein.w.n]

Note: I'm happy the drop this if you feel its not worth the complexity. Its mostly because D148593 causes a regression in the test without this change.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

 
 /// Materialize an extraction of bits from an integer in IR.
 static Value *extractIntPart(const IntPart &P, IRBuilderBase &Builder) {
   Value *V = P.From;
   if (P.StartBit)
     V = Builder.CreateLShr(V, P.StartBit);
   Type *TruncTy = V->getType()->getWithNewBitWidth(P.NumBits);
   if (TruncTy != V->getType())
-    V = Builder.CreateTrunc(V, TruncTy);
+    V = Builder.CreateZExtOrTrunc(V, TruncTy);
   return V;
 }
 
 /// (icmp eq X0, Y0) & (icmp eq X1, Y1) -> icmp eq X01, Y01
 /// (icmp ne X0, Y0) | (icmp ne X1, Y1) -> icmp ne X01, Y01
 /// where X0, X1 and Y0, Y1 are adjacent parts extracted from an integer.
 Value *InstCombinerImpl::foldEqOfParts(ICmpInst *Cmp0, ICmpInst *Cmp1,
                                        bool IsAnd) {
   if (!Cmp0->hasOneUse() || !Cmp1->hasOneUse())
     return nullptr;
 
   CmpInst::Predicate Pred = IsAnd ? CmpInst::ICMP_EQ : CmpInst::ICMP_NE;
-  if (Cmp0->getPredicate() != Pred || Cmp1->getPredicate() != Pred)
+  auto MatchPred = [&](ICmpInst *Cmp) -> std::pair<bool, const APInt *> {
+    if (Pred == Cmp->getPredicate())
+      return {true, nullptr};
+
+    const APInt *C;
+    // (icmp eq (lshr x, C), (lshr y, C)) gets optimized to:
+    // (icmp ult (xor x, y), 1 << C) so also look for that.
+    if (Pred == CmpInst::ICMP_EQ && Cmp->getPredicate() == CmpInst::ICMP_ULT)
+      return {match(Cmp->getOperand(1), m_APInt(C)) && C->isPowerOf2() &&
+                  match(Cmp->getOperand(0), m_Xor(m_Value(), m_Value())),
+              C};
+
+    // (icmp ne (lshr x, C), (lshr y, C)) gets optimized to:
+    // (icmp uge (xor x, y), (1 << C) - 1) so also look for that.
+    if (Pred == CmpInst::ICMP_NE && Cmp->getPredicate() == CmpInst::ICMP_UGT)
+      return {match(Cmp->getOperand(1), m_APInt(C)) && C->isMask() &&
+                  match(Cmp->getOperand(0), m_Xor(m_Value(), m_Value())),
+              C};
+
+    return {false, nullptr};
+  };
+
+  auto GetMatchPart = [&](std::pair<bool, const APInt *> MatchResult,
+                          ICmpInst *Cmp,
+                          unsigned OpNo) -> std::optional<IntPart> {
+    // Normal IntPart
+    if (MatchResult.second == nullptr)
+      return matchIntPart(Cmp->getOperand(OpNo));
+
+    // We have one of the ult/ugt patterns.
+    unsigned From;
+    const APInt *C = MatchResult.second;
+    if (Pred == CmpInst::ICMP_NE)
+      From = C->popcount();
+    else
+      From = (*C - 1).popcount();
+    Instruction *I = cast<Instruction>(Cmp->getOperand(0));
+    return {{I->getOperand(OpNo), From,
+             Cmp->getOperand(0)->getType()->getScalarSizeInBits()}};
+  };
+
+  auto Cmp0Match = MatchPred(Cmp0);
+  auto Cmp1Match = MatchPred(Cmp1);
+  if (!Cmp0Match.first || !Cmp1Match.first)
     return nullptr;
 
-  std::optional<IntPart> L0 = matchIntPart(Cmp0->getOperand(0));
-  std::optional<IntPart> R0 = matchIntPart(Cmp0->getOperand(1));
-  std::optional<IntPart> L1 = matchIntPart(Cmp1->getOperand(0));
-  std::optional<IntPart> R1 = matchIntPart(Cmp1->getOperand(1));
+  std::optional<IntPart> L0 = GetMatchPart(Cmp0Match, Cmp0, 0);
+  std::optional<IntPart> R0 = GetMatchPart(Cmp0Match, Cmp0, 1);
+  std::optional<IntPart> L1 = GetMatchPart(Cmp1Match, Cmp1, 0);
+  std::optional<IntPart> R1 = GetMatchPart(Cmp1Match, Cmp1, 1);
+
   if (!L0 || !R0 || !L1 || !R1)
     return nullptr;
 
   // Make sure the LHS/RHS compare a part of the same value, possibly after
   // an operand swap.
   if (L0->From != L1->From || R0->From != R1->From) {
     if (L0->From != R1->From || R0->From != L1->From)
       return nullptr;

llvm/test/Transforms/InstCombine/eq-of-parts.ll

;
   %c.1 = icmp eq i8 %x.1, %y.1
   %c.2 = icmp eq i8 %x.2, %y.2
   %c.210 = or i1 %c.2, %c.1
   ret i1 %c.210
 }
 
 define i1 @eq_optimized_highbits_cmp(i32 %x, i32 %y) {
 ; CHECK-LABEL: @eq_optimized_highbits_cmp(
-; CHECK-NEXT:    [[XOR:%.*]] = xor i32 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[CMP_HI:%.*]] = icmp ult i32 [[XOR]], 33554432
-; CHECK-NEXT:    [[TX:%.*]] = trunc i32 [[X]] to i25
-; CHECK-NEXT:    [[TY:%.*]] = trunc i32 [[Y]] to i25
-; CHECK-NEXT:    [[CMP_LO:%.*]] = icmp eq i25 [[TX]], [[TY]]
-; CHECK-NEXT:    [[R:%.*]] = and i1 [[CMP_HI]], [[CMP_LO]]
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %xor = xor i32 %y, %x
   %cmp_hi = icmp ult i32 %xor, 33554432
   %tx = trunc i32 %x to i25
   %ty = trunc i32 %y to i25
   %cmp_lo = icmp eq i25 %tx, %ty
   %r = and i1 %cmp_hi, %cmp_lo
;
   %ty = trunc i32 %y to i24
   %cmp_lo = icmp eq i24 %tx, %ty
   %r = and i1 %cmp_hi, %cmp_lo
   ret i1 %r
 }
 
 define i1 @ne_optimized_highbits_cmp(i32 %x, i32 %y) {
 ; CHECK-LABEL: @ne_optimized_highbits_cmp(
-; CHECK-NEXT:    [[XOR:%.*]] = xor i32 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[CMP_HI:%.*]] = icmp ugt i32 [[XOR]], 16777215
-; CHECK-NEXT:    [[TX:%.*]] = trunc i32 [[X]] to i24
-; CHECK-NEXT:    [[TY:%.*]] = trunc i32 [[Y]] to i24
-; CHECK-NEXT:    [[CMP_LO:%.*]] = icmp ne i24 [[TX]], [[TY]]
-; CHECK-NEXT:    [[R:%.*]] = or i1 [[CMP_HI]], [[CMP_LO]]
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i32 [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %xor = xor i32 %y, %x
   %cmp_hi = icmp ugt i32 %xor, 16777215
   %tx = trunc i32 %x to i24
   %ty = trunc i32 %y to i24
   %cmp_lo = icmp ne i24 %tx, %ty
   %r = or i1 %cmp_hi, %cmp_lo