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

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4629

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

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4629

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

78

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

92

precommit negative tests here

[spatel]

LGTM - see inline for some possible minor improvements.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4621

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

4642

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.

4659

Another transform that can be updated...

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4621

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_APIntAllowUndef(AP1)))) &&
+       match(Op1, m_OneUse(m_LShr(m_Value(B), m_APIntAllowUndef(AP2))))) ||
+      (match(Op0, m_OneUse(m_AShr(m_Value(A), m_APIntAllowUndef(AP1)))) &&
+       match(Op1, m_OneUse(m_AShr(m_Value(B), m_APIntAllowUndef(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));
+      return new ICmpInst(NewPred, Xor, ConstantInt::get(A->getType(), 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:    [[T0:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[B:%.*]], <i32 31, i32 31>
-; CHECK-NEXT:    [[T2:%.*]] = icmp eq <2 x i32> [[T0]], [[T1]]
+; 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
 }
 
 define <2 x i32> @test3vec_undef1(<2 x i32> %a, <2 x i32> %b) nounwind readnone {
 ; CHECK-LABEL: @test3vec_undef1(
-; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 24, i32 undef>
-; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[B:%.*]], <i32 24, i32 24>
-; CHECK-NEXT:    [[T2:%.*]] = icmp eq <2 x i32> [[T0]], [[T1]]
+; CHECK-NEXT:    [[T2_UNSHIFTED:%.*]] = xor <2 x i32> [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[T2:%.*]] = icmp ult <2 x i32> [[T2_UNSHIFTED]], <i32 16777216, i32 16777216>
 ; 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 24, i32 undef>
   %t1 = lshr <2 x i32> %b, <i32 24, i32 24>
   %t2 = icmp eq <2 x i32> %t0, %t1
   %t3 = zext <2 x i1> %t2 to <2 x i32>
   ret <2 x i32> %t3
 }
 
 define <2 x i32> @test3vec_undef2(<2 x i32> %a, <2 x i32> %b) nounwind readnone {
 ; CHECK-LABEL: @test3vec_undef2(
-; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[A:%.*]], <i32 undef, i32 17>
-; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[B:%.*]], <i32 undef, i32 17>
-; CHECK-NEXT:    [[T2:%.*]] = icmp eq <2 x i32> [[T0]], [[T1]]
+; CHECK-NEXT:    [[T2_UNSHIFTED:%.*]] = xor <2 x i32> [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[T2:%.*]] = icmp ult <2 x i32> [[T2_UNSHIFTED]], <i32 131072, i32 131072>
 ; 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 undef, i32 17>
   %t1 = lshr <2 x i32> %b, <i32 undef, i32 17>
   %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> @test3vec_diff(<2 x i32> %a, <2 x i32> %b) nounwind readnone {
 ; CHECK-LABEL: @test3vec_diff(
 ; CHECK-NEXT:    [[T0:%.*]] = lshr <2 x i32> [[A:%.*]], <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]]