[ValueTracking] Use m_LogicalAnd/Or to look into conditions

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

See llvm.org/pr48353 and D93065 for more context

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D93845

lib/Analysis/ValueTracking.cpp

@@ -2131,11 +2131,10 @@ static bool isKnownNonNullFromDominatingCondition(const Value *V,
         // 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 (match(Curr, m_LogicalAnd(m_Value(), m_Value()))) {
-            if (BO->getOpcode() == Instruction::And) {
+            for (auto *CurrU : Curr->users())
-              for (auto *BOU : BO->users())
+              if (Visited.insert(CurrU).second)
-                if (Visited.insert(BOU).second)
+                WorkList.push_back(CurrU);
-                  WorkList.push_back(BOU);
             continue;
           }
 
@@ -6156,25 +6155,25 @@ static Optional<bool> isImpliedCondICmps(const ICmpInst *LHS,
 
 /// 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) &&
+          LHS->getOpcode() == Instruction::Or ||
-         "Expected LHS to be 'and' or '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;
+  const Value *ALHS, *ARHS;
-  if ((!LHSIsTrue && match(LHS, m_Or(m_Value(ALHS), m_Value(ARHS)))) ||
+  if ((!LHSIsTrue && match(LHS, m_LogicalOr(m_Value(ALHS), m_Value(ARHS)))) ||
-      (LHSIsTrue && match(LHS, m_And(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))
@@ -6215,13 +6214,14 @@ llvm::isImpliedCondition(const Value *LHS, CmpInst::Predicate RHSPred,
     return isImpliedCondICmps(LHSCmp, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue,
                               Depth);
 
-  /// The LHS should be an 'or' or an 'and' instruction.  We expect the RHS to
+  /// The LHS should be an 'or', 'and', or a 'select' instruction.  We expect
-  /// be / an icmp. FIXME: Add support for and/or on the RHS.
+  /// the RHS to be an icmp.
-  const BinaryOperator *LHSBO = dyn_cast<BinaryOperator>(LHS);
+  /// FIXME: Add support for and/or/select on the RHS.
-  if (LHSBO) {
+  if (const Instruction *LHSI = dyn_cast<Instruction>(LHS)) {
-    if ((LHSBO->getOpcode() == Instruction::And ||
+    if ((LHSI->getOpcode() == Instruction::And ||
-         LHSBO->getOpcode() == Instruction::Or))
+         LHSI->getOpcode() == Instruction::Or ||
-      return isImpliedCondAndOr(LHSBO, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue,
+         LHSI->getOpcode() == Instruction::Select))
+      return isImpliedCondAndOr(LHSI, RHSPred, RHSOp0, RHSOp1, DL, LHSIsTrue,
                                 Depth);
   }
   return None;

unittests/Analysis/ValueTrackingTest.cpp

@@ -1033,6 +1033,30 @@ TEST_F(ValueTrackingTest, KnownNonZeroFromDomCond) {
   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) {
@@ -1052,6 +1076,25 @@ TEST_F(ValueTrackingTest, IsImpliedConditionAnd) {
   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) {
@@ -1071,6 +1114,25 @@ TEST_F(ValueTrackingTest, IsImpliedConditionOr) {
   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.