[llvm][NFC] Replace CallSite with CallBase in Inliner

[mtrofin]

Summary

*Almost* all uses are replaced. Left FIXMEs for the two sites that
require refactoring outside of Inliner, to scope this patch.

[dblaikie]

Looks good!

llvm/lib/Transforms/IPO/Inliner.cpp

948–949

Use 'auto *' here, now that CS is a pointer type

1156–1157

Side note (not related to this patch (saw some other instances while reviewing this too, though) - but might be handy for reference otherwise) - generally for a locally used lambda like this (especially a lambda only used within its expression) I'd encourage using [&] rather than an explicit capture list - like any other nested scope, it should have access to all the enclosing variables implicitly.

[mtrofin]

fix Harbormaster lint failure

llvm/lib/Transforms/IPO/Inliner.cpp

1156–1157

Ack - subsequent change. There are also naming inconsistencies (for (int i, instead of int I, etc))

[dblaikie]

llvm/lib/Transforms/IPO/Inliner.cpp

519

Yeah - this looks like Harbormaster just couldn't tell you didn't introduce this problem (the 'message' name not following LLVM naming conventions) - so I'd generally advise ignoring such a suggestion, or fixing it in an independent commit - no big deal/wouldn't bother fussing about it if you'd rather just commit it together this time.

[davidxl]

What is the convention here? Pointer CallBase* vs CallBase& ?

[dblaikie]

In D77817#1972599, @davidxl wrote:

What is the convention here? Pointer CallBase* vs CallBase& ?

I'd say "Same as other llvm::Instructions" - and I /assume/ the answer there is that they're still predominantly handled as pointers, rather than references?

Seems for /literally/ Instruction*/Instruction& (as opposed to all possible types of Instructions - harder to search/gather data for all of them) it's about 7:1 in favor of *, but that's hard to evaluate/compare - some of those may need to be pointers because null is an option.

Evidently references to Instructions aren't totally unheard of, so the preference "use references when something can't be null" could be observed here. A pretty clear endorsement of this is that "instructions(F)" gives you an iterator over Instruction&, not over Instruction*.

But yeah - given the current codebase I think it's dealer's choice & not something I'd push strongly either way in review personally. (but I don't do lots of Instruction handling, so maybe some other folks have stronger ideas about which direction this sort of thing should be going in)

[mtrofin]

In D77817#1972613, @dblaikie wrote:

In D77817#1972599, @davidxl wrote:

What is the convention here? Pointer CallBase* vs CallBase& ?

I'd say "Same as other llvm::Instructions" - and I /assume/ the answer there is that they're still predominantly handled as pointers, rather than references?

Seems for /literally/ Instruction*/Instruction& (as opposed to all possible types of Instructions - harder to search/gather data for all of them) it's about 7:1 in favor of *, but that's hard to evaluate/compare - some of those may need to be pointers because null is an option.

Evidently references to Instructions aren't totally unheard of, so the preference "use references when something can't be null" could be observed here. A pretty clear endorsement of this is that "instructions(F)" gives you an iterator over Instruction&, not over Instruction*.

This seems like a reasonable guideline, though - i.e. using references in function signatures where the contract is the parameter must be non-null. Looking at the refactored sites, there are a few spots where the caller is assumed to pass a valid CallSite. While that wasn't checked, now may be a good time to have that signal.

Example: setInlineRemark

But yeah - given the current codebase I think it's dealer's choice & not something I'd push strongly either way in review personally. (but I don't do lots of Instruction handling, so maybe some other folks have stronger ideas about which direction this sort of thing should be going in)

[mtrofin]

CallBase* -> CallBase& in places where the contract is that the value can't be null.

[davidxl]

The old CallSite based interface is a mixed bag -- it is passed by value, but behaves like a reference sometimes but also supports 'null/boolean' check which is pointer like.

looks good

llvm/include/llvm/Transforms/IPO/Inliner.h

struct LegacyInlinerBase : public CallGraphSCCPass {
 
   /// Remove now-dead linkonce functions at the end of processing to avoid
   /// breaking the SCC traversal.
   bool doFinalization(CallGraph &CG) override;
 
