[CodeGen] Peel off the dominant case in switch statement in lowering

[xur]

Summary

This patch peels off the top case in switch statement into a branch if the probability exceeds a threshold.
This will help the branch prediction and avoids the extra compares when lowering into chain of branches.
A previous patch of the same idea was on CFG simplification pass:
https://reviews.llvm.org/D37940
This patch is implemented in SelectionDAG.

[hans]

Thanks! I think doing it here rather than in SimplifyCFG is much better.

I think this can be made both more powerful and simpler, though.

First, I think we want to peel a little later, at least after adjacent cases have been merged into clusters. For example if we have case 0: case 1: case 2: foo; and the three together dominate the switch, we probably want to peel that whole cluster instead of only peeling one case. We could even do this after clustering for jump tables, but maybe we don't want to do that. I.e. if one case dominates heavily, maybe it's worth peeling the case separately rather than peeling the jump table out of the switch. On the other hand, if it dominates completely, that jump table should be well predicted...

In any case, to do the peeling I think we should re-use the current lowering mechanism. This should be done at some point before creating the WorkList. Loop through Clusters to find one that dominates. If we find it, call lowerWorkItem with a single work item which represents that cluster to lower it, update the blocks etc. Then remove the peeled cluster from Clusters and proceed as usual.

What do you think?

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

139

Could we have just one flag instead of two separate ones? I.e. one flag that enables peeling and sets the threshold (and setting it to 0 would turn it off or something)?

141

How was the value 66 chosen?

[xur]

Integrated the review comments and suggestion from Hans.

Thanks,

-Rong

[hans]

Thanks! I think that's looking better.

Sorry for the slow reply; I was travelling the last couple of days. Some comments inline.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

138

Perhaps the option name could make it more clear that this is about switch peeling. Maybe "-switch-peel-threshold"?

141

I'm still curious about this.

9839

No need for 'inline'. Can you add a comment that explains what this does?

9847

Please comment on what the return value is, that is updates Clusters, and what the PelledProb value is.

9849

Maybe call it peelDominantCaseCluster instead, since it's operating on the CaseClusters.

9865

Why are you skipping clusters where Low != High?

9866

This requires CC.Prob so be strictly greater than the DominantCasePercentThreshold, not just greater or equal. Is that intentional? I'd assume that if the threshold is 80% probability, and we have a case with 80% probability, we want to peel it.

9880

Is PeeledSwitchMBB the MBB that contains the peeled switch (i.e. the switch minus the most dominant case), or the one containing the peeled case? I wonder if we could find better names, or add comments to make this clear.

9883

Maybe move this down to where it's used (when creating W).

9917

Could we just update SwitchMBB instead, if that's the one that has the rest of the switch? (I'm a little confused by what SwitchMBB actually is at this point).

9994

When would DefaultMBB get replaced? If it can, it needs a test case.

test/CodeGen/Generic/switch-lower-peel-top-case.ll

13

Checking the debug printouts doesn't seem very nice. Can't we check the resulting machine-IR and make sure it has the correct probabilities?

[xur]

Hans, Thanks for the detailed reviews. Please check if the updated patch addresses your reviews.

One change not in the reviews is I move the code after rangeify clusters. This is together with the suggested change that removing Low != High check.

I also change the code a bit so that if the single case is removed, generate BR to the default directly.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

141

We have a case of ~70% that we would like to peel in our internal test. I lower a few percent as the buffering.

9847

Also changed the parameter for better clarity.

9865

I was concerning that this would create two jumps for the range. But
this turns out not to be the case. I moved this condition.

9866

you are right! changed it to "<"

9880

Containing the peeled switch statement. I added a comment here.

9917

My first try was to update SwitchMBB directly, but it did not work. It seems that we need SwitchMBB as the placehold to decide if we defer the case handling (i.e visitSwitchCase() or push_back to SwitchCases.)

9994

We replace the unreachable default with most popular destination. This actually exposed by an existing test case.

[xur]

Integrated with Hans's review comments.

[hans]

Starting to look really good.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9882

nit: "If there is only case is peeled." sentence looks weird.

9891

