[InstCombine] Handle (iszero(A & K1) | iszero(A & K2)) -> (A & (K1 | K2)) != (K1 | K2) when the one of the Ands is commuted relative to the other

[craig.topper]

Summary

Currently we expect A to be on the same side in both Ands but nothing guarantees that.

While there also switch to using matchers for some of the code.

[spatel]

LGTM. See inline for a possible follow-up and nit.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1599–1603 ↗

(On Diff #102632)

Could turn this into a matcher too; using m_Zero() would work with vectors.

test/Transforms/InstCombine/onehot_merge.ll

49 ↗

(On Diff #102632)

Remove the label to minimize?

[spatel]

test/Transforms/InstCombine/onehot_merge.ll

37 ↗

(On Diff #102632)

Remove the TODO.

llvm/trunk/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Value *InstCombiner::foldOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
                                    Instruction *CxtI) {
   ICmpInst::Predicate PredL = LHS->getPredicate(), PredR = RHS->getPredicate();
 
   // Fold (iszero(A & K1) | iszero(A & K2)) ->  (A & (K1 | K2)) != (K1 | K2)
   // if K1 and K2 are a one-bit mask.
   ConstantInt *LHSC = dyn_cast<ConstantInt>(LHS->getOperand(1));
   ConstantInt *RHSC = dyn_cast<ConstantInt>(RHS->getOperand(1));
 
+  // TODO support vector splats
   if (LHS->getPredicate() == ICmpInst::ICMP_EQ && LHSC && LHSC->isZero() &&
       RHS->getPredicate() == ICmpInst::ICMP_EQ && RHSC && RHSC->isZero()) {
 
-    BinaryOperator *LAnd = dyn_cast<BinaryOperator>(LHS->getOperand(0));
-    BinaryOperator *RAnd = dyn_cast<BinaryOperator>(RHS->getOperand(0));
-    if (LAnd && RAnd && LAnd->hasOneUse() && RHS->hasOneUse() &&
-        LAnd->getOpcode() == Instruction::And &&
-        RAnd->getOpcode() == Instruction::And) {
-
-      Value *Mask = nullptr;
-      Value *Masked = nullptr;
-      if (LAnd->getOperand(0) == RAnd->getOperand(0) &&
-          isKnownToBeAPowerOfTwo(LAnd->getOperand(1), false, 0, CxtI) &&
-          isKnownToBeAPowerOfTwo(RAnd->getOperand(1), false, 0, CxtI)) {
-        Mask = Builder->CreateOr(LAnd->getOperand(1), RAnd->getOperand(1));
-        Masked = Builder->CreateAnd(LAnd->getOperand(0), Mask);
-      } else if (LAnd->getOperand(1) == RAnd->getOperand(1) &&
-                 isKnownToBeAPowerOfTwo(LAnd->getOperand(0), false, 0, CxtI) &&
-                 isKnownToBeAPowerOfTwo(RAnd->getOperand(0), false, 0, CxtI)) {
-        Mask = Builder->CreateOr(LAnd->getOperand(0), RAnd->getOperand(0));
-        Masked = Builder->CreateAnd(LAnd->getOperand(1), Mask);
-      }
+    Value *A, *B, *C, *D;
+    if (match(LHS->getOperand(0), m_And(m_Value(A), m_Value(B))) &&
+        match(RHS->getOperand(0), m_And(m_Value(C), m_Value(D)))) {
+      if (A == D || B == D)
+        std::swap(C, D);
+      if (B == C)
+        std::swap(A, B);
 
-      if (Masked)
+      if (A == C &&
+          isKnownToBeAPowerOfTwo(B, false, 0, CxtI) &&
+          isKnownToBeAPowerOfTwo(D, false, 0, CxtI)) {
+        Value *Mask = Builder->CreateOr(B, D);
+        Value *Masked = Builder->CreateAnd(A, Mask);
         return Builder->CreateICmp(ICmpInst::ICMP_NE, Masked, Mask);
       }
     }
+  }
 
   // Fold (icmp ult/ule (A + C1), C3) | (icmp ult/ule (A + C2), C3)
   //                   -->  (icmp ult/ule ((A & ~(C1 ^ C2)) + max(C1, C2)), C3)
   // The original condition actually refers to the following two ranges:
   // [MAX_UINT-C1+1, MAX_UINT-C1+1+C3] and [MAX_UINT-C2+1, MAX_UINT-C2+1+C3]
   // We can fold these two ranges if:
   // 1) C1 and C2 is unsigned greater than C3.
   // 2) The two ranges are separated.

llvm/trunk/test/Transforms/InstCombine/onehot_merge.ll

bb:
   %tmp2 = icmp eq i32 %tmp1, 0
   %tmp5 = and i32 %tmp4, %k
   %tmp6 = icmp eq i32 %tmp5, 0
   %or = or i1 %tmp2, %tmp6
   ret i1 %or
 }
 
 ; Same as above but with operands commuted one of the ands, but not the other.
-; TODO handle this form correctly
 define i1 @foo1_and_commuted(i32 %k, i32 %c1, i32 %c2) {
 ; CHECK-LABEL: @foo1_and_commuted(
-; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[K2:%.*]] = mul i32 [[K:%.*]], [[K]]
 ; CHECK-NEXT:    [[TMP:%.*]] = shl i32 1, [[C1:%.*]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = lshr i32 -2147483648, [[C2:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[K2]], [[TMP]]
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], 0
-; CHECK-NEXT:    [[TMP5:%.*]] = and i32 [[TMP4]], [[K2]]
-; CHECK-NEXT:    [[TMP6:%.*]] = icmp eq i32 [[TMP5]], 0
-; CHECK-NEXT:    [[OR:%.*]] = or i1 [[TMP2]], [[TMP6]]
-; CHECK-NEXT:    ret i1 [[OR]]
+; CHECK-NEXT:    [[TMP0:%.*]] = or i32 [[TMP]], [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[K2]], [[TMP0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[TMP1]], [[TMP0]]
+; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
-bb:
   %k2 = mul i32 %k, %k ; to trick the complexity sorting
   %tmp = shl i32 1, %c1
   %tmp4 = lshr i32 -2147483648, %c2
   %tmp1 = and i32 %k2, %tmp
   %tmp2 = icmp eq i32 %tmp1, 0
   %tmp5 = and i32 %tmp4, %k2
   %tmp6 = icmp eq i32 %tmp5, 0
   %or = or i1 %tmp2, %tmp6
   ret i1 %or
 }