Switch lowering: omit range check for bit tests when default is unreachable (PR43129)

[hans]

Summary

This is modeled after the same functionality for jump tables, which was added in r357067.

[spatel]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2622

Please pre-commit this and other cosmetic diffs.

llvm/test/CodeGen/X86/switch-bt.ll

161–179

Please add this test with baseline CHECK lines as a pre-commit for this patch, so we show the functional difference from this patch.

Does this test cover both sides of the 'if' in the new code?

(MBB != NextBlock(SwitchBB))

and

(MBB == NextBlock(SwitchBB))

If not, can you add another test to provide that coverage?

[nickdesaulniers]

Thanks for the patch! Just some small thoughts on style.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2673

Please prefer:

if true:
  foo()
else:
  bar()

to

if !true:
  bar()
else
  foo()

2677

If you declared a SDValue before the conditional, you could simply initialize it in these branches, and share the setRoot assignment after the conditional.

[hans]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2622

r373191.

llvm/test/CodeGen/X86/switch-bt.ll

161–179

Pre-committed a test in r373190.

I think constructing a test where MBB != NextBlock(SwitchBB) would be very tricky.

[hans]

Rebasing and addressing comments.

Please take another look.

[spatel]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2624–2626

Can we sink RangeCmp into the !OmitRangeCheck block below? As written, we are potentially creating a node with no uses?

2678–2679

Is the optional branch to MBB identical on both sides of this if/else? If so, sink it after the if/else (and that simplifies the if/else to just !OmitRangeCheck)?

[nickdesaulniers]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2686

This statement and conditional matches L2667. I think we could flip the conditionals/check for this first? DRY

[hans]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2624–2626

It's a bit annoying, because by the time we're down here, the code above has changed the value of Sub which we would use in RangeCmp. But that can be fixed of course.. doing that.

[hans]

Please take another look.

[xbolva00]

llvm/test/CodeGen/X86/switch-bt.ll

165

Can you (re)generate CHECKs with update_llc_test_checks.py?

[hans]

llvm/test/CodeGen/X86/switch-bt.ll

165

Why? It's a hand-written test, and the point of the patch is specifically to omit the cmp here.

[nickdesaulniers]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2650

So it seems to me that CopyTo might not even end up as the root. It looks like we have 3 cases that affect the root:

1. !B.OmitRangeCheck
2. MBB != NextBlock(SwitchBB)
3. Neither 1 or 2

For cases 1 and 2, we do all this wasteful work to create CopyTo and Sub, just to throw it away. It seems like RangeSub isn't needed for case 2. I think it would be clearer and do less work for each case to have a simple 3 arm conditional, rather than successive conditionals that may just throw away previous work to reassign to Root.

[hans]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2650

Note that Root below isn't just assigned to, it's also passed as an argument when creating BRCOND or the BR node.

So in the cases where CopyTo doesn't end up being the root, it's been passed as the chain element to one of those instructions. In all cases, it is on the chain -- it's never thrown away. (And neither is RangeSub since it's an input to CopyTo.)

I realize that's maybe not super clear, but I think this style of updating the selection dag root by using the old root as the chain of the new root is common.

[nickdesaulniers]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

2650

yep, sorry I missed that.

[spatel]

LGTM

[RKSimon]

@hans This is causing EXPENSIVE_CHECKS failures, please can you take a look? http://lab.llvm.org:8011/builders/llvm-clang-x86_64-expensive-checks-win/builds/19967

[hans]

In D68131#1691127, @RKSimon wrote:

@hans This is causing EXPENSIVE_CHECKS failures, please can you take a look? http://lab.llvm.org:8011/builders/llvm-clang-x86_64-expensive-checks-win/builds/19967

Sorry about that. Reverted in r373454 until I can investigate it properly. Hopefully it's just a simple fix :-)

[hans]

Re-landed with fix in r373477.

llvm/include/llvm/CodeGen/SwitchLoweringUtils.h

