[SCCP] Remove dead switch cases based on range information

Determine whether switch edges are feasible based on range information,
and remove non-feasible edges lateron.

This does not try to determine whether the default edge is dead,
as we'd have to determine that the range is fully covered by the
cases for that.

Another limitation here is that we don't remove dead cases that
have the same successor as a live case. I'm not handling this
because I wanted to keep the edge removal based on feasible edges
only, rather than inspecting ranges again there -- this does not
seem like a particularly useful case to handle.

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

Author: nikic

llvm/lib/Transforms/Scalar/SCCP.cpp

@@ -649,17 +649,30 @@ void SCCPSolver::getFeasibleSuccessors(Instruction &TI,
       Succs[0] = true;
       return;
     }
-    ValueLatticeElement SCValue = getValueState(SI->getCondition());
-    ConstantInt *CI = getConstantInt(SCValue);
+    const ValueLatticeElement &SCValue = getValueState(SI->getCondition());
+    if (ConstantInt *CI = getConstantInt(SCValue)) {
+      Succs[SI->findCaseValue(CI)->getSuccessorIndex()] = true;
+      return;
+    }
 
-    if (!CI) {   // Overdefined or unknown condition?
-      // All destinations are executable!
-      if (!SCValue.isUnknownOrUndef())
-        Succs.assign(TI.getNumSuccessors(), true);
+    // TODO: Switch on undef is UB. Stop passing false once the rest of LLVM
+    // is ready.
+    if (SCValue.isConstantRange(/*UndefAllowed=*/false)) {
+      const ConstantRange &Range = SCValue.getConstantRange();
+      for (const auto &Case : SI->cases()) {
+        const APInt &CaseValue = Case.getCaseValue()->getValue();
+        if (Range.contains(CaseValue))
+          Succs[Case.getSuccessorIndex()] = true;
+      }
+
+      // TODO: Determine whether default case is reachable.
+      Succs[SI->case_default()->getSuccessorIndex()] = true;
       return;
     }
 
-    Succs[SI->findCaseValue(CI)->getSuccessorIndex()] = true;
+    // Overdefined or unknown condition? All destinations are executable!
+    if (!SCValue.isUnknownOrUndef())
+      Succs.assign(TI.getNumSuccessors(), true);
     return;
   }
 
@@ -1847,9 +1860,26 @@ static bool removeNonFeasibleEdges(const SCCPSolver &Solver, BasicBlock *BB,
 
     BranchInst::Create(OnlyFeasibleSuccessor, BB);
     TI->eraseFromParent();
+    DTU.applyUpdatesPermissive(Updates);
+  } else if (FeasibleSuccessors.size() > 1) {
+    SwitchInstProfUpdateWrapper SI(*cast<SwitchInst>(TI));
+    SmallVector<DominatorTree::UpdateType, 8> Updates;
+    for (auto CI = SI->case_begin(); CI != SI->case_end();) {
+      if (FeasibleSuccessors.contains(CI->getCaseSuccessor())) {
+        ++CI;
+        continue;
+      }
+
+      BasicBlock *Succ = CI->getCaseSuccessor();
+      Succ->removePredecessor(BB);
+      Updates.push_back({DominatorTree::Delete, BB, Succ});
+      SI.removeCase(CI);
+      // Don't increment CI, as we removed a case.
+    }
+
     DTU.applyUpdatesPermissive(Updates);
   } else {
-    llvm_unreachable("Either all successors are feasible, or exactly one is");
+    llvm_unreachable("Must have at least one feasible successor");
   }
   return true;
 }

llvm/test/Transforms/SCCP/switch.ll

@@ -73,34 +73,29 @@ end:
   ret i32 %phi
 }
 
-define i32 @test_duplicate_successors_phi_3(i1 %c1, i32 %x) {
+define i32 @test_duplicate_successors_phi_3(i1 %c1, i32* %p, i32 %y) {
 ; CHECK-LABEL: @test_duplicate_successors_phi_3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[C1:%.*]], label [[SWITCH:%.*]], label [[SWITCH_1:%.*]]
 ; CHECK:       switch:
-; CHECK-NEXT:    [[C2:%.*]] = icmp ult i32 [[X:%.*]], 3
-; CHECK-NEXT:    call void @llvm.assume(i1 [[C2]])
+; CHECK-NEXT:    [[X:%.*]] = load i32, i32* [[P:%.*]], align 4, !range !0
 ; CHECK-NEXT:    switch i32 [[X]], label [[SWITCH_DEFAULT:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[SWITCH_DEFAULT]]
 ; CHECK-NEXT:    i32 1, label [[SWITCH_0:%.*]]
 ; CHECK-NEXT:    i32 2, label [[SWITCH_0]]
-; CHECK-NEXT:    i32 3, label [[SWITCH_1]]
-; CHECK-NEXT:    i32 4, label [[SWITCH_1]]
 ; CHECK-NEXT:    ]
 ; CHECK:       switch.default:
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       switch.0:
 ; CHECK-NEXT:    ret i32 0
 ; CHECK:       switch.1:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[X]], [[ENTRY:%.*]] ], [ 0, [[SWITCH]] ], [ 0, [[SWITCH]] ]
-; CHECK-NEXT:    ret i32 [[PHI]]
+; CHECK-NEXT:    ret i32 [[Y:%.*]]
 ;
 entry:
   br i1 %c1, label %switch, label %switch.1
 
 switch:
