Move the implementation logic to AsmPrinter.

Author: rlavaee

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -135,6 +135,12 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
   /// default, this is equal to CurrentFnSym.
   MCSymbol *CurrentFnSymForSize = nullptr;
 
+  /// Vector of symbols marking the position of callsites in each basic block.
+  /// The callsite symbols of each block are stored in the order they appear
+  /// in that block.
+  DenseMap<const MachineBasicBlock *, SmallVector<MCSymbol *, 1>>
+      CurrentFnCallsiteSymbols;
+
   /// Provides the profile information for constants.
   const StaticDataProfileInfo *SDPI = nullptr;
 
@@ -295,6 +301,10 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
   /// to emit them as well, return the whole set.
   ArrayRef<MCSymbol *> getAddrLabelSymbolToEmit(const BasicBlock *BB);
 
+  /// Creates a new symbol to be used for the beginning of a callsite at th
+  /// specified basic block.
+  MCSymbol *createCallsiteSymbol(const MachineBasicBlock &MBB);
+
   /// If the specified function has had any references to address-taken blocks
   /// generated, but the block got deleted, return the symbol now so we can
   /// emit it.  This prevents emitting a reference to a symbol that has no
@@ -426,7 +436,7 @@ class LLVM_ABI AsmPrinter : public MachineFunctionPass {
 
   void emitStackUsage(const MachineFunction &MF);
 
-  void emitBBAddrMapSection(const MachineFunction &MF, bool HasCalls);
+  void emitBBAddrMapSection(const MachineFunction &MF);
 
   void emitKCFITrapEntry(const MachineFunction &MF, const MCSymbol *Symbol);
   virtual void emitKCFITypeId(const MachineFunction &MF);

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -242,10 +242,6 @@ class MachineBasicBlock
   /// calculate the size of the basic block, or the BB section ending with it.
   mutable MCSymbol *CachedEndMCSymbol = nullptr;
 
-  /// Vector of symbols marking the position of callsites in the basic
-  /// block, stored in the order they appear in the basic block.
-  mutable SmallVector<MCSymbol *, 1> CallsiteSymbols;
-
   // Intrusive list support
   MachineBasicBlock() = default;
 
@@ -329,12 +325,6 @@ class MachineBasicBlock
   /// its label be emitted.
   void setLabelMustBeEmitted() { LabelMustBeEmitted = true; }
 
-  /// Returns the symbols marking callsites in the order they appear in the
-  /// basic block.
-  const SmallVectorImpl<MCSymbol *> &getCallsiteSymbols() const {
-    return CallsiteSymbols;
-  }
-
   /// Return the MachineFunction containing this basic block.
   const MachineFunction *getParent() const { return xParent; }
   MachineFunction *getParent() { return xParent; }

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1428,11 +1428,11 @@ getBBAddrMapFeature(const MachineFunction &MF, int NumMBBSectionRanges,
           HasCalls};
 }
 
-void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF,
-                                      bool HasCalls) {
+void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
   MCSection *BBAddrMapSection =
       getObjFileLowering().getBBAddrMapSection(*MF.getSection());
   assert(BBAddrMapSection && ".llvm_bb_addr_map section is not initialized.");
+  bool HasCalls = !CurrentFnCallsiteSymbols.empty();
 
   const MCSymbol *FunctionSymbol = getFunctionBegin();
 
@@ -1498,7 +1498,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF,
       const MCSymbol *CurrentLabel = MBBSymbol;
       if (HasCalls) {
         const SmallVectorImpl<MCSymbol *> &CallsiteSymbols =
-            MBB.getCallsiteSymbols();
+            CurrentFnCallsiteSymbols.lookup(&MBB);
         OutStreamer->AddComment("number of callsites");
         OutStreamer->emitULEB128IntValue(CallsiteSymbols.size());
         for (const MCSymbol *CallsiteSymbol : CallsiteSymbols) {
@@ -1815,7 +1815,6 @@ void AsmPrinter::emitFunctionBody() {
 
   // Print out code for the function.
   bool HasAnyRealCode = false;
-  bool HasCalls = false;
   int NumInstsInFunction = 0;
   bool IsEHa = MMI->getModule()->getModuleFlag("eh-asynch");
 
@@ -1842,11 +1841,8 @@ void AsmPrinter::emitFunctionBody() {
           !MI.isDebugInstr()) {
         HasAnyRealCode = true;
       }
-      if (MI.isCall()) {
-        HasCalls = true;
-        if (MF->getTarget().Options.BBAddrMap)
-          OutStreamer->emitLabel(MBB.createCallsiteSymbol());
-      }
+      if (MI.isCall() && MF->getTarget().Options.BBAddrMap)
+        OutStreamer->emitLabel(createCallsiteSymbol(MBB));
 
       // If there is a pre-instruction symbol, emit a label for it here.
       if (MCSymbol *S = MI.getPreInstrSymbol())
@@ -2133,7 +2129,7 @@ void AsmPrinter::emitFunctionBody() {
   // BB labels are requested for this function. Skip empty functions.
   if (HasAnyRealCode) {
     if (MF->getTarget().Options.BBAddrMap)
-      emitBBAddrMapSection(*MF, HasCalls);
+      emitBBAddrMapSection(*MF);
     else if (PgoAnalysisMapFeatures.getBits() != 0)
       MF->getContext().reportWarning(
           SMLoc(), "pgo-analysis-map is enabled for function " + MF->getName() +
@@ -2794,6 +2790,14 @@ MCSymbol *AsmPrinter::getMBBExceptionSym(const MachineBasicBlock &MBB) {
   return Res.first->second;
 }
 
+MCSymbol *AsmPrinter::createCallsiteSymbol(const MachineBasicBlock &MBB) {
+  MCContext &Ctx = MF->getContext();
+  MCSymbol *Sym = Ctx.createTempSymbol("BB" + Twine(MF->getFunctionNumber()) +
+                                       "_" + Twine(MBB.getNumber()) + "_CS");
+  CurrentFnCallsiteSymbols[&MBB].push_back(Sym);
+  return Sym;
+}
+
 void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
   this->MF = &MF;
   const Function &F = MF.getFunction();
@@ -2828,6 +2832,7 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
   CurrentFnBegin = nullptr;
   CurrentFnBeginLocal = nullptr;
   CurrentSectionBeginSym = nullptr;
+  CurrentFnCallsiteSymbols.clear();
   MBBSectionRanges.clear();
   MBBSectionExceptionSyms.clear();
   bool NeedsLocalForSize = MAI->needsLocalForSize();

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -112,15 +112,6 @@ MCSymbol *MachineBasicBlock::getEndSymbol() const {
   return CachedEndMCSymbol;
 }
 
-MCSymbol *MachineBasicBlock::createCallsiteSymbol() const {
-  const MachineFunction *MF = getParent();
-  MCContext &Ctx = MF->getContext();
-  CallsiteSymbols.push_back(
-      Ctx.createTempSymbol("BB" + Twine(MF->getFunctionNumber()) + "_" +
-                           Twine(getNumber()) + "_CS"));
-  return CallsiteSymbols.back();
-}
-
 raw_ostream &llvm::operator<<(raw_ostream &OS, const MachineBasicBlock &MBB) {
   MBB.print(OS);
   return OS;