[CVP] Keep ReachableCaseCount in sync with range of condition (llvm#142302)

llvm#79993 assumes that a reachable
case must be contained by `CR`. However, it doesn't hold for some edge
cases. This patch adds additional checks to ensure `ReachableCaseCount`
is correct.

Note: Similar optimization in SCCP isn't affected by this bug because it
uses `CR` to compute `ReachableCaseCount`.

Closes llvm#142286.

Author: dtcxzyw

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

@@ -370,15 +370,30 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
   { // Scope for SwitchInstProfUpdateWrapper. It must not live during
     // ConstantFoldTerminator() as the underlying SwitchInst can be changed.
     SwitchInstProfUpdateWrapper SI(*I);
+    ConstantRange CR =
+        LVI->getConstantRangeAtUse(I->getOperandUse(0), /*UndefAllowed=*/false);
     unsigned ReachableCaseCount = 0;
 
     for (auto CI = SI->case_begin(), CE = SI->case_end(); CI != CE;) {
       ConstantInt *Case = CI->getCaseValue();
-      auto *Res = dyn_cast_or_null<ConstantInt>(
-          LVI->getPredicateAt(CmpInst::ICMP_EQ, Cond, Case, I,
-                              /* UseBlockValue */ true));
+      std::optional<bool> Predicate = std::nullopt;
+      if (!CR.contains(Case->getValue()))
+        Predicate = false;
+      else if (CR.isSingleElement() &&
+               *CR.getSingleElement() == Case->getValue())
+        Predicate = true;
+      if (!Predicate) {
+        // Handle missing cases, e.g., the range has a hole.
+        auto *Res = dyn_cast_or_null<ConstantInt>(
+            LVI->getPredicateAt(CmpInst::ICMP_EQ, Cond, Case, I,
+                                /* UseBlockValue=*/true));
+        if (Res && Res->isZero())
+          Predicate = false;
+        else if (Res && Res->isOne())
+          Predicate = true;
+      }
 
-      if (Res && Res->isZero()) {
+      if (Predicate && !*Predicate) {
         // This case never fires - remove it.
         BasicBlock *Succ = CI->getCaseSuccessor();
         Succ->removePredecessor(BB);
@@ -395,7 +410,7 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
           DTU.applyUpdatesPermissive({{DominatorTree::Delete, BB, Succ}});
         continue;
       }
-      if (Res && Res->isOne()) {
+      if (Predicate && *Predicate) {
         // This case always fires.  Arrange for the switch to be turned into an
         // unconditional branch by replacing the switch condition with the case
         // value.
@@ -410,28 +425,24 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
       ++ReachableCaseCount;
     }
 
-    BasicBlock *DefaultDest = SI->getDefaultDest();
-    if (ReachableCaseCount > 1 &&
-        !isa<UnreachableInst>(DefaultDest->getFirstNonPHIOrDbg())) {
-      ConstantRange CR = LVI->getConstantRangeAtUse(I->getOperandUse(0),
-                                                    /*UndefAllowed*/ false);
-      // The default dest is unreachable if all cases are covered.
-      if (!CR.isSizeLargerThan(ReachableCaseCount)) {
-        BasicBlock *NewUnreachableBB =
-            BasicBlock::Create(BB->getContext(), "default.unreachable",
-                               BB->getParent(), DefaultDest);
-        new UnreachableInst(BB->getContext(), NewUnreachableBB);
+    // The default dest is unreachable if all cases are covered.
+    if (!SI->defaultDestUndefined() &&
+        !CR.isSizeLargerThan(ReachableCaseCount)) {
+      BasicBlock *DefaultDest = SI->getDefaultDest();
+      BasicBlock *NewUnreachableBB =
+          BasicBlock::Create(BB->getContext(), "default.unreachable",
+                             BB->getParent(), DefaultDest);
+      new UnreachableInst(BB->getContext(), NewUnreachableBB);
 
-        DefaultDest->removePredecessor(BB);
-        SI->setDefaultDest(NewUnreachableBB);
+      DefaultDest->removePredecessor(BB);
+      SI->setDefaultDest(NewUnreachableBB);
 
-        if (SuccessorsCount[DefaultDest] == 1)
-          DTU.applyUpdates({{DominatorTree::Delete, BB, DefaultDest}});
-        DTU.applyUpdates({{DominatorTree::Insert, BB, NewUnreachableBB}});
+      if (SuccessorsCount[DefaultDest] == 1)
+        DTU.applyUpdates({{DominatorTree::Delete, BB, DefaultDest}});
+      DTU.applyUpdates({{DominatorTree::Insert, BB, NewUnreachableBB}});
 
-        ++NumDeadCases;
-        Changed = true;
-      }
+      ++NumDeadCases;
+      Changed = true;
     }
   }
 

llvm/test/Transforms/CorrelatedValuePropagation/switch.ll

@@ -294,6 +294,42 @@ cleanup:
   ret i32 %retval.0
 }
 
+; Make sure that we don't branch into unreachable.
+
+define void @pr142286() {
+; CHECK-LABEL: define void @pr142286() {
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br label [[LOOP2:%.*]]
+; CHECK:       loop2:
+; CHECK-NEXT:    br label [[LOOP3:%.*]]
+; CHECK:       loop3:
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+start:
+  br label %loop
+
+loop:
+  %phi = phi i8 [ -1, %start ], [ 0, %loop3 ]
+  br label %loop2
+
+loop2:
+  br label %loop3
+
+loop3:
+  switch i8 %phi, label %exit [
+  i8 0, label %loop3
+  i8 1, label %loop2
+  i8 2, label %loop
+  ]
+
+exit:
+  ret void
+}
+
 declare i32 @call0()
 declare i32 @call1()
 declare i32 @call2()