Actually, if there's only one cluster, maybe we shouldn't be peeling. I mean, it doesn't make much sense right? And maybe that would help downstream code that could then assume that there are still cases left after peeling.

test/CodeGen/Generic/switch-lower-peel-top-case.ll

13

This could probably be simplified a bit. If %n were an i32 instad of i16 you don't need the sext. You could drop local_unnamed_addr.

The case labels look pretty random; maybe chose cases with more purpose, and rename the target blocks to match. It would also be good to test the case where two adjacent cases are merged into a cluster and peeled.

[xur]

Here is the patch that Integrated with the review comments from Hans.

Thanks!

-Rong

[hans]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9994

I still don't understand why it matters here if DefaultMBB has changed?

test/CodeGen/Generic/switch-lower-peel-top-case.ll

13

The part about case values, and test for peeling two adjacent cases still applies.

[xur]

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9882

No need with the new patch.

9891

Ah. Right. Current logic handles it well. I will have a check that bailout earlier.

9994

When the DefaultMBB is changed (replaced with the most popular case), the branch probability not longer consistent. If I scale the original default case branch probability, it could create a insane probability (greater than 100%).

test/CodeGen/Generic/switch-lower-peel-top-case.ll

13

This is extracted from a real program. But I agree that it should be simplified by removing the unrelated details. I will change.

13

Sorry I might misread your earlier comments. Do you mean I should change the case value?

[hans]

test/CodeGen/Generic/switch-lower-peel-top-case.ll

13

Yes.

"The case labels look pretty random; maybe chose cases with more purpose, and rename the target blocks to match. It would also be good to test the case where two adjacent cases are merged into a cluster and peeled."

[xur]

Add a test case that tests merging adjacent cases.

Thanks!

[hans]

test/CodeGen/Generic/switch-lower-peel-top-case.ll

73

Actually, I think this test would pass even without switch peeling.
Since there are only three clusters, it will do straight-line comparisons, with the most probable one (i.e. 85-86) first.
I'd suggest adding a few more cases, so that the remaining cases are lowered with a binary tree. That will show more clearly that the switch is getting peeled.

[xur]

Added more switch cases to make the lowering generate binary search compares.
Thanks to Hans's suggestion.

I also moved the test case to from test/CodeGenGeneric/ test/CodeGen/X86/ as it now checks the IR for x86.

[hans]

lgtm

Thanks for working on this!

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h

private:
   void splitWorkItem(SwitchWorkList &WorkList, const SwitchWorkListItem &W,
                      Value *Cond, MachineBasicBlock *SwitchMBB);
 
   /// Lower W.
   void lowerWorkItem(SwitchWorkListItem W, Value *Cond,
                      MachineBasicBlock *SwitchMBB,
                      MachineBasicBlock *DefaultMBB);
 
+  /// Peel the top probability case if it exceeds the threshold
+  MachineBasicBlock *peelOneDominantWorkItem(const SwitchInst &SI,
+                                             CaseClusterVector &Clusters,
+                                             BranchProbability &PeeledProb);
 
   /// A class which encapsulates all of the information needed to generate a
   /// stack protector check and signals to isel via its state being initialized
   /// that a stack protector needs to be generated.
   ///
   /// *NOTE* The following is a high level documentation of SelectionDAG Stack
   /// Protector Generation. The reason that it is placed here is for a lack of
   /// other good places to stick it.

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 
 static cl::opt<unsigned, true>
 LimitFPPrecision("limit-float-precision",
                  cl::desc("Generate low-precision inline sequences "
                           "for some float libcalls"),
                  cl::location(LimitFloatPrecision),
                  cl::init(0));
 
+static cl::opt<unsigned> DominantCasePercentThreshold(
+    "dominant-case-percent-threshold", cl::Hidden, cl::init(66),
+    cl::desc("Set the case probability threshold for peeling the case from a "
+             "switch statement. A value greater than 100 will void this "
+             "optimization"));
+
 // Limit the width of DAG chains. This is important in general to prevent
 // DAG-based analysis from blowing up. For example, alias analysis and
 // load clustering may not complete in reasonable time. It is difficult to
 // recognize and avoid this situation within each individual analysis, and
 // future analyses are likely to have the same behavior. Limiting DAG width is
 // the safe approach and will be especially important with global DAGs.
 //
 // MaxParallelChains default is arbitrarily high to avoid affecting