-  %c2 = icmp ult i32 %x, 3
-  call void @llvm.assume(i1 %c2)
+  %x = load i32, i32* %p, !range !{i32 0, i32 3}
   switch i32 %x, label %switch.default [
   i32 0, label %switch.default
   i32 1, label %switch.0
@@ -116,19 +111,18 @@ switch.0:
   ret i32 0
 
 switch.1:
-  %phi = phi i32 [ %x, %entry ], [ 0, %switch ], [ 0, %switch ]
+  %phi = phi i32 [ %y, %entry ], [ 0, %switch ], [ 0, %switch ]
   ret i32 %phi
 }
 
-define i32 @test_local_range(i32 %x) {
+; TODO: Determine that the default destination is dead.
+define i32 @test_local_range(i32* %p) {
 ; CHECK-LABEL: @test_local_range(
-; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[X:%.*]], 3
-; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
+; CHECK-NEXT:    [[X:%.*]] = load i32, i32* [[P:%.*]], align 4, !range !0
 ; CHECK-NEXT:    switch i32 [[X]], label [[SWITCH_DEFAULT:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[SWITCH_0:%.*]]
 ; CHECK-NEXT:    i32 1, label [[SWITCH_1:%.*]]
 ; CHECK-NEXT:    i32 2, label [[SWITCH_2:%.*]]
-; CHECK-NEXT:    i32 3, label [[SWITCH_3:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       switch.default:
 ; CHECK-NEXT:    ret i32 -1
@@ -138,11 +132,8 @@ define i32 @test_local_range(i32 %x) {
 ; CHECK-NEXT:    ret i32 1
 ; CHECK:       switch.2:
 ; CHECK-NEXT:    ret i32 2
-; CHECK:       switch.3:
-; CHECK-NEXT:    ret i32 3
 ;
-  %c = icmp ult i32 %x, 3
-  call void @llvm.assume(i1 %c)
+  %x = load i32, i32* %p, !range !{i32 0, i32 3}
   switch i32 %x, label %switch.default [
   i32 0, label %switch.0
   i32 1, label %switch.1
@@ -166,29 +157,24 @@ switch.3:
   ret i32 3
 }
 
-define i32 @test_duplicate_successors(i32 %x) {
+; TODO: Determine that case i3 is dead, even though the edge is shared?
+define i32 @test_duplicate_successors(i32* %p) {
 ; CHECK-LABEL: @test_duplicate_successors(
-; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[X:%.*]], 3
-; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
+; CHECK-NEXT:    [[X:%.*]] = load i32, i32* [[P:%.*]], align 4, !range !0
 ; CHECK-NEXT:    switch i32 [[X]], label [[SWITCH_DEFAULT:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[SWITCH_0:%.*]]
 ; CHECK-NEXT:    i32 1, label [[SWITCH_0]]
 ; CHECK-NEXT:    i32 2, label [[SWITCH_1:%.*]]
 ; CHECK-NEXT:    i32 3, label [[SWITCH_1]]
-; CHECK-NEXT:    i32 4, label [[SWITCH_2:%.*]]
-; CHECK-NEXT:    i32 5, label [[SWITCH_2]]
 ; CHECK-NEXT:    ]
 ; CHECK:       switch.default:
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       switch.0:
 ; CHECK-NEXT:    ret i32 0
 ; CHECK:       switch.1:
 ; CHECK-NEXT:    ret i32 1
-; CHECK:       switch.2:
-; CHECK-NEXT:    ret i32 2
 ;
-  %c = icmp ult i32 %x, 3
-  call void @llvm.assume(i1 %c)
+  %x = load i32, i32* %p, !range !{i32 0, i32 3}
   switch i32 %x, label %switch.default [
   i32 0, label %switch.0
   i32 1, label %switch.0
@@ -211,18 +197,17 @@ switch.2:
   ret i32 2
 }
 
+; Case i32 2 is dead as well, but this cannot be determined based on
+; range information.
 define internal i32 @test_ip_range(i32 %x) {
 ; CHECK-LABEL: @test_ip_range(
 ; CHECK-NEXT:    switch i32 [[X:%.*]], label [[SWITCH_DEFAULT:%.*]] [
-; CHECK-NEXT:    i32 0, label [[SWITCH_0:%.*]]
+; CHECK-NEXT:    i32 3, label [[SWITCH_3:%.*]]
 ; CHECK-NEXT:    i32 1, label [[SWITCH_1:%.*]]
 ; CHECK-NEXT:    i32 2, label [[SWITCH_2:%.*]]
-; CHECK-NEXT:    i32 3, label [[SWITCH_3:%.*]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    ], !prof !1
 ; CHECK:       switch.default:
 ; CHECK-NEXT:    ret i32 -1
-; CHECK:       switch.0:
-; CHECK-NEXT:    ret i32 0
 ; CHECK:       switch.1:
 ; CHECK-NEXT:    ret i32 1
 ; CHECK:       switch.2:
@@ -235,7 +220,7 @@ define internal i32 @test_ip_range(i32 %x) {
   i32 1, label %switch.1
   i32 2, label %switch.2
   i32 3, label %switch.3
-  ]
+  ], !prof !{!"branch_weights", i32 1, i32 2, i32 3, i32 4, i32 5}
 
 switch.default:
   ret i32 -1
@@ -265,3 +250,5 @@ define void @call_test_ip_range() {
 }
 
 declare void @llvm.assume(i1)
+
+; CHECK: !1 = !{!"branch_weights", i32 1, i32 5, i32 3, i32 4}