[FuncSpec] Split the specialization bonus into CodeSize and Latency.

[labrinea]

Summary

Currently we use a combined metric TargetTransformInfo::TCK_SizeAndLatency
when estimating the specialization bonus. This is suboptimal, and in some
cases erroneous. For example we shouldn't be weighting the codesize decrease
attributed to constant propagation by the block frequency of the dead code.
Instead only the latency savings should be weighted by block frequency. The
total codesize savings from all the specialization arguments should be
deducted from the specialization cost.

[ChuanqiXu]

Sounds good.

[labrinea]

I run some experiments to measure compilation time. It seems if getUserBonus returns std::pair<Cost,Cost> instead of caching CodeSize and Latency to the InstCostVisitor it is slightly faster (perhaps a litle uglier too).
This patch improves geomean of instruction count for llvm-test-suite by -0.016% at O3 and regresses it by +0.07% at LTO. The alternative is -0.036 and +0.064% respectively.

[labrinea]

This revision is better in compile times (instruction count) for llvm test suite: Geomean with O3 is -0.049%, with LTO is +0.062%.

[labrinea]

Improved the description of estimateBasicBlocks

[labrinea]

Ping

[ChuanqiXu]

LGTM.

llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h

struct Spec {
 
   // Cloned function, a specialized version of the original one.
   Function *Clone = nullptr;
 
   // Specialization signature.
   SpecSig Sig;
 
   // Profitability of the specialization.
-  Cost Score;
+  unsigned Score;
 
   // List of call sites, matching this specialization.
   SmallVector<CallBase *> CallSites;
 
-  Spec(Function *F, const SpecSig &S, Cost Score)
+  Spec(Function *F, const SpecSig &S, unsigned Score)
       : F(F), Sig(S), Score(Score) {}
-  Spec(Function *F, const SpecSig &&S, Cost Score)
+  Spec(Function *F, const SpecSig &&S, unsigned Score)
       : F(F), Sig(S), Score(Score) {}
 };
 
+struct Bonus {
+  unsigned CodeSize = 0;
+  unsigned Latency = 0;
+
+  Bonus() = default;
+
+  Bonus(Cost CodeSize, Cost Latency) {
+    int64_t Sz = *CodeSize.getValue();
+    int64_t Ltc = *Latency.getValue();
+
+    assert(Sz >= 0 && Ltc >= 0 && "CodeSize and Latency cannot be negative");
+    // It is safe to down cast since we know the arguments
+    // cannot be negative and Cost is of type int64_t.
+    this->CodeSize = static_cast<unsigned>(Sz);
+    this->Latency = static_cast<unsigned>(Ltc);
+  }
+
+  Bonus &operator+=(const Bonus RHS) {
+    CodeSize += RHS.CodeSize;
+    Latency += RHS.Latency;
+    return *this;
+  }
+
+  Bonus operator+(const Bonus RHS) const {
+    return Bonus(CodeSize + RHS.CodeSize, Latency + RHS.Latency);
+  }
+
+  bool operator==(const Bonus RHS) const {
+    return CodeSize == RHS.CodeSize && Latency == RHS.Latency;
+  }
+};
+
 class InstCostVisitor : public InstVisitor<InstCostVisitor, Constant *> {
   const DataLayout &DL;
   BlockFrequencyInfo &BFI;
   TargetTransformInfo &TTI;
   SCCPSolver &Solver;
 
   ConstMap KnownConstants;
   // Basic blocks known to be unreachable after constant propagation.
   DenseSet<BasicBlock *> DeadBlocks;
   // PHI nodes we have visited before.
   DenseSet<Instruction *> VisitedPHIs;
   // PHI nodes we have visited once without successfully constant folding them.
   // Once the InstCostVisitor has processed all the specialization arguments,
   // it should be possible to determine whether those PHIs can be folded
   // (some of their incoming values may have become constant or dead).
   SmallVector<Instruction *> PendingPHIs;
 
   ConstMap::iterator LastVisited;
 
 public:
   InstCostVisitor(const DataLayout &DL, BlockFrequencyInfo &BFI,
                   TargetTransformInfo &TTI, SCCPSolver &Solver)
       : DL(DL), BFI(BFI), TTI(TTI), Solver(Solver) {}
 
-  Cost getUserBonus(Instruction *User, Value *Use = nullptr,
+  Bonus getUserBonus(Instruction *User, Value *Use = nullptr,
                      Constant *C = nullptr);
 
-  Cost getBonusFromPendingPHIs();
+  Bonus getBonusFromPendingPHIs();
 
 private:
   friend class InstVisitor<InstCostVisitor, Constant *>;
 
   Cost estimateSwitchInst(SwitchInst &I);
   Cost estimateBranchInst(BranchInst &I);
 
   Constant *visitInstruction(Instruction &I) { return nullptr; }
public:
 
   InstCostVisitor getInstCostVisitorFor(Function *F) {
     auto &BFI = GetBFI(*F);
     auto &TTI = GetTTI(*F);
     return InstCostVisitor(M.getDataLayout(), BFI, TTI, Solver);
   }
 
   /// Compute a bonus for replacing argument \p A with constant \p C.
-  Cost getSpecializationBonus(Argument *A, Constant *C,
+  Bonus getSpecializationBonus(Argument *A, Constant *C,
                                InstCostVisitor &Visitor);
 
 private:
   Constant *getPromotableAlloca(AllocaInst *Alloca, CallInst *Call);
 
   /// A constant stack value is an AllocaInst that has a single constant
   /// value stored to it. Return this constant if such an alloca stack value
   /// is a function argument.
   Constant *getConstantStackValue(CallInst *Call, Value *Val);
private:
 
   /// @brief  Find potential specialization opportunities.
   /// @param F Function to specialize
   /// @param SpecCost Cost of specializing a function. Final score is benefit
   /// minus this cost.
   /// @param AllSpecs A vector to add potential specializations to.
   /// @param SM  A map for a function's specialisation range
   /// @return True, if any potential specializations were found
-  bool findSpecializations(Function *F, Cost SpecCost,
+  bool findSpecializations(Function *F, unsigned SpecCost,
                            SmallVectorImpl<Spec> &AllSpecs, SpecMap &SM);
 
   bool isCandidateFunction(Function *F);
 
   /// @brief Create a specialization of \p F and prime the SCCPSolver
   /// @param F Function to specialize
   /// @param S Which specialization to create
   /// @return The new, cloned function

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 //
 // https://llvm-compile-time-tracker.com
 // https://github.com/nikic/llvm-compile-time-tracker
 static cl::opt<bool> SpecializeLiteralConstant(
     "funcspec-for-literal-constant", cl::init(false), cl::Hidden, cl::desc(
     "Enable specialization of functions that take a literal constant as an "
     "argument"));
 
-// Estimates the instruction cost of all the basic blocks in \p WorkList.
-// The successors of such blocks are added to the list as long as they are
-// executable and they have a unique predecessor. \p WorkList represents
-// the basic blocks of a specialization which become dead once we replace
-// instructions that are known to be constants. The aim here is to estimate
-// the combination of size and latency savings in comparison to the non
-// specialized version of the function.
+// Estimates the codesize savings due to dead code after constant propagation.
+// \p WorkList represents the basic blocks of a specialization which will
+// eventually become dead once we replace instructions that are known to be
+// constants. The successors of such blocks are added to the list as long as
+// the \p Solver found they were executable prior to specialization, and only
+// if they have a unique predecessor.
 static Cost estimateBasicBlocks(SmallVectorImpl<BasicBlock *> &WorkList,
                                 DenseSet<BasicBlock *> &DeadBlocks,
                                 ConstMap &KnownConstants, SCCPSolver &Solver,
-                                BlockFrequencyInfo &BFI,
                                 TargetTransformInfo &TTI) {
-  Cost Bonus = 0;
-
+  Cost CodeSize = 0;
   // Accumulate the instruction cost of each basic block weighted by frequency.
   while (!WorkList.empty()) {
     BasicBlock *BB = WorkList.pop_back_val();
 
-    uint64_t Weight = BFI.getBlockFreq(BB).getFrequency() /
-                      BFI.getEntryFreq();
-    if (!Weight)
-      continue;
-
     // These blocks are considered dead as far as the InstCostVisitor is
     // concerned. They haven't been proven dead yet by the Solver, but
     // may become if we propagate the constant specialization arguments.
     if (!DeadBlocks.insert(BB).second)
       continue;
 
     for (Instruction &I : *BB) {
       // Disregard SSA copies.
       if (auto *II = dyn_cast<IntrinsicInst>(&I))
         if (II->getIntrinsicID() == Intrinsic::ssa_copy)
           continue;
       // If it's a known constant we have already accounted for it.
       if (KnownConstants.contains(&I))
         continue;
 
-      Bonus += Weight *
-          TTI.getInstructionCost(&I, TargetTransformInfo::TCK_SizeAndLatency);
+      Cost C = TTI.getInstructionCost(&I, TargetTransformInfo::TCK_CodeSize);
 
-      LLVM_DEBUG(dbgs() << "FnSpecialization:     Bonus " << Bonus
-                        << " after user " << I << "\n");
+      LLVM_DEBUG(dbgs() << "FnSpecialization:     CodeSize " << C
+                        << " for user " << I << "\n");
+      CodeSize += C;
     }
 
     // Keep adding dead successors to the list as long as they are
     // executable and they have a unique predecessor.
     for (BasicBlock *SuccBB : successors(BB))
       if (Solver.isBlockExecutable(SuccBB) &&
           SuccBB->getUniquePredecessor() == BB)
         WorkList.push_back(SuccBB);
   }
-  return Bonus;
+  return CodeSize;
 }
 
 static Constant *findConstantFor(Value *V, ConstMap &KnownConstants) {
   if (auto *C = dyn_cast<Constant>(V))
     return C;
   if (auto It = KnownConstants.find(V); It != KnownConstants.end())
     return It->second;
   return nullptr;
 }
 
-Cost InstCostVisitor::getBonusFromPendingPHIs() {
-  Cost Bonus = 0;
+Bonus InstCostVisitor::getBonusFromPendingPHIs() {
+  Bonus B;
   while (!PendingPHIs.empty()) {
     Instruction *Phi = PendingPHIs.pop_back_val();
-    Bonus += getUserBonus(Phi);
+    B += getUserBonus(Phi);
   }
-  return Bonus;
+  return B;
 }
 
-Cost InstCostVisitor::getUserBonus(Instruction *User, Value *Use, Constant *C) {
+Bonus InstCostVisitor::getUserBonus(Instruction *User, Value *Use, Constant *C) {
   // Cache the iterator before visiting.
   LastVisited = Use ? KnownConstants.insert({Use, C}).first
                     : KnownConstants.end();
 
-  if (auto *I = dyn_cast<SwitchInst>(User))
-    return estimateSwitchInst(*I);
-
-  if (auto *I = dyn_cast<BranchInst>(User))
-    return estimateBranchInst(*I);
-
+  Cost CodeSize = 0;
+  if (auto *I = dyn_cast<SwitchInst>(User)) {
+    CodeSize = estimateSwitchInst(*I);
+  } else if (auto *I = dyn_cast<BranchInst>(User)) {
+    CodeSize = estimateBranchInst(*I);
+  } else {
     C = visit(*User);
     if (!C)
-    return 0;
-
+      return {0, 0};
     KnownConstants.insert({User, C});
+  }
+
+  CodeSize += TTI.getInstructionCost(User, TargetTransformInfo::TCK_CodeSize);
 
   uint64_t Weight = BFI.getBlockFreq(User->getParent()).getFrequency() /
                     BFI.getEntryFreq();
-  if (!Weight)
-    return 0;
 
-  Cost Bonus = Weight *
-      TTI.getInstructionCost(User, TargetTransformInfo::TCK_SizeAndLatency);
+  Cost Latency = Weight *
+      TTI.getInstructionCost(User, TargetTransformInfo::TCK_Latency);
 
-  LLVM_DEBUG(dbgs() << "FnSpecialization:     Bonus " << Bonus
-                    << " for user " << *User << "\n");
+  LLVM_DEBUG(dbgs() << "FnSpecialization:     {CodeSize = " << CodeSize
+                    << ", Latency = " << Latency << "} for user "
+                    << *User << "\n");
 
+  Bonus B(CodeSize, Latency);
   for (auto *U : User->users())
     if (auto *UI = dyn_cast<Instruction>(U))
       if (UI != User && Solver.isBlockExecutable(UI->getParent()))
-        Bonus += getUserBonus(UI, User, C);
+        B += getUserBonus(UI, User, C);
 
-  return Bonus;
+  return B;
 }
 
 Cost InstCostVisitor::estimateSwitchInst(SwitchInst &I) {
   assert(LastVisited != KnownConstants.end() && "Invalid iterator!");
 
   if (I.getCondition() != LastVisited->first)
     return 0;
 
Cost InstCostVisitor::estimateSwitchInst(SwitchInst &I) {
   for (const auto &Case : I.cases()) {
     BasicBlock *BB = Case.getCaseSuccessor();
     if (BB == Succ || !Solver.isBlockExecutable(BB) ||
         BB->getUniquePredecessor() != I.getParent())
       continue;
     WorkList.push_back(BB);
   }
 
-  return estimateBasicBlocks(WorkList, DeadBlocks, KnownConstants, Solver, BFI,
-                             TTI);
+  return estimateBasicBlocks(WorkList, DeadBlocks, KnownConstants, Solver, TTI);
 }
 
 Cost InstCostVisitor::estimateBranchInst(BranchInst &I) {
   assert(LastVisited != KnownConstants.end() && "Invalid iterator!");
 
   if (I.getCondition() != LastVisited->first)
     return 0;
 
   BasicBlock *Succ = I.getSuccessor(LastVisited->second->isOneValue());
   // Initialize the worklist with the dead successor as long as
   // it is executable and has a unique predecessor.
   SmallVector<BasicBlock *> WorkList;
   if (Solver.isBlockExecutable(Succ) &&
       Succ->getUniquePredecessor() == I.getParent())
     WorkList.push_back(Succ);
 
-  return estimateBasicBlocks(WorkList, DeadBlocks, KnownConstants, Solver, BFI,
-                             TTI);
+  return estimateBasicBlocks(WorkList, DeadBlocks, KnownConstants, Solver, TTI);
 }
 
 Constant *InstCostVisitor::visitPHINode(PHINode &I) {
   if (I.getNumIncomingValues() > MaxIncomingPhiValues)
     return nullptr;
 
   bool Inserted = VisitedPHIs.insert(&I).second;
   Constant *Const = nullptr;
if (Metrics.notDuplicatable || !Metrics.NumInsts.isValid() ||
       continue;
 
     // TODO: For now only consider recursive functions when running multiple
     // times. This should change if specialization on literal constants gets
     // enabled.
     if (!Inserted && !Metrics.isRecursive && !SpecializeLiteralConstant)
       continue;
 
+    int64_t Sz = *Metrics.NumInsts.getValue();
+    assert(Sz > 0 && "CodeSize should be positive");
+    // It is safe to down cast from int64_t, NumInsts is always positive.
+    unsigned SpecCost = static_cast<unsigned>(Sz);
+
     LLVM_DEBUG(dbgs() << "FnSpecialization: Specialization cost for "
-                      << F.getName() << " is " << Metrics.NumInsts << "\n");
+                      << F.getName() << " is " << SpecCost << "\n");
 
     if (Inserted && Metrics.isRecursive)
       promoteConstantStackValues(&F);
 
-    if (!findSpecializations(&F, Metrics.NumInsts, AllSpecs, SM)) {
+    if (!findSpecializations(&F, SpecCost, AllSpecs, SM)) {
       LLVM_DEBUG(
           dbgs() << "FnSpecialization: No possible specializations found for "
                  << F.getName() << "\n");
       continue;
     }
 
     ++NumCandidates;
   }

 /// the SCCPSolver in the cloned version.
 static Function *cloneCandidateFunction(Function *F) {
   ValueToValueMapTy Mappings;
   Function *Clone = CloneFunction(F, Mappings);
   removeSSACopy(*Clone);
   return Clone;
 }
 
-bool FunctionSpecializer::findSpecializations(Function *F, Cost SpecCost,
+bool FunctionSpecializer::findSpecializations(Function *F, unsigned SpecCost,
                                               SmallVectorImpl<Spec> &AllSpecs,
                                               SpecMap &SM) {
   // A mapping from a specialisation signature to the index of the respective
   // entry in the all specialisation array. Used to ensure uniqueness of
   // specialisations.
   DenseMap<SpecSig, unsigned> UniqueSpecs;
 
   // Get a list of interesting arguments.
if (auto It = UniqueSpecs.find(S); It != UniqueSpecs.end()) {
       // functions. Hence we don't rewrite the call directly, but match it with
       // the best specialisation once all specialisations are known.
       if (CS.getFunction() == F)
         continue;
       const unsigned Index = It->second;
       AllSpecs[Index].CallSites.push_back(&CS);
     } else {
       // Calculate the specialisation gain.
-      Cost Score = 0;
+      Bonus B;
       InstCostVisitor Visitor = getInstCostVisitorFor(F);
       for (ArgInfo &A : S.Args)
-        Score += getSpecializationBonus(A.Formal, A.Actual, Visitor);
-      Score += Visitor.getBonusFromPendingPHIs();
+        B += getSpecializationBonus(A.Formal, A.Actual, Visitor);
+      B += Visitor.getBonusFromPendingPHIs();
 
-      LLVM_DEBUG(dbgs() << "FnSpecialization: Specialization score = "
-                        << Score << "\n");
+      LLVM_DEBUG(dbgs() << "FnSpecialization: Specialization score {CodeSize = "
+                        << B.CodeSize << ", Latency = " << B.Latency
+                        << "}\n");
 
       // Discard unprofitable specialisations.
-      if (!ForceSpecialization && Score <= SpecCost)
+      if (!ForceSpecialization && B.Latency <= SpecCost - B.CodeSize)
         continue;
 
       // Create a new specialisation entry.
-      auto &Spec = AllSpecs.emplace_back(F, S, Score);
+      auto &Spec = AllSpecs.emplace_back(F, S, B.Latency);
       if (CS.getFunction() != F)
         Spec.CallSites.push_back(&CS);
       const unsigned Index = AllSpecs.size() - 1;
       UniqueSpecs[S] = Index;
       if (auto [It, Inserted] = SM.try_emplace(F, Index, Index + 1); !Inserted)
         It->second.second = Index + 1;
     }
   }
Function *FunctionSpecializer::createSpecialization(Function *F,
   // Mark all the specialized functions
   Specializations.insert(Clone);
   ++NumSpecsCreated;
 
   return Clone;
 }
 
 /// Compute a bonus for replacing argument \p A with constant \p C.
-Cost FunctionSpecializer::getSpecializationBonus(Argument *A, Constant *C,
+Bonus FunctionSpecializer::getSpecializationBonus(Argument *A, Constant *C,
                                                  InstCostVisitor &Visitor) {
   LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing bonus for constant: "
                     << C->getNameOrAsOperand() << "\n");
 
-  Cost TotalCost = 0;
+  Bonus B;
   for (auto *U : A->users())
     if (auto *UI = dyn_cast<Instruction>(U))
       if (Solver.isBlockExecutable(UI->getParent()))
-        TotalCost += Visitor.getUserBonus(UI, A, C);
+        B += Visitor.getUserBonus(UI, A, C);
 
-  LLVM_DEBUG(dbgs() << "FnSpecialization:   Accumulated user bonus "
-                    << TotalCost << " for argument " << *A << "\n");
+  LLVM_DEBUG(dbgs() << "FnSpecialization:   Accumulated bonus {CodeSize = "
+                    << B.CodeSize << ", Latency = " << B.Latency
+                    << "} for argument " << *A << "\n");
 
   // The below heuristic is only concerned with exposing inlining
   // opportunities via indirect call promotion. If the argument is not a
   // (potentially casted) function pointer, give up.
   //
   // TODO: Perhaps we should consider checking such inlining opportunities
   // while traversing the users of the specialization arguments ?
   Function *CalledFunction = dyn_cast<Function>(C->stripPointerCasts());
   if (!CalledFunction)
-    return TotalCost;
+    return B;
 
   // Get TTI for the called function (used for the inline cost).
   auto &CalleeTTI = (GetTTI)(*CalledFunction);
 
   // Look at all the call sites whose called value is the argument.
   // Specializing the function on the argument would allow these indirect
   // calls to be promoted to direct calls. If the indirect call promotion
   // would likely enable the called function to be inlined, specializing is a
   // good idea.
-  int Bonus = 0;
+  int InliningBonus = 0;
   for (User *U : A->users()) {
     if (!isa<CallInst>(U) && !isa<InvokeInst>(U))
       continue;
     auto *CS = cast<CallBase>(U);
     if (CS->getCalledOperand() != A)
       continue;
     if (CS->getFunctionType() != CalledFunction->getFunctionType())
       continue;
for (User *U : A->users()) {
     auto Params = getInlineParams();
     Params.DefaultThreshold += InlineConstants::IndirectCallThreshold;
     InlineCost IC =
         getInlineCost(*CS, CalledFunction, Params, CalleeTTI, GetAC, GetTLI);
 
     // We clamp the bonus for this call to be between zero and the default
     // threshold.
     if (IC.isAlways())
-      Bonus += Params.DefaultThreshold;
+      InliningBonus += Params.DefaultThreshold;
     else if (IC.isVariable() && IC.getCostDelta() > 0)
-      Bonus += IC.getCostDelta();
+      InliningBonus += IC.getCostDelta();
 
-    LLVM_DEBUG(dbgs() << "FnSpecialization:   Inlining bonus " << Bonus
+    LLVM_DEBUG(dbgs() << "FnSpecialization:   Inlining bonus " << InliningBonus
                       << " for user " << *U << "\n");
   }
 
-  return TotalCost + Bonus;
+  return B += {0, InliningBonus};
 }
 
 /// Determine if it is possible to specialise the function for constant values
 /// of the formal parameter \p A.
 bool FunctionSpecializer::isArgumentInteresting(Argument *A) {
   // No point in specialization if the argument is unused.
   if (A->user_empty())
     return false;

llvm/unittests/Transforms/IPO/FunctionSpecializationTest.cpp

FunctionSpecializer getSpecializerFor(Function *F) {
     for (Argument &Arg : F->args())
       Solver->markOverdefined(&Arg);
     Solver->solveWhileResolvedUndefsIn(*M);
 
     return FunctionSpecializer(*Solver, *M, &FAM, GetBFI, GetTLI, GetTTI,
                                GetAC);
   }
 
-  Cost getInstCost(Instruction &I) {
+  Bonus getInstCost(Instruction &I, bool SizeOnly = false) {
     auto &TTI = FAM.getResult<TargetIRAnalysis>(*I.getFunction());
     auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(*I.getFunction());
 
-    return BFI.getBlockFreq(I.getParent()).getFrequency() / BFI.getEntryFreq() *
-         TTI.getInstructionCost(&I, TargetTransformInfo::TCK_SizeAndLatency);
+    Cost CodeSize =
+        TTI.getInstructionCost(&I, TargetTransformInfo::TCK_CodeSize);
+
+    Cost Latency = SizeOnly ? 0 :
+        BFI.getBlockFreq(I.getParent()).getFrequency() / BFI.getEntryFreq() *
+        TTI.getInstructionCost(&I, TargetTransformInfo::TCK_Latency);
+
+    return {CodeSize, Latency};
   }
 };
 
 } // namespace llvm
 
 using namespace llvm;
 
 TEST_F(FunctionSpecializationTest, SwitchInst) {
TEST_F(FunctionSpecializationTest, SwitchInst) {
   Module &M = parseModule(ModuleString);
   Function *F = M.getFunction("foo");
   FunctionSpecializer Specializer = getSpecializerFor(F);
   InstCostVisitor Visitor = Specializer.getInstCostVisitorFor(F);
 
   Constant *One = ConstantInt::get(IntegerType::getInt32Ty(M.getContext()), 1);
 
   auto FuncIter = F->begin();
-  ++FuncIter;
+  BasicBlock &Loop = *++FuncIter;
   BasicBlock &Case1 = *++FuncIter;
   BasicBlock &Case2 = *++FuncIter;
   BasicBlock &BB1 = *++FuncIter;
   BasicBlock &BB2 = *++FuncIter;
 
+  Instruction &Switch = Loop.front();
   Instruction &Mul = Case1.front();
   Instruction &And = Case2.front();
   Instruction &Sdiv = *++Case2.begin();
   Instruction &BrBB2 = Case2.back();
   Instruction &Add = BB1.front();
   Instruction &Or = BB2.front();
   Instruction &BrLoop = BB2.back();
 
   // mul
-  Cost Ref = getInstCost(Mul);
-  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Bonus Ref = getInstCost(Mul);
+  Bonus Test = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
   // and + or + add
   Ref = getInstCost(And) + getInstCost(Or) + getInstCost(Add);
-  Bonus = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
-
-  // sdiv + br + br
-  Ref = getInstCost(Sdiv) + getInstCost(BrBB2) + getInstCost(BrLoop);
-  Bonus = Specializer.getSpecializationBonus(F->getArg(2), One, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Test = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
+
+  // switch + sdiv + br + br
+  Ref = getInstCost(Switch) +
+        getInstCost(Sdiv, /*SizeOnly =*/ true) +
+        getInstCost(BrBB2, /*SizeOnly =*/ true) +
+        getInstCost(BrLoop, /*SizeOnly =*/ true);
+  Test = Specializer.getSpecializationBonus(F->getArg(2), One, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 }
 
 TEST_F(FunctionSpecializationTest, BranchInst) {
   const char *ModuleString = R"(
     define void @foo(i32 %a, i32 %b, i1 %cond) {
     entry:
       br label %loop
     loop:
TEST_F(FunctionSpecializationTest, BranchInst) {
   Function *F = M.getFunction("foo");
   FunctionSpecializer Specializer = getSpecializerFor(F);
   InstCostVisitor Visitor = Specializer.getInstCostVisitorFor(F);
 
   Constant *One = ConstantInt::get(IntegerType::getInt32Ty(M.getContext()), 1);
   Constant *False = ConstantInt::getFalse(M.getContext());
 
   auto FuncIter = F->begin();
-  ++FuncIter;
+  BasicBlock &Loop = *++FuncIter;
   BasicBlock &BB0 = *++FuncIter;
   BasicBlock &BB1 = *++FuncIter;
 
+  Instruction &Branch = Loop.front();
   Instruction &Mul = BB0.front();
   Instruction &Sub = *++BB0.begin();
   Instruction &BrBB1 = BB0.back();
   Instruction &Add = BB1.front();
   Instruction &Sdiv = *++BB1.begin();
   Instruction &BrLoop = BB1.back();
 
   // mul
-  Cost Ref = getInstCost(Mul);
-  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Bonus Ref = getInstCost(Mul);
+  Bonus Test = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
   // add
   Ref = getInstCost(Add);
-  Bonus = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
-
-  // sub + br + sdiv + br
-  Ref = getInstCost(Sub) + getInstCost(BrBB1) + getInstCost(Sdiv) +
-        getInstCost(BrLoop);
-  Bonus = Specializer.getSpecializationBonus(F->getArg(2), False, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Test = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
+
+  // branch + sub + br + sdiv + br
+  Ref = getInstCost(Branch) +
+        getInstCost(Sub, /*SizeOnly =*/ true) +
+        getInstCost(BrBB1, /*SizeOnly =*/ true) +
+        getInstCost(Sdiv, /*SizeOnly =*/ true) +
+        getInstCost(BrLoop, /*SizeOnly =*/ true);
+  Test = Specializer.getSpecializationBonus(F->getArg(2), False, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 }
 
 TEST_F(FunctionSpecializationTest, Misc) {
   const char *ModuleString = R"(
     %struct_t = type { [8 x i16], [8 x i16], i32, i32, i32, ptr, [8 x i8] }
     @g = constant %struct_t zeroinitializer, align 16
 
     declare i32 @llvm.smax.i32(i32, i32)
TEST_F(FunctionSpecializationTest, Misc) {
   Instruction &Zext = *BlockIter++;
   Instruction &Select = *BlockIter++;
   Instruction &Gep = *BlockIter++;
   Instruction &Load = *BlockIter++;
   Instruction &Freeze = *BlockIter++;
   Instruction &Smax = *BlockIter++;
 
   // icmp + zext
-  Cost Ref = getInstCost(Icmp) + getInstCost(Zext);
-  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Bonus Ref = getInstCost(Icmp) + getInstCost(Zext);
+  Bonus Test = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
   // select
   Ref = getInstCost(Select);
-  Bonus = Specializer.getSpecializationBonus(F->getArg(1), True, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Test = Specializer.getSpecializationBonus(F->getArg(1), True, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
   // gep + load + freeze + smax
   Ref = getInstCost(Gep) + getInstCost(Load) + getInstCost(Freeze) +
         getInstCost(Smax);
-  Bonus = Specializer.getSpecializationBonus(F->getArg(2), GV, Visitor);
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  Test = Specializer.getSpecializationBonus(F->getArg(2), GV, Visitor);
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
-  Bonus = Specializer.getSpecializationBonus(F->getArg(3), Undef, Visitor);
-  EXPECT_TRUE(Bonus == 0);
+  Test = Specializer.getSpecializationBonus(F->getArg(3), Undef, Visitor);
+  EXPECT_TRUE(Test.CodeSize == 0 && Test.Latency == 0);
 }
 
 TEST_F(FunctionSpecializationTest, PhiNode) {
   const char *ModuleString = R"(
     define void @foo(i32 %a, i32 %b, i32 %i) {
     entry:
       br label %loop
     loop:
TEST_F(FunctionSpecializationTest, PhiNode) {
   auto FuncIter = F->begin();
   for (int I = 0; I < 4; ++I)
     ++FuncIter;
 
   BasicBlock &BB = *FuncIter;
 
   Instruction &Phi = BB.front();
   Instruction &Icmp = *++BB.begin();
+  Instruction &Branch = BB.back();
 
-  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor) +
+  Bonus Test = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor) +
                Specializer.getSpecializationBonus(F->getArg(1), One, Visitor) +
                Specializer.getSpecializationBonus(F->getArg(2), One, Visitor);
-  EXPECT_TRUE(Bonus > 0);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
-  // phi + icmp
-  Cost Ref = getInstCost(Phi) + getInstCost(Icmp);
-  Bonus = Visitor.getBonusFromPendingPHIs();
-  EXPECT_EQ(Bonus, Ref);
-  EXPECT_TRUE(Bonus > 0);
+  // phi + icmp + branch
+  Bonus Ref = getInstCost(Phi) + getInstCost(Icmp) + getInstCost(Branch);
+  Test = Visitor.getBonusFromPendingPHIs();
+  EXPECT_EQ(Test, Ref);
+  EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 }