CaseBlock CB(ISD::SETLT, Cond, Pivot, nullptr, LeftMBB, RightMBB, W.MBB,
                getCurSDLoc(), LeftProb, RightProb);
 
   if (W.MBB == SwitchMBB)
     visitSwitchCase(CB, SwitchMBB);
   else
     SwitchCases.push_back(CB);
 }
 
+inline static BranchProbability
+scaleCaseProbality(BranchProbability CaseProb, BranchProbability SwitchProb) {
+  if (SwitchProb == BranchProbability::getZero())
+    return BranchProbability::getZero();
+  return BranchProbability(CaseProb.getNumerator(),
+                           SwitchProb.scale(CaseProb.getDenominator()));
+}
+
+// Peel the top probability case if it exceeds the threshold.
+MachineBasicBlock *
+SelectionDAGBuilder::peelOneDominantWorkItem(const SwitchInst &SI,
+                                             CaseClusterVector &Clusters,
+                                             BranchProbability &PeeledProb) {
+  MachineBasicBlock *SwitchMBB = FuncInfo.MBB;
+  if (DominantCasePercentThreshold > 100 || !FuncInfo.BPI ||
+      TM.getOptLevel() == CodeGenOpt::None ||
+      SwitchMBB->getParent()->getFunction()->optForMinSize())
+    return SwitchMBB;
+
+  BranchProbability TopCaseProb =
+      BranchProbability(DominantCasePercentThreshold, 100);
+  unsigned PeeledCaseIndex = 0;
+  bool SwitchPeeled = false;
+  for (unsigned Index = 0; Index < Clusters.size(); ++Index) {
+    CaseCluster &CC = Clusters[Index];
+    if (CC.Low != CC.High)
+      continue;
+    if (CC.Prob <= TopCaseProb)
+      continue;
+    TopCaseProb = CC.Prob;
+    PeeledCaseIndex = Index;
+    SwitchPeeled = true;
+  }
+  if (!SwitchPeeled)
+    return SwitchMBB;
+
+  DEBUG(dbgs() << "Peeled one top case in switch stmt, prob: " << TopCaseProb
+               << "\n");
+  PeeledProb = TopCaseProb.getCompl();
+  MachineFunction::iterator BBI(SwitchMBB);
+  ++BBI;
+  MachineBasicBlock *PeeledSwitchMBB =
+      FuncInfo.MF->CreateMachineBasicBlock(SwitchMBB->getBasicBlock());
+  FuncInfo.MF->insert(BBI, PeeledSwitchMBB);
+  auto PeeledCaseIt = Clusters.begin() + PeeledCaseIndex;
+
+  ExportFromCurrentBlock(SI.getCondition());
+  SwitchWorkListItem W = {SwitchMBB, PeeledCaseIt, PeeledCaseIt,
+                          nullptr,   nullptr,      PeeledProb};
+  lowerWorkItem(W, SI.getCondition(), SwitchMBB, PeeledSwitchMBB);
+
+  Clusters.erase(PeeledCaseIt);
+  for (CaseCluster &CC : Clusters) {
+    DEBUG(dbgs() << "Scale the probablity for one cluster, before scaling: "
+                 << CC.Prob << "\n");
+    CC.Prob = scaleCaseProbality(CC.Prob, PeeledProb);
+    DEBUG(dbgs() << "After scaling: " << CC.Prob << "\n");
+  }
+  return PeeledSwitchMBB;
+}
+
 void SelectionDAGBuilder::visitSwitch(const SwitchInst &SI) {
   // Extract cases from the switch.
   BranchProbabilityInfo *BPI = FuncInfo.BPI;
   CaseClusterVector Clusters;
   Clusters.reserve(SI.getNumCases());
   for (auto I : SI.cases()) {
     MachineBasicBlock *Succ = FuncInfo.MBBMap[I.getCaseSuccessor()];
     const ConstantInt *CaseVal = I.getCaseValue();
     BranchProbability Prob =
         BPI ? BPI->getEdgeProbability(SI.getParent(), I.getSuccessorIndex())
             : BranchProbability(1, SI.getNumCases() + 1);
     Clusters.push_back(CaseCluster::range(CaseVal, CaseVal, Succ, Prob));
   }
 
+  // The probablity of the switch statement after peeling out the top case.
+  BranchProbability PeeledProb = BranchProbability::getOne();
+  MachineBasicBlock *PeeledSwitchMBB =
+      peelOneDominantWorkItem(SI, Clusters, PeeledProb);
+
   MachineBasicBlock *DefaultMBB = FuncInfo.MBBMap[SI.getDefaultDest()];
 
   // Cluster adjacent cases with the same destination. We do this at all
   // optimization levels because it's cheap to do and will make codegen faster
   // if there are many clusters.
   sortAndRangeify(Clusters);
 
   if (TM.getOptLevel() != CodeGenOpt::None) {
DEBUG({
     }
     dbgs() << '\n';
   });
 
   assert(!Clusters.empty());
   SwitchWorkList WorkList;
   CaseClusterIt First = Clusters.begin();
   CaseClusterIt Last = Clusters.end() - 1;
-  auto DefaultProb = getEdgeProbability(SwitchMBB, DefaultMBB);
-  WorkList.push_back({SwitchMBB, First, Last, nullptr, nullptr, DefaultProb});
+  auto DefaultProb = getEdgeProbability(PeeledSwitchMBB, DefaultMBB);
+  // Scale the branchprobability for DefaultMBB if the peel occurs and
+  // DefaultMBB is not replaced.
+  if (PeeledProb != BranchProbability::getOne() &&
+      DefaultMBB == FuncInfo.MBBMap[SI.getDefaultDest()])
+    DefaultProb = scaleCaseProbality(DefaultProb, PeeledProb);
+  WorkList.push_back(
+      {PeeledSwitchMBB, First, Last, nullptr, nullptr, DefaultProb});
 
   while (!WorkList.empty()) {
     SwitchWorkListItem W = WorkList.back();
     WorkList.pop_back();
     unsigned NumClusters = W.LastCluster - W.FirstCluster + 1;
 
     if (NumClusters > 3 && TM.getOptLevel() != CodeGenOpt::None &&
         !DefaultMBB->getParent()->getFunction()->optForMinSize()) {
       // For optimized builds, lower large range as a balanced binary tree.
       splitWorkItem(WorkList, W, SI.getCondition(), SwitchMBB);
       continue;
     }
 
     lowerWorkItem(W, SI.getCondition(), SwitchMBB, DefaultMBB);
   }
 }

