[Inliner] Inline through indirect call sites having !callees metadata

[mssimpso]

Summary

This patch allows the inliner to inspect !callees metadata attached to an indirect call site, and then inline the subset of callees for which inlining is profitable. The inlined code is conditionally executed depending on the run-time value of the call site's function pointer. !callees metadata is attached to an indirect call by the called value propagation pass and indicates the set of functions the call site might call.

Functionally, the patch versions an indirect call site for each possible callee that can be inlined, and then promotes those call sites to direct calls. If inlining is unsuccessful for any reason, a promoted call site is demoted and made indirect once again. This ensures that no direct call sites remain after inlining that weren't already in the program. Although call site promotion based on !callees metadata could conceivably be performed as a stand-alone pass, performing it in the inliner avoids extra work, since inline cost is used as the determining measure. Moreover, knowing if a callee will be inlined is probably one of the most reasonable static estimates of whether a call site should be promoted.

A dependent revision (D39339) modifies the call graph analyses to consider !callees metadata. It ensures the inliner will visit all of the functions listed in a given !callees metadata node prior to the functions containing call sites annotated by that node. This means the inliner's normal bottom-up traversal of the call graph is valid when promoting and then inlining the subset of possible callees.

[mssimpso]

Ping.

[davidxl]

A high level comment: I think the right direction is to share the implementation with the existing IndirectCallPromotion.cpp pass. The analyzer can synthesize the MD_prof profile data for the indirect targets, though the actual target count can be synthesized with some heuristics.

[mssimpso]

In D39869#936138, @davidxl wrote:

A high level comment: I think the right direction is to share the implementation with the existing IndirectCallPromotion.cpp pass. The analyzer can synthesize the MD_prof profile data for the indirect targets, though the actual target count can be synthesized with some heuristics.

Thanks very much for the feedback, David. Just to clarify, you mean both the indirect call promotion pass and the inliner patch here should use a common interface (like llvm::promoteIndirectCall, et al.) to perform the actual transformation? This makes sense to me, since there's a lot of overlap. There are a few differences in what I've implemented, though. For example, in this patch, we can do the promotion for a single callee without creating an if-then-else structure. I think sharing the implementation might require refactoring the existing code in the indirect call promotion pass somewhat. Are you OK with me doing that as a first step?

[davidxl]

Yes, I think it is good to do the necessary refactoring work first.

[mssimpso]

In D39869#936338, @davidxl wrote:

Yes, I think it is good to do the necessary refactoring work first.

OK, sounds good. I'll submit the refactoring patch for review, and then update this one afterwards. Thanks!

[mssimpso]

Rebased on top of D40658. We now use a common indirect call promotion interface.

[mssimpso]

Moved demoteCall, previously in D40751, to this patch.

[mssimpso]

Rebased.

[mssimpso]

Ping.

[davidxl]

I think 'tryToPromote' method should be moved to the indirect call promotion pass so that the callback to inline cost/benefit is exposed there (a more general cost/benefit model needs to be developed for indirect call promotion). Initially, I think we can limit the use of the getInlineCost callback to cases where profile data is not available, this will achieve what this patch is doing without affecting existing promotion heuristics.

[mssimpso]

In D39869#969989, @davidxl wrote:

I think 'tryToPromote' method should be moved to the indirect call promotion pass so that the callback to inline cost/benefit is exposed there (a more general cost/benefit model needs to be developed for indirect call promotion). Initially, I think we can limit the use of the getInlineCost callback to cases where profile data is not available, this will achieve what this patch is doing without affecting existing promotion heuristics.

Thanks for taking a look at this again, David. This makes sense to me -- I'll update the patch shortly.

[tejohnson]

Ping

[tejohnson]

In D39869#970309, @tejohnson wrote:

Ping

Wrong Patch - please ignore! Sorry.

[mssimpso]

Addressed David's comments.

• Made tryToPromote visible to the indirect call promotion pass
• Added a check in tryToPromote that gives up if a call site has profile metadata to avoid conflicting heuristics.

David,

I hope I've at least addressed the spirit of you comments. First, I've moved tryToPromote to Transforms/Utils/CallPromotionUtils.cpp. IndirectCallPromotion.cpp already includes the necessary header, so the function is now visible there.

Second, regarding the inline cost callback, I made the profitability measure completely generic. So the client needs to provide tryToPromote with a function to determine if promotion is profitable, given an indirect call site and possible direct callee. I think it makes sense to implement it in this generic way if we leave the function in CallPromotionUtils.cpp. When called from within the inliner, the profitability measure is a wrapper around shouldInline that takes care of caching the inline cost to avoid recomputing it, etc. If we want to reuse the code in the indirect call promotion pass, we would need to define provide a similar profitability function.

I was unsure if you intended that I completely scrap all of the inliner changes, and make tryToPromote only available to the indirect call promotion pass. I didn't take your comment to mean this, but I could be wrong. Just in case, here are a few words in defense of the current patch.

In the current approach, where promotion is directed by the inliner, we limit promotion to the callees we know will definitely be inlined. So promotion is really just an intermediate step. We know the direct callee will be eliminated by inlining, so we're actually just inlining code through the indirect call. We also know the profitability measure we're trying to satisfy, inline cost, is accurate (but note that we need D39339 to get this right). Although a useful indicator, I think if we use inline cost outside the inliner, or at least outside a CallGraphSCCPass, it may end up being wrong in some cases. Suppose, for example, we use inline cost in a module pass that promotes an indirect call site to directly call function F. Later, the inliner may decide to inline a lot of code into F, making F a bad candidate for inlining into its caller. So the benefit we thought we were getting by promotion (the code simplification, measured by inline cost, due to inlining F into its caller) turns out to be nothing. This kind of situation is avoided by letting the inliner direct the promotion.

Thanks again!

[davidxl]

I actually mean the former.

Regarding your comment "Although a useful indicator, I think if we use inline cost outside the inliner, or at least outside a CallGraphSCCPass, it may end up being wrong in some cases. Suppose, for example, we use inline cost in a module pass that promotes an indirect call site to directly call function F. Later, the inliner may decide to inline a lot of code into F, making F a bad candidate for inlining into its caller. So the benefit we thought we were getting by promotion (the code simplification, measured by inline cost, due to inlining F into its caller) turns out to be nothing. This kind of situation is avoided by letting the inliner direct the promotion",

I understand what you are saying -- due to the bottom up order (limitation) of the current inliner, doing inlining cost analysis at indirect call promootion may not reflect what is actually happening -- the indirect callsite that deem to be beneficial to be inlined at indirect call promotion pass may not be inlinable anymore because it may get too large.

I need to think about this a little.

[mssimpso]

Hi David,

I'd like to get back to this patch soon. Have you given the previous comments anymore thought? We had started to discuss whether the transformation should be done inside the inliner or the indirect call promotion pass. An advantage of letting the inliner do the promotion is that the inline cost, which is used to determine if an indirect call should be promoted, is known. If promotion is done outside the inliner, the cost would only be an estimate. An advantage of moving the transformation to the indirect call promotion pass may be that it is a better conceptual fit within that pass. On the other hand, it is not a pgo-driven transformation. A third possibility would be a separate pass. Would do you think?

[mssimpso]

Ping.

[davidxl]

I added Easwaran to get a closer look at the change.

[mssimpso]

In D39869#1042160, @davidxl wrote:

I added Easwaran to get a closer look at the change.

Thanks David.

[eraman]

lib/Transforms/IPO/Inliner.cpp

457

CallSite has a getCaller method.

1173

This same code repeats twice. This can be moved to a static function. Actually, most of the logic used is common to legacy and new PM (except for the callgraph updating bit) and hence could be consolidated into a separate function.

1194

There are two other checks (starting at line 1051) that could prevent a callsite from inlining before we call shouldInline. This means it is possible that the promoted call doesn't get inlined.

