From 84145dcd0872dba90bc280c08e39e1fe506b8ec0 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans@hanshq.net>
Date: Mon, 27 Apr 2015 20:21:17 +0000
Subject: [PATCH] Switch lowering: order bit tests by branch weight.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@235912 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../SelectionDAG/SelectionDAGBuilder.cpp      |  5 +-
 test/CodeGen/X86/switch.ll                    | 46 +++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index fd59d57580b..d7edd2c1966 100644
--- 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;
   });
 
diff --git a/test/CodeGen/X86/switch.ll b/test/CodeGen/X86/switch.ll
index 9baee19a2a6..fd386e1f821 100644
--- a/test/CodeGen/X86/switch.ll
+++ b/test/CodeGen/X86/switch.ll
@@ -367,3 +367,49 @@ return: ret void
 ; CHECK-LABEL: int_max_table_cluster
 ; CHECK: jmpq *.LJTI
 }
+
+
+define void @bt_order_by_weight(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 5, label %bb2
+    i32 8, label %bb2
+    i32 9, label %bb2
+  ], !prof !1
+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.
+
+; 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
+}
+
+!1 = !{!"branch_weights",
+       ; Default:
+       i32 1,
+       ; Cases 0,3,6:
+       i32 4, i32 4, i32 4,
+       ; Cases 1,4,7:
+       i32 4294967295, i32 2, i32 4294967295,
+       ; Cases 2,5,8,9:
+       i32 3, i32 3, i32 3, i32 3}
-- 
2.34.1