test/CodeGen/Generic/MachineBranchProb.ll

 ; CHECK-LABEL: Machine code for function test2:
 entry:
   %conv = sext i32 %x to i64
   switch i64 %conv, label %return [
     i64 0, label %sw.bb
     i64 1, label %sw.bb
     i64 4, label %sw.bb
     i64 5, label %sw.bb1
+    i64 15, label %sw.bb
   ], !prof !0
 ; CHECK: BB#0: derived from LLVM BB %entry
-; CHECK: Successors according to CFG: BB#2({{[0-9a-fx/= ]+}}75.29%) BB#4({{[0-9a-fx/= ]+}}24.71%)
+; CHECK: Successors according to CFG: BB#1({{[0-9a-fx/= ]+}}92.17%) BB#4({{[0-9a-fx/= ]+}}7.83%)
 ; CHECK: BB#4: derived from LLVM BB %entry
-; CHECK: Successors according to CFG: BB#1({{[0-9a-fx/= ]+}}47.62%) BB#5({{[0-9a-fx/= ]+}}52.38%)
+; CHECK: Successors according to CFG: BB#2({{[0-9a-fx/= ]+}}75.29%) BB#5({{[0-9a-fx/= ]+}}24.71%)
 ; CHECK: BB#5: derived from LLVM BB %entry
