Peel off the dominant case in switch statement

[xur]

Summary

This patch peels off the dominant case in switch statement into a branch. This avoids the extra compares when lowering the switch statement into chain of branches in binary search tree style.

For a switch statement like the following:

switch (i) {
case     8: s1; break;
case -8826: s2; break;
case 18312: s3; break;
case 18568: s4; break;
case   129: s5; break;
}

The lowering generates the following coe

1. BB#0: # %entry cmpl $128, %edi jg .LBB0_4
2. BB#1: # %entry movl $1, %eax cmpl $8, %edi jne .LBB0_2
3. BB#10: # %return retq

.LBB0_4: # %entry

cmpl    $129, %edi

Here the coml to $128 is extra. This is happening even if i==8 target has only non-zero branch-weights. Like !{!"branch_weights", i32 0, i32 100000, i32 0, i32 0, i32 0, i32 0}

This patch would peel i==8 into a branch and keep the rest in the switch statement.

[hans]

lib/Transforms/Utils/SimplifyCFG.cpp

5680

I don't think we should do this if SwitchToLookupTable (the call below) can turn the swich into a lookup instead.

In fact, I don't think this should be done at the IR level at all, since it's more of a switch lowering issue. SelectionDAGBuilder::visitSwitch() would be a better place.

That code already takes case weights into account, and when lowering to a binary search tree, it will balance it based on weight, favoring cases that are hot. Do you find that that's not sufficient? (I'm willing to believe that's the case, but I'd like to see it argued.)

[davidxl]

lib/Transforms/Utils/SimplifyCFG.cpp

5680

Peeling the dominating case avoids a memory read (table lookup). Besides it trades a indirect branch for a a highly biased direct branch (which is usually highly predictable by the branch predictor).

Like looping which happens at higher level, switch peeling can also enable more possible surrounding optimizations (eg like jump threading). Peeling it out also probably also makes it easier to for better code layout.

On the other hand, this should not be done when size optimization is on.

[• dberlin]

We've already discovered a slew of problem from over-optimizing things in simplifycfg, which is also being used as a canonicalization pass.
Either it needs a set of flags controlling what it does, or it needs to not be trying to do "performance improvement" and "canonicalization" as a single pass.
See, e.g, https://bugs.llvm.org/show_bug.cgi?id=34603

[hans]

lib/Transforms/Utils/SimplifyCFG.cpp

5680

SwitchToLookupTable creates lookup tables for switches that are used to select from a set of constant values (this is different from jump tables, sorry the names here are confusing), so there is no indirect branch. I think peeling off the common case is probably not a good idea for those kinds of switches.

I can see that peeling early might enable other high-level optimizations. But that's true in some sense for switch lowering in general. I'd suggest to at least look into how peeling in the regular switch lowering code at SelectionDAGBuilder::visitSwitch() would look like. It seems to me like the natural place to do it, and I suspect it's easier to do there.

lib/Transforms/Utils/SimplifyCFG.cpp

static cl::opt<bool> SpeculateOneExpensiveInst(
     cl::desc("Allow exactly one expensive instruction to be speculatively "
              "executed"));
 
 static cl::opt<unsigned> MaxSpeculationDepth(
     "max-speculation-depth", cl::Hidden, cl::init(10),
     cl::desc("Limit maximum recursion depth when calculating costs of "
              "speculatively executed instructions"));
 
+static cl::opt<bool> PeelDominantSwitchCase(
+    "peel-dominant-switch-case", cl::Hidden, cl::init(true),
+    cl::desc("peel off the dominant case in switch statement"));
+
+static cl::opt<unsigned> DominantSwitchCaseThreshold(
+    "dominant-switch-case-threshold", cl::Hidden, cl::init(80),
+    cl::desc("The case is considered to be dominant if its block weight is "
+             "greater than this percent of the total block weights"));
+
 STATISTIC(NumBitMaps, "Number of switch instructions turned into bitmaps");
 STATISTIC(NumLinearMaps,
           "Number of switch instructions turned into linear mapping");
 STATISTIC(NumLookupTables,
           "Number of switch instructions turned into lookup tables");
 STATISTIC(
     NumLookupTablesHoles,
     "Number of switch instructions turned into lookup tables (holes checked)");
 STATISTIC(NumTableCmpReuses, "Number of reused switch table lookup compares");
 STATISTIC(NumSinkCommons,
           "Number of common instructions sunk down to the end block");
 STATISTIC(NumSpeculations, "Number of speculative executed instructions");
