[SCCP] Remove dead switch cases based on range information

[nikic]

Summary

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.

[nikic]

Rebase

[nikic]

Rebase

[fhahn]

llvm/lib/Transforms/Scalar/SCCP.cpp

658

Do we have to exclude constant ranges that may be undef here? Should all successors be feasible in that case?

1866

Recently c++20 style contains members have been added to various data types. Might be good to use instead of count here. (https://en.cppreference.com/w/cpp/container/set/contains).

Also, maybe consider using an early continue/increment if the current case is a feasible successor.

[nikic]

llvm/lib/Transforms/Scalar/SCCP.cpp

658

Depends on how conservative we want to be :) Per langref, switch on undef is UB. In practical terms, I would expect the situation for switch to be less problematic than for br, because we generally don't create switches "out of thin air", e.g. loop unswitch does not create switch statements.

[fhahn]

llvm/lib/Transforms/Scalar/SCCP.cpp

658

I think it would be more consistent to follow to same logic as we do for branches here and flip the switch together. It might not be as big of a problem as for branches and there's no clear line, but it might be surprising for readers of the code. And we should really work towards flipping globally :)

[nikic]

llvm/lib/Transforms/Scalar/SCCP.cpp

658

Makes sense. In that case I'll wait on the assume change in D83643 (or I could split that off into a separate revision), to allow easy testing.

[nikic]

Set AllowUndef=false, switch tests to use !range instead of comparisons.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Scalar/SCCP.cpp

if (TI.isExceptionalTerminator()) {
     return;
   }
 
   if (auto *SI = dyn_cast<SwitchInst>(&TI)) {
     if (!SI->getNumCases()) {
       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);
+    if (SCValue.isConstantRange()) {
+      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;
   }
 
   // In case of indirect branch and its address is a blockaddress, we mark
   // the target as executable.
   if (auto *IBR = dyn_cast<IndirectBrInst>(&TI)) {
     // Casts are folded by visitCastInst.
     ValueLatticeElement IBRValue = getValueState(IBR->getAddress());
for (BasicBlock *Succ : successors(BB)) {
 
       Succ->removePredecessor(BB);
       Updates.push_back({DominatorTree::Delete, BB, Succ});
     }
 
     DTU.applyUpdatesPermissive(Updates);
     BranchInst::Create(OnlyFeasibleSuccessor, BB);
     TI->eraseFromParent();
+  } 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.count(CI->getCaseSuccessor())) {
+        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.
+        continue;
+      }
+
+      ++CI;
+    }
+
+    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;
 }
 
 bool llvm::runIPSCCP(
     Module &M, const DataLayout &DL,
     std::function<const TargetLibraryInfo &(Function &)> GetTLI,
     function_ref<AnalysisResultsForFn(Function &)> getAnalysis) {

llvm/test/Transforms/SCCP/switch.ll

 ; 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:    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 [[X]]
 ;
 entry:
   br i1 %c1, label %switch, label %switch.1
 
 switch:
   %c2 = icmp ult i32 %x, 3
   call void @llvm.assume(i1 %c2)
   switch i32 %x, label %switch.default [

 switch.0:
   ret i32 0
 
 switch.1:
   %phi = phi i32 [ %x, %entry ], [ 0, %switch ], [ 0, %switch ]
   ret i32 %phi
 }
 
+; TODO: Determine that the default destination is dead.
 define i32 @test_local_range(i32 %x) {
 ; CHECK-LABEL: @test_local_range(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[X:%.*]], 3
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
 ; 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
 ; CHECK:       switch.0:
 ; CHECK-NEXT:    ret i32 0
 ; CHECK:       switch.1:
 ; 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)
   switch i32 %x, label %switch.default [
   i32 0, label %switch.0
   i32 1, label %switch.1
   i32 2, label %switch.2
   i32 3, label %switch.3

 
 switch.2:
   ret i32 2
 
 switch.3:
   ret i32 3
 }
 
+; TODO: Determine that case i3 is dead, even though the edge is shared?
 define i32 @test_duplicate_successors(i32 %x) {
 ; CHECK-LABEL: @test_duplicate_successors(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[X:%.*]], 3
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
 ; 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)
   switch i32 %x, label %switch.default [
   i32 0, label %switch.0
   i32 1, label %switch.0
   i32 2, label %switch.1
   i32 3, label %switch.1

 
 switch.1:
   ret i32 1
 
 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 !0
 ; 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
 ; CHECK:       switch.3:
 ; CHECK-NEXT:    ret i32 3
 ;
   switch i32 %x, label %switch.default [
   i32 0, label %switch.0
   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
 
 switch.0:
   ret i32 0
 
 switch.1:

 ; CHECK-NEXT:    ret void
 ;
   call i32 @test_ip_range(i32 1)
   call i32 @test_ip_range(i32 3)
   ret void
 }
 
 declare void @llvm.assume(i1)
+
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 5, i32 3, i32 4}