Remove the final bit test during lowering switch statement if all cases in bit test cover a contiguous range.

[congh]

Summary

When lowering switch statement, if bit tests are used then LLVM will always generates a jump to the default statement in the last bit test. However, this is not necessary when all cases in bit tests cover a contiguous range. This is because when generating the bit tests header MBB, there is a range check that guarantees cases in bit tests won't go outside of [low, high], where low and high are minimum and maximum case values in the bit tests. This patch checks if this is the case and then doesn't emit jump to default statement and hence saves a bit test and a branch.

For example, given the following IR:

define void @foo(i32 %x) {
entry:
  switch i32 %x, label %return [
    i32 0, label %bb0
    i32 2, label %bb0
    i32 4, label %bb0
    i32 1, label %bb1
    i32 3, label %bb1
    i32 5, label %bb1
  ]
bb0: tail call void @a() br label %return
bb1: tail call void @b() br label %return
return: ret void
}
declare void @a()
declare void @b()

Previously LLVM generates the following MC:

BB#0: derived from LLVM BB %entry
    Live Ins: %EDI
	CMP32ri8 %EDI, 5, %EFLAGS<imp-def>
	JA_1 <BB#3>, %EFLAGS<imp-use>
    Successors according to CFG: BB#3(16) BB#1(16)

BB#1: derived from LLVM BB %entry
    Live Ins: %EDI
    Predecessors according to CFG: BB#0
	%EAX<def> = MOV32ri 21
	BT32rr %EAX<kill>, %EDI, %EFLAGS<imp-def>
	JAE_1 <BB#2>, %EFLAGS<imp-use>
    Successors according to CFG: BB#4(48) BB#2(48)

BB#4: derived from LLVM BB %bb0
    Predecessors according to CFG: BB#1
	TAILJMPd64 <ga:@a>, <regmask>, %RSP<imp-use>, %RSP<imp-use>

BB#2: derived from LLVM BB %entry
    Live Ins: %EDI
    Predecessors according to CFG: BB#1
	%EAX<def> = MOV32ri 42
	BT32rr %EAX<kill>, %EDI<kill>, %EFLAGS<imp-def>
	JAE_1 <BB#3>, %EFLAGS<imp-use>
    Successors according to CFG: BB#5(48) BB#3(16)

BB#5: derived from LLVM BB %bb1
    Predecessors according to CFG: BB#2
	TAILJMPd64 <ga:@b>, <regmask>, %RSP<imp-use>, %RSP<imp-use>

BB#3: derived from LLVM BB %return
    Predecessors according to CFG: BB#0 BB#2
	RETQ

With this patch the MC generated is:

BB#0: derived from LLVM BB %entry
    Live Ins: %EDI
	CMP32ri8 %EDI, 5, %EFLAGS<imp-def>
	JBE_1 <BB#1>, %EFLAGS<imp-use>
    Successors according to CFG: BB#3(16) BB#1(16)

BB#3: derived from LLVM BB %return
    Predecessors according to CFG: BB#0
	RETQ

BB#1: derived from LLVM BB %entry
    Live Ins: %EDI
    Predecessors according to CFG: BB#0
	%EAX<def> = MOV32ri 21
	BT32rr %EAX<kill>, %EDI<kill>, %EFLAGS<imp-def>
	JAE_1 <BB#2>, %EFLAGS<imp-use>
    Successors according to CFG: BB#4(48) BB#2(48)

BB#4: derived from LLVM BB %bb0
    Predecessors according to CFG: BB#1
	TAILJMPd64 <ga:@a>, <regmask>, %RSP<imp-use>, %RSP<imp-use>

BB#2: derived from LLVM BB %bb1
    Predecessors according to CFG: BB#1
	TAILJMPd64 <ga:@b>, <regmask>, %RSP<imp-use>, %RSP<imp-use>

We can save one basic block with 3 instructions.

[davidxl]

