[FuncSpec] Extract the analysis part as an analysis pass

[ChuanqiXu]

Summary

I want the ThinLTO to support Function Specialization Pass.
Now the ThinLTO importer would only import functions by inline heuristics.
And if we want to make ThinLTO importer knows the information used by function specialization, we need to extract the analysis part as an analysis pass to make other pass get the information without transforming the IR.

There are much works to do since the function specialization pass itself is not strong enough first. The first step should be harmless.

This patch intended to be a NFC patch so here is nothing to factor the cost model. We could do that in following patches.

Test Plan: check-llvm

[ChuanqiXu]

Format codes.

[SjoerdMeijer]

I don't know much/anything about ThinLTO, which is why I don't really get what problem we need to solve:

And if we want to make ThinLTO importer knows the information used by function specialization, we need to extract the analysis part

But before we go down the path of making a lot of changes, I think it would be best if we are reasonably confident this is the right approach, which we can achieve by adding reviewers familiar with the ThinLTO challenges and by posting follow up patches using this work. Like I said, I don't know much about the ThinLTO pipeline, perhaps @fhahn can recommend someone.

[SjoerdMeijer]

Ah, I remembered your post to the mailing list about this which provides more context:

https://lists.llvm.org/pipermail/llvm-dev/2021-May/150443.html

Looks like you're raising fundamental ThinLTO questions that need answering first. Perhaps ping that thread, perhaps update it with some more information, so that hopefully someone with ThinLTO knowledge can shed a light on this approach.

[ChuanqiXu]

In D105524#2861297, @SjoerdMeijer wrote:

Ah, I remembered your post to the mailing list about this which provides more context:

https://lists.llvm.org/pipermail/llvm-dev/2021-May/150443.html

Looks like you're raising fundamental ThinLTO questions that need answering first. Perhaps ping that thread, perhaps update it with some more information, so that hopefully someone with ThinLTO knowledge can shed a light on this approach.

Yeah, I posted it before.

Then I made it in the downstream to see the impact I raised in the mailing list. And the impact looks good from my experiments. So I want to see if it is OK to contribute them to the upstream.

Since I want to make things happen in order, so I updated this patch first. Now it looks like I'd better to update the remained patch and discuss in the thread.

I would do that. Thanks for suggestion.

[ChuanqiXu]

Waiting to get in consensus that we would import functions in ThinLTO for sepcialization.

llvm/include/llvm/Analysis/FuncSpecCost.h

+//===- FuncSpecCost.h - Cost analysis for function specialization -*- C++
+//-*-===//
+//
+// 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
+//
+//===------------------------------------------------------------------------===//
+//
+// This file implements heuristics for function specialization decisions.
+//
+//===------------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_FUNCSPECCOST_H
+#define LLVM_ANALYSIS_FUNCSPECCOST_H
+
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/Argument.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/Value.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/InstructionCost.h"
+
+namespace llvm {
+
+class AssumptionCache;
+class TargetTransformInfo;
+class TargetLibraryInfo;
+
+/// Cost/Bonus information for specialize a function with
+/// each argument.
+class FuncSpecCostInfo {
+  /// The cost to specialize the function.
+  InstructionCost Cost;
+  /// Map from the number of the argument, to the base bonuss for specialize it.
+  /// Base bonus stands for the bonus we could get by the function body.
+  MapVector<unsigned, unsigned> SpecBonusBaseMap;
+
+  unsigned getBonusBase(Argument *Arg) const;
+  unsigned getBonusBase(unsigned Index) const;
+
+public:
+  InstructionCost getCost() const { return Cost; }
+  unsigned
+  getBonus(Argument *Arg, Constant *C,
+           function_ref<TargetTransformInfo &(Function &)> GetTTI,
+           function_ref<AssumptionCache &(Function &)> GetAC,
+           function_ref<const TargetLibraryInfo &(Function &)> GetTLI) const;
+
+  FuncSpecCostInfo() {}
+  FuncSpecCostInfo(FuncSpecCostInfo &&Other)
+      : FuncSpecCostInfo(std::move(Other.Cost),
+                         std::move(Other.SpecBonusBaseMap)) {}
+  FuncSpecCostInfo(InstructionCost Cost,
+                   MapVector<unsigned, unsigned> &&BonusBaseMap)
+      : Cost(std::move(Cost)) {
+    SpecBonusBaseMap.swap(BonusBaseMap);
+  }
+};
+
+/// Analysis pass which computes a \c FuncSpecCostInfo.
+class FunctionSpecializationAnalysis
+    : public AnalysisInfoMixin<FunctionSpecializationAnalysis> {
+  friend AnalysisInfoMixin<FunctionSpecializationAnalysis>;
+  static AnalysisKey Key;
+
+public:
+  using Result = FuncSpecCostInfo;
+
+  Result run(Function &F, FunctionAnalysisManager &);
+};
+
+/// Legacy analysis pass which computes a \c FuncSpecCostInfo.
+class FunctionSpecializationWrapperPass : public FunctionPass {
+  Optional<FuncSpecCostInfo> Info;
+
+public:
+  static char ID;
+
+  FunctionSpecializationWrapperPass();
+
+  FuncSpecCostInfo &getFuncSpecCost() { return *Info; }
+  const FuncSpecCostInfo &getFuncSpecCost() const { return *Info; }
+
+  bool runOnFunction(Function &F) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+  StringRef getPassName() const override;
+};
+
+} // namespace llvm
+
+#endif

