Switch lowering: Take branch weight into account when ordering for fall-through

[oota-llvm.git] / lib / CodeGen / SelectionDAG / SelectionDAGBuilder.cpp

diff --git a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index fd59d57580bcfc74457208a44feca1e4b0833457..985a4dfd49d76adc0412b6df6ebf07538e4ebfcf 100644 (file)
--- a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -7482,12 +7482,15 @@ bool SelectionDAGBuilder::buildBitTests(CaseClusterVector &Clusters,
       CB->Bits++;
     }
     CB->ExtraWeight += Clusters[i].Weight;
+    assert(CB->ExtraWeight >= Clusters[i].Weight && "Weight sum overflowed!");
     TotalWeight += Clusters[i].Weight;
   }
 
   BitTestInfo BTI;
   std::sort(CBV.begin(), CBV.end(), [](const CaseBits &a, const CaseBits &b) {
-    // FIXME: Sort by weight.
+    // Sort by weight first, number of bits second.
+    if (a.ExtraWeight != b.ExtraWeight)
+      return a.ExtraWeight > b.ExtraWeight;
     return a.Bits > b.Bits;
   });
 
@@ -7675,11 +7678,12 @@ void SelectionDAGBuilder::lowerWorkItem(SwitchWorkListItem W, Value *Cond,
       return a.Weight > b.Weight;
     });
 
-    // Rearrange the case blocks so that the last one falls through if possible.
-    // Start at the bottom as that's the case with the lowest weight.
-    // FIXME: Take branch probability into account.
+    // Rearrange the case blocks so that the last one falls through if possible
+    // without without changing the order of weights.
     for (CaseClusterIt I = W.LastCluster; I > W.FirstCluster; ) {
       --I;
+      if (I->Weight > W.LastCluster->Weight)
+        break;
       if (I->Kind == CC_Range && I->MBB == NextMBB) {
         std::swap(*I, *W.LastCluster);
         break;