I have not looked at the details. Does your patch work for the case when a min value is not zero:

define void @foo(i32 %x) {
entry:

switch i32 %x, label %return [
  i32 1, label %bb0
  i32 3, label %bb0
  i32 5, label %bb0
  i32 2, label %bb1
  i32 4, label %bb1
  i32 6, label %bb1
]

bb0: tail call void @a() br label %return
bb1: tail call void @b() br label %return
return: ret void
}
declare void @a()
declare void @b()

[congh]

In D12249#230466, @davidxl wrote:

I have not looked at the details. Does your patch work for the case when a min value is not zero:

define void @foo(i32 %x) {
entry:

switch i32 %x, label %return [
  i32 1, label %bb0
  i32 3, label %bb0
  i32 5, label %bb0
  i32 2, label %bb1
  i32 4, label %bb1
  i32 6, label %bb1
]

bb0: tail call void @a() br label %return
bb1: tail call void @b() br label %return
return: ret void
}
declare void @a()
declare void @b()

Yes. In visitBitTestHeader(), a subtraction is always generated to subtract the min value, turning the min value into zero.

[davidxl]

The subtraction is generated with your patch? Without your patch, the subtraction is not generated unless the range is outside 64 bit.

[congh]

No, the subtraction is not from this patch. Originally, if the value can fit in 64-bit, the min value will be set to zero before visitBitTestHeader() is called, so a subtraction by zero will be eliminated. In this patch, when the value can fit in 64-bit, I check if the value range is contiguous. If it is, I assume we can save more than a subtraction with the optimization done in this patch. Therefore I still do the subtraction in this case, and the branch to the default statement will be eliminated. Otherwise, the subtraction will not be done as there is always a branch to the default statement.

[davidxl]

If subtraction is always done (when it is contiguous), it is always a win? There are cases the last test is rarely taken.

[congh]

In D12249#230525, @davidxl wrote:

If subtraction is always done (when it is contiguous), it is always a win? There are cases the last test is rarely taken.

OK. So should I calculated the potential win by taking the probability of the last test into consideration? Or if the value can fit in 64-bit we don't to this optimization?

[davidxl]

I think it should be an optimization that always wins -- the following conditions need to be satisfied:

1. value range is continuous AND
2. min == 0 OR the value range is outside 64 (or whatever largest width supported by the target)

[congh]

In D12249#230549, @davidxl wrote:

I think it should be an optimization that always wins -- the following conditions need to be satisfied:

1. value range is continuous AND
2. min == 0 OR the value range is outside 64 (or whatever largest width supported by the target)

Agree. I will update the patch later.

[congh]

Update the patch by disabling the optimization in this patch when the case value can fit in a word and the minimum value of the value range is not zero (in which case we can save a subtraction).

[davidxl]

looks fine to me. Needs to be examined by other reviewers. (suggest adding hans)

[congh]

In D12249#232250, @davidxl wrote:

looks fine to me. Needs to be examined by other reviewers. (suggest adding hans)

I have added Hans as a reviewer. Thanks for the review!

[hans]

Looks great! I like this :-)

Please add tests for the cases where this doesn't fire too.

test/CodeGen/X86/switch.ll

419 ↗

