[InstCombine] add vector support for (A >> C) == (B >> C) --> (A^B) u< (1 << C)

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4616

Does returning here affect any later combines that could have folded?

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4616

This is the only fold of (A >> AP1) == (B >> AP2) so far, and I think it should be moved here when other same pattern appear.

[RKSimon]

llvm/test/Transforms/InstCombine/compare-signs.ll

62

Please can add non-uniform vector tests as well?

[Chenbing.Zheng]

rebase precommit tests, and add support for vector's undef

[Chenbing.Zheng]

llvm/test/Transforms/InstCombine/compare-signs.ll

80–81

precommit negative tests here

80–81

undef is allowed
https://alive2.llvm.org/ce/z/gq9YWm

[spatel]

LGTM - see inline for some possible minor improvements.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4608

"C" in the result is not the same as "C" in the source - update this comment to make the transform clearer?

4629

Do you plan to change this one too? It will be easier to read the code and see the early return if we make a smaller helper function to handle the related patterns.

4646

Another transform that can be updated...

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4608

Sorry - disregard that comment. I misread the transform. This looks fine.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

if ((Op0->hasOneUse() && match(Op0, m_ZExt(m_Value(A))) &&
       (Op1->hasOneUse() && match(Op0, m_And(m_Value(B), m_ConstantInt(Cst1))) &&
        match(Op1, m_ZExt(m_Value(A))))) {
     APInt Pow2 = Cst1->getValue() + 1;
     if (Pow2.isPowerOf2() && isa<IntegerType>(A->getType()) &&
         Pow2.logBase2() == cast<IntegerType>(A->getType())->getBitWidth())
       return new ICmpInst(Pred, A, Builder.CreateTrunc(B, A->getType()));
   }
 
   // (A >> C) == (B >> C) --> (A^B) u< (1 << C)
   // For lshr and ashr pairs.
-  if ((match(Op0, m_OneUse(m_LShr(m_Value(A), m_ConstantInt(Cst1)))) &&
-       match(Op1, m_OneUse(m_LShr(m_Value(B), m_Specific(Cst1))))) ||
-      (match(Op0, m_OneUse(m_AShr(m_Value(A), m_ConstantInt(Cst1)))) &&
-       match(Op1, m_OneUse(m_AShr(m_Value(B), m_Specific(Cst1)))))) {
-    unsigned TypeBits = Cst1->getBitWidth();
-    unsigned ShAmt = (unsigned)Cst1->getLimitedValue(TypeBits);
+  const APInt *AP1, *AP2;
+  if ((match(Op0, m_OneUse(m_LShr(m_Value(A), m_APInt(AP1)))) &&
+       match(Op1, m_OneUse(m_LShr(m_Value(B), m_APInt(AP2))))) ||
+      (match(Op0, m_OneUse(m_AShr(m_Value(A), m_APInt(AP1)))) &&
+       match(Op1, m_OneUse(m_AShr(m_Value(B), m_APInt(AP2)))))) {
+    if (AP1 != AP2)
+      return nullptr;
+    unsigned TypeBits = AP1->getBitWidth();
+    unsigned ShAmt = AP1->getLimitedValue(TypeBits);
     if (ShAmt < TypeBits && ShAmt != 0) {
       ICmpInst::Predicate NewPred =
           Pred == ICmpInst::ICMP_NE ? ICmpInst::ICMP_UGE : ICmpInst::ICMP_ULT;
       Value *Xor = Builder.CreateXor(A, B, I.getName() + ".unshifted");
       APInt CmpVal = APInt::getOneBitSet(TypeBits, ShAmt);
       return new ICmpInst(NewPred, Xor, Builder.getInt(CmpVal));
     }
   }
 
   // (A << C) == (B << C) --> ((A^B) & (~0U >> C)) == 0
   if (match(Op0, m_OneUse(m_Shl(m_Value(A), m_ConstantInt(Cst1)))) &&
       match(Op1, m_OneUse(m_Shl(m_Value(B), m_Specific(Cst1))))) {
     unsigned TypeBits = Cst1->getBitWidth();
     unsigned ShAmt = (unsigned)Cst1->getLimitedValue(TypeBits);
     if (ShAmt < TypeBits && ShAmt != 0) {
       Value *Xor = Builder.CreateXor(A, B, I.getName() + ".unshifted");
       APInt AndVal = APInt::getLowBitsSet(TypeBits, TypeBits - ShAmt);
       Value *And = Builder.CreateAnd(Xor, Builder.getInt(AndVal),
                                       I.getName() + ".mask");
       return new ICmpInst(Pred, And, Constant::getNullValue(Cst1->getType()));
     }
   }
 
   // Transform "icmp eq (trunc (lshr(X, cst1)), cst" to
   // "icmp (and X, mask), cst"
   uint64_t ShAmt = 0;
   if (Op0->hasOneUse() &&
       match(Op0, m_Trunc(m_OneUse(m_LShr(m_Value(A), m_ConstantInt(ShAmt))))) &&
       match(Op1, m_ConstantInt(Cst1)) &&
       // Only do this when A has multiple uses.  This is most important to do
       // when it exposes other optimizations.
       !A->hasOneUse()) {
     unsigned ASize = cast<IntegerType>(A->getType())->getPrimitiveSizeInBits();
 
     if (ShAmt < ASize) {
       APInt MaskV =

llvm/test/Transforms/InstCombine/compare-signs.ll

 ;
   %t0 = lshr i32 %a, 31
   %t1 = lshr i32 %b, 31
   %t2 = icmp eq i32 %t0, %t1
   %t3 = zext i1 %t2 to i32
   ret i32 %t3
 }
 
-; TODO this should optimize but doesn't due to missing vector support in InstCombiner::foldICmpEquality.
 define <2 x i32> @test3vec(<2 x i32> %a, <2 x i32> %b) nounwind readnone {
 ; CHECK-LABEL: @test3vec(
+; CHECK-NEXT:    [[T2_UNSHIFTED:%.*]] = xor <2 x i32> [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[T2:%.*]] = icmp sgt <2 x i32> [[T2_UNSHIFTED]], <i32 -1, i32 -1>
+; CHECK-NEXT:    [[T3:%.*]] = zext <2 x i1> [[T2]] to <2 x i32>
+; CHECK-NEXT:    ret <2 x i32> [[T3]]
+;
+  %t0 = lshr <2 x i32> %a, <i32 31, i32 31>
+  %t1 = lshr <2 x i32> %b, <i32 31, i32 31>
+  %t2 = icmp eq <2 x i32> %t0, %t1
+  %t3 = zext <2 x i1> %t2 to <2 x i32>
+  ret <2 x i32> %t3
+}
+
+; negative test
+
+define <2 x i32> @test4vec(<2 x i32> %a, <2 x i32> %b) nounwind readnone {
+; CHECK-LABEL: @test4vec(
 ; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[B:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[B:%.*]], <i32 30, i32 30>
 ; CHECK-NEXT:    [[T2:%.*]] = icmp eq <2 x i32> [[T0]], [[T1]]
 ; CHECK-NEXT:    [[T3:%.*]] = zext <2 x i1> [[T2]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[T3]]
 ;
   %t0 = lshr <2 x i32> %a, <i32 31, i32 31>
-  %t1 = lshr <2 x i32> %b, <i32 31, i32 31>
+  %t1 = lshr <2 x i32> %b, <i32 30, i32 30>
   %t2 = icmp eq <2 x i32> %t0, %t1
   %t3 = zext <2 x i1> %t2 to <2 x i32>
   ret <2 x i32> %t3
 }
 
 ; Variation on @test3: checking the 2nd bit in a situation where the 5th bit
 ; is one, not zero.
 define i32 @test3i(i32 %a, i32 %b) nounwind readnone {
 ; CHECK-LABEL: @test3i(
 ; CHECK-NEXT:    [[T01:%.*]] = xor i32 [[A:%.*]], [[B:%.*]]
 ; CHECK-NEXT:    [[T4:%.*]] = icmp sgt i32 [[T01]], -1
 ; CHECK-NEXT:    [[T5:%.*]] = zext i1 [[T4]] to i32
 ; CHECK-NEXT:    ret i32 [[T5]]
 ;