[InstCombine] Extend #125676 to handle variable power of 2 (#125855)

Alive2: https://alive2.llvm.org/ce/z/dJehZ8

Author: dtcxzyw

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -4200,14 +4200,14 @@ Value *InstCombinerImpl::foldXorOfICmps(ICmpInst *LHS, ICmpInst *RHS,
 
     // Fold (icmp eq/ne (X & Pow2), 0) ^ (icmp eq/ne (Y & Pow2), 0) into
     // (icmp eq/ne ((X ^ Y) & Pow2), 0)
-    Value *X, *Y;
-    const APInt *Mask;
+    Value *X, *Y, *Pow2;
     if (ICmpInst::isEquality(PredL) && ICmpInst::isEquality(PredR) &&
         LC->isZero() && RC->isZero() && LHS->hasOneUse() && RHS->hasOneUse() &&
-        match(LHS0, m_And(m_Value(X), m_Power2(Mask))) &&
-        match(RHS0, m_And(m_Value(Y), m_SpecificInt(*Mask)))) {
+        match(LHS0, m_And(m_Value(X), m_Value(Pow2))) &&
+        match(RHS0, m_And(m_Value(Y), m_Specific(Pow2))) &&
+        isKnownToBeAPowerOfTwo(Pow2, /*OrZero=*/true, /*Depth=*/0, &I)) {
       Value *Xor = Builder.CreateXor(X, Y);
-      Value *And = Builder.CreateAnd(Xor, *Mask);
+      Value *And = Builder.CreateAnd(Xor, Pow2);
       return Builder.CreateICmp(PredL == PredR ? ICmpInst::ICMP_NE
                                                : ICmpInst::ICMP_EQ,
                                 And, ConstantInt::getNullValue(Xor->getType()));

llvm/test/Transforms/InstCombine/xor-icmps.ll

@@ -335,6 +335,42 @@ define i1 @test_xor_of_bittest_ne_ne(i8 %x, i8 %y) {
   ret i1 %xor
 }
 
+define i1 @test_xor_of_bittest_ne_ne_var_pow2(i8 %x, i8 %y, i8 %shamt) {
+; CHECK-LABEL: @test_xor_of_bittest_ne_ne_var_pow2(
+; CHECK-NEXT:    [[POW2:%.*]] = shl nuw i8 1, [[SHAMT:%.*]]
+; CHECK-NEXT:    [[Y:%.*]] = xor i8 [[X:%.*]], [[Y1:%.*]]
+; CHECK-NEXT:    [[MASK2:%.*]] = and i8 [[Y]], [[POW2]]
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp ne i8 [[MASK2]], 0
+; CHECK-NEXT:    ret i1 [[CMP2]]
+;
+  %pow2 = shl nuw i8 1, %shamt
+  %mask1 = and i8 %x, %pow2
+  %cmp1 = icmp ne i8 %mask1, 0
+  %mask2 = and i8 %y, %pow2
+  %cmp2 = icmp ne i8 %mask2, 0
+  %xor = xor i1 %cmp1, %cmp2
+  ret i1 %xor
+}
+
+define i1 @test_xor_of_bittest_ne_ne_var_pow2_or_zero(i8 %x, i8 %y, i8 %z) {
+; CHECK-LABEL: @test_xor_of_bittest_ne_ne_var_pow2_or_zero(
+; CHECK-NEXT:    [[NZ:%.*]] = sub i8 0, [[Z:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[TMP1]], [[NZ]]
+; CHECK-NEXT:    [[MASK2:%.*]] = and i8 [[TMP2]], [[Z]]
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp ne i8 [[MASK2]], 0
+; CHECK-NEXT:    ret i1 [[CMP2]]
+;
+  %nz = sub i8 0, %z
+  %pow2 = and i8 %z, %nz
+  %mask1 = and i8 %x, %pow2
+  %cmp1 = icmp ne i8 %mask1, 0
+  %mask2 = and i8 %y, %pow2
+  %cmp2 = icmp ne i8 %mask2, 0
+  %xor = xor i1 %cmp1, %cmp2
+  ret i1 %xor
+}
+
 define i1 @test_xor_of_bittest_eq_eq(i8 %x, i8 %y) {
 ; CHECK-LABEL: @test_xor_of_bittest_eq_eq(
 ; CHECK-NEXT:    [[Y:%.*]] = xor i8 [[X:%.*]], [[Y1:%.*]]