[ValueTracking] Handle recursive select/PHI in ComputeKnownBits (llvm#114689)

goldsteinn · web-flow · commit a56ba1fab07b · 2025-01-22T11:51:18.000-06:00
Finish porting llvm#114008 to `KnownBits` (Follow up to llvm#113707).
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
@@ -592,6 +592,36 @@ static bool cmpExcludesZero(CmpInst::Predicate Pred, const Value *RHS) {
   return true;
 }
 
+static void breakSelfRecursivePHI(const Use *U, const PHINode *PHI,
+                                  Value *&ValOut, Instruction *&CtxIOut) {
+  ValOut = U->get();
+  if (ValOut == PHI)
+    return;
+  CtxIOut = PHI->getIncomingBlock(*U)->getTerminator();
+  Value *V;
+  // If the Use is a select of this phi, compute analysis on other arm to break
+  // recursion.
+  // TODO: Min/Max
+  if (match(ValOut, m_Select(m_Value(), m_Specific(PHI), m_Value(V))) ||
+      match(ValOut, m_Select(m_Value(), m_Value(V), m_Specific(PHI))))
+    ValOut = V;
+
+  // Same for select, if this phi is 2-operand phi, compute analysis on other
+  // incoming value to break recursion.
+  // TODO: We could handle any number of incoming edges as long as we only have
+  // two unique values.
+  else if (auto *IncPhi = dyn_cast<PHINode>(ValOut);
+           IncPhi && IncPhi->getNumIncomingValues() == 2) {
+    for (int Idx = 0; Idx < 2; ++Idx) {
+      if (IncPhi->getIncomingValue(Idx) == PHI) {
+        ValOut = IncPhi->getIncomingValue(1 - Idx);
+        CtxIOut = IncPhi->getIncomingBlock(1 - Idx)->getTerminator();
+        break;
+      }
+    }
+  }
+}
+
 static bool isKnownNonZeroFromAssume(const Value *V, const SimplifyQuery &Q) {
   // Use of assumptions is context-sensitive. If we don't have a context, we
   // cannot use them!
@@ -1641,25 +1671,19 @@ static void computeKnownBitsFromOperator(const Operator *I,
 
       Known.Zero.setAllBits();
       Known.One.setAllBits();
-      for (unsigned u = 0, e = P->getNumIncomingValues(); u < e; ++u) {
-        Value *IncValue = P->getIncomingValue(u);
+      for (const Use &U : P->operands()) {
+        Value *IncValue;
+        Instruction *CxtI;
+        breakSelfRecursivePHI(&U, P, IncValue, CxtI);
         // Skip direct self references.
-        if (IncValue == P) continue;
-
-        // If the Use is a select of this phi, use the knownbit of the other
-        // operand to break the recursion.
-        if (auto *SI = dyn_cast<SelectInst>(IncValue)) {
-          if (SI->getTrueValue() == P || SI->getFalseValue() == P)
-            IncValue = SI->getTrueValue() == P ? SI->getFalseValue()
-                                               : SI->getTrueValue();
-        }
+        if (IncValue == P)
+          continue;
 
         // Change the context instruction to the "edge" that flows into the
         // phi. This is important because that is where the value is actually
         // "evaluated" even though it is used later somewhere else. (see also
         // D69571).
-        SimplifyQuery RecQ = Q.getWithoutCondContext();
-        RecQ.CxtI = P->getIncomingBlock(u)->getTerminator();
+        SimplifyQuery RecQ = Q.getWithoutCondContext().getWithInstruction(CxtI);
 
         Known2 = KnownBits(BitWidth);
 
@@ -6053,30 +6077,13 @@ void computeKnownFPClass(const Value *V, const APInt &DemandedElts,
       bool First = true;
 
       for (const Use &U : P->operands()) {
-        Value *IncValue = U.get();
+        Value *IncValue;
+        Instruction *CxtI;
+        breakSelfRecursivePHI(&U, P, IncValue, CxtI);
         // Skip direct self references.
         if (IncValue == P)
           continue;
 
-        Instruction *CxtI = P->getIncomingBlock(U)->getTerminator();
-
-        // If the Use is a select of this phi, use the fp class of the other
-        // operand to break the recursion. Same around 2-operand phi nodes
-        Value *V;
-        if (match(IncValue, m_Select(m_Value(), m_Specific(P), m_Value(V))) ||
-            match(IncValue, m_Select(m_Value(), m_Value(V), m_Specific(P)))) {
-          IncValue = V;
-        } else if (auto *IncPhi = dyn_cast<PHINode>(IncValue);
-                   IncPhi && IncPhi->getNumIncomingValues() == 2) {
-          for (int Idx = 0; Idx < 2; ++Idx) {
-            if (IncPhi->getIncomingValue(Idx) == P) {
-              IncValue = IncPhi->getIncomingValue(1 - Idx);
-              CxtI = IncPhi->getIncomingBlock(1 - Idx)->getTerminator();
-              break;
-            }
-          }
-        }
-
         KnownFPClass KnownSrc;
         // Recurse, but cap the recursion to two levels, because we don't want
         // to waste time spinning around in loops. We need at least depth 2 to
diff --git a/llvm/test/Analysis/ScalarEvolution/cycled_phis.ll b/llvm/test/Analysis/ScalarEvolution/cycled_phis.ll
@@ -8,9 +8,9 @@ define void @test_01() {
 ; CHECK-LABEL: 'test_01'
 ; CHECK-NEXT:  Classifying expressions for: @test_01
 ; CHECK-NEXT:    %phi_1 = phi i32 [ 10, %entry ], [ %phi_2, %loop ]
-; CHECK-NEXT:    --> %phi_1 U: full-set S: full-set Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %phi_1 U: [0,31) S: [0,31) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %phi_2 = phi i32 [ 20, %entry ], [ %phi_1, %loop ]
-; CHECK-NEXT:    --> %phi_2 U: full-set S: full-set Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %phi_2 U: [0,31) S: [0,31) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %cond = call i1 @cond()
 ; CHECK-NEXT:    --> %cond U: full-set S: full-set Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_01
diff --git a/llvm/test/Analysis/ScalarEvolution/unknown_phis.ll b/llvm/test/Analysis/ScalarEvolution/unknown_phis.ll
@@ -39,9 +39,9 @@ define void @merge_values_with_ranges_looped(ptr %a_len_ptr, ptr %b_len_ptr) {
 ; CHECK-NEXT:    %len_b = load i32, ptr %b_len_ptr, align 4, !range !0
 ; CHECK-NEXT:    --> %len_b U: [0,2147483647) S: [0,2147483647)
 ; CHECK-NEXT:    %p1 = phi i32 [ %len_a, %entry ], [ %p2, %loop ]
-; CHECK-NEXT:    --> %p1 U: full-set S: full-set Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %p1 U: [0,-2147483648) S: [0,-2147483648) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %p2 = phi i32 [ %len_b, %entry ], [ %p1, %loop ]
-; CHECK-NEXT:    --> %p2 U: full-set S: full-set Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %p2 U: [0,-2147483648) S: [0,-2147483648) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<nuw><nsw><%loop> U: [0,100) S: [0,100) Exits: 99 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.next = add i32 %iv, 1
diff --git a/llvm/test/Transforms/InstCombine/known-phi-recurse.ll b/llvm/test/Transforms/InstCombine/known-phi-recurse.ll
@@ -256,4 +256,142 @@ exit:
   ret i8 %bool
 }
 
+define i8 @knownbits_umax_select_test() {
+; CHECK-LABEL: @knownbits_umax_select_test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDVAR:%.*]] = phi i8 [ 0, [[ENTRY:%.*]] ], [ [[CONTAIN:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[COND0:%.*]] = call i1 @cond()
+; CHECK-NEXT:    [[CONTAIN]] = call i8 @llvm.umax.i8(i8 [[INDVAR]], i8 1)
+; CHECK-NEXT:    [[COND1:%.*]] = call i1 @cond()
+; CHECK-NEXT:    br i1 [[COND1]], label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[BOOL:%.*]] = and i8 [[CONTAIN]], 1
+; CHECK-NEXT:    ret i8 [[BOOL]]
+;
+entry:
+  br label %loop
+
+loop:
+  %indvar = phi i8 [ 0, %entry ], [ %contain, %loop ]
+  %cond0 = call i1 @cond()
+  %contain = call i8 @llvm.umax.i8(i8 1, i8 %indvar)
+  %cond1 = call i1 @cond()
+  br i1 %cond1, label %exit, label %loop
+
+exit:
+  %bool = and i8 %contain, 1
+  ret i8 %bool
+}
+
+define i8 @knownbits_phi_phi_test() {
+; CHECK-LABEL: @knownbits_phi_phi_test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDVAR:%.*]] = phi i8 [ 0, [[ENTRY:%.*]] ], [ [[CONTAIN:%.*]], [[LOOP_BB1:%.*]] ]
+; CHECK-NEXT:    [[COND0:%.*]] = call i1 @cond()
+; 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:    ret i8 [[CONTAIN]]
+;
+entry:
+  br label %loop
+
+loop:
+  %indvar = phi i8 [ 0, %entry ], [ %contain, %loop.bb1 ]
+  %cond0 = call i1 @cond()
+  br i1 %cond0, label %loop.bb0, label %loop.bb1
+loop.bb0:
+  call void @side.effect()
+  br label %loop.bb1
+loop.bb1:
+  %contain = phi i8 [ 1, %loop.bb0 ], [ %indvar, %loop ]
+  %cond1 = call i1 @cond()
+  br i1 %cond1, label %exit, label %loop
+
+exit:
+  %bool = and i8 %contain, 1
+  ret i8 %bool
+}
+
+
+define i1 @known_non_zero_phi_phi_test() {
+; CHECK-LABEL: @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:       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]]
+;
+entry:
+  br label %loop
+
+loop:
+  %indvar = phi i8 [ 2, %entry ], [ %contain, %loop.bb1 ]
+  %cond0 = call i1 @cond()
+  br i1 %cond0, label %loop.bb0, label %loop.bb1
+loop.bb0:
+  call void @side.effect()
+  br label %loop.bb1
+loop.bb1:
+  %contain = phi i8 [ 1, %loop.bb0 ], [ %indvar, %loop ]
+  %cond1 = call i1 @cond()
+  br i1 %cond1, label %exit, label %loop
+
+exit:
+  %bool = icmp eq i8 %contain, 0
+  ret i1 %bool
+}
+
+define i1 @known_non_zero_phi_select_test() {
+; CHECK-LABEL: @known_non_zero_phi_select_test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDVAR:%.*]] = phi i8 [ 2, [[ENTRY:%.*]] ], [ [[CONTAIN:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[COND0:%.*]] = call i1 @cond()
+; CHECK-NEXT:    [[CONTAIN]] = select i1 [[COND0]], i8 1, i8 [[INDVAR]]
+; 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]]
+;
+entry:
+  br label %loop
+
+loop:
+  %indvar = phi i8 [ 2, %entry ], [ %contain, %loop ]
+  %cond0 = call i1 @cond()
+  %contain = select i1 %cond0, i8 1, i8 %indvar
+  %cond1 = call i1 @cond()
+  br i1 %cond1, label %exit, label %loop
+
+exit:
+  %bool = icmp eq i8 %contain, 0
+  ret i1 %bool
+}
+
 declare i1 @cond()
+declare void @side.effect()
+
diff --git a/llvm/test/Transforms/SimplifyCFG/switch-branch-fold-indirectbr-102351.ll b/llvm/test/Transforms/SimplifyCFG/switch-branch-fold-indirectbr-102351.ll
@@ -6,35 +6,29 @@ define i32 @foo.1(i32 %arg, ptr %arg1) {
 ; CHECK-SAME: i32 [[ARG:%.*]], ptr [[ARG1:%.*]]) {
 ; CHECK-NEXT:  [[BB:.*]]:
 ; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [2 x ptr], align 16
-; CHECK-NEXT:    store ptr blockaddress(@foo.1, %[[BB8:.*]]), ptr [[ALLOCA]], align 16
+; CHECK-NEXT:    store ptr blockaddress(@foo.1, %[[BB2:.*]]), ptr [[ALLOCA]], align 16
 ; CHECK-NEXT:    [[GETELEMENTPTR:%.*]] = getelementptr inbounds [2 x ptr], ptr [[ALLOCA]], i64 0, i64 1
 ; CHECK-NEXT:    store ptr blockaddress(@foo.1, %[[BB16:.*]]), ptr [[GETELEMENTPTR]], align 8
-; CHECK-NEXT:    br label %[[PREFBB2:.*]]
-; CHECK:       [[PREFBB2]]:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, %[[BB]] ], [ [[PHI14:%.*]], %[[BB13:.*]] ]
-; CHECK-NEXT:    [[PHI3:%.*]] = phi i32 [ 0, %[[BB]] ], [ [[PHI15:%.*]], %[[BB13]] ]
-; CHECK-NEXT:    switch i32 [[PHI]], label %[[BB13]] [
-; CHECK-NEXT:      i32 0, label %[[PREFBB18:.*]]
-; CHECK-NEXT:      i32 1, label %[[BB8]]
+; CHECK-NEXT:    br label %[[BB2]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, %[[BB]] ], [ 2, %[[BB18:.*]] ]
+; CHECK-NEXT:    [[PHI3:%.*]] = phi i32 [ 0, %[[BB]] ], [ [[ARG]], %[[BB18]] ]
+; CHECK-NEXT:    switch i32 [[PHI]], label %[[BB2_UNREACHABLEDEFAULT:.*]] [
+; CHECK-NEXT:      i32 0, label %[[BB18]]
 ; CHECK-NEXT:      i32 2, label %[[PREFBB11:.*]]
 ; CHECK-NEXT:    ]
-; CHECK:       [[BB8]]:
-; CHECK-NEXT:    [[PHI10:%.*]] = phi i32 [ [[ARG]], %[[PREFBB18]] ], [ [[PHI3]], %[[PREFBB2]] ]
-; CHECK-NEXT:    br label %[[BB13]]
 ; CHECK:       [[PREFBB11]]:
 ; CHECK-NEXT:    [[CALL:%.*]] = call i32 @wombat(i32 noundef [[PHI3]])
 ; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[PHI3]], 1
-; CHECK-NEXT:    br label %[[PREFBB18]]
-; CHECK:       [[BB13]]:
-; CHECK-NEXT:    [[PHI14]] = phi i32 [ [[PHI]], %[[PREFBB2]] ], [ 2, %[[BB8]] ]
-; CHECK-NEXT:    [[PHI15]] = phi i32 [ [[PHI3]], %[[PREFBB2]] ], [ [[PHI10]], %[[BB8]] ]
-; CHECK-NEXT:    br label %[[PREFBB2]]
+; CHECK-NEXT:    br label %[[BB18]]
+; CHECK:       [[BB2_UNREACHABLEDEFAULT]]:
+; CHECK-NEXT:    unreachable
 ; CHECK:       [[BB16]]:
 ; CHECK-NEXT:    [[CALL17:%.*]] = call i32 @wombat(i32 noundef [[ARG]])
 ; CHECK-NEXT:    ret i32 0
-; CHECK:       [[PREFBB18]]:
+; CHECK:       [[BB18]]:
 ; CHECK-NEXT:    [[LOAD:%.*]] = load ptr, ptr [[ARG1]], align 8
-; CHECK-NEXT:    indirectbr ptr [[LOAD]], [label %[[BB8]], label %bb16]
+; CHECK-NEXT:    indirectbr ptr [[LOAD]], [label %[[BB2]], label %bb16]
 ;
 bb:
   %alloca = alloca [2 x ptr], align 16