   /// This method must be implemented by the subclass to determine the cost of
   /// inlining the specified call site.  If the cost returned is greater than
   /// the current inline threshold, the call site is not inlined.
+  // FIXME(mtrofin): remove this in favor of the CallBase-based one
   virtual InlineCost getInlineCost(CallSite CS) = 0;
 
+  /// This method must be implemented by the subclass to determine the cost of
+  /// inlining the specified call site.  If the cost returned is greater than
+  /// the current inline threshold, the call site is not inlined.
+  virtual InlineCost getInlineCost(CallBase &CB) {
+    return getInlineCost(CallSite(&CB));
+  }
+
   /// Remove dead functions.
   ///
   /// This also includes a hack in the form of the 'AlwaysInlineOnly' flag
   /// which restricts it to deleting functions with an 'AlwaysInline'
   /// attribute. This is useful for the InlineAlways pass that only wants to
   /// deal with that subset of the functions.
   bool removeDeadFunctions(CallGraph &CG, bool AlwaysInlineOnly = false);
 

llvm/lib/Transforms/IPO/Inliner.cpp

 /// do so and update the CallGraph for this operation.
 ///
 /// This function also does some basic book-keeping to update the IR.  The
 /// InlinedArrayAllocas map keeps track of any allocas that are already
 /// available from other functions inlined into the caller.  If we are able to
 /// inline this call site we attempt to reuse already available allocas or add
 /// any new allocas to the set if not possible.
 static InlineResult InlineCallIfPossible(
-    CallSite CS, InlineFunctionInfo &IFI,
+    CallBase &CS, InlineFunctionInfo &IFI,
     InlinedArrayAllocasTy &InlinedArrayAllocas, int InlineHistory,
     bool InsertLifetime, function_ref<AAResults &(Function &)> &AARGetter,
     ImportedFunctionsInliningStatistics &ImportedFunctionsStats) {
   Function *Callee = CS.getCalledFunction();
   Function *Caller = CS.getCaller();
 
   AAResults &AAR = AARGetter(*Callee);
 
   // Try to inline the function.  Get the list of static allocas that were
   // inlined.
-  InlineResult IR = InlineFunction(CS, IFI, &AAR, InsertLifetime);
+  InlineResult IR = InlineFunction(&CS, IFI, &AAR, InsertLifetime);
   if (!IR.isSuccess())
     return IR;
 
   if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
     ImportedFunctionsStats.recordInline(*Caller, *Callee);
 
   AttributeFuncs::mergeAttributesForInlining(*Caller, *Callee);
 
   if (!DisableInlinedAllocaMerging)
     mergeInlinedArrayAllocas(Caller, IFI, InlinedArrayAllocas, InlineHistory);
 
   return IR; // success
 }
 