+STATISTIC(NumPeeledCases,
+          "Number of dominant cases peeled in switch statements");
 
 namespace {
 
 // The first field contains the value that the switch produces when a certain
 // case group is selected, and the second field is a vector containing the
 // cases composing the case group.
 typedef SmallVector<std::pair<Constant *, SmallVector<ConstantInt *, 4>>, 2>
     SwitchCaseResultVectorTy;
for (unsigned i = 1, e = MD->getNumOperands(); i < e; ++i) {
     Weights.push_back(CI->getValue().getZExtValue());
   }
 
   // If TI is a conditional eq, the default case is the false case,
   // and the corresponding branch-weight data is at index 2. We swap the
   // default weight to be the first entry.
   if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
     assert(Weights.size() == 2);
-    ICmpInst *ICI = cast<ICmpInst>(BI->getCondition());
+    if (ICmpInst *ICI = dyn_cast<ICmpInst>(BI->getCondition())) {
       if (ICI->getPredicate() == ICmpInst::ICMP_EQ)
         std::swap(Weights.front(), Weights.back());
     }
   }
+}
 
 /// Keep halving the weights until all can fit in uint32_t.
 static void FitWeights(MutableArrayRef<uint64_t> Weights) {
   uint64_t Max = *std::max_element(Weights.begin(), Weights.end());
   if (Max > UINT_MAX) {
     unsigned Offset = 32 - countLeadingZeros(Max);
     for (uint64_t &I : Weights)
       I >>= Offset;
   }
 }
 