llvm/include/llvm/InitializePasses.h

 void initializeFixupStatepointCallerSavedPass(PassRegistry&);
 void initializeFlattenCFGPassPass(PassRegistry&);
 void initializeFloat2IntLegacyPassPass(PassRegistry&);
 void initializeForceFunctionAttrsLegacyPassPass(PassRegistry&);
 void initializeForwardControlFlowIntegrityPass(PassRegistry&);
 void initializeFuncletLayoutPass(PassRegistry&);
 void initializeFunctionImportLegacyPassPass(PassRegistry&);
 void initializeFunctionSpecializationLegacyPassPass(PassRegistry &);
+void initializeFunctionSpecializationWrapperPassPass(PassRegistry &);
 void initializeGCMachineCodeAnalysisPass(PassRegistry&);
 void initializeGCModuleInfoPass(PassRegistry&);
 void initializeGCOVProfilerLegacyPassPass(PassRegistry&);
 void initializeGVNHoistLegacyPassPass(PassRegistry&);
 void initializeGVNLegacyPassPass(PassRegistry&);
 void initializeGVNSinkLegacyPassPass(PassRegistry&);
 void initializeGlobalDCELegacyPassPass(PassRegistry&);
 void initializeGlobalMergePass(PassRegistry&);

llvm/include/llvm/Transforms/Scalar/SCCP.h

 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Transforms/Utils/PredicateInfo.h"
 #include "llvm/Transforms/Utils/SCCPSolver.h"
 
 namespace llvm {
 
 class PostDominatorTree;
+class FuncSpecCostInfo;
 
 /// This pass performs function-level constant propagation and merging.
 class SCCPPass : public PassInfoMixin<SCCPPass> {
 public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 bool runIPSCCP(Module &M, const DataLayout &DL,
                std::function<const TargetLibraryInfo &(Function &)> GetTLI,
                function_ref<AnalysisResultsForFn(Function &)> getAnalysis);
 
 bool runFunctionSpecialization(
     Module &M, const DataLayout &DL,
     std::function<TargetLibraryInfo &(Function &)> GetTLI,
     std::function<TargetTransformInfo &(Function &)> GetTTI,
     std::function<AssumptionCache &(Function &)> GetAC,
+    std::function<FuncSpecCostInfo &(Function &)> GetFSCI,
     function_ref<AnalysisResultsForFn(Function &)> GetAnalysis);
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_SCCP_H

llvm/lib/Analysis/Analysis.cpp

void llvm::initializeAnalysis(PassRegistry &Registry) {
   initializeCFLSteensAAWrapperPassPass(Registry);
   initializeDependenceAnalysisWrapperPassPass(Registry);
   initializeDelinearizationPass(Registry);
   initializeDemandedBitsWrapperPassPass(Registry);
   initializeDominanceFrontierWrapperPassPass(Registry);
   initializeDomViewerPass(Registry);
   initializeDomPrinterPass(Registry);
   initializeDomOnlyViewerPass(Registry);
+  initializeFunctionSpecializationWrapperPassPass(Registry);
   initializePostDomViewerPass(Registry);
   initializeDomOnlyPrinterPass(Registry);
   initializePostDomPrinterPass(Registry);
   initializePostDomOnlyViewerPass(Registry);
   initializePostDomOnlyPrinterPass(Registry);
   initializeAAResultsWrapperPassPass(Registry);
   initializeGlobalsAAWrapperPassPass(Registry);
   initializeIVUsersWrapperPassPass(Registry);

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   DependenceAnalysis.cpp
   DependenceGraphBuilder.cpp
   DevelopmentModeInlineAdvisor.cpp
   DivergenceAnalysis.cpp
   DomPrinter.cpp
   DomTreeUpdater.cpp
   DominanceFrontier.cpp
   EHPersonalities.cpp
+  FuncSpecCost.cpp
   FunctionPropertiesAnalysis.cpp
   GlobalsModRef.cpp
   GuardUtils.cpp
   HeatUtils.cpp
   IRSimilarityIdentifier.cpp
   IVDescriptors.cpp
   IVUsers.cpp
   ImportedFunctionsInliningStatistics.cpp

llvm/lib/Analysis/FuncSpecCost.cpp

+//===- FuncSpecCost.cpp - Cost analysis for function specializer
+//-*---------===//
+//-*-===//
+//
+// 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
+//
+//===-----------------------------------------------------------------------===//
+//
+// This file implements heuristics for specializing function decisions.
+//
+//===-----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/FuncSpecCost.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Support/CommandLine.h"
+#include <cassert>
+#include <cmath>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "func-spec-cost"
+
+static cl::opt<unsigned>
+    AvgLoopIterationCount("func-specialization-avg-iters-cost", cl::Hidden,
+                          cl::desc("Average loop iteration count cost"),
+                          cl::init(10));
+
+static Function *getFunction(Value *CalledValue) {
+  // Since the argument is a function pointer, its incoming constant values
+  // should be functions or constant expressions. The code below attempts to
+  // look through cast expressions to find the function that will be called.
+  if (!isa<PointerType>(CalledValue->getType()) ||
+      !isa<FunctionType>(CalledValue->getType()->getPointerElementType()))
+    return nullptr;
+
+  while (isa<ConstantExpr>(CalledValue) &&
+         cast<ConstantExpr>(CalledValue)->isCast())
+    CalledValue = cast<User>(CalledValue)->getOperand(0);
+  return dyn_cast<Function>(CalledValue);
+}
+
+/// Bonus equals to basebonus plus extrabonus.
+/// The extrabonus means bonus we could find at the callsites.
+/// For example, if the callsite passes a function to the specializing function,
+/// we could compute the bonus for the callsite by the bonus of inline
+/// the function into the specializing function.
+unsigned FuncSpecCostInfo::getBonus(
+    Argument *Arg, Constant *C,
+    function_ref<TargetTransformInfo &(Function &)> GetTTI,
+    function_ref<AssumptionCache &(Function &)> GetAC,
+    function_ref<const TargetLibraryInfo &(Function &)> GetTLI) const {
+  LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing bonus for: " << *Arg
+                    << "\n");
+
+  unsigned BonusBase = getBonusBase(Arg);
+
+  // The below heuristic is only concerned with exposing inlining
+  // opportunities via indirect call promotion. If the argument is not a
+  // function pointer, give up.
+  Function *CalledFunction = getFunction(C);
+  if (!CalledFunction)
+    return BonusBase;
+
+  // Get TTI for the called function (used for the inline cost).
+  auto &CalleeTTI = (GetTTI)(*CalledFunction);
+
+  int Bonus = 0;
+  for (User *U : Arg->users()) {
+    if (!isa<CallInst>(U) && !isa<InvokeInst>(U))
+      continue;
+    auto *CS = cast<CallBase>(U);
+    if (CS->getCalledOperand() != Arg)
+      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();
+  }
+
+  return BonusBase + Bonus;
+}
+
+unsigned FuncSpecCostInfo::getBonusBase(Argument *Arg) const {
+  return getBonusBase(Arg->getArgNo());
+}
+
+unsigned FuncSpecCostInfo::getBonusBase(unsigned Index) const {
+  if (!SpecBonusBaseMap.count(Index))
+    return 0;
+  return SpecBonusBaseMap.lookup(Index);
+}
+
+static InstructionCost getSpecializationCost(Function &F, AssumptionCache &AC,
+                                             TargetTransformInfo &TTI) {
+  // Compute the code metrics for the function.
+  SmallPtrSet<const Value *, 32> EphValues;
+  CodeMetrics::collectEphemeralValues(&F, &AC, EphValues);
+  CodeMetrics Metrics;
+  for (BasicBlock &BB : F)
+    Metrics.analyzeBasicBlock(&BB, TTI, EphValues);
+
+  // If the code metrics reveal that we shouldn't duplicate the function, we
+  // shouldn't specialize it. Set the specialization cost to Invalid.
+  if (Metrics.notDuplicatable) {
+    InstructionCost C{};
+    C.setInvalid();
+    return C;
+  }
+
+  return Metrics.NumInsts * InlineConstants::InstrCost;
+}
+
+static InstructionCost getUserBonus(User *U, TargetTransformInfo &TTI,
+                                    LoopInfo &LI) {
+  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 std::numeric_limits<unsigned>::min();
+
+  InstructionCost Cost =
+      TTI.getUserCost(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())
+      Cost += getUserBonus(User, TTI, LI);
+
+  // Increase the cost if it is inside the loop.
+  unsigned LoopDepth = LI.getLoopDepth(I->getParent());
+  Cost *= std::pow((long double)AvgLoopIterationCount, LoopDepth);
+  return Cost;
+}
+
+static FuncSpecCostInfo CreateFuncSpecCostInfo(Function &F, AssumptionCache &AC,
+                                               TargetTransformInfo &TTI,
+                                               LoopInfo &LI) {
+  InstructionCost SpecializeCost = getSpecializationCost(F, AC, TTI);
+  MapVector<unsigned, unsigned> SpecBonusBaseMap;
+  for (Argument &A : F.args()) {
+    InstructionCost Cost = 0;
+    for (auto *U : A.users())
+      Cost += getUserBonus(U, TTI, LI);
+
+    if (Cost.isValid() && Cost > 0)
+      SpecBonusBaseMap.insert({A.getArgNo(), *Cost.getValue()});
+  }
+  return FuncSpecCostInfo(SpecializeCost, std::move(SpecBonusBaseMap));
+}
+
+AnalysisKey FunctionSpecializationAnalysis::Key;
+
+FuncSpecCostInfo
+FunctionSpecializationAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
+  TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
+  LoopInfo &LI = FAM.getResult<LoopAnalysis>(F);
+  AssumptionCache &AC = FAM.getResult<AssumptionAnalysis>(F);
+  return CreateFuncSpecCostInfo(F, AC, TTI, LI);
+}
+
+char FunctionSpecializationWrapperPass::ID = 0;
+INITIALIZE_PASS_BEGIN(FunctionSpecializationWrapperPass, "func-spec-cost",
+                      "Function Specialization Cost Analysis", false, true)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
+INITIALIZE_PASS_END(FunctionSpecializationWrapperPass, "func-spec-cost",
+                    "Function Specialization Cost Analysis", false, true)
+
+FunctionSpecializationWrapperPass::FunctionSpecializationWrapperPass()
+    : FunctionPass(ID) {
+  initializeFunctionSpecializationWrapperPassPass(
+      *PassRegistry::getPassRegistry());
+}
+
+void FunctionSpecializationWrapperPass::getAnalysisUsage(
+    AnalysisUsage &AU) const {
+  AU.setPreservesAll();
+  AU.addRequired<AssumptionCacheTracker>();
+  AU.addRequired<LoopInfoWrapperPass>();
+  AU.addRequired<TargetTransformInfoWrapperPass>();
+}
+
+bool FunctionSpecializationWrapperPass::runOnFunction(Function &F) {
+  TargetTransformInfo &TTI =
+      getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+  AssumptionCache &AC =
+      getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+  Info.emplace(CreateFuncSpecCostInfo(F, AC, TTI, LI));
+  return false;
+}
+
+StringRef FunctionSpecializationWrapperPass::getPassName() const {
+  return "Function Specialization Cost Analysis";
+}
 No newline at end of file

llvm/lib/Passes/PassBuilder.cpp

 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/DDG.h"
 #include "llvm/Analysis/DDGPrinter.h"
 #include "llvm/Analysis/Delinearization.h"
 #include "llvm/Analysis/DemandedBits.h"
 #include "llvm/Analysis/DependenceAnalysis.h"
 #include "llvm/Analysis/DivergenceAnalysis.h"
 #include "llvm/Analysis/DominanceFrontier.h"
+#include "llvm/Analysis/FuncSpecCost.h"
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/IRSimilarityIdentifier.h"
 #include "llvm/Analysis/IVUsers.h"
 #include "llvm/Analysis/InlineAdvisor.h"
 #include "llvm/Analysis/InlineSizeEstimatorAnalysis.h"
 #include "llvm/Analysis/InstCount.h"
 #include "llvm/Analysis/LazyCallGraph.h"

llvm/lib/Passes/PassRegistry.def

 FUNCTION_ANALYSIS("assumptions", AssumptionAnalysis())
 FUNCTION_ANALYSIS("block-freq", BlockFrequencyAnalysis())
 FUNCTION_ANALYSIS("branch-prob", BranchProbabilityAnalysis())
 FUNCTION_ANALYSIS("domtree", DominatorTreeAnalysis())
 FUNCTION_ANALYSIS("postdomtree", PostDominatorTreeAnalysis())
 FUNCTION_ANALYSIS("demanded-bits", DemandedBitsAnalysis())
 FUNCTION_ANALYSIS("domfrontier", DominanceFrontierAnalysis())
 FUNCTION_ANALYSIS("func-properties", FunctionPropertiesAnalysis())
+FUNCTION_ANALYSIS("func-spec-analysis", FunctionSpecializationAnalysis())
 FUNCTION_ANALYSIS("loops", LoopAnalysis())
 FUNCTION_ANALYSIS("lazy-value-info", LazyValueAnalysis())
 FUNCTION_ANALYSIS("da", DependenceAnalysis())
 FUNCTION_ANALYSIS("inliner-size-estimator", InlineSizeEstimatorAnalysis())
 FUNCTION_ANALYSIS("memdep", MemoryDependenceAnalysis())
 FUNCTION_ANALYSIS("memoryssa", MemorySSAAnalysis())
 FUNCTION_ANALYSIS("phi-values", PhiValuesAnalysis())
 FUNCTION_ANALYSIS("regions", RegionInfoAnalysis())

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 //   and thus control compile-times / code-size.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/FuncSpecCost.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Transforms/Scalar/SCCP.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 #include <cmath>
static cl::opt<unsigned> FuncSpecializationMaxIters(
     cl::init(1));
 
 static cl::opt<unsigned> MaxConstantsThreshold(
     "func-specialization-max-constants", cl::Hidden,
     cl::desc("The maximum number of clones allowed for a single function "
              "specialization"),
     cl::init(3));
 
-static cl::opt<unsigned>
-    AvgLoopIterationCount("func-specialization-avg-iters-cost", cl::Hidden,
-                          cl::desc("Average loop iteration count cost"),
-                          cl::init(10));
-
 static cl::opt<bool> EnableSpecializationForLiteralConstant(
     "function-specialization-for-literal-constant", cl::init(false), cl::Hidden,
     cl::desc("Make function specialization available for literal constant."));
 
 // Helper to check if \p LV is either overdefined or a constant int.
 static bool isOverdefined(const ValueLatticeElement &LV) {
   return !LV.isUnknownOrUndef() && !LV.isConstant();
 }
 
 class FunctionSpecializer {
 
   /// The IPSCCP Solver.
   SCCPSolver &Solver;
 
   /// Analyses used to help determine if a function should be specialized.
   std::function<AssumptionCache &(Function &)> GetAC;
   std::function<TargetTransformInfo &(Function &)> GetTTI;
   std::function<TargetLibraryInfo &(Function &)> GetTLI;
+  std::function<FuncSpecCostInfo &(Function &)> GetFSCI;
 
   SmallPtrSet<Function *, 2> SpecializedFuncs;
 
 public:
   FunctionSpecializer(SCCPSolver &Solver,
                       std::function<AssumptionCache &(Function &)> GetAC,
                       std::function<TargetTransformInfo &(Function &)> GetTTI,
-                      std::function<TargetLibraryInfo &(Function &)> GetTLI)
-      : Solver(Solver), GetAC(GetAC), GetTTI(GetTTI), GetTLI(GetTLI) {}
+                      std::function<TargetLibraryInfo &(Function &)> GetTLI,
+                      std::function<FuncSpecCostInfo &(Function &)> GetFSCI)
+      : Solver(Solver), GetAC(GetAC), GetTTI(GetTTI), GetTLI(GetTLI),
+        GetFSCI(GetFSCI) {}
 
   /// Attempt to specialize functions in the module to enable constant
   /// propagation across function boundaries.
   ///
   /// \returns true if at least one function is specialized.
   bool
   specializeFunctions(SmallVectorImpl<Function *> &FuncDecls,
                       SmallVectorImpl<Function *> &CurrentSpecializations) {
for (Argument &A : F->args()) {
       return true;
     }
 
     return false;
   }
 
   /// Compute the cost of specializing function \p F.
   InstructionCost getSpecializationCost(Function *F) {
-    // Compute the code metrics for the function.
-    SmallPtrSet<const Value *, 32> EphValues;
-    CodeMetrics::collectEphemeralValues(F, &(GetAC)(*F), EphValues);
-    CodeMetrics Metrics;
-    for (BasicBlock &BB : *F)
-      Metrics.analyzeBasicBlock(&BB, (GetTTI)(*F), EphValues);
-
-    // If the code metrics reveal that we shouldn't duplicate the function, we
-    // shouldn't specialize it. Set the specialization cost to Invalid.
-    if (Metrics.notDuplicatable) {
-      InstructionCost C{};
-      C.setInvalid();
-      return C;
-    }
-
-    // Otherwise, set the specialization cost to be the cost of all the
-    // instructions in the function and penalty for specializing more functions.
-    unsigned Penalty = NbFunctionsSpecialized + 1;
-    return Metrics.NumInsts * InlineConstants::InstrCost * Penalty;
-  }
-
-  InstructionCost getUserBonus(User *U, llvm::TargetTransformInfo &TTI,
-                               LoopInfo &LI) {
-    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 std::numeric_limits<unsigned>::min();
-
-    auto Cost = TTI.getUserCost(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())
-        Cost += getUserBonus(User, TTI, LI);
-
-    // Increase the cost if it is inside the loop.
-    auto LoopDepth = LI.getLoopDepth(I->getParent());
-    Cost *= std::pow((double)AvgLoopIterationCount, LoopDepth);
-    return Cost;
+    FuncSpecCostInfo &FSCI = GetFSCI(*F);
+    // Penalty for specializing more functions.
+    unsigned Penalty = (NumFuncSpecialized + 1);
+    return FSCI.getCost() * Penalty;
   }
 
   /// Compute a bonus for replacing argument \p A with constant \p C.
   InstructionCost getSpecializationBonus(Argument *A, Constant *C) {
-    Function *F = A->getParent();
-    DominatorTree DT(*F);
-    LoopInfo LI(DT);
-    auto &TTI = (GetTTI)(*F);
-    LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing bonus for: " << *A
-                      << "\n");
-
-    InstructionCost TotalCost = 0;
-    for (auto *U : A->users()) {
-      TotalCost += getUserBonus(U, TTI, LI);
-      LLVM_DEBUG(dbgs() << "FnSpecialization: User cost ";
-                 TotalCost.print(dbgs()); dbgs() << " for: " << *U << "\n");
-    }
-
-    // The below heuristic is only concerned with exposing inlining
-    // opportunities via indirect call promotion. If the argument is not a
-    // function pointer, give up.
-    if (!isa<PointerType>(A->getType()) ||
-        !isa<FunctionType>(A->getType()->getPointerElementType()))
-      return TotalCost;
-
-    // Since the argument is a function pointer, its incoming constant values
-    // should be functions or constant expressions. The code below attempts to
-    // look through cast expressions to find the function that will be called.
-    Value *CalledValue = C;
-    while (isa<ConstantExpr>(CalledValue) &&
-           cast<ConstantExpr>(CalledValue)->isCast())
-      CalledValue = cast<User>(CalledValue)->getOperand(0);
-    Function *CalledFunction = dyn_cast<Function>(CalledValue);
-    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;
-
-      // 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();
-    }
-
-    return TotalCost + Bonus;
+    FuncSpecCostInfo &FSCI = GetFSCI(*A->getParent());
+    return FSCI.getBonus(A, C, GetTTI, GetAC, GetTLI);
   }
 
   /// Determine if we should specialize a function based on the incoming values
   /// of the given argument.
   ///
   /// This function implements the goal-directed heuristic. It determines if
   /// specializing the function based on the incoming values of argument \p A
   /// would result in any significant optimization opportunities. If
