Re-commit r355490 "[CodeGen] Omit range checks from jump tables when lowering switches with unreachable default"

Original commit by Ayonam Ray.

This commit adds a regression test for the issue discovered in the
previous commit: that the range check for the jump table can only be
omitted if the fall-through destination of the jump table is
unreachable, which isn't necessarily true just because the default of
the switch is unreachable.

This addresses the missing optimization in PR41242.

> During the lowering of a switch that would result in the generation of a
> jump table, a range check is performed before indexing into the jump
> table, for the switch value being outside the jump table range and a
> conditional branch is inserted to jump to the default block. In case the
> default block is unreachable, this conditional jump can be omitted. This
> patch implements omitting this conditional branch for unreachable
> defaults.
>
> Differential Revision: https://reviews.llvm.org/D52002
> Reviewers: Hans Wennborg, Eli Freidman, Roman Lebedev

llvm-svn: 357067

5 files changed, 153 insertions, 72 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 7ccc371f548..68994ae8e04 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -2366,9 +2366,7 @@ void SelectionDAGBuilder::visitJumpTableHeader(JumpTable &JT,
                                                MachineBasicBlock *SwitchBB) {
   SDLoc dl = getCurSDLoc();
 
-  // Subtract the lowest switch case value from the value being switched on and
-  // conditional branch to default mbb if the result is greater than the
-  // difference between smallest and largest cases.
+  // Subtract the lowest switch case value from the value being switched on.
   SDValue SwitchOp = getValue(JTH.SValue);
   EVT VT = SwitchOp.getValueType();
   SDValue Sub = DAG.getNode(ISD::SUB, dl, VT, SwitchOp,
@@ -2388,24 +2386,33 @@ void SelectionDAGBuilder::visitJumpTableHeader(JumpTable &JT,
                                     JumpTableReg, SwitchOp);
   JT.Reg = JumpTableReg;
 
-  // Emit the range check for the jump table, and branch to the default block
-  // for the switch statement if the value being switched on exceeds the largest
-  // case in the switch.
-  SDValue CMP = DAG.getSetCC(
-      dl, TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(),
-                                 Sub.getValueType()),
-      Sub, DAG.getConstant(JTH.Last - JTH.First, dl, VT), ISD::SETUGT);
-
-  SDValue BrCond = DAG.getNode(ISD::BRCOND, dl,
-                               MVT::Other, CopyTo, CMP,
-                               DAG.getBasicBlock(JT.Default));
-
-  // Avoid emitting unnecessary branches to the next block.
-  if (JT.MBB != NextBlock(SwitchBB))
-    BrCond = DAG.getNode(ISD::BR, dl, MVT::Other, BrCond,
-                         DAG.getBasicBlock(JT.MBB));
-
-  DAG.setRoot(BrCond);
+  if (!JTH.OmitRangeCheck) {
+    // Emit the range check for the jump table, and branch to the default block
+    // for the switch statement if the value being switched on exceeds the
+    // largest case in the switch.
+    SDValue CMP = DAG.getSetCC(
+        dl, TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(),
+                                   Sub.getValueType()),
+        Sub, DAG.getConstant(JTH.Last - JTH.First, dl, VT), ISD::SETUGT);
+
+    SDValue BrCond = DAG.getNode(ISD::BRCOND, dl,
+                                 MVT::Other, CopyTo, CMP,
+                                 DAG.getBasicBlock(JT.Default));
+
+    // Avoid emitting unnecessary branches to the next block.
+    if (JT.MBB != NextBlock(SwitchBB))
+      BrCond = DAG.getNode(ISD::BR, dl, MVT::Other, BrCond,
+                           DAG.getBasicBlock(JT.MBB));
+
+    DAG.setRoot(BrCond);
+  } else {
+    // Avoid emitting unnecessary branches to the next block.
+    if (JT.MBB != NextBlock(SwitchBB))
+      DAG.setRoot(DAG.getNode(ISD::BR, dl, MVT::Other, CopyTo,
+                              DAG.getBasicBlock(JT.MBB)));
+    else
+      DAG.setRoot(CopyTo);
+  }
 }
 
 /// Create a LOAD_STACK_GUARD node, and let it carry the target specific global
@@ -10335,7 +10342,15 @@ void SelectionDAGBuilder::lowerWorkItem(SwitchWorkListItem W, Value *Cond,
           }
         }
 
-        addSuccessorWithProb(CurMBB, Fallthrough, FallthroughProb);
+        if (Fallthrough == DefaultMBB &&
+            isa<UnreachableInst>(
+                DefaultMBB->getBasicBlock()->getFirstNonPHIOrDbg())) {
+          // Skip the range check if the fallthrough block is unreachable.
+          JTH->OmitRangeCheck = true;
+        }
+
+        if (!JTH->OmitRangeCheck)
+          addSuccessorWithProb(CurMBB, Fallthrough, FallthroughProb);
         addSuccessorWithProb(CurMBB, JumpMBB, JumpProb);
         CurMBB->normalizeSuccProbs();
 