[mssimpso]

Addressed Easwaran's comments.

• Reused more code between the new and legacy pass managers.
• Ensured that promotion only happens when inlining is successful.

Thanks for your comments, Easwaran. This a fairly large update, primarily meant to address your concern that a promoted call site doesn't end up getting inlined. Ensuring this situation doesn't happen is somewhat difficult because the inlining algorithm itself can abort, even after the inline cost has been computed and inlining found to be profitable. The solution I've implemented here is to perform the promotion before the inline cost is computed (as in the previous version of the patch), but if inlining fails for any reason, the call site is demoted and made indirect again. So there shouldn't be any new direct calls in the program that weren't there before.

[eraman]

This is looking good. Main comment I have is about the number of promotions+inlining that can take place is unbounded. One could argue that if there were many small direct calls, we would inline all of them, but still I feel it is better to limit the number of promotions.

lib/Transforms/IPO/Inliner.cpp

574

Is there any limit on the size of the callees metadata? Otherwise you might potentially end up promoting a large number of calls and inline them

881

Nit: The call is not necessarily inlined. More precisely, the call is deleted (either due to inlining or because it is dead). Perhaps just modify the comment above or rename the function.

lib/Transforms/Utils/CallPromotionUtils.cpp

615

You could set the CPI->Versioned = true here instead of passing an additional parameter to versionCallSite.

649

It looks like isLegalToPromote doesn't check if the callee's body is available. If F->isDeclaration() is true, is there any reason to promote? The inliner is going to not inline the callee and ends up unnecessarily promote and demote the callee. This is even more important in the versioned case since it invalidates the analysis.

test/Transforms/Inline/callees-metadata.ll

23

This is relying on how instructions get simplified after inlining. Why not check that there are not call instructions and followed by a ret?

[mssimpso]

Addressed Easwaran's comments. Thanks again!

• Added a command line option to limit the number of callees that are inlined for a given indirect call site.
• Changed tests to not run simplification after inlining.
• Renamed setInlined to setCallSiteRemoved.
• Addressed the remaining comments for CallPromotionUtils.

lib/Transforms/IPO/Inliner.cpp

574

The CalledValuePropagation pass limits the size of the callees set to 4 functions. But it's probably a good idea to have a limit in the inliner as well. I'll add a command line option for limiting the number of callees we inline for a given call site.

lib/Transforms/Utils/CallPromotionUtils.cpp

649

Nice catch! I think this check was present in a previous version of this patch, but it must have gotten lost in an update. I'll add it back.

[mssimpso]

Ping.

[eraman]

LGTM

lib/Transforms/IPO/Inliner.cpp

801

Typo

[mssimpso]

Thanks very much for the review Easwaran! As this patch still depends on D39339 to get the function iteration order right in the inliner, would you mind taking at look at that patch if you have a chance, or recommending someone for the review? Thanks again.

include/llvm/Transforms/Utils/CallPromotionUtils.h

 /// This function ensures that the number and type of the call site's arguments
 /// and return value match those of the given function. If the types do not
 /// match exactly, they must at least be bitcast compatible. If \p FailureReason
 /// is non-null and the indirect call cannot be promoted, the failure reason
 /// will be stored in it.
 bool isLegalToPromote(CallSite CS, Function *Callee,
                       const char **FailureReason = nullptr);
 
+/// CallPromotionInfo holds information created during the promotion and
+/// possible versioning of a call site. This information is needed for the call
+/// site to be demoted and made to be indirect once again. If demotion is not
+/// required, a CallPromotionInfo object need not be created.
+struct CallPromotionInfo {
+
+  /// Indicates whether the call site was versioned during promotion or not.
+  /// Versioning creates an if-then-else structure where the new direct call
+  /// site is in the "then" block and the original indirect call site is in the
+  /// "else" block.
+  bool Versioned = false;
+
+  /// The called value of the original indirect call site.
+  Value *CalledValue = nullptr;
+
+  /// The !callees metadata attached to the original indirect call.
+  MDNode *Callees = nullptr;
+
+  /// A cast instruction created for the return value of the promoted call
+  /// site, if required.
+  CastInst *RetCast = nullptr;
+
+  /// Holds cast instructions created for the arguments of the promoted call
+  /// site, if required.
+  SmallVector<CastInst *, 8> ArgCasts;
+};
+
 /// Promote the given indirect call site to unconditionally call \p Callee.
 ///
 /// This function promotes the given call site, returning the direct call or
-/// invoke instruction. If the function type of the call site doesn't match that
-/// of the callee, bitcast instructions are inserted where appropriate. If \p
-/// RetBitCast is non-null, it will be used to store the return value bitcast,
-/// if created.
+/// invoke instruction. If the function type of the call site doesn't match
+/// that of the callee, bitcast instructions are inserted where appropriate. If
+/// \p CPI is non-null, these bitcasts are collected in the provided call
+/// promotion information object.
 Instruction *promoteCall(CallSite CS, Function *Callee,
-                         CastInst **RetBitCast = nullptr);
+                         CallPromotionInfo *CPI = nullptr);
 
 /// Promote the given indirect call site to conditionally call \p Callee.
 ///
 /// This function creates an if-then-else structure at the location of the call
 /// site. The original call site is moved into the "else" block. A clone of the
-/// indirect call site is promoted, placed in the "then" block, and returned. If
-/// \p BranchWeights is non-null, it will be used to set !prof metadata on the
-/// new conditional branch.
+/// indirect call site is promoted, placed in the "then" block, and returned.
+/// If \p BranchWeights is non-null, it will be used to set !prof metadata on
+/// the new conditional branch. If \p CPI is non-null, information will be
+/// stored in the provided call promotion information object that will allow
+/// the call site to be demoted, if required.
 Instruction *promoteCallWithIfThenElse(CallSite CS, Function *Callee,
-                                       MDNode *BranchWeights = nullptr);
+                                       MDNode *BranchWeights = nullptr,
+                                       CallPromotionInfo *CPI = nullptr);
+
+/// Demote the given direct call site.
+///
+/// This function is meant to be used in conjunction with \p promoteCall and
+/// \promoteCallWithIfThenElse. As such, the call promotion information created
+/// by these functions will be used to demote a call that was previously
+/// promoted. Together, these functions can be used to temporarily promote a
+/// call in order to facilitate analyzing the program as if the call were
+/// direct. The indirect call or invoke instruction is returned.
+Instruction *demoteCall(CallSite CS, CallPromotionInfo &CPI);
 
+/// Try to promote an indirect call site.
+///
+/// This function attempts to promote the given indirect call site. Promotion
+/// is attempted for call sites having !callees metadata. For indirect call
+/// sites having a small number of possible callees, !callees metadata
+/// indicates what those functions are.
+///
+/// The approach begins by inspecting the !callees metadata attached to a call
+/// site. There are a few cases to consider based on the number of possible
+/// callees. The first two are trivial. First, if the metadata indicates that
+/// there are zero possible callees, the call site must be dead. Such a call
+/// site is not touched. Second, if the metadata indicates that there is only
+/// one possible callee, it is promoted to a direct call site in place. The
+/// last and most interesting case is when the metadata indicates that there is
+/// more than one possible callee.
+///
+/// When a call site has more than one possible callee, the call site is
+/// promoted inside an if-then-else structure, where the "then" block contains
+/// a new direct call to a possible callee, and the "else" block contains the
+/// original indirect call site. The possible callee selected for the direct
+/// call is the first callee given in the indirect call site's !callees
+/// metadata that is not also in \p CalleesToIgnore. This container can be used
+/// to record callees that are already known to be not profitable for
+/// promotion.
+///
+/// The function returns a pointer to the direct call or invoke instruction if
+/// the given call site was promoted. Otherwise, it returns a null pointer. If
+/// \p CPI is non-null and promotion is successful, information will be stored
+/// in the provided call promotion information object that will allow the call
+/// site to be demoted, if required. Demotion also removes the if-then-else
+/// structure if if the call site was versioned.
+Instruction *tryToPromoteCallSite(CallSite CS,
+                                  SmallPtrSetImpl<Function *> &CalleesToIgnore,
+                                  CallPromotionInfo *CPI = nullptr);
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_CALLPROMOTIONUTILS_H