 /// Return true if inlining of CS can block the caller from being
 /// inlined which is proved to be more beneficial. \p IC is the
 /// estimated inline cost associated with callsite \p CS.
 /// \p TotalSecondaryCost will be set to the estimated cost of inlining the
 /// caller if \p CS is suppressed for inlining.
 static bool
-shouldBeDeferred(Function *Caller, CallSite CS, InlineCost IC,
-                 int &TotalSecondaryCost,
-                 function_ref<InlineCost(CallSite CS)> GetInlineCost) {
+shouldBeDeferred(Function *Caller, InlineCost IC, int &TotalSecondaryCost,
+                 function_ref<InlineCost(CallBase &CS)> GetInlineCost) {
   // For now we only handle local or inline functions.
   if (!Caller->hasLocalLinkage() && !Caller->hasLinkOnceODRLinkage())
     return false;
   // If the cost of inlining CS is non-positive, it is not going to prevent the
   // caller from being inlined into its callers and hence we don't need to
   // defer.
   if (IC.getCost() <= 0)
     return false;
shouldBeDeferred(Function *Caller, InlineCost IC, int &TotalSecondaryCost,
   // This bool tracks what happens if we DO inline C into B.
   bool inliningPreventsSomeOuterInline = false;
   for (User *U : Caller->users()) {
     // If the caller will not be removed (either because it does not have a
     // local linkage or because the LastCallToStaticBonus has been already
     // applied), then we can exit the loop early.
     if (!ApplyLastCallBonus && TotalSecondaryCost >= IC.getCost())
       return false;
-    CallSite CS2(U);
+    CallBase *CS2 = dyn_cast<CallBase>(U);
 
     // If this isn't a call to Caller (it could be some other sort
     // of reference) skip it.  Such references will prevent the caller
     // from being removed.
-    if (!CS2 || CS2.getCalledFunction() != Caller) {
+    if (!CS2 || CS2->getCalledFunction() != Caller) {
       ApplyLastCallBonus = false;
       continue;
     }
 
-    InlineCost IC2 = GetInlineCost(CS2);
+    InlineCost IC2 = GetInlineCost(*CS2);
     ++NumCallerCallersAnalyzed;
     if (!IC2) {
       ApplyLastCallBonus = false;
       continue;
     }
     if (IC2.isAlways())
       continue;
 
static std::string inlineCostStr(const InlineCost &IC) {
   Remark << IC;
   return Remark.str();
 }
 
 /// Return the cost only if the inliner should attempt to inline at the given
 /// CallSite. If we return the cost, we will emit an optimisation remark later
 /// using that cost, so we won't do so from this function.
 static Optional<InlineCost>
-shouldInline(CallSite CS, function_ref<InlineCost(CallSite CS)> GetInlineCost,
+shouldInline(CallBase &CS, function_ref<InlineCost(CallBase &CS)> GetInlineCost,
              OptimizationRemarkEmitter &ORE) {
   using namespace ore;
 
   InlineCost IC = GetInlineCost(CS);
-  Instruction *Call = CS.getInstruction();
+  Instruction *Call = &CS;
   Function *Callee = CS.getCalledFunction();
   Function *Caller = CS.getCaller();
 
   if (IC.isAlways()) {
     LLVM_DEBUG(dbgs() << "    Inlining " << inlineCostStr(IC)
-                      << ", Call: " << *CS.getInstruction() << "\n");
+                      << ", Call: " << CS << "\n");
     return IC;
   }
 
   if (IC.isNever()) {
     LLVM_DEBUG(dbgs() << "    NOT Inlining " << inlineCostStr(IC)
-                      << ", Call: " << *CS.getInstruction() << "\n");
+                      << ", Call: " << CS << "\n");
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "NeverInline", Call)
              << NV("Callee", Callee) << " not inlined into "
              << NV("Caller", Caller) << " because it should never be inlined "
              << IC;
     });
     return IC;
   }
 
   if (!IC) {
     LLVM_DEBUG(dbgs() << "    NOT Inlining " << inlineCostStr(IC)
-                      << ", Call: " << *CS.getInstruction() << "\n");
+                      << ", Call: " << CS << "\n");
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "TooCostly", Call)
              << NV("Callee", Callee) << " not inlined into "
              << NV("Caller", Caller) << " because too costly to inline " << IC;
     });
     return IC;
   }
 
   int TotalSecondaryCost = 0;
-  if (shouldBeDeferred(Caller, CS, IC, TotalSecondaryCost, GetInlineCost)) {
-    LLVM_DEBUG(dbgs() << "    NOT Inlining: " << *CS.getInstruction()
+  if (shouldBeDeferred(Caller, IC, TotalSecondaryCost, GetInlineCost)) {
+    LLVM_DEBUG(dbgs() << "    NOT Inlining: " << CS
                       << " Cost = " << IC.getCost()
                       << ", outer Cost = " << TotalSecondaryCost << '\n');
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "IncreaseCostInOtherContexts",
                                       Call)
              << "Not inlining. Cost of inlining " << NV("Callee", Callee)
              << " increases the cost of inlining " << NV("Caller", Caller)
              << " in other contexts";
     });
 
     // IC does not bool() to false, so get an InlineCost that will.
     // This will not be inspected to make an error message.
     return None;
   }
 
