[ValueTracking] Use m_LogicalAnd/Or to look into conditions

[aqjune]

Summary

This patch updates isImpliedCondition/isKnownNonZero to look into select form of
and/or as well.

See llvm.org/pr48353 and D93065 for more context

[nikic]

LGTM

llvm/lib/Analysis/ValueTracking.cpp

for (auto *CmpU : U->users()) {
       while (!WorkList.empty()) {
         auto *Curr = WorkList.pop_back_val();
 
         // If a user is an AND, add all its users to the work list. We only
         // propagate "pred != null" condition through AND because it is only
         // correct to assume that all conditions of AND are met in true branch.
         // TODO: Support similar logic of OR and EQ predicate?
         if (NonNullIfTrue)
-          if (auto *BO = dyn_cast<BinaryOperator>(Curr))
-            if (BO->getOpcode() == Instruction::And) {
-              for (auto *BOU : BO->users())
-                if (Visited.insert(BOU).second)
-                  WorkList.push_back(BOU);
-              continue;
-            }
+          if (match(Curr, m_LogicalAnd(m_Value(), m_Value()))) {
+            for (auto *CurrU : Curr->users())
+              if (Visited.insert(CurrU).second)
+                WorkList.push_back(CurrU);
 
         if (const BranchInst *BI = dyn_cast<BranchInst>(Curr)) {
           assert(BI->isConditional() && "uses a comparison!");
 
           BasicBlock *NonNullSuccessor =
               BI->getSuccessor(NonNullIfTrue ? 0 : 1);
           BasicBlockEdge Edge(BI->getParent(), NonNullSuccessor);
           if (Edge.isSingleEdge() && DT->dominates(Edge, CtxI->getParent()))
static Optional<bool> isImpliedCondICmps(const ICmpInst *LHS,
 
   if (APred == BPred)
     return isImpliedCondOperands(APred, ALHS, ARHS, BLHS, BRHS, DL, Depth);
   return None;
 }
 
 /// Return true if LHS implies RHS is true.  Return false if LHS implies RHS is
 /// false.  Otherwise, return None if we can't infer anything.  We expect the
-/// RHS to be an icmp and the LHS to be an 'and' or an 'or' instruction.
+/// RHS to be an icmp and the LHS to be an 'and', 'or', or a 'select' instruction.
 static Optional<bool>
-isImpliedCondAndOr(const BinaryOperator *LHS, CmpInst::Predicate RHSPred,
+isImpliedCondAndOr(const Instruction *LHS, CmpInst::Predicate RHSPred,
                    const Value *RHSOp0, const Value *RHSOp1,
-
                    const DataLayout &DL, bool LHSIsTrue, unsigned Depth) {
-  // The LHS must be an 'or' or an 'and' instruction.
+  // The LHS must be an 'or', 'and', or a 'select' instruction.
   assert((LHS->getOpcode() == Instruction::And ||
-          LHS->getOpcode() == Instruction::Or) &&
-         "Expected LHS to be 'and' or 'or'.");
+          LHS->getOpcode() == Instruction::Or ||
+          LHS->getOpcode() == Instruction::Select) &&
+         "Expected LHS to be 'and', 'or', or 'select'.");
 
   assert(Depth <= MaxAnalysisRecursionDepth && "Hit recursion limit");
 
   // If the result of an 'or' is false, then we know both legs of the 'or' are
   // false.  Similarly, if the result of an 'and' is true, then we know both
   // legs of the 'and' are true.
-  Value *ALHS, *ARHS;
-  if ((!LHSIsTrue && match(LHS, m_Or(m_Value(ALHS), m_Value(ARHS)))) ||
-      (LHSIsTrue && match(LHS, m_And(m_Value(ALHS), m_Value(ARHS))))) {
+  const Value *ALHS, *ARHS;
+  if ((!LHSIsTrue && match(LHS, m_LogicalOr(m_Value(ALHS), m_Value(ARHS)))) ||
+      (LHSIsTrue && match(LHS, m_LogicalAnd(m_Value(ALHS), m_Value(ARHS))))) {
     // FIXME: Make this non-recursion.
     if (Optional<bool> Implication = isImpliedCondition(
             ALHS, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue, Depth + 1))
       return Implication;
     if (Optional<bool> Implication = isImpliedCondition(
             ARHS, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue, Depth + 1))
       return Implication;
     return None;
llvm::isImpliedCondition(const Value *LHS, CmpInst::Predicate RHSPred,
   assert(OpTy->isIntegerTy(1) && "implied by above");
 
   // Both LHS and RHS are icmps.
   const ICmpInst *LHSCmp = dyn_cast<ICmpInst>(LHS);
   if (LHSCmp)
     return isImpliedCondICmps(LHSCmp, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue,
                               Depth);
 
-  /// The LHS should be an 'or' or an 'and' instruction.  We expect the RHS to
-  /// be / an icmp. FIXME: Add support for and/or on the RHS.
-  const BinaryOperator *LHSBO = dyn_cast<BinaryOperator>(LHS);
-  if (LHSBO) {
-    if ((LHSBO->getOpcode() == Instruction::And ||
-         LHSBO->getOpcode() == Instruction::Or))
-      return isImpliedCondAndOr(LHSBO, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue,
+  /// The LHS should be an 'or', 'and', or a 'select' instruction.  We expect
+  /// the RHS to be an icmp.
+  /// FIXME: Add support for and/or/select on the RHS.
+  if (const Instruction *LHSI = dyn_cast<Instruction>(LHS)) {
+    if ((LHSI->getOpcode() == Instruction::And ||
+         LHSI->getOpcode() == Instruction::Or ||
+         LHSI->getOpcode() == Instruction::Select))
+      return isImpliedCondAndOr(LHSI, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue,
                                 Depth);
   }
   return None;
 }
 
 Optional<bool> llvm::isImpliedCondition(const Value *LHS, const Value *RHS,
                                         const DataLayout &DL, bool LHSIsTrue,
                                         unsigned Depth) {

llvm/unittests/Analysis/ValueTrackingTest.cpp

TEST_F(ValueTrackingTest, KnownNonZeroFromDomCond) {
   )");
   AssumptionCache AC(*F);
   DominatorTree DT(*F);
   DataLayout DL = M->getDataLayout();
   EXPECT_EQ(isKnownNonZero(A, DL, 0, &AC, CxtI, &DT), true);
   EXPECT_EQ(isKnownNonZero(A, DL, 0, &AC, CxtI2, &DT), false);
 }
 
+TEST_F(ValueTrackingTest, KnownNonZeroFromDomCond2) {
+  parseAssembly(R"(
+    declare i8* @f_i8()
+    define void @test(i1 %c) {
+      %A = call i8* @f_i8()
+      %B = call i8* @f_i8()
+      %c1 = icmp ne i8* %A, null
+      %cond = select i1 %c, i1 %c1, i1 false
+      br i1 %cond, label %T, label %Q
+    T:
+      %CxtI = add i32 0, 0
+      ret void
+    Q:
+      %CxtI2 = add i32 0, 0
+      ret void
+    }
+  )");
+  AssumptionCache AC(*F);
+  DominatorTree DT(*F);
+  DataLayout DL = M->getDataLayout();
+  EXPECT_EQ(isKnownNonZero(A, DL, 0, &AC, CxtI, &DT), true);
+  EXPECT_EQ(isKnownNonZero(A, DL, 0, &AC, CxtI2, &DT), false);
+}
+
 TEST_F(ValueTrackingTest, IsImpliedConditionAnd) {
   parseAssembly(R"(
     define void @test(i32 %x, i32 %y) {
       %c1 = icmp ult i32 %x, 10
       %c2 = icmp ult i32 %y, 15
       %A = and i1 %c1, %c2
       ; x < 10 /\ y < 15
       %A2 = icmp ult i32 %x, 20
       %A3 = icmp uge i32 %y, 20
       %A4 = icmp ult i32 %x, 5
       ret void
     }
   )");
   DataLayout DL = M->getDataLayout();
   EXPECT_EQ(isImpliedCondition(A, A2, DL), true);
   EXPECT_EQ(isImpliedCondition(A, A3, DL), false);
   EXPECT_EQ(isImpliedCondition(A, A4, DL), None);
 }
 
+TEST_F(ValueTrackingTest, IsImpliedConditionAnd2) {
+  parseAssembly(R"(
+    define void @test(i32 %x, i32 %y) {
+      %c1 = icmp ult i32 %x, 10
+      %c2 = icmp ult i32 %y, 15
+      %A = select i1 %c1, i1 %c2, i1 false
+      ; x < 10 /\ y < 15
+      %A2 = icmp ult i32 %x, 20
+      %A3 = icmp uge i32 %y, 20
+      %A4 = icmp ult i32 %x, 5
+      ret void
+    }
+  )");
+  DataLayout DL = M->getDataLayout();
+  EXPECT_EQ(isImpliedCondition(A, A2, DL), true);
+  EXPECT_EQ(isImpliedCondition(A, A3, DL), false);
+  EXPECT_EQ(isImpliedCondition(A, A4, DL), None);
+}
+
 TEST_F(ValueTrackingTest, IsImpliedConditionOr) {
   parseAssembly(R"(
     define void @test(i32 %x, i32 %y) {
       %c1 = icmp ult i32 %x, 10
       %c2 = icmp ult i32 %y, 15
       %A = or i1 %c1, %c2 ; negated
       ; x >= 10 /\ y >= 15
       %A2 = icmp ult i32 %x, 5
       %A3 = icmp uge i32 %y, 10
       %A4 = icmp ult i32 %x, 15
       ret void
     }
   )");
   DataLayout DL = M->getDataLayout();
   EXPECT_EQ(isImpliedCondition(A, A2, DL, false), false);
   EXPECT_EQ(isImpliedCondition(A, A3, DL, false), true);
   EXPECT_EQ(isImpliedCondition(A, A4, DL, false), None);
 }
 
+TEST_F(ValueTrackingTest, IsImpliedConditionOr2) {
+  parseAssembly(R"(
+    define void @test(i32 %x, i32 %y) {
+      %c1 = icmp ult i32 %x, 10
+      %c2 = icmp ult i32 %y, 15
+      %A = select i1 %c1, i1 true, i1 %c2 ; negated
+      ; x >= 10 /\ y >= 15
+      %A2 = icmp ult i32 %x, 5
+      %A3 = icmp uge i32 %y, 10
+      %A4 = icmp ult i32 %x, 15
+      ret void
+    }
+  )");
+  DataLayout DL = M->getDataLayout();
+  EXPECT_EQ(isImpliedCondition(A, A2, DL, false), false);
+  EXPECT_EQ(isImpliedCondition(A, A3, DL, false), true);
+  EXPECT_EQ(isImpliedCondition(A, A4, DL, false), None);
+}
+
 TEST_F(ComputeKnownBitsTest, KnownNonZeroShift) {
   // %q is known nonzero without known bits.
   // Because %q is nonzero, %A[0] is known to be zero.
   parseAssembly(
       "define i8 @test(i8 %p, i8* %pq) {\n"
       "  %q = load i8, i8* %pq, !range !0\n"
       "  %A = shl i8 %p, %q\n"
       "  ret i8 %A\n"