[llvm][NFC] Refactor APIs operating on CallBase

[mtrofin]

Summary

Refactored the parameter and return type where they are too generally
typed as Instruction.

[dblaikie]

Looks good!

llvm/include/llvm/Analysis/IndirectCallVisitor.h

 #define LLVM_ANALYSIS_INDIRECTCALLVISITOR_H
 
 #include "llvm/IR/InstVisitor.h"
 #include <vector>
 
 namespace llvm {
 // Visitor class that finds all indirect call.
 struct PGOIndirectCallVisitor : public InstVisitor<PGOIndirectCallVisitor> {
-  std::vector<Instruction *> IndirectCalls;
+  std::vector<CallBase *> IndirectCalls;
   PGOIndirectCallVisitor() {}
 
   void visitCallBase(CallBase &Call) {
     if (Call.isIndirectCall())
       IndirectCalls.push_back(&Call);
   }
 };
 
 // Helper function that finds all indirect call sites.
-inline std::vector<Instruction *> findIndirectCalls(Function &F) {
+inline std::vector<CallBase *> findIndirectCalls(Function &F) {
   PGOIndirectCallVisitor ICV;
   ICV.visit(F);
   return ICV.IndirectCalls;
 }
 } // namespace llvm
 
 #endif

llvm/include/llvm/Transforms/Instrumentation.h

 
 namespace llvm {
 
 class Triple;
 class FunctionPass;
 class ModulePass;
 class OptimizationRemarkEmitter;
 class Comdat;
+class CallBase;
 
 /// Instrumentation passes often insert conditional checks into entry blocks.
 /// Call this function before splitting the entry block to move instructions
 /// that must remain in the entry block up before the split point. Static
 /// allocas and llvm.localescape calls, for example, must remain in the entry
 /// block.
 BasicBlock::iterator PrepareToSplitEntryBlock(BasicBlock &BB,
                                               BasicBlock::iterator IP);

 
 // The pgo-specific indirect call promotion function declared below is used by
 // the pgo-driven indirect call promotion and sample profile passes. It's a
 // wrapper around llvm::promoteCall, et al. that additionally computes !prof
 // metadata. We place it in a pgo namespace so it's not confused with the
 // generic utilities.
 namespace pgo {
 
-// Helper function that transforms Inst (either an indirect-call instruction, or
+// Helper function that transforms CB (either an indirect-call instruction, or
 // an invoke instruction , to a conditional call to F. This is like:
 //     if (Inst.CalledValue == F)
 //        F(...);
 //     else
 //        Inst(...);
 //     end
 // TotalCount is the profile count value that the instruction executes.
 // Count is the profile count value that F is the target function.
 // These two values are used to update the branch weight.
 // If \p AttachProfToDirectCall is true, a prof metadata is attached to the
 // new direct call to contain \p Count.
 // Returns the promoted direct call instruction.
-Instruction *promoteIndirectCall(Instruction *Inst, Function *F, uint64_t Count,
-                                 uint64_t TotalCount,
-                                 bool AttachProfToDirectCall,
-                                 OptimizationRemarkEmitter *ORE);
+CallBase &promoteIndirectCall(CallBase &CB, Function *F, uint64_t Count,
+                              uint64_t TotalCount, bool AttachProfToDirectCall,
 } // namespace pgo
 
 /// Options for the frontend instrumentation based profiling pass.
 struct InstrProfOptions {
   // Add the 'noredzone' attribute to added runtime library calls.
   bool NoRedZone = false;
 
   // Do counter register promotion

llvm/lib/Transforms/IPO/SampleProfile.cpp

for (CallBase *I : CIS) {
 
           const char *Reason = "Callee function not available";
           auto R = SymbolMap.find(CalleeFunctionName);
           if (R != SymbolMap.end() && R->getValue() &&
               !R->getValue()->isDeclaration() &&
               R->getValue()->getSubprogram() &&
               isLegalToPromote(*I, R->getValue(), &Reason)) {
             uint64_t C = FS->getEntrySamples();
-            Instruction *DI =
-                pgo::promoteIndirectCall(I, R->getValue(), C, Sum, false, ORE);
+            auto &DI =
+                pgo::promoteIndirectCall(*I, R->getValue(), C, Sum, false, ORE);
             Sum -= C;
             PromotedInsns.insert(I);
             // If profile mismatches, we should not attempt to inline DI.
             if ((isa<CallInst>(DI) || isa<InvokeInst>(DI)) &&
-                inlineCallInstruction(*cast<CallBase>(DI))) {
+                inlineCallInstruction(cast<CallBase>(DI))) {
               localNotInlinedCallSites.erase(I);
               LocalChanged = true;
               ++NumCSInlined;
             }
           } else {
             LLVM_DEBUG(dbgs()
                        << "\nFailed to promote indirect call to "
                        << CalleeFunctionName << " because " << Reason << "\n");

llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

private:
   // TotalCount is the total profiled count of call executions, and
   // NumCandidates is the number of candidate entries in ValueDataRef.
   std::vector<PromotionCandidate> getPromotionCandidatesForCallSite(
       Instruction *Inst, const ArrayRef<InstrProfValueData> &ValueDataRef,
       uint64_t TotalCount, uint32_t NumCandidates);
 
   // Promote a list of targets for one indirect-call callsite. Return
   // the number of promotions.
-  uint32_t tryToPromote(Instruction *Inst,
+  uint32_t tryToPromote(CallBase &CB,
                         const std::vector<PromotionCandidate> &Candidates,
                         uint64_t &TotalCount);
 
 public:
   ICallPromotionFunc(Function &Func, Module *Modu, InstrProfSymtab *Symtab,
                      bool SamplePGO, OptimizationRemarkEmitter &ORE)
       : F(Func), M(Modu), Symtab(Symtab), SamplePGO(SamplePGO), ORE(ORE) {}
   ICallPromotionFunc(const ICallPromotionFunc &) = delete;
for (uint32_t I = 0; I < NumCandidates; I++) {
     }
 
     Ret.push_back(PromotionCandidate(TargetFunction, Count));
     TotalCount -= Count;
   }
   return Ret;
 }
 
-// FIXME(callsite): the Instruction* parameter and return can be changed to
-// CallBase
-Instruction *llvm::pgo::promoteIndirectCall(Instruction *Inst,
-                                            Function *DirectCallee,
-                                            uint64_t Count, uint64_t TotalCount,
-                                            bool AttachProfToDirectCall,
-                                            OptimizationRemarkEmitter *ORE) {
+CallBase &llvm::pgo::promoteIndirectCall(CallBase &CB, Function *DirectCallee,
 
   uint64_t ElseCount = TotalCount - Count;
   uint64_t MaxCount = (Count >= ElseCount ? Count : ElseCount);
   uint64_t Scale = calculateCountScale(MaxCount);
-  MDBuilder MDB(Inst->getContext());
+  MDBuilder MDB(CB.getContext());
   MDNode *BranchWeights = MDB.createBranchWeights(
       scaleBranchCount(Count, Scale), scaleBranchCount(ElseCount, Scale));
 
-  CallBase &NewInst = promoteCallWithIfThenElse(*cast<CallBase>(Inst),
-                                                DirectCallee, BranchWeights);
+  CallBase &NewInst =
+      promoteCallWithIfThenElse(CB, DirectCallee, BranchWeights);
 
   if (AttachProfToDirectCall) {
     MDBuilder MDB(NewInst.getContext());
     NewInst.setMetadata(
         LLVMContext::MD_prof,
         MDB.createBranchWeights({static_cast<uint32_t>(Count)}));
   }
 
   using namespace ore;
 
   if (ORE)
     ORE->emit([&]() {
-      return OptimizationRemark(DEBUG_TYPE, "Promoted", Inst)
+      return OptimizationRemark(DEBUG_TYPE, "Promoted", &CB)
              << "Promote indirect call to " << NV("DirectCallee", DirectCallee)
              << " with count " << NV("Count", Count) << " out of "
              << NV("TotalCount", TotalCount);
     });
-  return &NewInst;
+  return NewInst;
 }
 
 // Promote indirect-call to conditional direct-call for one callsite.
 uint32_t ICallPromotionFunc::tryToPromote(
-    Instruction *Inst, const std::vector<PromotionCandidate> &Candidates,
+    CallBase &CB, const std::vector<PromotionCandidate> &Candidates,
     uint64_t &TotalCount) {
   uint32_t NumPromoted = 0;
 
   for (auto &C : Candidates) {
     uint64_t Count = C.Count;
-    pgo::promoteIndirectCall(Inst, C.TargetFunction, Count, TotalCount,
-                             SamplePGO, &ORE);
+    pgo::promoteIndirectCall(CB, C.TargetFunction, Count, TotalCount, SamplePGO,
+                             &ORE);
     assert(TotalCount >= Count);
     TotalCount -= Count;
     NumOfPGOICallPromotion++;
     NumPromoted++;
   }
   return NumPromoted;
 }
 
 // Traverse all the indirect-call callsite and get the value profile
 // annotation to perform indirect-call promotion.
 bool ICallPromotionFunc::processFunction(ProfileSummaryInfo *PSI) {
   bool Changed = false;
   ICallPromotionAnalysis ICallAnalysis;
-  for (auto &I : findIndirectCalls(F)) {
+  for (auto *CB : findIndirectCalls(F)) {
     uint32_t NumVals, NumCandidates;
     uint64_t TotalCount;
     auto ICallProfDataRef = ICallAnalysis.getPromotionCandidatesForInstruction(
-        I, NumVals, TotalCount, NumCandidates);
+        CB, NumVals, TotalCount, NumCandidates);
     if (!NumCandidates ||
         (PSI && PSI->hasProfileSummary() && !PSI->isHotCount(TotalCount)))
       continue;
     auto PromotionCandidates = getPromotionCandidatesForCallSite(
-        I, ICallProfDataRef, TotalCount, NumCandidates);
-    uint32_t NumPromoted = tryToPromote(I, PromotionCandidates, TotalCount);
+        CB, ICallProfDataRef, TotalCount, NumCandidates);
+    uint32_t NumPromoted = tryToPromote(*CB, PromotionCandidates, TotalCount);
     if (NumPromoted == 0)
       continue;
 
     Changed = true;
     // Adjust the MD.prof metadata. First delete the old one.
-    I->setMetadata(LLVMContext::MD_prof, nullptr);
+    CB->setMetadata(LLVMContext::MD_prof, nullptr);
     // If all promoted, we don't need the MD.prof metadata.
     if (TotalCount == 0 || NumPromoted == NumVals)
       continue;
     // Otherwise we need update with the un-promoted records back.
-    annotateValueSite(*M, *I, ICallProfDataRef.slice(NumPromoted), TotalCount,
+    annotateValueSite(*M, *CB, ICallProfDataRef.slice(NumPromoted), TotalCount,
                       IPVK_IndirectCallTarget, NumCandidates);
   }
   return Changed;
 }
 
 // A wrapper function that does the actual work.
 static bool promoteIndirectCalls(Module &M, ProfileSummaryInfo *PSI,
                                  bool InLTO, bool SamplePGO,

llvm/lib/Transforms/Instrumentation/ValueProfilePlugins.inc

class IndirectCallPromotionPlugin {
   Function &F;
 
 public:
   static constexpr InstrProfValueKind Kind = IPVK_IndirectCallTarget;
 
   IndirectCallPromotionPlugin(Function &Fn) : F(Fn) {}
 
   void run(std::vector<CandidateInfo> &Candidates) {
-    std::vector<Instruction *> Result = findIndirectCalls(F);
+    std::vector<CallBase *> Result = findIndirectCalls(F);
     for (Instruction *I : Result) {
       Value *Callee = cast<CallBase>(I)->getCalledOperand();
       Instruction *InsertPt = I;
       Instruction *AnnotatedInst = I;
       Candidates.emplace_back(CandidateInfo{Callee, InsertPt, AnnotatedInst});
     }
   }
 };
 
 ///----------------------- Registration of the plugins -------------------------
 /// For now, registering a plugin with the ValueProfileCollector is done by
 /// adding the plugin type to the VP_PLUGIN_LIST macro.
 #define VP_PLUGIN_LIST           \
     MemIntrinsicPlugin,          \
     IndirectCallPromotionPlugin