static void cleanup(Module &M) {
   }
 }
 
 bool llvm::runFunctionSpecialization(
     Module &M, const DataLayout &DL,
     std::function<TargetLibraryInfo &(Function &)> GetTLI,
     std::function<TargetTransformInfo &(Function &)> GetTTI,
     std::function<AssumptionCache &(Function &)> GetAC,
+    std::function<FuncSpecCostInfo &(Function &)> GetFSCI,
     function_ref<AnalysisResultsForFn(Function &)> GetAnalysis) {
   SCCPSolver Solver(DL, GetTLI, M.getContext());
-  FunctionSpecializer FS(Solver, GetAC, GetTTI, GetTLI);
+  FunctionSpecializer FS(Solver, GetAC, GetTTI, GetTLI, GetFSCI);
   bool Changed = false;
 
   // Loop over all functions, marking arguments to those with their addresses
   // taken or that are external as overdefined.
   for (Function &F : M) {
     if (F.isDeclaration())
       continue;
     if (F.hasFnAttribute(Attribute::NoDuplicate))

llvm/lib/Transforms/IPO/SCCP.cpp

 //===-- SCCP.cpp ----------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements Interprocedural Sparse Conditional Constant Propagation.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/SCCP.h"
 #include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/FuncSpecCost.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Scalar/SCCP.h"
 
 using namespace llvm;
auto GetTLI = [&FAM](Function &F) -> TargetLibraryInfo & {
     return FAM.getResult<TargetLibraryAnalysis>(F);
   };
   auto GetTTI = [&FAM](Function &F) -> TargetTransformInfo & {
     return FAM.getResult<TargetIRAnalysis>(F);
   };
   auto GetAC = [&FAM](Function &F) -> AssumptionCache & {
     return FAM.getResult<AssumptionAnalysis>(F);
   };
+  auto GetFSCI = [&FAM](Function &F) -> FuncSpecCostInfo & {
+    return FAM.getResult<FunctionSpecializationAnalysis>(F);
+  };
   auto GetAnalysis = [&FAM](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)};
   };
 