+; CHECK: Successors according to CFG: BB#1({{[0-9a-fx/= ]+}}47.62%) BB#6({{[0-9a-fx/= ]+}}52.38%)
+; CHECK: BB#6: derived from LLVM BB %entry
 ; CHECK: Successors according to CFG: BB#1({{[0-9a-fx/= ]+}}36.36%) BB#3({{[0-9a-fx/= ]+}}63.64%)
 
 sw.bb:
 ; this call will prevent simplifyCFG from optimizing the block away in ARM/AArch64.
   tail call void @foo()
   br label %return
 
 sw.bb1:
   br label %return
 
 return:
   %retval.0 = phi i32 [ 5, %sw.bb1 ], [ 1, %sw.bb ], [ 0, %entry ]
   ret i32 %retval.0
 }
 
-!0 = !{!"branch_weights", i32 7, i32 6, i32 4, i32 4, i32 64}
+!0 = !{!"branch_weights", i32 7, i32 6, i32 4, i32 4, i32 64, i21 1000}
 
 
 declare void @g(i32)
 define void @left_leaning_weight_balanced_tree(i32 %x) {
 entry:
   switch i32 %x, label %return [
     i32 0,  label %bb0
     i32 100, label %bb1

test/CodeGen/Generic/switch-lower-peel-top-case.ll

+; RUN: llc -debug-only=isel < %s  -o /dev/null 2>&1 | FileCheck %s
+
+define i32 @foo(i16 signext %n) local_unnamed_addr !prof !1 {
+entry:
+  %conv = sext i16 %n to i32
+  switch i32 %conv, label %sw.epilog [
+    i32 8, label %return
+    i32 -8826, label %sw.bb1
+    i32 18312, label %sw.bb3
+    i32 18568, label %sw.bb5
+    i32 129, label %sw.bb7
+  ], !prof !2
+; CHECK: Peeled one top case in switch stmt, prob: 0x5999999a / 0x80000000 = 70.00%
+; CHECK: Scale the probablity for one cluster, before scaling: 0x0020c49c / 0x80000000 = 0.10%
+; CHECK: After scaling: 0x006d3a08 / 0x80000000 = 0.33%
+; CHECK: Scale the probablity for one cluster, before scaling: 0x00418937 / 0x80000000 = 0.20%
+; CHECK: After scaling: 0x00da740d / 0x80000000 = 0.67%
+; CHECK: Scale the probablity for one cluster, before scaling: 0x25c28f5c / 0x80000000 = 29.50%
+; CHECK: After scaling: 0x7ddddddf / 0x80000000 = 98.33%
+; CHECK: Scale the probablity for one cluster, before scaling: 0x003126e9 / 0x80000000 = 0.15%
+; CHECK: After scaling: 0x00a3d709 / 0x80000000 = 0.50%
+; CHECK: Case clusters: -8826 8 129 18312
+
+sw.bb1:
+  br label %return
+
+sw.bb3:
+  br label %return
+
+sw.bb5:
+  br label %return
+
+sw.bb7:
+  br label %return
+
+sw.epilog:
+  br label %return
+
+return:
+  %retval = phi i32 [ 0, %sw.epilog ], [ 5, %sw.bb7 ], [ 4, %sw.bb5 ], [ 3, %sw.bb3 ], [ 2, %sw.bb1 ], [ 1, %entry ]
+  ret i32 %retval
+}
+
+!1 = !{!"function_entry_count", i64 100000}
+!2 = !{!"branch_weights", i32 50, i32 100, i32 200, i32 29500, i32 70000, i32 150}
+

test/CodeGen/X86/switch.ll

-; RUN: llc -mtriple=x86_64-linux-gnu %s -o - -jump-table-density=40 -verify-machineinstrs | FileCheck %s
+; RUN: llc -mtriple=x86_64-linux-gnu %s -o - -jump-table-density=40 -dominant-case-percent-threshold=101 -verify-machineinstrs | FileCheck %s
 ; RUN: llc -mtriple=x86_64-linux-gnu %s -o - -O0 -jump-table-density=40 -verify-machineinstrs | FileCheck --check-prefix=NOOPT %s
 
 declare void @g(i32)
 
 define void @basic(i32 %x) {
 entry:
   switch i32 %x, label %return [
     i32 3, label %bb0