@@ -10649,38 +10664,6 @@ void SelectionDAGBuilder::visitSwitch(const SwitchInst &SI) {
   // if there are many clusters.
   sortAndRangeify(Clusters);
 
-  if (TM.getOptLevel() != CodeGenOpt::None) {
-    // Replace an unreachable default with the most popular destination.
-    // FIXME: Exploit unreachable default more aggressively.
-    bool UnreachableDefault =
-        isa<UnreachableInst>(SI.getDefaultDest()->getFirstNonPHIOrDbg());
-    if (UnreachableDefault && !Clusters.empty()) {
-      DenseMap<const BasicBlock *, unsigned> Popularity;
-      unsigned MaxPop = 0;
-      const BasicBlock *MaxBB = nullptr;
-      for (auto I : SI.cases()) {
-        const BasicBlock *BB = I.getCaseSuccessor();
-        if (++Popularity[BB] > MaxPop) {
-          MaxPop = Popularity[BB];
-          MaxBB = BB;
-        }
-      }
-      // Set new default.
-      assert(MaxPop > 0 && MaxBB);
-      DefaultMBB = FuncInfo.MBBMap[MaxBB];
-
-      // Remove cases that were pointing to the destination that is now the
-      // default.
-      CaseClusterVector New;
-      New.reserve(Clusters.size());
-      for (CaseCluster &CC : Clusters) {
-        if (CC.MBB != DefaultMBB)
-          New.push_back(CC);
-      }
-      Clusters = std::move(New);
-    }
-  }
-
   // The branch probablity of the peeled case.
   BranchProbability PeeledCaseProb = BranchProbability::getZero();
   MachineBasicBlock *PeeledSwitchMBB =
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h
index e0d397c2d41..3ed1b64adbe 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h
@@ -278,11 +278,12 @@ private:
     const Value *SValue;
     MachineBasicBlock *HeaderBB;
     bool Emitted;
+    bool OmitRangeCheck;
 
     JumpTableHeader(APInt F, APInt L, const Value *SV, MachineBasicBlock *H,
                     bool E = false)
         : First(std::move(F)), Last(std::move(L)), SValue(SV), HeaderBB(H),
-          Emitted(E) {}
+          Emitted(E), OmitRangeCheck(false) {}
   };
   using JumpTableBlock = std::pair<JumpTableHeader, JumpTable>;
 
diff --git a/llvm/test/CodeGen/AArch64/switch-unreachable-default.ll b/llvm/test/CodeGen/AArch64/switch-unreachable-default.ll
new file mode 100644
index 00000000000..9db3cee14e5
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/switch-unreachable-default.ll
@@ -0,0 +1,99 @@
+; RUN: llc -O3 -o - %s | FileCheck %s
+
+; Test that the output in the presence of an unreachable default does not have
+; a compare and branch at the top of the switch to handle the default case.
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Function Attrs: nounwind
+define void @fn(i4) {
+  switch i4 %0, label %default [
+    i4 0, label %case_0
+    i4 1, label %case_1
+    i4 2, label %case_2
+    i4 3, label %case_3
+    i4 4, label %case_4
+    i4 5, label %case_5
+  ]
+
+; CHECK-LABEL: fn:
+; CHECK-NOT:    sub
+; CHECK-NOT:    cmp
+; CHECK-NOT:    b.hi
+; CHECK:        ldrb {{w[0-9]+}}, [{{x[0-9]+}}, {{x[0-9]+}}]
+; CHECK:        add {{x[0-9]+}}, {{x[0-9]+}}, {{x[0-9]+}}, lsl #2
+; CHECK:        br {{x[0-9]+}}
+
+default:
+  unreachable
+
+case_0:
+  tail call void @handle_case_00(i4 %0) #2
+  br label %return_label
+
+case_1:
+  tail call void @handle_case_01(i4 %0) #2
+  br label %return_label
+
+case_2:
+  tail call void @handle_case_02(i4 %0) #2
+  br label %return_label
+
+case_3:
+  tail call void @handle_case_03(i4 %0) #2
+  br label %return_label
+
+case_4:
+  tail call void @handle_case_04(i4 %0) #2
+  br label %return_label
+
+case_5:
+  tail call void @handle_case_05(i4 %0) #2
+  br label %return_label
+
+return_label:
+  ret void
+}
+
+declare  void @handle_case_00(i4)
+declare  void @handle_case_01(i4)
+declare  void @handle_case_02(i4)
+declare  void @handle_case_03(i4)
+declare  void @handle_case_04(i4)
+declare  void @handle_case_05(i4)
+
+
+
+define i32 @reachable_fallthrough(i32 %x) {
+entry:
+  switch i32 %x, label %def [
+    i32 1,  label %bb1
+    i32 8,  label %bb2
+    i32 16, label %bb3
+    i32 32, label %bb4
+    i32 64, label %bb5
+  ]
+
+; The switch is lowered with a jump table for cases 1--32 and case 64 checked
+; separately. Even though the default of switch is unreachable, the fall-through
+; for the jump table *is* reachable so the range check must be emitted.
+;
+; CHECK-LABEL: reachable_fallthrough
+; CHECK: sub
+; CHECK: cmp
+; CHECK: b.hi
+;
+; TODO: Drop the cmp for the 64 case since its fallthrough is unreachable.
+
+
+def: unreachable
+bb1: br label %return
+bb2: br label %return
+bb3: br label %return
+bb4: br label %return
+bb5: br label %return
+
+return:
+  %p = phi i32 [ 3, %bb1 ], [ 2, %bb2 ], [ 1, %bb3 ], [ 0, %bb4 ], [ 0, %bb5 ]
+  ret i32 %p
+}
diff --git a/llvm/test/CodeGen/X86/pr38743.ll b/llvm/test/CodeGen/X86/pr38743.ll
index 8154baf0804..bb8cbb13f7b 100644
--- a/llvm/test/CodeGen/X86/pr38743.ll
+++ b/llvm/test/CodeGen/X86/pr38743.ll
@@ -19,26 +19,24 @@ declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture writeonly, i8* nocapture r
 define void @pr38743(i32 %a0) #1 align 2 {
 ; CHECK-LABEL: pr38743:
 ; CHECK:       # %bb.0: # %bb
