From 7953fd5824b6ebc743cb7296e20ba453d0cc6b9e Mon Sep 17 00:00:00 2001
From: Dan Gohman <gohman@apple.com>
Date: Thu, 24 Jun 2010 02:06:24 +0000
Subject: [PATCH] Optimize the "bit test" code path for switch lowering in the
 case where the bit mask has exactly one bit.

llvm-svn: 106716
---
 .../SelectionDAG/SelectionDAGBuilder.cpp      | 40 ++++++++++-----
 test/CodeGen/X86/switch-bt.ll                 | 51 +++++++++++++++++++
 2 files changed, 77 insertions(+), 14 deletions(-)
 create mode 100644 test/CodeGen/X86/switch-bt.ll

diff --git a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index f43c4e4e7e4..fd8fe22f61d 100644
--- a/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -1557,29 +1557,41 @@ void SelectionDAGBuilder::visitBitTestCase(MachineBasicBlock* NextMBB,
                                            unsigned Reg,
                                            BitTestCase &B,
                                            MachineBasicBlock *SwitchBB) {
-  // Make desired shift
   SDValue ShiftOp = DAG.getCopyFromReg(getControlRoot(), getCurDebugLoc(), Reg,
                                        TLI.getPointerTy());
-  SDValue SwitchVal = DAG.getNode(ISD::SHL, getCurDebugLoc(),
-                                  TLI.getPointerTy(),
-                                  DAG.getConstant(1, TLI.getPointerTy()),
-                                  ShiftOp);
+  SDValue Cmp;
+  if (CountPopulation_64(B.Mask) == 1) {
+    // Testing for a single bit; just compare the shift count with what it
+    // would need to be to shift a 1 bit in that position.
+    Cmp = DAG.getSetCC(getCurDebugLoc(),
+                       TLI.getSetCCResultType(ShiftOp.getValueType()),
+                       ShiftOp,
+                       DAG.getConstant(CountTrailingZeros_64(B.Mask),
+                                       TLI.getPointerTy()),
+                       ISD::SETEQ);
+  } else {
+    // Make desired shift
+    SDValue SwitchVal = DAG.getNode(ISD::SHL, getCurDebugLoc(),
+                                    TLI.getPointerTy(),
+                                    DAG.getConstant(1, TLI.getPointerTy()),
+                                    ShiftOp);
 
-  // Emit bit tests and jumps
-  SDValue AndOp = DAG.getNode(ISD::AND, getCurDebugLoc(),
-                              TLI.getPointerTy(), SwitchVal,
-                              DAG.getConstant(B.Mask, TLI.getPointerTy()));
-  SDValue AndCmp = DAG.getSetCC(getCurDebugLoc(),
-                                TLI.getSetCCResultType(AndOp.getValueType()),
-                                AndOp, DAG.getConstant(0, TLI.getPointerTy()),
-                                ISD::SETNE);
+    // Emit bit tests and jumps
+    SDValue AndOp = DAG.getNode(ISD::AND, getCurDebugLoc(),
+                                TLI.getPointerTy(), SwitchVal,
+                                DAG.getConstant(B.Mask, TLI.getPointerTy()));
+    Cmp = DAG.getSetCC(getCurDebugLoc(),
+                       TLI.getSetCCResultType(AndOp.getValueType()),
+                       AndOp, DAG.getConstant(0, TLI.getPointerTy()),
+                       ISD::SETNE);
+  }
 
   SwitchBB->addSuccessor(B.TargetBB);
   SwitchBB->addSuccessor(NextMBB);
 
   SDValue BrAnd = DAG.getNode(ISD::BRCOND, getCurDebugLoc(),
                               MVT::Other, getControlRoot(),
-                              AndCmp, DAG.getBasicBlock(B.TargetBB));
+                              Cmp, DAG.getBasicBlock(B.TargetBB));
 
   // Set NextBlock to be the MBB immediately after the current one, if any.
   // This is used to avoid emitting unnecessary branches to the next block.
diff --git a/test/CodeGen/X86/switch-bt.ll b/test/CodeGen/X86/switch-bt.ll
new file mode 100644
index 00000000000..ed3266ec422
--- /dev/null
+++ b/test/CodeGen/X86/switch-bt.ll
@@ -0,0 +1,51 @@
+; RUN: llc -march=x86-64 -asm-verbose=false < %s | FileCheck %s
+
+; This switch should use bit tests, and the third bit test case is just
+; testing for one possible value, so it doesn't need a bt.
+
+;      CHECK: movabsq $2305843009482129440, %r
+; CHECK-NEXT: btq %rax, %r
+; CHECK-NEXT: jb
+; CHECK-NEXT: movl  $671088640, %e
+; CHECK-NEXT: btq %rax, %r
+; CHECK-NEXT: jb
+; CHECK-NEXT: testq %rax, %r
+; CHECK-NEXT: j
+
+define void @test(i8* %l) nounwind {
+entry:
+  %l.addr = alloca i8*, align 8                   ; <i8**> [#uses=2]
+  store i8* %l, i8** %l.addr
+  %tmp = load i8** %l.addr                        ; <i8*> [#uses=1]
+  %tmp1 = load i8* %tmp                           ; <i8> [#uses=1]
+  %conv = sext i8 %tmp1 to i32                    ; <i32> [#uses=1]
+  switch i32 %conv, label %sw.default [
+    i32 62, label %sw.bb
+    i32 60, label %sw.bb
+    i32 38, label %sw.bb2
+    i32 94, label %sw.bb2
+    i32 61, label %sw.bb2
+    i32 33, label %sw.bb4
+  ]
+
+sw.bb:                                            ; preds = %entry, %entry
+  call void @foo(i32 0)
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry, %entry, %entry
+  call void @foo(i32 1)
+  br label %sw.epilog
+
+sw.bb4:                                           ; preds = %entry
+  call void @foo(i32 3)
+  br label %sw.epilog
+
+sw.default:                                       ; preds = %entry
+  call void @foo(i32 97)
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %sw.default, %sw.bb4, %sw.bb2, %sw.bb
+  ret void
+}
+
+declare void @foo(i32)