-  if (!runFunctionSpecialization(M, DL, GetTLI, GetTTI, GetAC, GetAnalysis))
+  if (!runFunctionSpecialization(M, DL, GetTLI, GetTTI, GetAC, GetFSCI,
+                                 GetAnalysis))
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserve<DominatorTreeAnalysis>();
   PA.preserve<PostDominatorTreeAnalysis>();
   PA.preserve<FunctionAnalysisManagerModuleProxy>();
   return PA;
 }
 
 struct FunctionSpecializationLegacyPass : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   FunctionSpecializationLegacyPass() : ModulePass(ID) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
+    AU.addRequired<FunctionSpecializationWrapperPass>();
   }
 
   virtual bool runOnModule(Module &M) override {
     if (skipModule(M))
       return false;
 
     const DataLayout &DL = M.getDataLayout();
     auto GetTLI = [this](Function &F) -> TargetLibraryInfo & {
       return this->getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
     };
     auto GetTTI = [this](Function &F) -> TargetTransformInfo & {
       return this->getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
     };
     auto GetAC = [this](Function &F) -> AssumptionCache & {
       return this->getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
     };
+    auto GetFSCI = [this](Function &F) -> FuncSpecCostInfo & {
+      return this->getAnalysis<FunctionSpecializationWrapperPass>(F)
+          .getFuncSpecCost();
+    };
 
     auto GetAnalysis = [this](Function &F) -> AnalysisResultsForFn {
       DominatorTree &DT =
           this->getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
       return {
           std::make_unique<PredicateInfo>(
               F, DT,
               this->getAnalysis<AssumptionCacheTracker>().getAssumptionCache(
                   F)),
           nullptr,  // We cannot preserve the DT or PDT with the legacy pass
           nullptr}; // manager, so set them to nullptr.
     };
-    return runFunctionSpecialization(M, DL, GetTLI, GetTTI, GetAC, GetAnalysis);
+    return runFunctionSpecialization(M, DL, GetTLI, GetTTI, GetAC, GetFSCI,
+                                     GetAnalysis);
   }
 };
 
 char FunctionSpecializationLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(
     FunctionSpecializationLegacyPass, "function-specialization",
     "Propagate constant arguments by specializing the function", false, false)