struct BitTestBlock {
   MVT RegVT;
   bool Emitted;
   bool ContiguousRange;
   MachineBasicBlock *Parent;
   MachineBasicBlock *Default;
   BitTestInfo Cases;
   BranchProbability Prob;
   BranchProbability DefaultProb;
+  bool OmitRangeCheck;
 
   BitTestBlock(APInt F, APInt R, const Value *SV, unsigned Rg, MVT RgVT, bool E,
                bool CR, MachineBasicBlock *P, MachineBasicBlock *D,
                BitTestInfo C, BranchProbability Pr)
       : First(std::move(F)), Range(std::move(R)), SValue(SV), Reg(Rg),
         RegVT(RgVT), Emitted(E), ContiguousRange(CR), Parent(P), Default(D),
-        Cases(std::move(C)), Prob(Pr) {}
+        Cases(std::move(C)), Prob(Pr), OmitRangeCheck(false) {}
 };
 
 /// Return the range of values within a range.
 uint64_t getJumpTableRange(const CaseClusterVector &Clusters, unsigned First,
                            unsigned Last);
 
 /// Return the number of cases within a range.
 uint64_t getJumpTableNumCases(const SmallVectorImpl<unsigned> &TotalCases,

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 }
 
 /// visitBitTestHeader - This function emits necessary code to produce value
 /// suitable for "bit tests"
 void SelectionDAGBuilder::visitBitTestHeader(BitTestBlock &B,
                                              MachineBasicBlock *SwitchBB) {
   SDLoc dl = getCurSDLoc();
 
   // Subtract the minimum value.
   SDValue SwitchOp = getValue(B.SValue);
   EVT VT = SwitchOp.getValueType();
-  SDValue Sub = DAG.getNode(ISD::SUB, dl, VT, SwitchOp,
-                            DAG.getConstant(B.First, dl, VT));
-
-  // Check range.
-  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
-  SDValue RangeCmp = DAG.getSetCC(
-      dl, TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(),
-                                 Sub.getValueType()),
-      Sub, DAG.getConstant(B.Range, dl, VT), ISD::SETUGT);
+  SDValue RangeSub =
+      DAG.getNode(ISD::SUB, dl, VT, SwitchOp, DAG.getConstant(B.First, dl, VT));
 
   // Determine the type of the test operands.
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   bool UsePtrType = false;
   if (!TLI.isTypeLegal(VT)) {
     UsePtrType = true;
   } else {
     for (unsigned i = 0, e = B.Cases.size(); i != e; ++i)
       if (!isUIntN(VT.getSizeInBits(), B.Cases[i].Mask)) {
         // Switch table case range are encoded into series of masks.
         // Just use pointer type, it's guaranteed to fit.
         UsePtrType = true;
         break;
       }
   }
+  SDValue Sub = RangeSub;
   if (UsePtrType) {
     VT = TLI.getPointerTy(DAG.getDataLayout());
     Sub = DAG.getZExtOrTrunc(Sub, dl, VT);
   }
 
   B.RegVT = VT.getSimpleVT();
   B.Reg = FuncInfo.CreateReg(B.RegVT);
   SDValue CopyTo = DAG.getCopyToReg(getControlRoot(), dl, B.Reg, Sub);
 
   MachineBasicBlock* MBB = B.Cases[0].ThisBB;
 
   addSuccessorWithProb(SwitchBB, B.Default, B.DefaultProb);
   addSuccessorWithProb(SwitchBB, MBB, B.Prob);
   SwitchBB->normalizeSuccProbs();
 
-  SDValue BrRange = DAG.getNode(ISD::BRCOND, dl,
-                                MVT::Other, CopyTo, RangeCmp,
-                                DAG.getBasicBlock(B.Default));
+  SDValue Root = CopyTo;
+  if (!B.OmitRangeCheck) {
+    // Conditional branch to the default block.
+    SDValue RangeCmp = DAG.getSetCC(dl,
+        TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(),
+                               RangeSub.getValueType()),
+        RangeSub, DAG.getConstant(B.Range, dl, RangeSub.getValueType()),
+        ISD::SETUGT);
+
+    Root = DAG.getNode(ISD::BRCOND, dl, MVT::Other, Root, RangeCmp,
+  }
 
   // Avoid emitting unnecessary branches to the next block.
   if (MBB != NextBlock(SwitchBB))