(On Diff #32887)

I think this means the optimization affects both bit-test tests in this file. It would be great to have tests where it doesn't fire too, either because the case range is not covered, or because of the "fits in machine word, so optimize away subtraction" case.

[congh]

Add two test cases on which the optimization in this patch won't be done.

[hans]

Thanks! lgtm

test/CodeGen/X86/switch.ll

145 ↗

(On Diff #33113)

nit: The newline seems unnecessary here. I guess it was just copied from line 110, so maybe remote it there too? Also in bt_is_better3.

[congh]

test/CodeGen/X86/switch.ll

488 ↗

(On Diff #33113)

I have added two test cases covering two cases you mentioned. Thank you for the suggestion!

[congh]

Remove three blank lines from switch.ll.

test/CodeGen/X86/switch.ll

144 ↗

(On Diff #33117)

Removed. Thanks!

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h

struct BitTestCase {
     MachineBasicBlock *TargetBB;
     uint32_t ExtraWeight;
   };
 
   typedef SmallVector<BitTestCase, 3> BitTestInfo;
 
   struct BitTestBlock {
     BitTestBlock(APInt F, APInt R, const Value* SV,
-                 unsigned Rg, MVT RgVT, bool E,
+                 unsigned Rg, MVT RgVT, bool E, bool CR,
                  MachineBasicBlock* P, MachineBasicBlock* D,
                  BitTestInfo C):
       First(F), Range(R), SValue(SV), Reg(Rg), RegVT(RgVT), Emitted(E),
-      Parent(P), Default(D), Cases(std::move(C)) { }
+      ContiguousRange(CR), Parent(P), Default(D), Cases(std::move(C)) { }
     APInt First;
     APInt Range;
     const Value *SValue;
     unsigned Reg;
     MVT RegVT;
     bool Emitted;
+    bool ContiguousRange;
     MachineBasicBlock *Parent;
     MachineBasicBlock *Default;
     BitTestInfo Cases;
   };
 
   /// Minimum jump table density, in percent.
   enum { MinJumpTableDensity = 40 };
 

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

bool SelectionDAGBuilder::buildBitTests(CaseClusterVector &Clusters,
   APInt LowBound;
   APInt CmpRange;
 
   const int BitWidth = DAG.getTargetLoweringInfo()
                            .getPointerTy(DAG.getDataLayout())
                            .getSizeInBits();
   assert(rangeFitsInWord(Low, High) && "Case range must fit in bit mask!");
 
-  if (Low.isNonNegative() && High.slt(BitWidth)) {
-    // Optimize the case where all the case values fit in a
-    // word without having to subtract minValue. In this case,
-    // we can optimize away the subtraction.
+  // Check if the clusters cover a contiguous range such that no value in the
+  // range will jump to the default statement.
+  bool ContiguousRange = true;
+  for (int64_t I = First + 1; I <= Last; ++I) {
+    if (Clusters[I].Low->getValue() != Clusters[I - 1].High->getValue() + 1) {
+      ContiguousRange = false;
+      break;
+    }
+  }
+
+  if (Low.isStrictlyPositive() && High.slt(BitWidth)) {
+    // Optimize the case where all the case values fit in a word without having
+    // to subtract minValue. In this case, we can optimize away the subtraction.
     LowBound = APInt::getNullValue(Low.getBitWidth());
     CmpRange = High;
+    ContiguousRange = false;
   } else {
     LowBound = Low;
     CmpRange = High - Low;
   }
 
   CaseBitsVector CBV;
   uint32_t TotalWeight = 0;
   for (unsigned i = First; i <= Last; ++i) {
bool SelectionDAGBuilder::buildBitTests(CaseClusterVector &Clusters,
   });
 
   for (auto &CB : CBV) {
     MachineBasicBlock *BitTestBB =
         FuncInfo.MF->CreateMachineBasicBlock(SI->getParent());
     BTI.push_back(BitTestCase(CB.Mask, BitTestBB, CB.BB, CB.ExtraWeight));
   }
   BitTestCases.emplace_back(std::move(LowBound), std::move(CmpRange),
-                            SI->getCondition(), -1U, MVT::Other, false, nullptr,
-                            nullptr, std::move(BTI));
+                            SI->getCondition(), -1U, MVT::Other, false,
+                            ContiguousRange, nullptr, nullptr, std::move(BTI));
 
   BTCluster = CaseCluster::bitTests(Clusters[First].Low, Clusters[Last].High,
                                     BitTestCases.size() - 1, TotalWeight);
   return true;
 }
 
 void SelectionDAGBuilder::findBitTestClusters(CaseClusterVector &Clusters,
                                               const SwitchInst *SI) {

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

for (unsigned j = 0, ej = SDB->BitTestCases[i].Cases.size(); j != ej; ++j)
       UnhandledWeight += SDB->BitTestCases[i].Cases[j].ExtraWeight;
 
     for (unsigned j = 0, ej = SDB->BitTestCases[i].Cases.size(); j != ej; ++j) {
       UnhandledWeight -= SDB->BitTestCases[i].Cases[j].ExtraWeight;
       // Set the current basic block to the mbb we wish to insert the code into
       FuncInfo->MBB = SDB->BitTestCases[i].Cases[j].ThisBB;
       FuncInfo->InsertPt = FuncInfo->MBB->end();
       // Emit the code
-      if (j+1 != ej)
-        SDB->visitBitTestCase(SDB->BitTestCases[i],
-                              SDB->BitTestCases[i].Cases[j+1].ThisBB,
-                              UnhandledWeight,
-                              SDB->BitTestCases[i].Reg,
-                              SDB->BitTestCases[i].Cases[j],
-                              FuncInfo->MBB);
+
+      // If all cases cover a contiguous range, it is not necessary to jump to
+      // the default block after the last bit test fails. This is because the
+      // range check during bit test header creation has guaranteed that every
+      // case here doesn't go outside the range.
+      MachineBasicBlock *NextMBB;
+      if (SDB->BitTestCases[i].ContiguousRange && j + 2 == ej)
+        NextMBB = SDB->BitTestCases[i].Cases[j + 1].TargetBB;
+      else if (j + 1 != ej)
+        NextMBB = SDB->BitTestCases[i].Cases[j + 1].ThisBB;
       else
-        SDB->visitBitTestCase(SDB->BitTestCases[i],
-                              SDB->BitTestCases[i].Default,
-                              UnhandledWeight,
-                              SDB->BitTestCases[i].Reg,
-                              SDB->BitTestCases[i].Cases[j],
-                              FuncInfo->MBB);
+        NextMBB = SDB->BitTestCases[i].Default;
+
+                            NextMBB,
 
-
       CurDAG->setRoot(SDB->getRoot());
       SDB->clear();
       CodeGenAndEmitDAG();
+
+      if (SDB->BitTestCases[i].ContiguousRange && j + 2 == ej)
+        break;
     }
 
     // Update PHI Nodes
     for (unsigned pi = 0, pe = FuncInfo->PHINodesToUpdate.size();
          pi != pe; ++pi) {
       MachineInstrBuilder PHI(*MF, FuncInfo->PHINodesToUpdate[pi].first);
       MachineBasicBlock *PHIBB = PHI->getParent();
       assert(PHI->isPHI() &&

llvm/trunk/test/CodeGen/X86/switch.ll

switch i32 %x, label %return [
     i32 3, label %bb0
     i32 6, label %bb0
     i32 1, label %bb1
     i32 4, label %bb1
     i32 7, label %bb1
     i32 2, label %bb2
     i32 5, label %bb2
     i32 8, label %bb2
-
   ]
 bb0: tail call void @g(i32 0) br label %return
 bb1: tail call void @g(i32 1) br label %return
 bb2: tail call void @g(i32 2) br label %return
 return: ret void
 
 ; This could be lowered as a jump table, but bit tests is more efficient.
 ; CHECK-LABEL: bt_is_better
+; The bit test on 2,5,8 is unnecessary as all cases cover the rage [0, 8].
+; The range check guarantees that cases other than 0,3,6 and 1,4,7 must be
+; in 2,5,8.
+;
 ; 73 = 2^0 + 2^3 + 2^6
 ; CHECK: movl $73
 ; CHECK: btl
 ; 146 = 2^1 + 2^4 + 2^7
 ; CHECK: movl $146
 ; CHECK: btl
 ; 292 = 2^2 + 2^5 + 2^8
-; CHECK: movl $292
+; CHECK-NOT: movl $292
+; CHECK-NOT: btl
+}
+
+define void @bt_is_better2(i32 %x) {
+entry:
+  switch i32 %x, label %return [
+    i32 0, label %bb0
+    i32 3, label %bb0
+    i32 6, label %bb0
+    i32 1, label %bb1
+    i32 4, label %bb1
+    i32 7, label %bb1
+    i32 2, label %bb2
+    i32 8, label %bb2
+  ]
+bb0: tail call void @g(i32 0) br label %return
+bb1: tail call void @g(i32 1) br label %return
+bb2: tail call void @g(i32 2) br label %return
+return: ret void
+
+; This will also be lowered as bit test, but as the range [0,8] is not fully
+; covered (5 missing), the default statement can be jumped to and we end up
+; with one more branch.
+; CHECK-LABEL: bt_is_better2
+;
+; 73 = 2^0 + 2^3 + 2^6
+; CHECK: movl $73
+; CHECK: btl
+; 146 = 2^1 + 2^4 + 2^7
+; CHECK: movl $146
+; CHECK: btl
+; 260 = 2^2 + 2^8
+; CHECK: movl $260
+; CHECK: btl
+}
+
+define void @bt_is_better3(i32 %x) {
+entry:
+  switch i32 %x, label %return [
+    i32 10, label %bb0
+    i32 13, label %bb0
+    i32 16, label %bb0
+    i32 11, label %bb1
+    i32 14, label %bb1
+    i32 17, label %bb1
+    i32 12, label %bb2
+    i32 18, label %bb2
+  ]
+bb0: tail call void @g(i32 0) br label %return
+bb1: tail call void @g(i32 1) br label %return
+bb2: tail call void @g(i32 2) br label %return
+return: ret void
+
+; We don't have to subtract 10 from the case value to let the range become
+; [0, 8], as each value in the range [10, 18] can be represented by bits in a
+; word. Then we still need a branch to jump to the default statement for the
+; range [0, 10).
+; CHECK-LABEL: bt_is_better3
+;
+; 74752 = 2^10 + 2^13 + 2^16
+; CHECK: movl $74752
+; CHECK: btl
+; 149504 = 2^11 + 2^14 + 2^17
+; CHECK: movl $149504
+; CHECK: btl
+; 266240 = 2^12 + 2^15 + 2^18
+; CHECK: movl $266240
 ; CHECK: btl
 }
 
 
 define void @optimal_pivot1(i32 %x) {
 entry:
   switch i32 %x, label %return [
     i32 100, label %bb0

 bb0: tail call void @g(i32 0) br label %return
 bb1: tail call void @g(i32 1) br label %return
 bb2: tail call void @g(i32 2) br label %return
 return: ret void
 
 ; Cases 1,4,7 have a very large branch weight (which shouldn't overflow), so
 ; their bit test should come first. 0,3,6 and 2,5,8,9 both have a weight of 12,
 ; but the latter set has more cases, so should be tested for earlier.
+; The bit test on 0,3,6 is unnecessary as all cases cover the rage [0, 9].
+; The range check guarantees that cases other than 1,4,7 and 2,5,8,9 must be
+; in 0,3,6.
 
 ; CHECK-LABEL: bt_order_by_weight
 ; 146 = 2^1 + 2^4 + 2^7
 ; CHECK: movl $146
 ; CHECK: btl
 ; 292 = 2^2 + 2^5 + 2^8 + 2^9
 ; CHECK: movl $804
 ; CHECK: btl
 ; 73 = 2^0 + 2^3 + 2^6
-; CHECK: movl $73
-; CHECK: btl
+; CHECK-NOT: movl $73
+; CHECK-NOT: btl
 }
 
 !1 = !{!"branch_weights",
        ; Default:
        i32 1,
        ; Cases 0,3,6:
        i32 4, i32 4, i32 4,
        ; Cases 1,4,7: