[FuncSpec] Improve the accuracy of the cost model.

[labrinea]

Summary

Instead of blindly traversing the use-def chain of constant arguments, compute known constants along the way. Stop as soon as a user cannot be replaced by a constant. Keep it light-weight by handling some basic instruction types.

Stats from the llvm testsuite (CTMark):
-O3 pipeline

Nspecs before	ClamAV : 5	SPASS : 2	Bullet : 1
Nspecs after	ClamAV : 5	consumer-typeset : 1	Bullet : 1
InstrCount delta (geomean) : -0.029%

LTO pipeline

Nspecs before	sqlite3 : 1	consumer-typeset : 1
Nspecs after	SPASS : 3	consumer-typeset : 1
InstrCount delta (geomean) : -0.015%

[ChuanqiXu]

Can we find some method to test this? If a lit test is not good enough, maybe we can try to look at unit tests?

[labrinea]

In D150464#4340945, @ChuanqiXu wrote:

Can we find some method to test this? If a lit test is not good enough, maybe we can try to look at unit tests?

Sounds like a good idea. I'll write some.

[labrinea]

I've added unittests but the results make more sense with D150649.

[ChuanqiXu]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

790–791

Is this change necessary? Or what is the intention? It looks not good to me at the first sight.

794

We prefer code style with shorter ident.

794

Maybe we can make it look prettier by using InstVisitor for such code patterns.

798–800

We don't like else-after-return pattern.

810–836

Oh, good idea. But I only understand what you're trying to do after reading the codes. A comment here will be helpful.

826

Constant expression is going to be deprecated: https://discourse.llvm.org/t/rfc-remove-most-constant-expressions/63179. Let's not add more constant expression.

832–834

Such code pattern looks a little bit dirty. Let's try to wrap it into a function.

841–843

This looks similar with estimateSwitchInst. Can we try to merge them?

845

Ditto.

876

[labrinea]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

790–791

Quoting a comment I left on the parent revision (D150375) in response to @mingmingl :

The Specialization cost estimation is wrong and I am planning to fix that in later patches. Metrics.NumInsts already contains TargetTransformInfo::TCK_CodeSize.

See https://github.com/llvm/llvm-project/blob/main/llvm/lib/Analysis/CodeMetrics.cpp#L180

794

Can you give an example?

794

will do

810–836

will do

826

What alternative do we have?

841–843

I can factor out common code to a static function perhaps.

[ChuanqiXu]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

790–791

Oh, sorry for disturbing.

794

There are multiple examples: InstructionAllowed and CallAnalyzer and so on. Maybe you would feel the code become longer after we adopt that. But I think it will be easier to read and understand. This is a suggestion and not a requirement.

826

I don't know if we have an off-the-shelf solution for this. But I'm sure that the resulting 'C' here shouldn't be ConstantExpr. If there is not an off-the-shelf solution, maybe we can add one for it. Or we can leave the case for later patches.

[labrinea]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

794

Ah, you meant the instruction visitor like the one used by the SCCPSolver. I was experimenting with this idea yesterday. One problem I see is that the visit methods take only one argument (the visited instruction) so we would have to cache the extra parameters inside the visitor class. I'll investigate more.

826

I think we can use the data-layout agnostic constant folding.

[labrinea]

Added an InstructionVisitor for the bonus calculation.

[labrinea]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

148–150

This is a bit ugly but I couldn't think of a better way. As I explained before the visit functions only take one parameter, the instruction.

[labrinea]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

148–150

Perhaps I could replace these two with ConstMap::iterator since we are inserting the {Use, C} pair on every visit.

[labrinea]

Cached the Map Iterator in the InstructionVisitor instead of the {Key, Value} pair.

[labrinea]

llvm/unittests/Transforms/IPO/FunctionSpecializationTest.cpp

2

I've spotted a problem. This unittest is relying on target specific costs, so it will fail on some targets. I need to either make it x86 specific or change the expected values based on some logic.

[labrinea]

Made the unittest more generic so that it succeeds on any target.

[ChuanqiXu]

LGTM with nit comments. Thanks.

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

100–103

Let's add a comment for this since its semantics are not clear from the signature.

210

234–236

Could we refactor it into this? Then we can remove AllOperandsAreConst.

[labrinea]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

100–103

Will do.

210

I don't understand. The worklist is populated by a chain of basic blocks. If 32 is too much perharps we could leave it to default? Surely 1 is too conservative and expanding the small vector at runtime seems expensive.

234–236

Oh yes, I had the old impementation (pre InstructionVisitor) in mind at the time of writing this.

[thevinster]

@labrinea This change hits a crash when trying to build lld with thinLTO (i.e. stage2 lld) and with assertions enabled. I verified that the revision before this change works. I get the following assertion failure

opt: <workspace>/llvm-project/llvm/include/llvm/Support/Casting.h:578: decltype(auto) llvm::cast(From *) [To = llvm::ConstantInt, From = llvm::Constant]: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.

I reduced the issue and this can be reproed by building opt on this revision and then running

$ opt reduced.ll -o /dev/null -O3

where reduced.ll is https://reviews.llvm.org/P8311

Could you please take a look and advise the path forward?

cc: @smeenai

[thevinster]

In D150464#4444716, @thevinster wrote:

@labrinea This change hits a crash when trying to build lld with thinLTO (i.e. stage2 lld) and with assertions enabled. I verified that the revision before this change works. I get the following assertion failure

opt: <workspace>/llvm-project/llvm/include/llvm/Support/Casting.h:578: decltype(auto) llvm::cast(From *) [To = llvm::ConstantInt, From = llvm::Constant]: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.

I reduced the issue and this can be reproed by building opt on this revision and then running

$ opt reduced.ll -o /dev/null -O3

where reduced.ll is https://reviews.llvm.org/P8311

Could you please take a look and advise the path forward?

cc: @smeenai