lib/Transforms/IPO/Inliner.cpp

 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/CallPromotionUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/ImportedFunctionsInliningStatistics.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <algorithm>
 #include <cassert>
 #include <functional>
 #include <tuple>
 #include <utility>

 /// prior to LLVM's code generator having support for stack coloring based on
 /// lifetime markers. It is now in the process of being removed. To experiment
 /// with disabling it and relying fully on lifetime marker based stack
 /// coloring, you can pass this flag to LLVM.
 static cl::opt<bool>
     DisableInlinedAllocaMerging("disable-inlined-alloca-merging",
                                 cl::init(false), cl::Hidden);
 
+/// Flag to limit indirect call promotion based on !callees metdata.
+///
+/// For indirect call sites having a small number of possible callees, !callees
+/// metadata indicates what those functions are. If such metadata is present,
+/// the inliner attempts to promote and inline profitable callees by versioning
+/// the call site. This flag can be used to limit the number of inlined callees
+/// for a given indirect call site.
+static cl::opt<unsigned> InlineMaxCalleesToPromote(
+    "inline-max-calls-to-promote", cl::init(4), cl::Hidden,
+    cl::desc("Limit the number of promotions for indirect call sites"));
+
 namespace {
 
 enum class InlinerFunctionImportStatsOpts {
   No = 0,
   Basic = 1,
   Verbose = 2,
 };
 

 /// 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) {
   while (InlineHistoryID != -1) {
     assert(unsigned(InlineHistoryID) < InlineHistory.size() &&
            "Invalid inline history ID");
     if (InlineHistory[InlineHistoryID].first == F)
       return true;
     InlineHistoryID = InlineHistory[InlineHistoryID].second;
   }
   return false;
 }
 
