Switch lowering: use profile info to build weight-balanced binary search trees

This will cause hot nodes to appear closer to the root.

The literature says building the tree like this makes it a near-optimal (in
terms of search time given key frequencies) binary search tree. In LLVM's case,
we can do up to 3 comparisons in each leaf node, so it might be better to opt
for lower tree height in some cases; that's something to look into in the
future.

Differential Revision: http://reviews.llvm.org/D9318

llvm-svn: 236192

1 files changed, 35 insertions, 7 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index db41c407b03..1c14d4d8d5f 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -7986,14 +7986,41 @@ void SelectionDAGBuilder::splitWorkItem(SwitchWorkList &WorkList,
   unsigned NumClusters = W.LastCluster - W.FirstCluster + 1;
   assert(NumClusters >= 2 && "Too small to split!");
 
-  // FIXME: When we have profile info, we might want to balance the tree based
-  // on weights instead of node count.
+  // Balance the tree based on branch weights to create a near-optimal (in terms
+  // of search time given key frequency) binary search tree. See e.g. Kurt
+  // Mehlhorn "Nearly Optimal Binary Search Trees" (1975).
+  CaseClusterIt LastLeft = W.FirstCluster;
+  CaseClusterIt FirstRight = W.LastCluster;
+  uint32_t LeftWeight = LastLeft->Weight;
+  uint32_t RightWeight = FirstRight->Weight;
+
+  // Move LastLeft and FirstRight towards each other from opposite directions to
+  // find a partitioning of the clusters which balances the weight on both
+  // sides.
+  while (LastLeft + 1 < FirstRight) {
+    // Zero-weight nodes would cause skewed trees since they don't affect
+    // LeftWeight or RightWeight.
+    assert(LastLeft->Weight != 0);
+    assert(FirstRight->Weight != 0);
+
+    if (LeftWeight < RightWeight)
+      LeftWeight += (++LastLeft)->Weight;
+    else
+      RightWeight += (--FirstRight)->Weight;
+  }
+  assert(LastLeft + 1 == FirstRight);
+  assert(LastLeft >= W.FirstCluster);
+  assert(FirstRight <= W.LastCluster);
+
+  // Use the first element on the right as pivot since we will make less-than
+  // comparisons against it.
+  CaseClusterIt PivotCluster = FirstRight;
+  assert(PivotCluster > W.FirstCluster);
+  assert(PivotCluster <= W.LastCluster);
 
-  CaseClusterIt PivotCluster = W.FirstCluster + NumClusters / 2;
   CaseClusterIt FirstLeft = W.FirstCluster;
-  CaseClusterIt LastLeft = PivotCluster - 1;
-  CaseClusterIt FirstRight = PivotCluster;
   CaseClusterIt LastRight = W.LastCluster;
+
   const ConstantInt *Pivot = PivotCluster->Low;
 
   // New blocks will be inserted immediately after the current one.
@@ -8032,7 +8059,8 @@ void SelectionDAGBuilder::splitWorkItem(SwitchWorkList &WorkList,
   }
 
   // Create the CaseBlock record that will be used to lower the branch.
-  CaseBlock CB(ISD::SETLT, Cond, Pivot, nullptr, LeftMBB, RightMBB, W.MBB);
+  CaseBlock CB(ISD::SETLT, Cond, Pivot, nullptr, LeftMBB, RightMBB, W.MBB,
+               LeftWeight, RightWeight);
 
   if (W.MBB == SwitchMBB)
     visitSwitchCase(CB, SwitchMBB);
@@ -8048,7 +8076,7 @@ void SelectionDAGBuilder::visitSwitch(const SwitchInst &SI) {
   for (auto I : SI.cases()) {
     MachineBasicBlock *Succ = FuncInfo.MBBMap[I.getCaseSuccessor()];
     const ConstantInt *CaseVal = I.getCaseValue();
-    uint32_t Weight = 0; // FIXME: Use 1 instead?
+    uint32_t Weight = 1;
     if (BPI) {
       Weight = BPI->getEdgeWeight(SI.getParent(), I.getSuccessorIndex());
       assert(Weight <= UINT32_MAX / SI.getNumSuccessors());