[ValueTracking] Handle recursive select/PHI in IsKnownNonZero

Author: goldsteinn

llvm/lib/Analysis/ValueTracking.cpp

@@ -3033,18 +3033,36 @@ static bool isKnownNonZeroFromOperator(const Operator *I,
       return true;
 
     // Check if all incoming values are non-zero using recursion.
-    SimplifyQuery RecQ = Q.getWithoutCondContext();
-    unsigned NewDepth = std::max(Depth, MaxAnalysisRecursionDepth - 1);
     return llvm::all_of(PN->operands(), [&](const Use &U) {
-      if (U.get() == PN)
+      Value *IncValue = U.get();
+      if (IncValue == PN)
         return true;
-      RecQ.CxtI = PN->getIncomingBlock(U)->getTerminator();
+
+      Instruction *CxtI = PN->getIncomingBlock(U)->getTerminator();
+      unsigned NewDepth = std::max(Depth, MaxAnalysisRecursionDepth - 1);
+      if (auto *SI = dyn_cast<SelectInst>(IncValue)) {
+        if (SI->getTrueValue() == PN || SI->getFalseValue() == PN) {
+          IncValue = SI->getTrueValue() == PN ? SI->getFalseValue()
+                                              : SI->getTrueValue();
+          NewDepth = Depth;
+        }
+      } else if (auto *IncPhi = dyn_cast<PHINode>(IncValue);
+                 IncPhi && IncPhi->getNumIncomingValues() == 2) {
+        for (int Idx = 0; Idx < 2; ++Idx) {
+          if (IncPhi->getIncomingValue(Idx) == PN) {
+            IncValue = IncPhi->getIncomingValue(1 - Idx);
+            CxtI = IncPhi->getIncomingBlock(1 - Idx)->getTerminator();
+            break;
+          }
+        }
+      }
+      SimplifyQuery RecQ = Q.getWithoutCondContext().getWithInstruction(CxtI);
       // Check if the branch on the phi excludes zero.
       ICmpInst::Predicate Pred;
       Value *X;
       BasicBlock *TrueSucc, *FalseSucc;
       if (match(RecQ.CxtI,
-                m_Br(m_c_ICmp(Pred, m_Specific(U.get()), m_Value(X)),
+                m_Br(m_c_ICmp(Pred, m_Specific(IncValue), m_Value(X)),
                      m_BasicBlock(TrueSucc), m_BasicBlock(FalseSucc)))) {
         // Check for cases of duplicate successors.
         if ((TrueSucc == PN->getParent()) != (FalseSucc == PN->getParent())) {
@@ -3056,7 +3074,7 @@ static bool isKnownNonZeroFromOperator(const Operator *I,
         }
       }
       // Finally recurse on the edge and check it directly.
-      return isKnownNonZero(U.get(), DemandedElts, RecQ, NewDepth);
+      return isKnownNonZero(IncValue, DemandedElts, RecQ, NewDepth);
     });
   }
   case Instruction::InsertElement: {

llvm/test/Transforms/InstCombine/known-phi-recurse.ll

@@ -300,19 +300,16 @@ define i1 @known_non_zero_phi_phi_test() {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[INDVAR:%.*]] = phi i8 [ 2, [[ENTRY:%.*]] ], [ [[CONTAIN:%.*]], [[LOOP_BB1:%.*]] ]
 ; CHECK-NEXT:    [[COND0:%.*]] = call i1 @cond()
-; CHECK-NEXT:    br i1 [[COND0]], label [[LOOP_BB0:%.*]], label [[LOOP_BB1]]
+; CHECK-NEXT:    br i1 [[COND0]], label [[LOOP_BB0:%.*]], label [[LOOP_BB1:%.*]]
 ; CHECK:       loop.bb0:
 ; CHECK-NEXT:    call void @side.effect()
 ; CHECK-NEXT:    br label [[LOOP_BB1]]
 ; CHECK:       loop.bb1:
-; CHECK-NEXT:    [[CONTAIN]] = phi i8 [ 1, [[LOOP_BB0]] ], [ [[INDVAR]], [[LOOP]] ]
 ; CHECK-NEXT:    [[COND1:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[COND1]], label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[BOOL:%.*]] = icmp eq i8 [[CONTAIN]], 0
-; CHECK-NEXT:    ret i1 [[BOOL]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   br label %loop
@@ -345,8 +342,7 @@ define i1 @known_non_zero_phi_select_test() {
 ; CHECK-NEXT:    [[COND1:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[COND1]], label [[EXIT:%.*]], label [[LOOP]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[BOOL:%.*]] = icmp eq i8 [[CONTAIN]], 0
-; CHECK-NEXT:    ret i1 [[BOOL]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   br label %loop