+/// A class used to manage call promotion and demotion.
+///
+/// This class facilitates inlining through indirect call sites. During
+/// inlining, if an InlinerCallPromotionGuard object is constructed for a given
+/// indirect call site, an attempt is made to promote the indirect call site.
+/// If successful, the inliner will attempt to inline the resulting direct call
+/// site's callee. If inlining fails for any reason, the promoted call site is
+/// demoted back to an indirect call site.
+///
+/// Promotion and demotion are handled by the guard's constructor and
+/// destructor, respectively. If promotion is successful, and the inliner is
+/// able to inline or otherwise remove the direct call site, the inliner
+/// informs the guard of this transformation using the \p setCallSiteRemoved
+/// function. When the guard is deleted or goes out of scope (e.g., when
+/// inlining aborts), if the guard's corresponding call site has been promoted,
+/// but its callee hasn't been removed, the call site will be demoted back to
+/// the original indirect call site.
+class InlinerCallPromotionGuard {
+public:
+  InlinerCallPromotionGuard(SmallVectorImpl<std::pair<CallSite, int>> &Calls,
+                            int &Index, SmallPtrSetImpl<Function *> &Visited,
+                            FunctionAnalysisManager *FAM = nullptr,
+                            CallGraph *CG = nullptr)
+      : Calls(Calls), Index(Index), Visited(Visited), FAM(FAM), CG(CG),
+        Caller(Calls[Index].first.getCaller()) {
+
+    // If the current call site is indirect, we can try to promote it to see if
+    // we can inline its callee.
+    if (!Calls[Index].first.getCalledFunction())
+      tryToPromoteCallSite();
+  }
+
+  ~InlinerCallPromotionGuard() {
+
+    // If the current call site was promoted but we weren't able to inline its
+    // callee, restore the indirect call.
+    if (!CallSiteRemoved && CallSitePromoted)
+      demoteCall();
+  }
+
+  /// Inform the call promotion guard that the direct call site was removed,
+  /// either due to inlining or because it is dead. If the call site was
+  /// promoted, this will prevent the guard from trying to demote it.
+  void setCallSiteRemoved() { CallSiteRemoved = true; };
+
+private:
+  /// The identified call sites in the caller. This list is updated when
+  /// promotion is successful.
+  SmallVectorImpl<std::pair<CallSite, int>> &Calls;
+
+  /// The index of the current call site in Calls. This is a reference so we
+  /// can decrement the index to revisit the current call site. We want to
+  /// revisit the current call site when it is demoted but has additional known
+  /// callees that haven't yet been visited.
+  int &Index;
+
+  /// The known callees of the current call site that we have already analyzed.
+  SmallPtrSetImpl<Function *> &Visited;
+
+  /// A function analysis manager for the new pass manager. Promotion with
+  /// versioning changes the control-flow graph, so we will have to invalidate
+  /// or update some analyses.
+  FunctionAnalysisManager *FAM;
+
+  /// The legacy call graph. For the legacy pass manager, we will need to
+  /// update the call graph when we promote or demote a call site.
+  CallGraph *CG;
+
+  /// The caller function.
+  Function *Caller;
+
+  /// Indicates if a call site was promoted.
+  bool CallSitePromoted = false;
+
+  /// Indicates if a call site was removed, either due to inlining or because
+  /// it is dead.
+  bool CallSiteRemoved = false;
+
+  /// Call promotion information to use for promotion and demotion. The
+  /// information is populated during promotion and then consumed during
+  /// demotion to restore the indirect call site.
+  CallPromotionInfo CPI;
+
+  void tryToPromoteCallSite() {
+
+    // The current indirect call site.
+    CallSite IndirectCS = Calls[Index].first;
+
+    // Compute the number of callees that have already been inlined or deleted
+    // for this call site. The number of callees is equal to the number of
+    // callees in the set difference between the set of visited callees and the
+    // set of known callees given by the !callees metadata attached to the call
+    // site. If the number of callees reaches InlineMaxCalleesToPromote, we end
+    // promotions for this call site to avoid code bloat.
+    unsigned NumPromotedCallees = Visited.size();
+    for (Function *F : IndirectCS.getKnownCallees())
+      if (Visited.count(F))
+        --NumPromotedCallees;
+    if (NumPromotedCallees == InlineMaxCalleesToPromote)
+      return;
+
+    // Try to promote the call site. If unsuccessful, there's nothing more to
+    // do. The call promotion information will be saved in CPI if successful.
+    Instruction *DirectCSInst =
+        llvm::tryToPromoteCallSite(IndirectCS, Visited, &CPI);
+    if (!DirectCSInst)
+      return;
+
+    // Promotion was successful. Create the direct call site from the returned
+    // instruction.
+    auto DirectCS = CallSite(DirectCSInst);
+    CallSitePromoted = true;
+
+    if (CPI.Versioned) {
+      // Unfortunately, since versioning the call site modifies the
+      // control-flow graph, we either need to update cached analyses for the
+      // caller or invalidate them. For now, we just invalidate them. Some of
+      // the analyses are needed by getInlineCost (e.g., BlockFreqencyInfo).
+      if (FAM)
+        FAM->invalidate(*Caller, PreservedAnalyses::none());
+
+      // The direct call site is a new instruction. Add it to the list of Calls
+      // and swap the current call site with the direct call site. The current
+      // call site will now be the direct one, and the indirect one will be
+      // visited again.
+      Calls.push_back({DirectCS, Calls[Index].second});
+      std::swap(Calls[Index], Calls.back());
+    }
+
+    // Update the legacy call graph to reflect the promotion.
+    if (CG) {
+      CallGraphNode *CallerN = (*CG)[Caller];
+      CallGraphNode *CalleeN = (*CG)[DirectCS.getCalledFunction()];
+      if (CPI.Versioned) {
+        CallerN->addCalledFunction(DirectCS, CalleeN);
+        CallerN->replaceCallEdge(IndirectCS, IndirectCS,
+                                 CG->getOrInsertNodeForCalleesMD(
+                                     IndirectCS.getInstruction()->getMetadata(
+                                         LLVMContext::MD_callees)));
+      } else {
+        CallerN->replaceCallEdge(IndirectCS, DirectCS, CalleeN);
+      }
+    }
+  }
+
+  void demoteCall() {
+
+    // If the call site was versioned and we need to update the legacy call
+    // graph, remove the edge for the direct call site. This instruction will
+    // be erased during demotion.
+    if (CG && CPI.Versioned) {
+      CallSite DirectCS = Calls[Index].first;
+      CallGraphNode *CallerN = (*CG)[Caller];
+      CallerN->removeCallEdgeFor(DirectCS);
+    }
+
+    // Demote the current call site using the call promotion information we
+    // created during promotion.
+    llvm::demoteCall(Calls[Index].first, CPI);
+
+    // If we versioned the call site, move the indirect call site back to its
+    // original location in the Calls list.
+    if (CPI.Versioned)
+      std::swap(Calls[Index], Calls.back());
+
+    // If we need to update the legacy call graph, make the edge from the
+    // indirect call site connect to the external node with the correct
+    // !callees metadata.
+    if (CG) {
+      CallSite IndirectCS = Calls[Index].first;
+      CallGraphNode *CallerN = (*CG)[Caller];
+      CallerN->replaceCallEdge(IndirectCS, IndirectCS,
+                               CG->getOrInsertNodeForCalleesMD(
+                                   IndirectCS.getInstruction()->getMetadata(
+                                       LLVMContext::MD_callees)));
+    }
+
+    if (!CPI.Versioned)
+      return;
+
+    // Unfortunately, since versioning the call site modifies the
+    // control-flow graph, we either need to update cached analyses for the
+    // caller or invalidate them. For now, we just invalidate them. Some of
+    // the analyses are needed by getInlineCost (e.g., BlockFreqencyInfo).
+    if (FAM)
+      FAM->invalidate(*Caller, PreservedAnalyses::none());
+
+    // If the current call site has known callees that we haven't yet tried to
+    // promote, we want to visit it again. Decrement the index.
+    if (any_of(Calls[Index].first.getKnownCallees(),
+               [&](Function *F) { return !Visited.count(F); }))
+      --Index;
+
+    // Remove the direct call site from the list.
+    Calls.pop_back();
+  }
+};
+
 bool LegacyInlinerBase::doInitialization(CallGraph &CG) {
   if (InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No)
     ImportedFunctionsStats.setModuleInfo(CG.getModule());
   return false; // No changes to CallGraph.
 }
 
 bool LegacyInlinerBase::runOnSCC(CallGraphSCC &SCC) {
   if (skipSCC(SCC))
inlineCallsImpl(CallGraphSCC &SCC, CallGraph &CG,
   for (unsigned i = 0; i < FirstCallInSCC; ++i)
     if (Function *F = CallSites[i].first.getCalledFunction())
       if (SCCFunctions.count(F))
         std::swap(CallSites[i--], CallSites[--FirstCallInSCC]);
 
   InlinedArrayAllocasTy InlinedArrayAllocas;
   InlineFunctionInfo InlineInfo(&CG, &GetAssumptionCache, PSI);
 
+  // A cache of callees per indirect call site that we've attempted to inline.
+  // We don't want to consider them again.
+  DenseMap<Instruction *, SmallPtrSet<Function *, 4>> AnalyzedCallees;
+
   // 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;
+    for (int CSi = 0; CSi != (int)CallSites.size(); ++CSi) {
+
+      // Construct a call promotion guard and attempt to inline through
+      // indirect call sites. If the call site is indirect and has a known set
+      // of possible callees (e.g., as specified by !callees metadata), it will
+      // be promoted (with possible versioning). If inlining fails, the call
+      // site will be demoted.
+      InlinerCallPromotionGuard Guard(
+          CallSites, CSi,
+          AnalyzedCallees[CallSites[CSi].first.getInstruction()], nullptr, &CG);
 
+      CallSite CS = CallSites[CSi].first;
+      int InlineHistoryID = CallSites[CSi].second;
       Function *Caller = CS.getCaller();
-      Function *Callee = CS.getCalledFunction();
+
+      // If the call site remains indirect, even after attempting to promote
+      // it, there's nothing we can do.
+      if (!CS.getCalledFunction())
+        continue;
 
       // We can only inline direct calls to non-declarations.
-      if (!Callee || Callee->isDeclaration())
+      Function *Callee = CS.getCalledFunction();
+      if (Callee->isDeclaration())
         continue;
 
       Instruction *Instr = CS.getInstruction();
 
       bool IsTriviallyDead = isInstructionTriviallyDead(Instr, &TLI);
 
-      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
         // itself.  If so, we'd be recursively inlining the same function,
         // which would provide the same callsites, which would cause us to
-        // infinitely inline.
-        InlineHistoryID = CallSites[CSi].second;
+        // infinitel inline.
         if (InlineHistoryID != -1 &&
             InlineHistoryIncludes(Callee, InlineHistoryID, InlineHistory))
           continue;
       }
 
       // FIXME for new PM: because of the old PM we currently generate ORE and
       // in turn BFI on demand.  With the new PM, the ORE dependency should
       // just become a regular analysis dependency.
for (int CSi = 0; CSi != (int)CallSites.size(); ++CSi) {
           int NewHistoryID = InlineHistory.size();
           InlineHistory.push_back(std::make_pair(Callee, InlineHistoryID));
 
           for (Value *Ptr : InlineInfo.InlinedCalls)
             CallSites.push_back(std::make_pair(CallSite(Ptr), NewHistoryID));
         }
       }
 
+      // Inform the call promotion guard that the call site was removed. If the
+      // call site was promoted, this will prevent the guard from trying to
+      // demote it.
+      Guard.setCallSiteRemoved();
+
       // 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.
           !SCCFunctions.count(Callee) &&
           // The function may be apparently dead, but if there are indirect
           // callgraph references to the node, we cannot delete it yet, this
           // could invalidate the CGSCC iterator.
PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   for (auto &N : InitialC) {
     // 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 (!Callee->isDeclaration())
-            Calls.push_back({CS, -1});
   }
   if (Calls.empty())
     return PreservedAnalyses::all();
 
   // Capture updatable variables for the current SCC and RefSCC.
   auto *C = &InitialC;
   auto *RC = &C->getOuterRefSCC();
 
for (int i = 0; i < (int)Calls.size(); ++i) {
 
     auto GetInlineCost = [&](CallSite CS) {
       Function &Callee = *CS.getCalledFunction();
       auto &CalleeTTI = FAM.getResult<TargetIRAnalysis>(Callee);
       return getInlineCost(CS, Params, CalleeTTI, GetAssumptionCache, {GetBFI},
                            PSI, &ORE);
     };
 
+    // A cache of callees per indirect call site that we've attempted to
+    // inline. We don't want to consider them again.
+    DenseMap<Instruction *, SmallPtrSet<Function *, 4>> AnalyzedCallees;
+
     // 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) {
+
+      // Construct a call promotion guard and attempt to inline through
+      // indirect call sites. If the call site is indirect and has a known set
+      // of possible callees (e.g., as specified by !callees metadata), it will
+      // be promoted (with possible versioning). If inlining fails, the call
+      // site will be demoted.
+      InlinerCallPromotionGuard Guard(
+          Calls, i, AnalyzedCallees[Calls[i].first.getInstruction()], &FAM);
+
       int InlineHistoryID;
       CallSite CS;
       std::tie(CS, InlineHistoryID) = Calls[i];
+
+      // If the call site remains indirect, even after attempting to promote
+      // it, there's nothing we can do.
+      if (!CS.getCalledFunction())
+        continue;
+
+      // We can only inline direct calls to non-declarations.
       Function &Callee = *CS.getCalledFunction();
+      if (Callee.isDeclaration())
+        continue;
 
       if (InlineHistoryID != -1 &&
           InlineHistoryIncludes(&Callee, InlineHistoryID, InlineHistory))
         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})) {
         DEBUG(dbgs() << "Skipping inlining internal SCC edge from a node "
                         "previously split out of this SCC by inlining: "
                      << F.getName() << " -> " << Callee.getName() << "\n");
         continue;
       }
 
       Optional<InlineCost> OIC = shouldInline(CS, GetInlineCost, ORE);
       // Check whether we want to inline this callsite.
       if (!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>(Callee));