+static bool WeightIsDominant(uint64_t Weight, uint64_t TotalWeight) {
+  assert(Weight <= TotalWeight);
+  return (TotalWeight / 100 < Weight / DominantSwitchCaseThreshold);
+}
+
 /// The specified terminator is a value equality comparison instruction
 /// (either a switch or a branch on "X == c").
 /// See if any of the predecessors of the terminator block are value comparisons
 /// on the same value.  If so, and if safe to do so, fold them together.
 bool SimplifyCFGOpt::FoldValueComparisonIntoPredecessors(TerminatorInst *TI,
                                                          IRBuilder<> &Builder) {
   BasicBlock *BB = TI->getParent();
   Value *CV = isValueEqualityComparison(TI); // CondVal
   assert(CV && "Not a comparison?");
   bool Changed = false;
 
   SmallVector<BasicBlock *, 16> Preds(pred_begin(BB), pred_end(BB));
   while (!Preds.empty()) {
     BasicBlock *Pred = Preds.pop_back_val();
 
-    // See if the predecessor is a comparison with the same value.
+    // Don't fold with predessors that having strongly biased branches. With
+    // PeelDominantSwitchCase on, we will peel them off.
     TerminatorInst *PTI = Pred->getTerminator();
-    Value *PCV = isValueEqualityComparison(PTI); // PredCondVal
+    SmallVector<uint64_t, 8> Weights;
+    bool PredHasWeights = HasBranchWeights(PTI);
+    if (PredHasWeights) {
+      GetBranchWeights(PTI, Weights);
+      if (PeelDominantSwitchCase && dyn_cast<BranchInst>(PTI)) {
+        assert(Weights.size() == 2);
+        uint64_t TotalWeight = Weights[0] + Weights[1];
+        if (WeightIsDominant(Weights[0], TotalWeight) ||
+            WeightIsDominant(Weights[1], TotalWeight))
+          return false;
+      }
+    }
 
+    // See if the predecessor is a comparison with the same value.
+    Value *PCV = isValueEqualityComparison(PTI); // PredCondVal
     if (PCV == CV && TI != PTI) {
       SmallSetVector<BasicBlock*, 4> FailBlocks;
       if (!SafeToMergeTerminators(TI, PTI, &FailBlocks)) {
         for (auto *Succ : FailBlocks) {
           if (!SplitBlockPredecessors(Succ, TI->getParent(), ".fold.split"))
             return false;
         }
       }
 
       // Figure out which 'cases' to copy from SI to PSI.
       std::vector<ValueEqualityComparisonCase> BBCases;
       BasicBlock *BBDefault = GetValueEqualityComparisonCases(TI, BBCases);
 
       std::vector<ValueEqualityComparisonCase> PredCases;
       BasicBlock *PredDefault = GetValueEqualityComparisonCases(PTI, PredCases);
 
       // Based on whether the default edge from PTI goes to BB or not, fill in
       // PredCases and PredDefault with the new switch cases we would like to
       // build.
       SmallVector<BasicBlock *, 8> NewSuccessors;
 
       // Update the branch weight metadata along the way
-      SmallVector<uint64_t, 8> Weights;
-      bool PredHasWeights = HasBranchWeights(PTI);
       bool SuccHasWeights = HasBranchWeights(TI);
 
       if (PredHasWeights) {
-        GetBranchWeights(PTI, Weights);
         // branch-weight metadata is inconsistent here.
         if (Weights.size() != 1 + PredCases.size())
           PredHasWeights = SuccHasWeights = false;
       } else if (SuccHasWeights)
         // If there are no predecessor weights but there are successor weights,
         // populate Weights with 1, which will later be scaled to the sum of
         // successor's weights
         Weights.assign(1 + PredCases.size(), 1);
for (ForwardingNodesMap::iterator I = ForwardingNodes.begin(),
     for (size_t I = 0, E = Indexes.size(); I != E; ++I)
       Phi->setIncomingValue(Indexes[I], SI->getCondition());
     Changed = true;
   }
 
   return Changed;
 }
 
+// Peel off the dominant case in the switch statement.
+// For a switch stament like the following:
+//   switch %cond, label %default [
+//     constant1: label %c1
+//     constant2: label %c2
+//     constant3: label %c3
+//     ...
+//   ],
+// If label_c2 is the domninant target, transform it into:
+//   %cmp = icmp eq %cond, constant2
+//   br %cmp, label %c2, label %peeled.switch
+//   peeled.switch:
+//     switch %cond, label %default [
+//       constant1: label %c1
+//       constant3: label %c3
+//       ...
+//     ],
+static bool PeelDominantCase(SwitchInst *SI, IRBuilder<> &Builder) {
+  if (!PeelDominantSwitchCase)
+    return false;
+  // Collect branch weights into a vector.
+  SmallVector<uint32_t, 8> Weights;
+  MDNode *MD = SI->getMetadata(LLVMContext::MD_prof);
+  bool HasWeight = MD && (MD->getNumOperands() == 2 + SI->getNumCases());
+  if (!HasWeight)
+    return false;
+
+  uint32_t MaxWeight = 0;
+  uint64_t TotalWeight = MaxWeight;
+  uint32_t MaxWeightIndex;
+  for (unsigned MD_i = 1, MD_e = MD->getNumOperands(); MD_i < MD_e; ++MD_i) {
+    ConstantInt *CI = mdconst::extract<ConstantInt>(MD->getOperand(MD_i));
+    uint32_t Weight = CI->getValue().getZExtValue();
+    TotalWeight += Weight;
+    if (Weight > MaxWeight) {
+      MaxWeight = Weight;
+      MaxWeightIndex = MD_i - 1;
+    }
+    Weights.push_back(CI->getValue().getZExtValue());
+  }
+
+  // Don't do anything if the default has the max weight.
+  if (MaxWeightIndex == 0)
+    return false;
+  // Don't perform if there is no dominent case.
+  if (!WeightIsDominant(MaxWeight, TotalWeight))
+    return false;
+  // Don't perform if this is already peeled.
+  BasicBlock *BB = SI->getParent();
+  if (BasicBlock *Pred = BB->getSinglePredecessor()) {
+    if (BranchInst *BI = dyn_cast<BranchInst>(Pred->getTerminator())) {
+      ICmpInst *ICI = cast<ICmpInst>(BI->getCondition());
+      if (ICI && ICI->getPredicate() == ICmpInst::ICMP_EQ &&
+          ICI->getOperand(0) == SI->getCondition() &&
+          isa<ConstantInt>(ICI->getOperand(1)))
+        return false;
+    }
+  }
+
+  DEBUG(dbgs() << "Peel DominantCase case.\n== Basic Block Before ==" << *BB);
+
+  Weights.erase(Weights.begin() + MaxWeightIndex);
+  uint32_t Index = 1;
+  ConstantInt *CaseVal = nullptr;
+  BasicBlock *CaseSuccessor = nullptr;
+  for (auto i = SI->case_begin(), e = SI->case_end(); i != e; ++i, ++Index) {
+    if (Index == MaxWeightIndex) {
+      CaseVal = i->getCaseValue();
+      CaseSuccessor = i->getCaseSuccessor();
+      SI->removeCase(i);
+      SI->setMetadata(LLVMContext::MD_prof,
+                      MDBuilder(SI->getParent()->getContext())
+                          .createBranchWeights(Weights));
+      break;
+    }
+  }
+  assert(CaseVal != nullptr);
+
+  Value *Cmp =
+      Builder.CreateICmpEQ(SI->getCondition(), CaseVal, "switch.peeledcmp");
+  MDNode *BranchWeights =
+      MDBuilder(SI->getContext())
+          .createBranchWeights(MaxWeight, TotalWeight - MaxWeight);
+  BasicBlock *NewSwitchBB = BB->splitBasicBlock(SI->getIterator());
+  NewSwitchBB->setName("peeled.switch");
+
+  TerminatorInst *OldTI = BB->getTerminator();
+  IRBuilder<> CmpBB_Builder(NewSwitchBB->getSinglePredecessor());
+  CmpBB_Builder.CreateCondBr(Cmp, CaseSuccessor, NewSwitchBB, BranchWeights);
+  OldTI->eraseFromParent();
+
+  DEBUG(dbgs() << "== Basic Block After ==" << *BB << *(SI->getParent())
+               << "\n");
+  NumPeeledCases++;
+  return true;
+}
+
 /// Return true if the backend will be able to handle
 /// initializing an array of constants like C.
 static bool ValidLookupTableConstant(Constant *C, const TargetTransformInfo &TTI) {
   if (C->isThreadDependent())
     return false;
   if (C->isDLLImportDependent())
     return false;
 
if (EliminateDeadSwitchCases(SI, AC, DL))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   if (SwitchToSelect(SI, Builder, AC, DL, TTI))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
   if (ForwardSwitchConditionToPHI(SI))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 
+  if (PeelDominantCase(SI, Builder))
+    return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
+
   // The conversion from switch to lookup tables results in difficult
   // to analyze code and makes pruning branches much harder.
   // This is a problem of the switch expression itself can still be
   // restricted as a result of inlining or CVP. There only apply this
   // transformation during late steps of the optimisation chain.
   if (LateSimplifyCFG && SwitchToLookupTable(SI, Builder, DL, TTI))
     return SimplifyCFG(BB, TTI, BonusInstThreshold, AC) | true;
 

test/Transforms/SimplifyCFG/peel_switch_case.ll

+; RUN: opt -simplifycfg -S < %s | FileCheck %s
+
+define i32 @foo(i16 signext %n) local_unnamed_addr !prof !1 {
+entry:
+  %conv = sext i16 %n to i32
+  switch i32 %conv, label %sw.epilog [
+    i32 8, label %return
+    i32 -8826, label %sw.bb1
+    i32 18312, label %sw.bb3
+    i32 18568, label %sw.bb5
+    i32 129, label %sw.bb7
+  ], !prof !2
+; CHECK: %switch.peeledcmp = icmp eq i32 %conv, 8
+; CHECK: br i1 %switch.peeledcmp, label %return, label %peeled.switch, !prof ![[BRANCH_COUNT:[0-9]+]]
+; CHECK: peeled.switch:
+; CHECK-NOT: i32 8, label %return
+; CHECK: ], !prof ![[SWITCH_COUNT:[0-9]+]]
+
+
+sw.bb1:
+  br label %return
+
+sw.bb3:
+  br label %return
+
+sw.bb5:
+  br label %return
+
+sw.bb7:
+  br label %return
+
+sw.epilog:
+  br label %return
+
+return:
+  %retval = phi i32 [ 0, %sw.epilog ], [ 5, %sw.bb7 ], [ 4, %sw.bb5 ], [ 3, %sw.bb3 ], [ 2, %sw.bb1 ], [ 1, %entry ]
+  ret i32 %retval
+}
+
+!1 = !{!"function_entry_count", i64 101000}
+!2 = !{!"branch_weights", i32 0, i32 100000, i32 0, i32 1000, i32 0, i32 0}
+; CHECK: ![[BRANCH_COUNT]] = !{!"branch_weights", i32 100000, i32 1000}
+; CHECK: ![[SWITCH_COUNT]] = !{!"branch_weights", i32 0, i32 0, i32 1000, i32 0, i32 0}
+

test/Transforms/SimplifyCFG/preserve-branchweights-switch-create.ll

-; RUN: opt -simplifycfg -S -o - < %s | FileCheck %s
+; RUN: opt -simplifycfg -peel-dominant-switch-case=false -S -o - < %s | FileCheck %s
 
 declare void @func2(i32)
 declare void @func4(i32)
 declare void @func6(i32)
 declare void @func8(i32)
 
 ;; test1 - create a switch with case 2 and case 4 from two branches: N == 2
 ;; and N == 4.