-  LLVM_DEBUG(dbgs() << "    Inlining " << inlineCostStr(IC)
-                    << ", Call: " << *CS.getInstruction() << '\n');
+  LLVM_DEBUG(dbgs() << "    Inlining " << inlineCostStr(IC) << ", Call: " << CS
+                    << '\n');
   return IC;
 }
 
 /// Return true if the specified inline history ID
 /// indicates an inline history that includes the specified function.
 static bool InlineHistoryIncludes(
     Function *F, int InlineHistoryID,
     const SmallVectorImpl<std::pair<Function *, int>> &InlineHistory) {
ORE.emit([&]() {
     bool AlwaysInline = IC.isAlways();
     StringRef RemarkName = AlwaysInline ? "AlwaysInline" : "Inlined";
     return OptimizationRemark(DEBUG_TYPE, RemarkName, DLoc, Block)
            << ore::NV("Callee", &Callee) << " inlined into "
            << ore::NV("Caller", &Caller) << " with " << IC;
   });
 }
 
-static void setInlineRemark(CallSite &CS, StringRef message) {
+static void setInlineRemark(CallBase &CS, StringRef Message) {
   if (!InlineRemarkAttribute)
     return;
 
-  Attribute attr = Attribute::get(CS->getContext(), "inline-remark", message);
+  Attribute attr = Attribute::get(CS.getContext(), "inline-remark", Message);
   CS.addAttribute(AttributeList::FunctionIndex, attr);
 }
 
 static bool
 inlineCallsImpl(CallGraphSCC &SCC, CallGraph &CG,
                 std::function<AssumptionCache &(Function &)> GetAssumptionCache,
                 ProfileSummaryInfo *PSI,
                 std::function<const TargetLibraryInfo &(Function &)> GetTLI,
                 bool InsertLifetime,
-                function_ref<InlineCost(CallSite CS)> GetInlineCost,
+                function_ref<InlineCost(CallBase &CS)> GetInlineCost,
                 function_ref<AAResults &(Function &)> AARGetter,
                 ImportedFunctionsInliningStatistics &ImportedFunctionsStats) {
   SmallPtrSet<Function *, 8> SCCFunctions;
   LLVM_DEBUG(dbgs() << "Inliner visiting SCC:");
   for (CallGraphNode *Node : SCC) {
     Function *F = Node->getFunction();
     if (F)
       SCCFunctions.insert(F);
     LLVM_DEBUG(dbgs() << " " << (F ? F->getName() : "INDIRECTNODE"));
   }
 
   // Scan through and identify all call sites ahead of time so that we only
   // inline call sites in the original functions, not call sites that result
   // from inlining other functions.
-  SmallVector<std::pair<CallSite, int>, 16> CallSites;
+  SmallVector<std::pair<CallBase *, int>, 16> CallSites;
 
   // When inlining a callee produces new call sites, we want to keep track of
   // the fact that they were inlined from the callee.  This allows us to avoid
   // infinite inlining in some obscure cases.  To represent this, we use an
   // index into the InlineHistory vector.
   SmallVector<std::pair<Function *, int>, 8> InlineHistory;
 
   for (CallGraphNode *Node : SCC) {
     Function *F = Node->getFunction();
     if (!F || F->isDeclaration())
       continue;
 
     OptimizationRemarkEmitter ORE(F);
     for (BasicBlock &BB : *F)
       for (Instruction &I : BB) {
-        CallSite CS(cast<Value>(&I));
+        auto *CS = dyn_cast<CallBase>(&I);
         // If this isn't a call, or it is a call to an intrinsic, it can
         // never be inlined.
         if (!CS || isa<IntrinsicInst>(I))
           continue;
 
         // If this is a direct call to an external function, we can never inline
         // it.  If it is an indirect call, inlining may resolve it to be a
         // direct call, so we keep it.
-        if (Function *Callee = CS.getCalledFunction())
+        if (Function *Callee = CS->getCalledFunction())
           if (Callee->isDeclaration()) {
             using namespace ore;
 
-            setInlineRemark(CS, "unavailable definition");
+            setInlineRemark(*CS, "unavailable definition");
             ORE.emit([&]() {
               return OptimizationRemarkMissed(DEBUG_TYPE, "NoDefinition", &I)
                      << NV("Callee", Callee) << " will not be inlined into "
-                     << NV("Caller", CS.getCaller())
+                     << NV("Caller", CS->getCaller())
                      << " because its definition is unavailable"
                      << setIsVerbose();
             });
             continue;
           }
 
         CallSites.push_back(std::make_pair(CS, -1));
       }
   }
 
   LLVM_DEBUG(dbgs() << ": " << CallSites.size() << " call sites.\n");
 
   // If there are no calls in this function, exit early.
   if (CallSites.empty())
     return false;
 
   // Now that we have all of the call sites, move the ones to functions in the
   // current SCC to the end of the list.
   unsigned FirstCallInSCC = CallSites.size();
   for (unsigned i = 0; i < FirstCallInSCC; ++i)
-    if (Function *F = CallSites[i].first.getCalledFunction())
+    if (Function *F = CallSites[i].first->getCalledFunction())
       if (SCCFunctions.count(F))
         std::swap(CallSites[i--], CallSites[--FirstCallInSCC]);
 
   InlinedArrayAllocasTy InlinedArrayAllocas;
   InlineFunctionInfo InlineInfo(&CG, &GetAssumptionCache, PSI);
 
   // Now that we have all of the call sites, loop over them and inline them if
   // it looks profitable to do so.
   bool Changed = false;
   bool LocalChange;
   do {
     LocalChange = false;
     // Iterate over the outer loop because inlining functions can cause indirect
     // calls to become direct calls.
     // CallSites may be modified inside so ranged for loop can not be used.
     for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
-      CallSite CS = CallSites[CSi].first;
+      CallBase &CS = *CallSites[CSi].first;
 
       Function *Caller = CS.getCaller();
       Function *Callee = CS.getCalledFunction();
 
       // We can only inline direct calls to non-declarations.
       if (!Callee || Callee->isDeclaration())
         continue;
 
-      Instruction *Instr = CS.getInstruction();
+      Instruction *Instr = &CS;
 
       bool IsTriviallyDead =
           isInstructionTriviallyDead(Instr, &GetTLI(*Caller));
 
       int InlineHistoryID;
       if (!IsTriviallyDead) {
         // If this call site was obtained by inlining another function, verify
         // that the include path for the function did not include the callee
for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
       // If this call site is dead and it is to a readonly function, we should
       // just delete the call instead of trying to inline it, regardless of
       // size.  This happens because IPSCCP propagates the result out of the
       // call and then we're left with the dead call.
       if (IsTriviallyDead) {
         LLVM_DEBUG(dbgs() << "    -> Deleting dead call: " << *Instr << "\n");
         // Update the call graph by deleting the edge from Callee to Caller.
         setInlineRemark(CS, "trivially dead");
-        CG[Caller]->removeCallEdgeFor(*cast<CallBase>(CS.getInstruction()));
+        CG[Caller]->removeCallEdgeFor(CS);
         Instr->eraseFromParent();
         ++NumCallsDeleted;
       } else {
         // Get DebugLoc to report. CS will be invalid after Inliner.
-        DebugLoc DLoc = CS->getDebugLoc();
+        DebugLoc DLoc = CS.getDebugLoc();
         BasicBlock *Block = CS.getParent();
 
         // Attempt to inline the function.
         using namespace ore;
 
         InlineResult IR = InlineCallIfPossible(
             CS, InlineInfo, InlinedArrayAllocas, InlineHistoryID,
             InsertLifetime, AARGetter, ImportedFunctionsStats);
for (unsigned CSi = 0; CSi != CallSites.size(); ++CSi) {
           int NewHistoryID = InlineHistory.size();
           InlineHistory.push_back(std::make_pair(Callee, InlineHistoryID));
 
 #ifndef NDEBUG
           // Make sure no dupplicates in the inline candidates. This could
           // happen when a callsite is simpilfied to reusing the return value
           // of another callsite during function cloning, thus the other
           // callsite will be reconsidered here.
-          DenseSet<CallSite> DbgCallSites;
+          DenseSet<CallBase *> DbgCallSites;
           for (auto &II : CallSites)
             DbgCallSites.insert(II.first);
 #endif
 
           for (Value *Ptr : InlineInfo.InlinedCalls) {
 #ifndef NDEBUG
-            assert(DbgCallSites.count(CallSite(Ptr)) == 0);
+            assert(DbgCallSites.count(dyn_cast<CallBase>(Ptr)) == 0);
 #endif
-            CallSites.push_back(std::make_pair(CallSite(Ptr), NewHistoryID));
+            CallSites.push_back(
+                std::make_pair(dyn_cast<CallBase>(Ptr), NewHistoryID));
           }
         }
       }
 
       // If we inlined or deleted the last possible call site to the function,
       // delete the function body now.
       if (Callee && Callee->use_empty() && Callee->hasLocalLinkage() &&
           // TODO: Can remove if in SCC now.
bool LegacyInlinerBase::inlineCalls(CallGraphSCC &SCC) {
   GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
     return getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   };
   auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
     return ACT->getAssumptionCache(F);
   };
   return inlineCallsImpl(
       SCC, CG, GetAssumptionCache, PSI, GetTLI, InsertLifetime,
-      [this](CallSite CS) { return getInlineCost(CS); }, LegacyAARGetter(*this),
-      ImportedFunctionsStats);
+      [this](CallBase &CS) { return getInlineCost(CS); },
+      LegacyAARGetter(*this), ImportedFunctionsStats);
 }
 
 /// Remove now-dead linkonce functions at the end of
 /// processing to avoid breaking the SCC traversal.
 bool LegacyInlinerBase::doFinalization(CallGraph &CG) {
   if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
     ImportedFunctionsStats.dump(InlinerFunctionImportStats ==
                                 InlinerFunctionImportStatsOpts::Verbose);
PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   // don't have a topological ordering to prevent explosions.
   //
   // To compensate for this, we don't process transitive edges made immediate
   // by inlining until we've done one pass of inlining across the entire SCC.
   // Large, highly connected SCCs still lead to some amount of code bloat in
   // this model, but it is uniformly spread across all the functions in the SCC
   // and eventually they all become too large to inline, rather than
   // incrementally maknig a single function grow in a super linear fashion.
-  SmallVector<std::pair<CallSite, int>, 16> Calls;
+  SmallVector<std::pair<CallBase *, int>, 16> Calls;
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(InitialC, CG)
           .getManager();
 
   // Populate the initial list of calls in this SCC.
   for (auto &N : InitialC) {
     auto &ORE =
         FAM.getResult<OptimizationRemarkEmitterAnalysis>(N.getFunction());
     // We want to generally process call sites top-down in order for
     // simplifications stemming from replacing the call with the returned value
     // after inlining to be visible to subsequent inlining decisions.
     // FIXME: Using instructions sequence is a really bad way to do this.
     // Instead we should do an actual RPO walk of the function body.
     for (Instruction &I : instructions(N.getFunction()))
-      if (auto CS = CallSite(&I))
-        if (Function *Callee = CS.getCalledFunction()) {
+      if (auto *CS = dyn_cast<CallBase>(&I))
+        if (Function *Callee = CS->getCalledFunction()) {
           if (!Callee->isDeclaration())
             Calls.push_back({CS, -1});
           else if (!isa<IntrinsicInst>(I)) {
             using namespace ore;
-            setInlineRemark(CS, "unavailable definition");
+            setInlineRemark(*CS, "unavailable definition");
             ORE.emit([&]() {
               return OptimizationRemarkMissed(DEBUG_TYPE, "NoDefinition", &I)
                      << NV("Callee", Callee) << " will not be inlined into "
-                     << NV("Caller", CS.getCaller())
+                     << NV("Caller", CS->getCaller())
                      << " because its definition is unavailable"
                      << setIsVerbose();
             });
           }
         }
   }
   if (Calls.empty())
     return PreservedAnalyses::all();
PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
 
   // Loop forward over all of the calls. Note that we cannot cache the size as
   // inlining can introduce new calls that need to be processed.
   for (int i = 0; i < (int)Calls.size(); ++i) {
     // We expect the calls to typically be batched with sequences of calls that
     // have the same caller, so we first set up some shared infrastructure for
     // this caller. We also do any pruning we can at this layer on the caller
     // alone.
-    Function &F = *Calls[i].first.getCaller();
+    Function &F = *Calls[i].first->getCaller();
     LazyCallGraph::Node &N = *CG.lookup(F);
     if (CG.lookupSCC(N) != C)
       continue;
     if (F.hasOptNone()) {
-      setInlineRemark(Calls[i].first, "optnone attribute");
+      setInlineRemark(*Calls[i].first, "optnone attribute");
       continue;
     }
 
     LLVM_DEBUG(dbgs() << "Inlining calls in: " << F.getName() << "\n");
 
     // Get a FunctionAnalysisManager via a proxy for this particular node. We
     // do this each time we visit a node as the SCC may have changed and as
     // we're going to mutate this particular function we want to make sure the
for (int i = 0; i < (int)Calls.size(); ++i) {
     };
     auto GetBFI = [&](Function &F) -> BlockFrequencyInfo & {
       return FAM.getResult<BlockFrequencyAnalysis>(F);
     };
     auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
       return FAM.getResult<TargetLibraryAnalysis>(F);
     };
 
-    auto GetInlineCost = [&](CallSite CS) {
+    auto GetInlineCost = [&](CallBase &CS) {
       Function &Callee = *CS.getCalledFunction();
       auto &CalleeTTI = FAM.getResult<TargetIRAnalysis>(Callee);
       bool RemarksEnabled =
           Callee.getContext().getDiagHandlerPtr()->isMissedOptRemarkEnabled(
               DEBUG_TYPE);
-      return getInlineCost(cast<CallBase>(*CS.getInstruction()), Params,
-                           CalleeTTI, GetAssumptionCache, {GetBFI}, GetTLI, PSI,
-                           RemarksEnabled ? &ORE : nullptr);
+      return getInlineCost(CS, Params, CalleeTTI, GetAssumptionCache, {GetBFI},
+                           GetTLI, PSI, RemarksEnabled ? &ORE : nullptr);
     };
 
     // Now process as many calls as we have within this caller in the sequnece.
     // We bail out as soon as the caller has to change so we can update the
     // call graph and prepare the context of that new caller.
     bool DidInline = false;
-    for (; i < (int)Calls.size() && Calls[i].first.getCaller() == &F; ++i) {
+    for (; i < (int)Calls.size() && Calls[i].first->getCaller() == &F; ++i) {
       int InlineHistoryID;
-      CallSite CS;
+      CallBase *CS = nullptr;
       std::tie(CS, InlineHistoryID) = Calls[i];
-      Function &Callee = *CS.getCalledFunction();
+      Function &Callee = *CS->getCalledFunction();
 
       if (InlineHistoryID != -1 &&
           InlineHistoryIncludes(&Callee, InlineHistoryID, InlineHistory)) {
-        setInlineRemark(CS, "recursive");
+        setInlineRemark(*CS, "recursive");
         continue;
       }
 
       // Check if this inlining may repeat breaking an SCC apart that has
       // already been split once before. In that case, inlining here may
       // trigger infinite inlining, much like is prevented within the inliner
       // itself by the InlineHistory above, but spread across CGSCC iterations
       // and thus hidden from the full inline history.
       if (CG.lookupSCC(*CG.lookup(Callee)) == C &&
           UR.InlinedInternalEdges.count({&N, C})) {
         LLVM_DEBUG(dbgs() << "Skipping inlining internal SCC edge from a node "
                              "previously split out of this SCC by inlining: "
                           << F.getName() << " -> " << Callee.getName() << "\n");
-        setInlineRemark(CS, "recursive SCC split");
+        setInlineRemark(*CS, "recursive SCC split");
         continue;
       }
 
-      Optional<InlineCost> OIC = shouldInline(CS, GetInlineCost, ORE);
+      Optional<InlineCost> OIC = shouldInline(*CS, GetInlineCost, ORE);
       // Check whether we want to inline this callsite.
       if (!OIC.hasValue()) {
-        setInlineRemark(CS, "deferred");
+        setInlineRemark(*CS, "deferred");
         continue;
       }
 
       if (!OIC.getValue()) {
         // shouldInline() call returned a negative inline cost that explains
         // why this callsite should not be inlined.
-        setInlineRemark(CS, inlineCostStr(*OIC));
+        setInlineRemark(*CS, inlineCostStr(*OIC));
         continue;
       }
 
       // Setup the data structure used to plumb customization into the
       // `InlineFunction` routine.
       InlineFunctionInfo IFI(
           /*cg=*/nullptr, &GetAssumptionCache, PSI,
-          &FAM.getResult<BlockFrequencyAnalysis>(*(CS.getCaller())),
+          &FAM.getResult<BlockFrequencyAnalysis>(*(CS->getCaller())),
           &FAM.getResult<BlockFrequencyAnalysis>(Callee));
 
       // Get DebugLoc to report. CS will be invalid after Inliner.
       DebugLoc DLoc = CS->getDebugLoc();
-      BasicBlock *Block = CS.getParent();
+      BasicBlock *Block = CS->getParent();
 
       using namespace ore;
 
       InlineResult IR = InlineFunction(CS, IFI);
       if (!IR.isSuccess()) {
-        setInlineRemark(CS, std::string(IR.getFailureReason()) + "; " +
+        setInlineRemark(*CS, std::string(IR.getFailureReason()) + "; " +
                                  inlineCostStr(*OIC));
         ORE.emit([&]() {
           return OptimizationRemarkMissed(DEBUG_TYPE, "NotInlined", DLoc, Block)
                  << NV("Callee", &Callee) << " will not be inlined into "
                  << NV("Caller", &F) << ": "
                  << NV("Reason", IR.getFailureReason());
         });
         continue;
       }
       DidInline = true;
       InlinedCallees.insert(&Callee);
 
       ++NumInlined;
 
       emit_inlined_into(ORE, DLoc, Block, Callee, F, *OIC);
 
       // Add any new callsites to defined functions to the worklist.
       if (!IFI.InlinedCallSites.empty()) {
         int NewHistoryID = InlineHistory.size();
         InlineHistory.push_back({&Callee, InlineHistoryID});
+
+        // FIXME(mtrofin): refactor IFI.InlinedCallSites to be CallBase-based
         for (CallSite &CS : reverse(IFI.InlinedCallSites)) {
           Function *NewCallee = CS.getCalledFunction();
           if (!NewCallee) {
             // Try to promote an indirect (virtual) call without waiting for the
             // post-inline cleanup and the next DevirtSCCRepeatedPass iteration
             // because the next iteration may not happen and we may miss
             // inlining it.
             if (tryPromoteCall(CS))
               NewCallee = CS.getCalledFunction();
           }
           if (NewCallee)
             if (!NewCallee->isDeclaration())
-              Calls.push_back({CS, NewHistoryID});
+              Calls.push_back(
+                  {cast<CallBase>(CS.getInstruction()), NewHistoryID});
         }
       }
 
       if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
         ImportedFunctionsStats->recordInline(F, Callee);
 
       // Merge the attributes based on the inlining.
       AttributeFuncs::mergeAttributesForInlining(F, Callee);
 
       // For local functions, check whether this makes the callee trivially
       // dead. In that case, we can drop the body of the function eagerly
       // which may reduce the number of callers of other functions to one,
       // changing inline cost thresholds.
       if (Callee.hasLocalLinkage()) {
         // To check this we also need to nuke any dead constant uses (perhaps
         // made dead by this operation on other functions).
         Callee.removeDeadConstantUsers();
         if (Callee.use_empty() && !CG.isLibFunction(Callee)) {
           Calls.erase(
               std::remove_if(Calls.begin() + i + 1, Calls.end(),
-                             [&Callee](const std::pair<CallSite, int> &Call) {
-                               return Call.first.getCaller() == &Callee;
+                             [&Callee](const std::pair<CallBase *, int> &Call) {
+                               return Call.first->getCaller() == &Callee;
                              }),
               Calls.end());
           // Clear the body and queue the function itself for deletion when we
           // finish inlining and call graph updates.
           // Note that after this point, it is an error to do anything other
           // than use the callee's address or delete it.
           Callee.dropAllReferences();
           assert(find(DeadFunctions, &Callee) == DeadFunctions.end() &&