[InstCombine] Fold xor of icmps of power of 2

[StephenFan]

Summary

https://alive2.llvm.org/ce/z/HKt_Rg

The motivation to do this is D146149 caused a code size regression in
our ci https://github.com/dtcxzyw/llvm-ci/issues/23.

[StephenFan]

Rename some variables.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

if (!Known.isNonNegative())
     return nullptr;
 
   if (Inverted)
     NewPred = ICmpInst::getInversePredicate(NewPred);
 
   return Builder.CreateICmp(NewPred, Input, RangeEnd);
 }
 
+// Fold (iszero(A & K1) ^ iszero(B & K2))
+//          => iszero(((lshr A, log2(K1)) ^ (lshr B, log2(K2))) & 1)
+Value *InstCombinerImpl::foldXorOfICmpsOfAndWithPow2(ICmpInst *LHS,
+                                                        ICmpInst *RHS,
+                                                        Instruction *CtxI) {
+  if (LHS->getPredicate() != CmpInst::ICMP_NE ||
+      RHS->getPredicate() != CmpInst::ICMP_NE)
+    return nullptr;
+
+  if (!match(LHS->getOperand(1), m_Zero()) ||
+      !match(RHS->getOperand(1), m_Zero()))
+    return nullptr;
+
+  if (LHS->getOperand(0)->getType() != RHS->getOperand(0)->getType())
+    return nullptr;
+
+  Value *L1, *L2, *R1, *R2;
+  if (match(LHS->getOperand(0), m_And(m_Value(L1), m_Value(L2))) &&
+      match(RHS->getOperand(0), m_And(m_Value(R1), m_Value(R2)))) {
+    if (L1 == R1)
+      return nullptr;
+    const APInt *L2Int, *R2Int;
+    if (match(L2, m_Power2(L2Int)) && match(R2, m_Power2(R2Int))) {
+      if (L2Int->eq(*R2Int)) {
+        Value *Xor = Builder.CreateXor(L1, R1);
+        Value *AndMask = Builder.CreateAnd(
+            Xor, ConstantInt::get(Xor->getType(), L2Int->getZExtValue()));
+        return Builder.CreateICmp(
+            CmpInst::ICMP_NE, AndMask,
+            ConstantInt::getNullValue(AndMask->getType()));
+      } else {
+        Value *L1Lshr = Builder.CreateLShr(
+            L1, ConstantInt::get(L1->getType(), L2Int->logBase2()));
+        Value *R1Lshr = Builder.CreateLShr(
+            R1, ConstantInt::get(R1->getType(), R2Int->logBase2()));
+        Value *Xor = Builder.CreateXor(L1Lshr, R1Lshr);
+        Value *And =
+            Builder.CreateAnd(Xor, ConstantInt::get(Xor->getType(), 1));
+        return Builder.CreateICmp(CmpInst::ICMP_NE, And,
+                                  ConstantInt::getNullValue(And->getType()));
+      }
+    }
+  }
+  return nullptr;
+}
+
 // Fold (iszero(A & K1) | iszero(A & K2)) -> (A & (K1 | K2)) != (K1 | K2)
 // Fold (!iszero(A & K1) & !iszero(A & K2)) -> (A & (K1 | K2)) == (K1 | K2)
 Value *InstCombinerImpl::foldAndOrOfICmpsOfAndWithPow2(ICmpInst *LHS,
                                                        ICmpInst *RHS,
                                                        Instruction *CxtI,
                                                        bool IsAnd,
                                                        bool IsLogical) {
   CmpInst::Predicate Pred = IsAnd ? CmpInst::ICMP_NE : CmpInst::ICMP_EQ;
Value *InstCombinerImpl::foldXorOfICmps(ICmpInst *LHS, ICmpInst *RHS,
                                         BinaryOperator &I) {
   assert(I.getOpcode() == Instruction::Xor && I.getOperand(0) == LHS &&
          I.getOperand(1) == RHS && "Should be 'xor' with these operands");
 
   ICmpInst::Predicate PredL = LHS->getPredicate(), PredR = RHS->getPredicate();
   Value *LHS0 = LHS->getOperand(0), *LHS1 = LHS->getOperand(1);
   Value *RHS0 = RHS->getOperand(0), *RHS1 = RHS->getOperand(1);
 
+  if (Value *V = foldXorOfICmpsOfAndWithPow2(LHS, RHS, &I))
+    return V;
+
   if (predicatesFoldable(PredL, PredR)) {
     if (LHS0 == RHS1 && LHS1 == RHS0) {
       std::swap(LHS0, LHS1);
       PredL = ICmpInst::getSwappedPredicate(PredL);
     }
     if (LHS0 == RHS0 && LHS1 == RHS1) {
       // (icmp1 A, B) ^ (icmp2 A, B) --> (icmp3 A, B)
       unsigned Code = getICmpCode(PredL) ^ getICmpCode(PredR);

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

Value *foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS, Instruction &I,
                           bool IsAnd, bool IsLogical = false);
   Value *foldXorOfICmps(ICmpInst *LHS, ICmpInst *RHS, BinaryOperator &Xor);
 
   Value *foldEqOfParts(ICmpInst *Cmp0, ICmpInst *Cmp1, bool IsAnd);
 
   Value *foldAndOrOfICmpsUsingRanges(ICmpInst *ICmp1, ICmpInst *ICmp2,
                                      bool IsAnd);
 
+  Value *foldXorOfICmpsOfAndWithPow2(ICmpInst *LHS, ICmpInst *RHS,
+                                        Instruction *CtxI);
+
   /// Optimize (fcmp)&(fcmp) or (fcmp)|(fcmp).
   /// NOTE: Unlike most of instcombine, this returns a Value which should
   /// already be inserted into the function.
   Value *foldLogicOfFCmps(FCmpInst *LHS, FCmpInst *RHS, bool IsAnd,
                           bool IsLogicalSelect = false);
 
   Instruction *foldLogicOfIsFPClass(BinaryOperator &Operator, Value *LHS,
                                     Value *RHS);

llvm/test/Transforms/InstCombine/bit-checks.ll

;
   %tobool3 = icmp sge i8 %trunc2, 0
   %or.cond = select i1 %tobool, i1 true, i1 %tobool3
   %retval.0 = select i1 %or.cond, i32 2, i32 1
   ret i32 %retval.0
 }
 
 define i1 @xor_of_icmps_of_pow2(i32 %a, i32 %b) {
 ; CHECK-LABEL: @xor_of_icmps_of_pow2(
-; CHECK-NEXT:    [[AND_A:%.*]] = and i32 [[A:%.*]], 2
-; CHECK-NEXT:    [[AND_A_ZERO:%.*]] = icmp ne i32 [[AND_A]], 0
-; CHECK-NEXT:    [[AND_B:%.*]] = and i32 [[B:%.*]], 8
-; CHECK-NEXT:    [[AND_B_ZERO:%.*]] = icmp ne i32 [[AND_B]], 0
-; CHECK-NEXT:    [[RET:%.*]] = xor i1 [[AND_A_ZERO]], [[AND_B_ZERO]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[A:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i32 [[B:%.*]], 3
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i32 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = and i32 [[TMP3]], 1
+; CHECK-NEXT:    [[RET:%.*]] = icmp ne i32 [[TMP4]], 0
 ; CHECK-NEXT:    ret i1 [[RET]]
 ;
   %and_a = and i32 %a, 2
   %and_a_zero = icmp ne i32 %and_a, 0
   %and_b = and i32 %b, 8
   %and_b_zero = icmp ne i32 %and_b, 0
   %ret = xor i1 %and_a_zero, %and_b_zero
   ret i1 %ret
 }
 
 define i1 @xor_of_icmps_of_pow2_1(i32 %a, i32 %b) {
 ; CHECK-LABEL: @xor_of_icmps_of_pow2_1(
-; CHECK-NEXT:    [[AND_A:%.*]] = and i32 [[A:%.*]], 1
-; CHECK-NEXT:    [[AND_A_ZERO:%.*]] = icmp ne i32 [[AND_A]], 0
-; CHECK-NEXT:    [[AND_B:%.*]] = and i32 [[B:%.*]], 8
-; CHECK-NEXT:    [[AND_B_ZERO:%.*]] = icmp ne i32 [[AND_B]], 0
-; CHECK-NEXT:    [[RET:%.*]] = xor i1 [[AND_A_ZERO]], [[AND_B_ZERO]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[B:%.*]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i32 [[TMP1]], [[A:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], 1
+; CHECK-NEXT:    [[RET:%.*]] = icmp ne i32 [[TMP3]], 0
 ; CHECK-NEXT:    ret i1 [[RET]]
 ;
   %and_a = and i32 %a, 1
   %and_a_zero = icmp ne i32 %and_a, 0
   %and_b = and i32 %b, 8
   %and_b_zero = icmp ne i32 %and_b, 0
   %ret = xor i1 %and_a_zero, %and_b_zero
   ret i1 %ret
;
   %and_b = and i16 %b, 8
   %and_b_zero = icmp ne i16 %and_b, 0
   %ret = xor i1 %and_a_zero, %and_b_zero
   ret i1 %ret
 }
 
 define i1 @xor_of_icmps_of_pow2_same_constant(i32 %a, i32 %b) {
 ; CHECK-LABEL: @xor_of_icmps_of_pow2_same_constant(
-; CHECK-NEXT:    [[AND_A:%.*]] = and i32 [[A:%.*]], 8
-; CHECK-NEXT:    [[AND_A_ZERO:%.*]] = icmp ne i32 [[AND_A]], 0
-; CHECK-NEXT:    [[AND_B:%.*]] = and i32 [[B:%.*]], 8
-; CHECK-NEXT:    [[AND_B_ZERO:%.*]] = icmp ne i32 [[AND_B]], 0
-; CHECK-NEXT:    [[RET:%.*]] = xor i1 [[AND_A_ZERO]], [[AND_B_ZERO]]
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 8
+; CHECK-NEXT:    [[RET:%.*]] = icmp ne i32 [[TMP2]], 0
 ; CHECK-NEXT:    ret i1 [[RET]]
 ;
   %and_a = and i32 %a, 8
   %and_a_zero = icmp ne i32 %and_a, 0
   %and_b = and i32 %b, 8
   %and_b_zero = icmp ne i32 %and_b, 0
   %ret = xor i1 %and_a_zero, %and_b_zero
   ret i1 %ret
 }