I put up D154159 that should address the crashing issue.

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

 
 #ifndef LLVM_TRANSFORMS_IPO_FUNCTIONSPECIALIZATION_H
 #define LLVM_TRANSFORMS_IPO_FUNCTIONSPECIALIZATION_H
 
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/IR/InstVisitor.h"
 #include "llvm/Transforms/Scalar/SCCP.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/SCCPSolver.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 
 using namespace llvm;
 
 namespace llvm {
 // Map of potential specializations for each function. The FunctionSpecializer
 // keeps the discovered specialisation opportunities for the module in a single
 // vector, where the specialisations of each function form a contiguous range.
 // This map's value is the beginning and the end of that range.
 using SpecMap = DenseMap<Function *, std::pair<unsigned, unsigned>>;
 
 // Just a shorter abbreviation to improve indentation.
 using Cost = InstructionCost;
 
+// Map of known constants found during the specialization bonus estimation.
+using ConstMap = DenseMap<Value *, Constant *>;
+
 // Specialization signature, used to uniquely designate a specialization within
 // a function.
 struct SpecSig {
   // Hashing support, used to distinguish between ordinary, empty, or tombstone
   // keys.
   unsigned Key = 0;
   SmallVector<ArgInfo, 4> Args;
 
struct Spec {
   SmallVector<CallBase *> CallSites;
 
   Spec(Function *F, const SpecSig &S, Cost Score)
       : F(F), Sig(S), Score(Score) {}
   Spec(Function *F, const SpecSig &&S, Cost Score)
       : F(F), Sig(S), Score(Score) {}
 };
 
+class InstCostVisitor : public InstVisitor<InstCostVisitor, Constant *> {
+  const DataLayout &DL;
+  BlockFrequencyInfo &BFI;
+  TargetTransformInfo &TTI;
+  SCCPSolver &Solver;
+
+  ConstMap KnownConstants;
+
+  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, Constant *C);
+
+private:
+  friend class InstVisitor<InstCostVisitor, Constant *>;
+
+  Cost estimateSwitchInst(SwitchInst &I);
+  Cost estimateBranchInst(BranchInst &I);
+
+  Constant *visitInstruction(Instruction &I) { return nullptr; }
+  Constant *visitLoadInst(LoadInst &I);
+  Constant *visitGetElementPtrInst(GetElementPtrInst &I);
+  Constant *visitSelectInst(SelectInst &I);
+  Constant *visitCastInst(CastInst &I);
+  Constant *visitCmpInst(CmpInst &I);
+  Constant *visitUnaryOperator(UnaryOperator &I);
+  Constant *visitBinaryOperator(BinaryOperator &I);
+};
+
 class FunctionSpecializer {
 
   /// The IPSCCP Solver.
   SCCPSolver &Solver;
 
   Module &M;
 
   /// Analysis manager, needed to invalidate analyses.
FunctionSpecializer(
         GetTTI(GetTTI), GetAC(GetAC) {}
 
   ~FunctionSpecializer();
 
   bool isClonedFunction(Function *F) { return Specializations.count(F); }
 
   bool run();
 
+  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,
+                              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:
   /// @param F Function to specialize
   /// @param S Which specialization to create
   /// @return The new, cloned function
   Function *createSpecialization(Function *F, const SpecSig &S);
 
   /// Compute and return the cost of specializing function \p F.
   Cost getSpecializationCost(Function *F);
 
-  /// Compute a bonus for replacing argument \p A with constant \p C.
-  Cost getSpecializationBonus(Argument *A, Constant *C);
-
   /// Determine if it is possible to specialise the function for constant values
   /// of the formal parameter \p A.
   bool isArgumentInteresting(Argument *A);
 
   /// Check if the value \p V  (an actual argument) is a constant or can only
   /// have a constant value. Return that constant.
   Constant *getCandidateConstant(Value *V);
 

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 // - 2021 LLVM Dev Mtg “Introducing function specialisation, and can we enable
 //   it by default?”, https://www.youtube.com/watch?v=zJiCjeXgV5Q
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/FunctionSpecialization.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueLattice.h"
 #include "llvm/Analysis/ValueLatticeUtils.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/ConstantFold.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Transforms/Scalar/SCCP.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/SCCPSolver.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 #include <cmath>
 
 using namespace llvm;

 //
 // 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.
+static Cost estimateBasicBlocks(SmallVectorImpl<BasicBlock *> &WorkList,
+                                ConstMap &KnownConstants, SCCPSolver &Solver,
+                                BlockFrequencyInfo &BFI,
+                                TargetTransformInfo &TTI) {
+  Cost Bonus = 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;
+
+    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);
+
+      LLVM_DEBUG(dbgs() << "FnSpecialization:     Bonus " << Bonus
+                        << " after user " << I << "\n");
+    }
+
+    // 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;
+}
+
+static Constant *findConstantFor(Value *V, ConstMap &KnownConstants) {
+  if (auto It = KnownConstants.find(V); It != KnownConstants.end())
+    return It->second;
+  return nullptr;
+}
+
+Cost InstCostVisitor::getUserBonus(Instruction *User, Value *Use, Constant *C) {
+  // Cache the iterator before visiting.
+  LastVisited = KnownConstants.insert({Use, C}).first;
+
+  if (auto *I = dyn_cast<SwitchInst>(User))
+    return estimateSwitchInst(*I);
+
+  if (auto *I = dyn_cast<BranchInst>(User))
+    return estimateBranchInst(*I);
+
+  C = visit(*User);
+  if (!C)
+    return 0;
+
+  KnownConstants.insert({User, C});
+
+  uint64_t Weight = BFI.getBlockFreq(User->getParent()).getFrequency() /
+                    BFI.getEntryFreq();
+  if (!Weight)
+    return 0;
+
+  Cost Bonus = Weight *
+      TTI.getInstructionCost(User, TargetTransformInfo::TCK_SizeAndLatency);
+
+  LLVM_DEBUG(dbgs() << "FnSpecialization:     Bonus " << Bonus
+                    << " for user " << *User << "\n");
+
+  for (auto *U : User->users())
+    if (auto *UI = dyn_cast<Instruction>(U))
+      if (Solver.isBlockExecutable(UI->getParent()))
+        Bonus += getUserBonus(UI, User, C);
+
+  return Bonus;
+}
+
+Cost InstCostVisitor::estimateSwitchInst(SwitchInst &I) {
+  if (I.getCondition() != LastVisited->first)
+    return 0;
+
+  auto *C = cast<ConstantInt>(LastVisited->second);
+  BasicBlock *Succ = I.findCaseValue(C)->getCaseSuccessor();
+  // Initialize the worklist with the dead basic blocks. These are the
+  // destination labels which are different from the one corresponding
+  // to \p C. They should be executable and have a unique predecessor.
+  SmallVector<BasicBlock *> WorkList;
+  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, KnownConstants, Solver, BFI, TTI);
+}
+
+Cost InstCostVisitor::estimateBranchInst(BranchInst &I) {
+  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, KnownConstants, Solver, BFI, TTI);
+}
+
+Constant *InstCostVisitor::visitLoadInst(LoadInst &I) {
+  if (isa<ConstantPointerNull>(LastVisited->second))
+    return nullptr;
+  return ConstantFoldLoadFromConstPtr(LastVisited->second, I.getType(), DL);
+}
+
+Constant *InstCostVisitor::visitGetElementPtrInst(GetElementPtrInst &I) {
+  SmallVector<Value *, 8> Operands;
+  Operands.reserve(I.getNumOperands());
+
+  for (unsigned Idx = 0, E = I.getNumOperands(); Idx != E; ++Idx) {
+    Value *V = I.getOperand(Idx);
+    auto *C = dyn_cast<Constant>(V);
+    if (!C)
+      C = findConstantFor(V, KnownConstants);
+    if (!C)
+      return nullptr;
+    Operands.push_back(C);
+  }
+
+  auto *Ptr = cast<Constant>(Operands[0]);
+  auto Ops = ArrayRef(Operands.begin() + 1, Operands.end());
+  return ConstantFoldGetElementPtr(I.getSourceElementType(), Ptr,
+                                   I.isInBounds(), std::nullopt, Ops);
+}
+
+Constant *InstCostVisitor::visitSelectInst(SelectInst &I) {
+  if (I.getCondition() != LastVisited->first)
+    return nullptr;
+
+  Value *V = LastVisited->second->isZeroValue() ? I.getFalseValue()
+                                                : I.getTrueValue();
+  auto *C = dyn_cast<Constant>(V);
+  if (!C)
+    C = findConstantFor(V, KnownConstants);
+  return C;
+}
+
+Constant *InstCostVisitor::visitCastInst(CastInst &I) {
+  return ConstantFoldCastOperand(I.getOpcode(), LastVisited->second,
+                                 I.getType(), DL);
+}
+
+Constant *InstCostVisitor::visitCmpInst(CmpInst &I) {
+  bool Swap = I.getOperand(1) == LastVisited->first;
+  Value *V = Swap ? I.getOperand(0) : I.getOperand(1);
+  auto *Other = dyn_cast<Constant>(V);
+  if (!Other)
+    Other = findConstantFor(V, KnownConstants);
+
+  if (!Other)
+    return nullptr;
+
+  Constant *Const = LastVisited->second;
+  return Swap ?
+        ConstantFoldCompareInstOperands(I.getPredicate(), Other, Const, DL)
+      : ConstantFoldCompareInstOperands(I.getPredicate(), Const, Other, DL);
+}
+
+Constant *InstCostVisitor::visitUnaryOperator(UnaryOperator &I) {
+  return ConstantFoldUnaryOpOperand(I.getOpcode(), LastVisited->second, DL);
+}
+
+Constant *InstCostVisitor::visitBinaryOperator(BinaryOperator &I) {
+  bool Swap = I.getOperand(1) == LastVisited->first;
+  Value *V = Swap ? I.getOperand(0) : I.getOperand(1);
+  auto *Other = dyn_cast<Constant>(V);
+  if (!Other)
+    Other = findConstantFor(V, KnownConstants);
+
+  if (!Other)
+    return nullptr;
+
+  Constant *Const = LastVisited->second;
+  return dyn_cast_or_null<Constant>(Swap ?
+        simplifyBinOp(I.getOpcode(), Other, Const, SimplifyQuery(DL))
+      : simplifyBinOp(I.getOpcode(), Const, Other, SimplifyQuery(DL)));
+}
+
 Constant *FunctionSpecializer::getPromotableAlloca(AllocaInst *Alloca,
                                                    CallInst *Call) {
   Value *StoreValue = nullptr;
   for (auto *User : Alloca->users()) {
     // We can't use llvm::isAllocaPromotable() as that would fail because of
     // the usage in the CallInst, which is what we check here.
     if (User == Call)
       continue;
CodeMetrics &FunctionSpecializer::analyzeFunction(Function *F) {
   auto I = FunctionMetrics.insert({F, CodeMetrics()});
   CodeMetrics &Metrics = I.first->second;
   if (I.second) {
     // The code metrics were not cached.
     SmallPtrSet<const Value *, 32> EphValues;
     CodeMetrics::collectEphemeralValues(F, &(GetAC)(*F), EphValues);
     for (BasicBlock &BB : *F)
       Metrics.analyzeBasicBlock(&BB, (GetTTI)(*F), EphValues);
-
-    LLVM_DEBUG(dbgs() << "FnSpecialization: Code size of function "
-                      << F->getName() << " is " << Metrics.NumInsts
-                      << " instructions\n");
   }
   return Metrics;
 }
 
 /// Clone the function \p F and remove the ssa_copy intrinsics added by
 /// the SCCPSolver in the cloned version.
 static Function *cloneCandidateFunction(Function *F) {
   ValueToValueMapTy Mappings;
if (auto It = UniqueSpecs.find(S); It != UniqueSpecs.end()) {
       // 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 - SpecCost;
+      InstCostVisitor Visitor = getInstCostVisitorFor(F);
       for (ArgInfo &A : S.Args)
-        Score += getSpecializationBonus(A.Formal, A.Actual);
+        Score += getSpecializationBonus(A.Formal, A.Actual, Visitor);
 
       // Discard unprofitable specialisations.
       if (!ForceSpecialization && Score <= 0)
         continue;
 
       // Create a new specialisation entry.
       auto &Spec = AllSpecs.emplace_back(F, S, Score);
       if (CS.getFunction() != F)
Cost FunctionSpecializer::getSpecializationCost(Function *F) {
   // Or if the lines of codes implies that this function is easy to get
   // inlined so that we shouldn't specialize it.
   if (Metrics.notDuplicatable || !Metrics.NumInsts.isValid() ||
       (!ForceSpecialization && !F->hasFnAttribute(Attribute::NoInline) &&
        Metrics.NumInsts < MinFunctionSize))
     return InstructionCost::getInvalid();
 
   // Otherwise, set the specialization cost to be the cost of all the
   // instructions in the function.
-  return Metrics.NumInsts * InlineConstants::getInstrCost();
-}
-
-static Cost getUserBonus(User *U, TargetTransformInfo &TTI,
-                         BlockFrequencyInfo &BFI) {
-  auto *I = dyn_cast_or_null<Instruction>(U);
-  // If not an instruction we do not know how to evaluate.
-  // Keep minimum possible cost for now so that it doesnt affect
-  // specialization.
-  if (!I)
-    return 0;
-
-  uint64_t Weight = BFI.getBlockFreq(I->getParent()).getFrequency() /
-                    BFI.getEntryFreq();
-  if (!Weight)
-    return 0;
-
-  Cost Bonus = Weight *
-      TTI.getInstructionCost(U, TargetTransformInfo::TCK_SizeAndLatency);
-
-  // Traverse recursively if there are more uses.
-  // TODO: Any other instructions to be added here?
-  if (I->mayReadFromMemory() || I->isCast())
-    for (auto *User : I->users())
-      Bonus += getUserBonus(User, TTI, BFI);
-
-  return Bonus;
+  return Metrics.NumInsts;
 }
 
 /// Compute a bonus for replacing argument \p A with constant \p C.
-Cost FunctionSpecializer::getSpecializationBonus(Argument *A, Constant *C) {
-  Function *F = A->getParent();
-  auto &TTI = (GetTTI)(*F);
-  auto &BFI = (GetBFI)(*F);
+Cost FunctionSpecializer::getSpecializationBonus(Argument *A, Constant *C,
+                                                 InstCostVisitor &Visitor) {
   LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing bonus for constant: "
                     << C->getNameOrAsOperand() << "\n");
 
   Cost TotalCost = 0;
-  for (auto *U : A->users()) {
-    TotalCost += getUserBonus(U, TTI, BFI);
-    LLVM_DEBUG(dbgs() << "FnSpecialization:   User cost ";
-               TotalCost.print(dbgs()); dbgs() << " for: " << *U << "\n");
-  }
+  for (auto *U : A->users())
+    if (auto *UI = dyn_cast<Instruction>(U))
+      if (Solver.isBlockExecutable(UI->getParent()))
+        TotalCost += Visitor.getUserBonus(UI, A, C);
+
+  LLVM_DEBUG(dbgs() << "FnSpecialization:   Accumulated user bonus "
+                    << TotalCost << " 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.
   Function *CalledFunction = dyn_cast<Function>(C->stripPointerCasts());
   if (!CalledFunction)
     return TotalCost;
 
   // 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;
   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;
 
     // Get the cost of inlining the called function at this call site. Note
     // that this is only an estimate. The called function may eventually
     // change in a way that leads to it not being inlined here, even though
     // inlining looks profitable now. For example, one of its called
     // functions may be inlined into it, making the called function too large
     // to be inlined into this call site.
     //
     // We apply a boost for performing indirect call promotion by increasing
     // the default threshold by the threshold for indirect calls.
     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;
     else if (IC.isVariable() && IC.getCostDelta() > 0)
       Bonus += IC.getCostDelta();
bool FunctionSpecializer::isArgumentInteresting(Argument *A) {
 
   Type *Ty = A->getType();
   if (!Ty->isPointerTy() && (!SpecializeLiteralConstant ||
       (!Ty->isIntegerTy() && !Ty->isFloatingPointTy() && !Ty->isStructTy())))
     return false;
 
   // SCCP solver does not record an argument that will be constructed on
   // stack.
   if (A->hasByValAttr() && !A->getParent()->onlyReadsMemory())
     return false;
 
   // For non-argument-tracked functions every argument is overdefined.
   if (!Solver.isArgumentTrackedFunction(A->getParent()))
     return true;
 
   // Check the lattice value and decide if we should attemt to specialize,
   // based on this argument. No point in specialization, if the lattice value

llvm/unittests/Transforms/IPO/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   IPO
   Support
   TargetParser
   TransformUtils
   )
 
 add_llvm_unittest(IPOTests
   LowerTypeTests.cpp
   WholeProgramDevirt.cpp
   AttributorTest.cpp
+  FunctionSpecializationTest.cpp
   )
 
 set_property(TARGET IPOTests PROPERTY FOLDER "Tests/UnitTests/TransformsTests")

llvm/unittests/Transforms/IPO/FunctionSpecializationTest.cpp

+//===- FunctionSpecializationTest.cpp - Cost model unit tests -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Transforms/IPO/FunctionSpecialization.h"
+#include "llvm/Transforms/Utils/SCCPSolver.h"
+#include "gtest/gtest.h"
+#include <memory>
+
+namespace llvm {
+
+class FunctionSpecializationTest : public testing::Test {
+protected:
+  LLVMContext Ctx;
+  FunctionAnalysisManager FAM;
+  std::unique_ptr<Module> M;
+  std::unique_ptr<SCCPSolver> Solver;
+
+  FunctionSpecializationTest() {
+    FAM.registerPass([&] { return TargetLibraryAnalysis(); });
+    FAM.registerPass([&] { return TargetIRAnalysis(); });
+    FAM.registerPass([&] { return BlockFrequencyAnalysis(); });
+    FAM.registerPass([&] { return BranchProbabilityAnalysis(); });
+    FAM.registerPass([&] { return LoopAnalysis(); });
+    FAM.registerPass([&] { return AssumptionAnalysis(); });
+    FAM.registerPass([&] { return DominatorTreeAnalysis(); });
+    FAM.registerPass([&] { return PostDominatorTreeAnalysis(); });
+    FAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+  }
+
+  Module &parseModule(const char *ModuleString) {
+    SMDiagnostic Err;
+    M = parseAssemblyString(ModuleString, Err, Ctx);
+    EXPECT_TRUE(M);
+    return *M;
+  }
+
+  FunctionSpecializer getSpecializerFor(Function *F) {
+    auto GetTLI = [this](Function &F) -> const TargetLibraryInfo & {
+      return FAM.getResult<TargetLibraryAnalysis>(F);
+    };
+    auto GetTTI = [this](Function &F) -> TargetTransformInfo & {
+      return FAM.getResult<TargetIRAnalysis>(F);
+    };
+    auto GetBFI = [this](Function &F) -> BlockFrequencyInfo & {
+      return FAM.getResult<BlockFrequencyAnalysis>(F);
+    };
+    auto GetAC = [this](Function &F) -> AssumptionCache & {
+      return FAM.getResult<AssumptionAnalysis>(F);
+    };
+    auto GetAnalysis = [this](Function &F) -> AnalysisResultsForFn {
+      DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
+      return { std::make_unique<PredicateInfo>(F, DT,
+                                FAM.getResult<AssumptionAnalysis>(F)),
+               &DT, FAM.getCachedResult<PostDominatorTreeAnalysis>(F) };
+    };
+
+    Solver = std::make_unique<SCCPSolver>(M->getDataLayout(), GetTLI, Ctx);
+
+    Solver->addAnalysis(*F, GetAnalysis(*F));
+    Solver->markBlockExecutable(&F->front());
+    for (Argument &Arg : F->args())
+      Solver->markOverdefined(&Arg);
+    Solver->solveWhileResolvedUndefsIn(*M);
+
+    return FunctionSpecializer(*Solver, *M, &FAM, GetBFI, GetTLI, GetTTI,
+                               GetAC);
+  }
+
+  Cost getInstCost(Instruction &I) {
+    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);
+  }
+};
+
+} // namespace llvm
+
+using namespace llvm;
+
+TEST_F(FunctionSpecializationTest, SwitchInst) {
+  const char *ModuleString = R"(
+    define void @foo(i32 %a, i32 %b, i32 %i) {
+    entry:
+      switch i32 %i, label %default
+      [ i32 1, label %case1
+        i32 2, label %case2 ]
+    case1:
+      %0 = mul i32 %a, 2
+      %1 = sub i32 6, 5
+      br label %bb1
+    case2:
+      %2 = and i32 %b, 3
+      %3 = sdiv i32 8, 2
+      br label %bb2
+    bb1:
+      %4 = add i32 %0, %b
+      br label %default
+    bb2:
+      %5 = or i32 %2, %a
+      br label %default
+    default:
+      ret void
+    }
+  )";
+
+  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();
+  BasicBlock &Case1 = *++FuncIter;
+  BasicBlock &Case2 = *++FuncIter;
+  BasicBlock &BB1 = *++FuncIter;
+  BasicBlock &BB2 = *++FuncIter;
+
+  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 &BrDefault = BB2.back();
+
+  // mul
+  Cost Ref = getInstCost(Mul);
+  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+
+  // and + or + add
+  Ref = getInstCost(And) + getInstCost(Or) + getInstCost(Add);
+  Bonus = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+
+  // sdiv + br + br
+  Ref = getInstCost(Sdiv) + getInstCost(BrBB2) + getInstCost(BrDefault);
+  Bonus = Specializer.getSpecializationBonus(F->getArg(2), One, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+}
+
+TEST_F(FunctionSpecializationTest, BranchInst) {
+  const char *ModuleString = R"(
+    define void @foo(i32 %a, i32 %b, i1 %cond) {
+    entry:
+      br i1 %cond, label %bb0, label %bb2
+    bb0:
+      %0 = mul i32 %a, 2
+      %1 = sub i32 6, 5
+      br label %bb1
+    bb1:
+      %2 = add i32 %0, %b
+      %3 = sdiv i32 8, 2
+      br label %bb2
+    bb2:
+      ret void
+    }
+  )";
+
+  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);
+  Constant *False = ConstantInt::getFalse(M.getContext());
+
+  auto FuncIter = F->begin();
+  BasicBlock &BB0 = *++FuncIter;
+  BasicBlock &BB1 = *++FuncIter;
+
+  Instruction &Mul = BB0.front();
+  Instruction &Sub = *++BB0.begin();
+  Instruction &BrBB1 = BB0.back();
+  Instruction &Add = BB1.front();
+  Instruction &Sdiv = *++BB1.begin();
+  Instruction &BrBB2 = BB1.back();
+
+  // mul
+  Cost Ref = getInstCost(Mul);
+  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+
+  // add
+  Ref = getInstCost(Add);
+  Bonus = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+
+  // sub + br + sdiv + br
+  Ref = getInstCost(Sub) + getInstCost(BrBB1) + getInstCost(Sdiv) +
+        getInstCost(BrBB2);
+  Bonus = Specializer.getSpecializationBonus(F->getArg(2), False, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+}
+
+TEST_F(FunctionSpecializationTest, Misc) {
+  const char *ModuleString = R"(
+    @g = constant [2 x i32] zeroinitializer, align 4
+
+    define i32 @foo(i8 %a, i1 %cond, ptr %b) {
+      %cmp = icmp eq i8 %a, 10
+      %ext = zext i1 %cmp to i32
+      %sel = select i1 %cond, i32 %ext, i32 1
+      %gep = getelementptr i32, ptr %b, i32 %sel
+      %ld = load i32, ptr %gep
+      ret i32 %ld
+    }
+  )";
+
+  Module &M = parseModule(ModuleString);
+  Function *F = M.getFunction("foo");
+  FunctionSpecializer Specializer = getSpecializerFor(F);
+  InstCostVisitor Visitor = Specializer.getInstCostVisitorFor(F);
+
+  GlobalVariable *GV = M.getGlobalVariable("g");
+  Constant *One = ConstantInt::get(IntegerType::getInt8Ty(M.getContext()), 1);
+  Constant *True = ConstantInt::getTrue(M.getContext());
+
+  auto BlockIter = F->front().begin();
+  Instruction &Icmp = *BlockIter++;
+  Instruction &Zext = *BlockIter++;
+  Instruction &Select = *BlockIter++;
+  Instruction &Gep = *BlockIter++;
+  Instruction &Load = *BlockIter++;
+
+  // icmp + zext
+  Cost Ref = getInstCost(Icmp) + getInstCost(Zext);
+  Cost Bonus = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+
+  // select
+  Ref = getInstCost(Select);
+  Bonus = Specializer.getSpecializationBonus(F->getArg(1), True, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+
+  // gep + load
+  Ref = getInstCost(Gep) + getInstCost(Load);
+  Bonus = Specializer.getSpecializationBonus(F->getArg(2), GV, Visitor);
+  EXPECT_EQ(Bonus, Ref);
+}