-; CHECK-NEXT:    cmpl $3, %edi
-; CHECK-NEXT:    je .LBB0_4
-; CHECK-NEXT:  # %bb.1: # %bb
-; CHECK-NEXT:    cmpl $1, %edi
-; CHECK-NEXT:    je .LBB0_2
-; CHECK-NEXT:  # %bb.3: # %bb5
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NEXT:    decl %edi
+; CHECK-NEXT:    jmpq *.LJTI0_0(,%rdi,8)
+; CHECK-NEXT:  .LBB0_2: # %bb5
 ; CHECK-NEXT:    movzwl .str.17+{{.*}}(%rip), %eax
 ; CHECK-NEXT:    movw %ax, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movq {{.*}}(%rip), %rax
-; CHECK-NEXT:    jmp .LBB0_5
-; CHECK-NEXT:  .LBB0_4: # %bb8
-; CHECK-NEXT:    movq .str.18+{{.*}}(%rip), %rax
+; CHECK-NEXT:    jmp .LBB0_4
+; CHECK-NEXT:  .LBB0_1: # %bb2
+; CHECK-NEXT:    movq .str.16+{{.*}}(%rip), %rax
 ; CHECK-NEXT:    movq %rax, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movq {{.*}}(%rip), %rax
-; CHECK-NEXT:    jmp .LBB0_5
-; CHECK-NEXT:  .LBB0_2: # %bb2
-; CHECK-NEXT:    movq .str.16+{{.*}}(%rip), %rax
+; CHECK-NEXT:    jmp .LBB0_4
+; CHECK-NEXT:  .LBB0_3: # %bb8
+; CHECK-NEXT:    movq .str.18+{{.*}}(%rip), %rax
 ; CHECK-NEXT:    movq %rax, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movq {{.*}}(%rip), %rax
-; CHECK-NEXT:  .LBB0_5: # %bb12
+; CHECK-NEXT:  .LBB0_4: # %bb12
 ; CHECK-NEXT:    movq %rax, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movq -{{[0-9]+}}(%rsp), %rax
 ; CHECK-NEXT:    movq %rax, (%rax)
diff --git a/llvm/test/CodeGen/X86/switch-jump-table.ll b/llvm/test/CodeGen/X86/switch-jump-table.ll
index 1e1f7c5cef0..4c7937078e8 100644
--- a/llvm/test/CodeGen/X86/switch-jump-table.ll
+++ b/llvm/test/CodeGen/X86/switch-jump-table.ll
@@ -2,14 +2,12 @@
 ; RUN: llc -mtriple=i686-pc-gnu-linux -print-machineinstrs=expand-isel-pseudos %s -o /dev/null 2>&1 | FileCheck %s -check-prefix=CHECK-JT-PROB
 
 
-; An unreachable default destination is replaced with the most popular case label.
+; An unreachable default destination is ignored and no compare and branch
+; is generated for the default values.
 
 define void @foo(i32 %x, i32* %to) {
 ; CHECK-LABEL: foo:
 ; CHECK: movl 4(%esp), [[REG:%e[a-z]{2}]]
-; CHECK: cmpl $3, [[REG]]
-; CHECK: ja .LBB0_6
-; CHECK-NEXT: # %bb.1:
 ; CHECK-NEXT: jmpl *.LJTI0_0(,[[REG]],4)
 ; CHECK: movl $4
 ; CHECK: retl
@@ -45,10 +43,12 @@ default:
 
 ; The jump table has four entries.
 ; CHECK-LABEL: .LJTI0_0:
+; CHECK-NEXT: .long  .LBB0_1
 ; CHECK-NEXT: .long  .LBB0_2
 ; CHECK-NEXT: .long  .LBB0_3
 ; CHECK-NEXT: .long  .LBB0_4
 ; CHECK-NEXT: .long  .LBB0_5
+; CHECK-NEXT: .long  .LBB0_5
 }
 
 ; Check if branch probabilities are correctly assigned to the jump table.