for (; i < (int)Calls.size() && Calls[i].first.getCaller() == &F; ++i) {
                  << NV("Callee", &Callee) << " will not be inlined into "
                  << NV("Caller", &F);
         });
         continue;
       }
       DidInline = true;
       InlinedCallees.insert(&Callee);
 
+      // Inform the call promotion guard that the call site was removed. If the
+      // call site was promoted, this will prevent the guard from trying to
+      // demote it.
+      Guard.setCallSiteRemoved();
+
       ORE.emit([&]() {
         bool AlwaysInline = OIC->isAlways();
         StringRef RemarkName = AlwaysInline ? "AlwaysInline" : "Inlined";
         OptimizationRemark R(DEBUG_TYPE, RemarkName, DLoc, Block);
         R << NV("Callee", &Callee) << " inlined into ";
         R << NV("Caller", &F);
         if (AlwaysInline)
           R << " with cost=always";

lib/Transforms/Utils/CallPromotionUtils.cpp

 //===- CallPromotionUtils.cpp - Utilities for call promotion ----*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities useful for promoting indirect call sites to
 // direct call sites.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/CallPromotionUtils.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "call-promotion-utils"
 
 /// Fix-up phi nodes in an invoke instruction's normal destination.
 ///

 ///   orig_bb:
 ///     %t0 = invoke i32 @func() to label %split_bb unwind label %unwind_dst
 ///
 ///   split_bb:
 ///     ; Uses of the original return value are replaced by uses of the bitcast.
 ///     %t1 = bitcast i32 %t0 to ...
 ///     br label %normal_dst
 ///
-static void createRetBitCast(CallSite CS, Type *RetTy, CastInst **RetBitCast) {
+static void createRetBitCast(CallSite CS, Type *RetTy, CallPromotionInfo *CPI) {
 
   // Save the users of the calling instruction. These uses will be changed to
   // use the bitcast after we create it.
   SmallVector<User *, 16> UsersToUpdate;
   for (User *U : CS.getInstruction()->users())
     UsersToUpdate.push_back(U);
 
   // Determine an appropriate location to create the bitcast for the return
   // value. The location depends on if we have a call or invoke instruction.
   Instruction *InsertBefore = nullptr;
   if (auto *Invoke = dyn_cast<InvokeInst>(CS.getInstruction()))
     InsertBefore =
         &SplitEdge(Invoke->getParent(), Invoke->getNormalDest())->front();
   else
     InsertBefore = &*std::next(CS.getInstruction()->getIterator());
 
   // Bitcast the return value to the correct type.
   auto *Cast = CastInst::Create(Instruction::BitCast, CS.getInstruction(),
                                 RetTy, "", InsertBefore);
-  if (RetBitCast)
-    *RetBitCast = Cast;
+  if (CPI)
+    CPI->RetCast = Cast;
 
   // Replace all the original uses of the calling instruction with the bitcast.
   for (User *U : UsersToUpdate)
     U->replaceUsesOfWith(CS.getInstruction(), Cast);
 }
 
 /// Predicate and clone the given call site.
 ///
if (!CastInst::isBitCastable(ActualTy, FormalTy)) {
       return false;
     }
   }
 
   return true;
 }
 
 Instruction *llvm::promoteCall(CallSite CS, Function *Callee,
-                               CastInst **RetBitCast) {
+                               CallPromotionInfo *CPI) {
   assert(!CS.getCalledFunction() && "Only indirect call sites can be promoted");
 
+  // If the client provided call promotion information, save the called value
+  // of the indirect call site and the !callees metata attached to it.
+  if (CPI) {
+    CPI->CalledValue = CS.getCalledValue();
+    if (CS.getInstruction()->getMetadata(LLVMContext::MD_callees))
+      CPI->Callees = CS.getInstruction()->getMetadata(LLVMContext::MD_callees);
+  }
+
   // Set the called function of the call site to be the given callee.
   CS.setCalledFunction(Callee);
 
   // Since the call site will no longer be direct, we must clear metadata that
   // is only appropriate for indirect calls. This includes !prof and !callees
   // metadata.
   CS.getInstruction()->setMetadata(LLVMContext::MD_prof, nullptr);
   CS.getInstruction()->setMetadata(LLVMContext::MD_callees, nullptr);
Instruction *llvm::promoteCall(CallSite CS, Function *Callee,
   for (Use &U : CS.args()) {
     unsigned ArgNo = CS.getArgumentNo(&U);
     Type *FormalTy = Callee->getFunctionType()->getParamType(ArgNo);
     Type *ActualTy = U.get()->getType();
     if (FormalTy != ActualTy) {
       auto *Cast = CastInst::Create(Instruction::BitCast, U.get(), FormalTy, "",
                                     CS.getInstruction());
       CS.setArgument(ArgNo, Cast);
+      if (CPI)
+        CPI->ArgCasts.push_back(Cast);
     }
   }
 
   // If the return type of the call site doesn't match that of the callee, cast
   // the returned value to the appropriate type.
   if (!CallSiteRetTy->isVoidTy() && CallSiteRetTy != CalleeRetTy)
-    createRetBitCast(CS, CallSiteRetTy, RetBitCast);
+    createRetBitCast(CS, CallSiteRetTy, CPI);
+
+  return CS.getInstruction();
+}
+
+Instruction *llvm::demoteCall(CallSite CS, CallPromotionInfo &CPI) {
+  assert(CS.getCalledFunction() && "Only direct call sites can be demoted");
+
+  // Handle argument casts. For each cast instruction, replace all of its uses
+  // with its source operand.
+  for (CastInst *Cast : CPI.ArgCasts) {
+    while (!Cast->user_empty())
+      Cast->user_back()->replaceUsesOfWith(Cast, Cast->getOperand(0));
+    Cast->eraseFromParent();
+  }
+
+  // Handle the return value cast. Similar to arguments casts, we replace all
+  // of uses of the cast with its source operand. If the call site is an invoke
+  // instruction, the return value cast is placed along the edge between the
+  // block containing the invoke and it's normal destination. The block
+  // containing this cast is also erased.
+  if (CPI.RetCast) {
+    while (!CPI.RetCast->user_empty())
+      CPI.RetCast->user_back()->replaceUsesOfWith(CPI.RetCast,
+                                                  CPI.RetCast->getOperand(0));
+    if (!isa<InvokeInst>(CS.getInstruction())) {
+      CPI.RetCast->eraseFromParent();
+    } else {
+      // For invoke instructions, the edge between the block containing the
+      // invoke and the invoke's normal destination has been split (see
+      // createRetBitCast):
+      //
+      //   orig_bb:
+      //     %t0 = invoke i32 @func()
+      //              to label %split_bb unwind label %unwind_dst
+      //
+      //   ; The following block will be erased.
+      //   split_bb:
+      //     %t1 = bitcast i32 %t0 to ...
+      //     br label %normal_dst
+      //
+      BasicBlock *SplitBlock = CPI.RetCast->getParent();
+      BasicBlock *OrigBlock = SplitBlock->getSinglePredecessor();
+      BasicBlock *NormalDest = SplitBlock->getSingleSuccessor();
+      auto *OrigInvoke = cast<InvokeInst>(OrigBlock->getTerminator());
+      OrigInvoke->setNormalDest(NormalDest);
+      SplitBlock->replaceAllUsesWith(OrigBlock);
+      SplitBlock->eraseFromParent();
+    }
+  }
+
+  if (!CPI.Versioned) {
+    // If the call site wasn't versioned, we demote the direct call site in
+    // place. For the versioning case, we simply erase the direct call site and
+    // move the indirect call site to the correct position. First, set the
+    // called value of the call site, and mutate its function type.
+    CS.setCalledFunction(CPI.CalledValue);
+    CS.mutateFunctionType(cast<FunctionType>(
+        cast<PointerType>(CPI.CalledValue->getType())->getElementType()));
+
+    // Attach !callees metadata to the now indirect call site. The metadata
+    // node should have been saved at the time the call site was promoted.
+    if (CPI.Callees)
+      CS.getInstruction()->setMetadata(LLVMContext::MD_callees, CPI.Callees);
+    return CS.getInstruction();
+  }
+
+  // A wrapper around llvm::MergeBlockIntoPredecessor that checks that the
+  // merging was successful. For our use, merging should not fail.
+  auto mergeBlockIntoPredecessor = [](BasicBlock *BB) {
+    bool Merged = MergeBlockIntoPredecessor(BB);
+    assert(Merged && "Block not merged into predecessor");
+    (void)Merged;
+  };
+
+  // Reconstruct the control-flow created during promotion. We have an
+  // if-then-else structure where a new direct call site is in the "then"
+  // block, and the original indirect call site is in the "else" block. For
+  // call instructions, this looks like:
+  //
+  //   orig_bb:
+  //     %cond = icmp eq i32 ()* %ptr, @func
+  //     br i1 %cond, %then_bb, %else_bb
+  //
+  //   then_bb:
+  //     %t1 = call i32 @func()
+  //     br merge_bb
+  //
+  //   else_bb:
+  //     %t0 = call i32 %ptr()
+  //     br merge_bb
+  //
+  //   merge_bb:
+  //     %t2 = phi i32 [ %t0, %else_bb ], [ %t1, %then_bb ]
+  //
+  // For invoke instructions, this looks like:
+  //
+  //   orig_bb:
+  //     %cond = icmp eq i32 ()* %ptr, @func
+  //     br i1 %cond, %then_bb, %else_bb
+  //
+  //   then_bb:
+  //     %t1 = invoke i32 @func() to label %merge_bb unwind label %unwind_dst
+  //
+  //   else_bb:
+  //     %t0 = invoke i32 %ptr() to label %merge_bb unwind label %unwind_dst
+  //
+  //   merge_bb:
+  //     %t2 = phi i32 [ %t0, %else_bb ], [ %t1, %then_bb ]
+  //     br %normal_dst
+  //
+  BasicBlock *ThenBlock = CS.getInstruction()->getParent();
+  BasicBlock *OrigBlock = ThenBlock->getSinglePredecessor();
+  TerminatorInst *OrigTerm = OrigBlock->getTerminator();
+  Value *Cond = cast<BranchInst>(OrigTerm)->getCondition();
+  BasicBlock *ElseBlock = OrigTerm->getSuccessor(1);
+
+  // The merge block depends on whether we have a call or an invoke
+  // instruction.
+  BasicBlock *MergeBlock = nullptr;
+  if (auto *PromotedInvoke = dyn_cast<InvokeInst>(CS.getInstruction()))
+    MergeBlock = PromotedInvoke->getNormalDest();
+  else
+    MergeBlock = ElseBlock->getSingleSuccessor();
+
+  // The indirect call site we want to keep is in the "else" block. We will
+  // erase the direct call in the "then" block.
+  CS = CallSite(&ElseBlock->front());
+
+  // Attach !callees metadata to the now indirect call site. The metadata node
+  // should have been saved at the time the call site was promoted.
+  if (CPI.Callees)
+    CS.getInstruction()->setMetadata(LLVMContext::MD_callees, CPI.Callees);
+
+  // Erase the condition and conditional branch from the end of the "original"
+  // block and replace it with an unconditional branch to the "else" block.
+  OrigTerm->eraseFromParent();
+  if (auto *I = dyn_cast<Instruction>(Cond))
+    I->eraseFromParent();
+  BranchInst::Create(ElseBlock, OrigBlock);
+
+  // We're going to delete the "then" block. First we need to fix-up the phi
+  // node in the "merge" block.
+  for (auto &Phi : MergeBlock->phis())
+    Phi.removeIncomingValue(ThenBlock);
+  FoldSingleEntryPHINodes(MergeBlock);
+
+  // Erase the terminator from the "then" block. For the invoke case, this will
+  // erase the direct call site. For the call case, this removes the branch so
+  // we can merge the "merge" block into a single predecessor.
+  ThenBlock->getTerminator()->eraseFromParent();
+
+  // Merge the "else" block into the "original" block.
+  mergeBlockIntoPredecessor(ElseBlock);
+
+  // For the call case, merge the "merge" block into the "original" block and
+  // erase the "then" block. No other work is required.
+  auto *PromotedInvoke = dyn_cast<InvokeInst>(OrigBlock->getTerminator());
+  if (!PromotedInvoke) {
+    mergeBlockIntoPredecessor(MergeBlock);
+    ThenBlock->eraseFromParent();
+    return CS.getInstruction();
+  }
+
+  // For the invoke case, we need to fix-up the phi nodes in the normal and
+  // unwind destination blocks. First, for the normal destination, values
+  // coming from the "merge" block now come from the "original" block.
+  PromotedInvoke->setNormalDest(MergeBlock->getSingleSuccessor());
+  for (auto &Phi : PromotedInvoke->getNormalDest()->phis()) {
+    int Idx = Phi.getBasicBlockIndex(MergeBlock);
+    if (Idx == -1)
+      continue;
+    Phi.setIncomingBlock(Idx, OrigBlock);
+  }
+
+  // For the unwind destination, there are now no longer any value coming from
+  // the "then" block.
+  for (auto &Phi : PromotedInvoke->getUnwindDest()->phis()) {
+    int Idx = Phi.getBasicBlockIndex(ThenBlock);
+    if (Idx == -1)
+      continue;
+    Phi.removeIncomingValue(Idx);
+  }
+
+  // Lastly, erase the "merge" and "then" blocks for the invoke case.
+  MergeBlock->eraseFromParent();
+  ThenBlock->eraseFromParent();
 
   return CS.getInstruction();
 }
 
 Instruction *llvm::promoteCallWithIfThenElse(CallSite CS, Function *Callee,
-                                             MDNode *BranchWeights) {
+                                             MDNode *BranchWeights,
+                                             CallPromotionInfo *CPI) {
 
   // Version the indirect call site. If the called value is equal to the given
   // callee, 'NewInst' will be executed, otherwise the original call site will
   // be executed.
   Instruction *NewInst = versionCallSite(CS, Callee, BranchWeights);
 
+  // If the client requested call promotion information, indicate that the call
+  // site was versioned.
+  if (CPI)
+    CPI->Versioned = true;
+
   // Promote 'NewInst' so that it directly calls the desired function.
-  return promoteCall(CallSite(NewInst), Callee);
+  return promoteCall(CallSite(NewInst), Callee, CPI);
+}
+
+Instruction *
+llvm::tryToPromoteCallSite(CallSite CS,
+                           SmallPtrSetImpl<Function *> &CalleesToIgnore,
+                           CallPromotionInfo *CPI) {
+  assert(!CS.getCalledFunction() && "Can only promote indirect calls");
+
+  // If this call site has profile metadata attached to it, don't attempt
+  // promotion. We don't want to disagree with profile-guideded heuristics.
+  if (CS.getInstruction()->getMetadata(LLVMContext::MD_prof))
+    return nullptr;
+
+  // Collect the functions the call site might call. If the call site doesn't
+  // have !callees metadata, we don't know its possible callees, so there's
+  // nothing to do.
+  if (!CS.getInstruction()->getMetadata(LLVMContext::MD_callees))
+    return nullptr;
+  SmallVector<Function *, 4> PossibleCallees = CS.getKnownCallees();
+
+  // If the metadata indicates that there is only one possible callee, we
+  // attempt to promote the call site. If successful, we clear the call site's
+  // !callees metadata.
+  if (PossibleCallees.size() == 1) {
+    Function *F = PossibleCallees[0];
+    if (F->isDeclaration() || CalleesToIgnore.count(F) ||
+        !isLegalToPromote(CS, F))
+      return nullptr;
+    CalleesToIgnore.insert(F);
+    Instruction *Direct = promoteCall(CS, F, CPI);
+    return Direct;
+  }
+
+  // If the metadata indicates that there is more than one possible callee, we
+  // will attempt to promote the call site by versioning it.
+  for (unsigned I = 0; I < PossibleCallees.size(); ++I) {
+    Function *F = PossibleCallees[I];
+
+    // Skip this function if we already know it to be unprofitable or it can't
+    // be promoted.
+    if (F->isDeclaration() || CalleesToIgnore.count(F) ||
+        !isLegalToPromote(CS, F))
+      continue;
+    CalleesToIgnore.insert(F);
+
+    // Promote the indirect call site such that we conditionally execute the
+    // direct call. The original instruction is moved to the "else" block, and
+    // the new direct call in the "then" block is returned.
+    Instruction *Direct = promoteCallWithIfThenElse(CS, F, nullptr, CPI);
+
+    // The original call site remains indirect, but we must remove the function
+    // from its !callees metadata. After the promotion, we know this call site
+    // cannot target it.
+    PossibleCallees.erase(PossibleCallees.begin() + I);
+    MDBuilder MDB(CS.getInstruction()->getContext());
+    CS.getInstruction()->setMetadata(LLVMContext::MD_callees,
+                                     MDB.createCallees(PossibleCallees));
+
+    return Direct;
+  }
+
+  return nullptr;
 }
 
 #undef DEBUG_TYPE

test/Other/new-pm-defaults.ll

 ; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ProfileSummaryAnalysis
 ; CHECK-O-NEXT: Running analysis: ProfileSummaryAnalysis
 ; CHECK-O-NEXT: Running pass: ModuleToPostOrderCGSCCPassAdaptor<{{.*}}LazyCallGraph{{.*}}>
 ; CHECK-O-NEXT: Running analysis: InnerAnalysisManagerProxy
 ; CHECK-O-NEXT: Running analysis: LazyCallGraphAnalysis
 ; CHECK-O-NEXT: Starting CGSCC pass manager run.
 ; CHECK-O-NEXT: Running pass: InlinerPass
 ; CHECK-O-NEXT: Running analysis: OuterAnalysisManagerProxy<{{.*}}LazyCallGraph{{.*}}>
-; CHECK-O-NEXT: Running pass: PostOrderFunctionAttrsPass
 ; CHECK-O-NEXT: Running analysis: FunctionAnalysisManagerCGSCCProxy
+; CHECK-O-NEXT: Running pass: PostOrderFunctionAttrsPass
 ; CHECK-O3-NEXT: Running pass: ArgumentPromotionPass
 ; CHECK-O-NEXT: Running pass: CGSCCToFunctionPassAdaptor<{{.*}}PassManager{{.*}}>
 ; CHECK-O-NEXT: Starting llvm::Function pass manager run.
 ; CHECK-O-NEXT: Running pass: SROA
 ; CHECK-O-NEXT: Running pass: EarlyCSEPass
 ; CHECK-O-NEXT: Running analysis: MemorySSAAnalysis
 ; CHECK-O-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O-NEXT: Running pass: JumpThreadingPass

test/Other/new-pm-thinlto-defaults.ll

 ; CHECK-O-NEXT: Running pass: RequireAnalysisPass<{{.*}}ProfileSummaryAnalysis
 ; CHECK-PRELINK-O-NEXT: Running analysis: ProfileSummaryAnalysis
 ; CHECK-O-NEXT: Running pass: ModuleToPostOrderCGSCCPassAdaptor<{{.*}}LazyCallGraph{{.*}}>
 ; CHECK-O-NEXT: Running analysis: InnerAnalysisManagerProxy
 ; CHECK-O-NEXT: Running analysis: LazyCallGraphAnalysis
 ; CHECK-O-NEXT: Starting CGSCC pass manager run.
 ; CHECK-O-NEXT: Running pass: InlinerPass
 ; CHECK-O-NEXT: Running analysis: OuterAnalysisManagerProxy<{{.*}}LazyCallGraph{{.*}}>
-; CHECK-O-NEXT: Running pass: PostOrderFunctionAttrsPass
 ; CHECK-O-NEXT: Running analysis: FunctionAnalysisManagerCGSCCProxy
+; CHECK-O-NEXT: Running pass: PostOrderFunctionAttrsPass
 ; CHECK-O3-NEXT: Running pass: ArgumentPromotionPass
 ; CHECK-O-NEXT: Running pass: CGSCCToFunctionPassAdaptor<{{.*}}PassManager{{.*}}>
 ; CHECK-O-NEXT: Starting llvm::Function pass manager run.
 ; CHECK-O-NEXT: Running pass: SROA
 ; CHECK-O-NEXT: Running pass: EarlyCSEPass
 ; CHECK-O-NEXT: Running analysis: MemorySSAAnalysis
 ; CHECK-O-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O-NEXT: Running pass: JumpThreadingPass

test/Transforms/Inline/callees-metadata.ll

+; RUN: opt < %s -inline -S | FileCheck %s
+; RUN: opt < %s -passes='cgscc(inline)' -S | FileCheck %s
+
+; A call instruction with empty !callees metadata. This is undefined behavior.
+; The inliner doesn't mess with it.
+;
+; CHECK-LABEL: @call.0.0(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void %fp()
+; CHECK-NEXT:    ret void
+;
+define void @call.0.0(void ()* %fp) {
+entry:
+  call void %fp(), !callees !{}
+  ret void
+}
+
+; A call instruction with two possible callees that can both inlined.
+;
+; CHECK-LABEL: @call.2.2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 (i64, i64)* %fp, @select_def_0
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[MERGE]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i64 [ %y, %[[ELSE]] ], [ %x, %[[THEN]] ]
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+define i64 @call.2.2(i64 %x, i64 %y, i64 (i64, i64)* %fp) {
+entry:
+  %tmp0 = call i64 %fp(i64 %x, i64 %y),
+     !callees !{i64 (i64, i64)* @select_def_0, i64 (i64, i64)* @select_def_1}
+  ret i64 %tmp0
+}
+
+; A call instruction with three possible callees, only one of which can be
+; inlined. Since more than one possible callee remains after promotion, the
+; cloned call is indirect.
+;
+; CHECK-LABEL: @call.3.1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 (i64, i64)* %fp, @select_def_0
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 %fp(i64 %x, i64 %y), !callees
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[MERGE]]:
+; CHECK-NEXT:    [[TMP2:%.*]] = phi i64 [ [[TMP1]], %[[ELSE]] ], [ %x, %[[THEN]] ]
+; CHECK-NEXT:    ret i64 [[TMP2]]
+;
+define i64 @call.3.1(i64 %x, i64 %y, i64 (i64, i64)* %fp) {
+entry:
+  %tmp0 = call i64 %fp(i64 %x, i64 %y),
+     !callees !{i64 (i64, i64)* @select_def_0, i64 (i64, i64)* @select_dec_1,
+                i64 (i64, i64)* @select_dec_0}
+  ret i64 %tmp0
+}
+
+; A call instruction with two possible callees, only one of which can be
+; inlined. Although only one possible callee remains after promotion, the
+; cloned call is not promoted since the callee cannot be inlined.
+;
+; CHECK-LABEL: @call.2.1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 (i64, i64)* %fp, @select_def_0
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 %fp(i64 %x, i64 %y), !callees
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[MERGE]]:
+; CHECK-NEXT:    [[TMP2:%.*]] = phi i64 [ [[TMP1]], %[[ELSE]] ], [ %x, %[[THEN]] ]
+; CHECK-NEXT:    ret i64 [[TMP2]]
+;
+define i64 @call.2.1(i64 %x, i64 %y, i64 (i64, i64)* %fp) {
+entry:
+  %tmp0 = call i64 %fp(i64 %x, i64 %y),
+     !callees !{i64 (i64, i64)* @select_def_0, i64 (i64, i64)* @select_dec_1}
+  ret i64 %tmp0
+}
+
+; A call instruction with two possible callees, only one of which can be
+; inlined. Although only one possible callee remains after promotion, the
+; cloned call is not promoted since the callee cannot be inlined. Bitcasts must
+; be inserted for the promoted call's arguments and return value since the
+; callee's function type doesn't match that of the function pointer.
+;
+; CHECK-LABEL: @call.2.1.bitcast(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32* (i32**)* %fp, bitcast (i64* (i64**)* @load_def to i32* (i32**)*)
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32** %x to i64**
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64*, i64** [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i64* [[TMP2]] to i32*
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32* %fp(i32** %x), !callees
+; CHECK-NEXT:    br label %[[MERGE]]
+; CHECK:       [[MERGE]]:
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i32* [ [[TMP4]], %[[ELSE]] ], [ [[TMP3]], %[[THEN]] ]
+; CHECK-NEXT:    ret i32* [[TMP5]]
+;
+define i32* @call.2.1.bitcast(i32** %x, i32* (i32**)* %fp) {
+entry:
+  %tmp0 = call i32* %fp(i32** %x),
+     !callees !{i64* (i64**)* @load_def, i64* (i64**)* @load_dec}
+  ret i32* %tmp0
+}
+
+; An invoke instruction with two possible callees, only one of which can be
+; inlined. Although only one possible callee remains after promotion, the
+; cloned invoke is not promoted since the callee cannot be inlined. The normal
+; destination has a phi node that must be fixed.
+;
+; CHECK-LABEL: @invoke.2.1(
+; CHECK-NEXT:  entry:
+; CHECK:         br i1 undef, label %[[TRY:.*]], label %[[NORMAL:.*]]
+; CHECK:       [[TRY]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64* (i64**)* %fp, @load_def
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = load i64*, i64** %x
+; CHECK-NEXT:    br label %[[MERGE]]
+; CHECK:       [[ELSE]]:
+; CHECK:         [[TMP2:%.*]] = invoke i64* %fp(i64** %x)
+; CHECK-NEXT:                      to label %[[MERGE]] unwind label %[[EXCEPT:.*]], !callees
+; CHECK:       [[MERGE]]:
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i64* [ [[TMP2]], %[[ELSE]] ], [ [[TMP1]], %[[THEN]] ]
+; CHECK-NEXT:    br label %[[NORMAL:.*]]
+; CHECK:       [[NORMAL]]:
+; CHECK-NEXT:    [[TMP4:%.*]] = phi i64* [ [[TMP3]], %[[MERGE]] ], [ null, %entry ]
+; CHECK-NEXT:    ret i64* [[TMP4]]
+; CHECK:       [[EXCEPT]]:
+; CHECK-NEXT:    [[TMP5:%.*]] = landingpad { i8*, i64 }
+; CHECK-NEXT:                   cleanup
+; CHECK-NEXT:    ret i64* null
+;
+define i64* @invoke.2.1(i64** %x, i64* (i64**)* %fp) personality i64 (...)* null {
+entry:
+  br i1 undef, label %try, label %normal
+try:
+  %tmp0 = invoke i64* %fp(i64 **%x)
+             to label %normal unwind label %except,
+                !callees !{i64* (i64**)* @load_def, i64* (i64**)* @load_dec}
+normal:
+  %tmp1 = phi i64* [ %tmp0, %try ], [ null, %entry ]
+  ret i64* %tmp1
+except:
+  %tmp2 = landingpad {i8*, i64}
+          cleanup
+  ret i64* null
+}
+
+; An invoke instruction with two possible callees, only one of which can be
+; inlined. Although only one possible callee remains after promotion, the
+; cloned invoke is not promoted since the callee cannot be inlined. The normal
+; destination has a phi node that must be fixed. Additionally, bitcasts must be
+; inserted for the promoted invoke's arguments and return value since the
+; callee's function type doesn't match that of the function pointer.
+;
+; CHECK-LABEL: @invoke.2.1.bitcast(
+; CHECK-NEXT:  entry:
+; CHECK:         br i1 undef, label %[[TRY:.*]], label %[[NORMAL:.*]]
+; CHECK:       [[TRY]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32* (i32**)* %fp, bitcast (i64* (i64**)* @load_def to i32* (i32**)*)
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[THEN:.*]], label %[[ELSE:.*]]
+; CHECK:       [[THEN]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32** %x to i64**
+; CHECK-NEXT:    [[TMP2:%.*]] = load i64*, i64** [[TMP1]]
+; CHECK-NEXT:    br label %[[SPLIT:.*]]
+; CHECK:       [[SPLIT]]:
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i64* [[TMP2]] to i32*
+; CHECK-NEXT:    br label %[[MERGE:.*]]
+; CHECK:       [[ELSE]]:
+; CHECK:         [[TMP4:%.*]] = invoke i32* %fp(i32** %x)
+; CHECK-NEXT:                      to label %[[MERGE]] unwind label %[[EXCEPT:.*]], !callees
+; CHECK:       [[MERGE]]:
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i32* [ [[TMP4]], %[[ELSE]] ], [ [[TMP3]], %[[SPLIT]] ]
+; CHECK-NEXT:    br label %[[NORMAL]]
+; CHECK:       [[NORMAL]]:
+; CHECK-NEXT:    [[TMP6:%.*]] = phi i32* [ [[TMP5]], %[[MERGE]] ], [ null, %entry ]
+; CHECK-NEXT:    ret i32* [[TMP6]]
+; CHECK:       [[EXCEPT]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = landingpad { i8*, i32 }
+; CHECK-NEXT:                   cleanup
+; CHECK-NEXT:    ret i32* null
+;
+define i32* @invoke.2.1.bitcast(i32** %x, i32* (i32**)* %fp) personality i32 (...)* null {
+entry:
+  br i1 undef, label %try, label %normal
+try:
+  %tmp0 = invoke i32* %fp(i32 **%x)
+             to label %normal unwind label %except,
+                !callees !{i64* (i64**)* @load_def, i64* (i64**)* @load_dec}
+normal:
+  %tmp1 = phi i32* [ %tmp0, %try ], [ null, %entry ]
+  ret i32* %tmp1
+except:
+  %tmp2 = landingpad {i8*, i32}
+          cleanup
+  ret i32* null
+}
+
+; The functions declarations and definitions below are used in the above tests
+; and are not interesting. The !callees metadata that references the functions
+; is expanded and attached to the call sites.
+;
+declare i64 @select_dec_0(i64 %x, i64 %y)
+declare i64 @select_dec_1(i64 %x, i64 %y)
+declare i64* @load_dec(i64** %x)
+
+define i64 @select_def_0(i64 %x, i64 %y) { ret i64 %x }
+define i64 @select_def_1(i64 %x, i64 %y) { ret i64 %y }
+define i64* @load_def(i64** %x) {
+  %tmp0 = load i64*, i64** %x
+  ret i64* %tmp0
+}