-    BrRange = DAG.getNode(ISD::BR, dl, MVT::Other, BrRange,
-                          DAG.getBasicBlock(MBB));
+    Root = DAG.getNode(ISD::BR, dl, MVT::Other, Root, DAG.getBasicBlock(MBB));
 
-  DAG.setRoot(BrRange);
+  DAG.setRoot(Root);
 }
 
 /// visitBitTestCase - this function produces one "bit test"
 void SelectionDAGBuilder::visitBitTestCase(BitTestBlock &BB,
                                            MachineBasicBlock* NextMBB,
                                            BranchProbability BranchProbToNext,
                                            unsigned Reg,
                                            BitTestCase &B,
                                            MachineBasicBlock *SwitchBB) {
   SDLoc dl = getCurSDLoc();
   MVT VT = BB.RegVT;
   SDValue ShiftOp = DAG.getCopyFromReg(getControlRoot(), dl, Reg, VT);
   SDValue Cmp;
   unsigned PopCount = countPopulation(B.Mask);
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   if (PopCount == 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(
switch (I->Kind) {
         // If we're in the right place, emit the jump table header right now.
         if (CurMBB == SwitchMBB) {
           visitJumpTableHeader(*JT, *JTH, SwitchMBB);
           JTH->Emitted = true;
         }
         break;
       }
       case CC_BitTests: {
-        // FIXME: If Fallthrough is unreachable, skip the range check.
-
         // FIXME: Optimize away range check based on pivot comparisons.
         BitTestBlock *BTB = &SL->BitTestCases[I->BTCasesIndex];
 
         // The bit test blocks haven't been inserted yet; insert them here.
         for (BitTestCase &BTC : BTB->Cases)
           CurMF->insert(BBI, BTC.ThisBB);
 
         // Fill in fields of the BitTestBlock.
         BTB->Parent = CurMBB;
         BTB->Default = Fallthrough;
 
         BTB->DefaultProb = UnhandledProbs;
         // If the cases in bit test don't form a contiguous range, we evenly
         // distribute the probability on the edge to Fallthrough to two
         // successors of CurMBB.
         if (!BTB->ContiguousRange) {
           BTB->Prob += DefaultProb / 2;
           BTB->DefaultProb -= DefaultProb / 2;
         }
 
+        if (FallthroughUnreachable) {
+          // Skip the range check if the fallthrough block is unreachable.
+          BTB->OmitRangeCheck = true;
+        }
+
         // If we're in the right place, emit the bit test header right now.
         if (CurMBB == SwitchMBB) {
           visitBitTestHeader(*BTB, SwitchMBB);
           BTB->Emitted = true;
         }
         break;
       }
       case CC_Range: {

llvm/test/CodeGen/X86/switch-bt.ll

 ; CHECK: je
 ; CHECK: cmpl $50
 ; CHECK: je
 ; CHECK: cmpl $60
 ; CHECK: jne
 }
 
 
-; TODO: Omit the range check when the default case is unreachable, see PR43129.
+; Omit the range check when the default case is unreachable, see PR43129.
 declare void @g(i32)
 define void @test5(i32 %x) {
 
 ; CHECK-LABEL: test5
-; CHECK: cmpl $8, %edi
-; CHECK: ja
+; CHECK-NOT: cmp
 
 ; 73 = 2^0 + 2^3 + 2^6
 ; CHECK:      movl $73
 ; CHECK-NEXT: btl
 ; CHECK-NEXT: jb
 
 ; 146 = 2^1 + 2^4 + 2^7
 ; CHECK-NEXT: movl $146
 ; CHECK-NEXT: btl
 ; CHECK-NEXT: jae
 
 
 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