[Propeller] Use Fixed MBB ID instead of volatile MachineBasicBlock::Number.

Let Propeller use specialized IDs for basic blocks, instead of MBB number.

This allows optimizations not just prior to asm-printer, but throughout the entire codegen.
This patch only implements the functionality under the new `LLVM_BB_ADDR_MAP` version, but the old version is still being used. A later patch will change the used version.

####Background
Today Propeller uses machine basic block (MBB) numbers, which already exist, to map native assembly to machine IR.  This is done as follows.
    - Basic block addresses are captured and dumped into the `LLVM_BB_ADDR_MAP` section just before the AsmPrinter pass which writes out object files. This ensures that we have a mapping that is close to assembly.
    - Profiling mapping works by taking a virtual address of an instruction and looking up the `LLVM_BB_ADDR_MAP` section to find the MBB number it corresponds to.
    - While this works well today, we need to do better when we scale Propeller to target other Machine IR optimizations like spill code optimization.  Register allocation happens earlier in the Machine IR pipeline and we need an annotation mechanism that is valid at that point.
    - The current scheme will not work in this scenario because the MBB number of a particular basic block is not fixed and changes over the course of codegen (via renumbering, adding, and removing the basic blocks).
    - In other words, the volatile MBB numbers do not provide a one-to-one correspondence throughout the lifetime of Machine IR.  Profile annotation using MBB numbers is restricted to a fixed point; only valid at the exact point where it was dumped.
    - Further, the object file can only be dumped before AsmPrinter and cannot be dumped at an arbitrary point in the Machine IR pass pipeline.  Hence, MBB numbers are not suitable and we need something else.
####Solution
We propose using fixed unique incremental MBB IDs for basic blocks instead of volatile MBB numbers. These IDs are assigned upon the creation of machine basic blocks. We modify `MachineFunction::CreateMachineBasicBlock` to assign the fixed ID to every newly created basic block.  It assigns `MachineFunction::NextMBBID` to the MBB ID and then increments it, which ensures having unique IDs.

 To ensure correct profile attribution, multiple equivalent compilations must generate the same Propeller IDs. This is guaranteed as long as the MachineFunction passes run in the same order. Since the `NextBBID` variable is scoped to `MachineFunction`, interleaving of codegen for different functions won't cause any inconsistencies.

The new encoding is generated under the new version number 2 and we keep backward-compatibility with older versions.

####Impact on Size of the `LLVM_BB_ADDR_MAP` Section
Emitting the Propeller ID results in a 23% increase in the size of the `LLVM_BB_ADDR_MAP` section for the clang binary.

Reviewed By: tmsriram

Differential Revision: https://reviews.llvm.org/D100808

Author: rlavaee

llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h

@@ -31,8 +31,8 @@ namespace llvm {
 
 // The cluster information for a machine basic block.
 struct BBClusterInfo {
-  // MachineBasicBlock ID.
-  unsigned MBBNumber;
+  // Unique ID for this basic block.
+  unsigned BBID;
   // Cluster ID this basic block belongs to.
   unsigned ClusterID;
   // Position of basic block within the cluster.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -169,6 +169,10 @@ class MachineBasicBlock
   /// Indicate that this basic block is the entry block of a cleanup funclet.
   bool IsCleanupFuncletEntry = false;
 
+  /// Fixed unique ID assigned to this basic block upon creation. Used with
+  /// basic block sections and basic block labels.
+  Optional<unsigned> BBID;
+
   /// With basic block sections, this stores the Section ID of the basic block.
   MBBSectionID SectionID{0};
 
@@ -620,6 +624,14 @@ class MachineBasicBlock
 
   void setIsEndSection(bool V = true) { IsEndSection = V; }
 
+  Optional<unsigned> getBBID() const { return BBID; }
+
+  /// Returns the BBID of the block when BBAddrMapVersion >= 2, otherwise
+  /// returns `MachineBasicBlock::Number`.
+  /// TODO: Remove this function when version 1 is deprecated and replace its
+  /// uses with `getBBID()`.
+  unsigned getBBIDOrNumber() const;
+
   /// Returns the section ID of this basic block.
   MBBSectionID getSectionID() const { return SectionID; }
 
@@ -629,6 +641,12 @@ class MachineBasicBlock
            ((unsigned)SectionID.Type) + SectionID.Number;
   }
 
+  /// Sets the fixed BBID of this basic block.
+  void setBBID(unsigned V) {
+    assert(!BBID.has_value() && "Cannot change BBID.");
+    BBID = V;
+  }
+
   /// Sets the section ID for this basic block.
   void setSectionID(MBBSectionID V) { SectionID = V; }
 

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -374,6 +374,9 @@ class LLVM_EXTERNAL_VISIBILITY MachineFunction {
   bool HasEHScopes = false;
   bool HasEHFunclets = false;
 
+  /// BBID to assign to the next basic block of this function.
+  unsigned NextBBID = 0;
+
   /// Section Type for basic blocks, only relevant with basic block sections.
   BasicBlockSection BBSectionsType = BasicBlockSection::None;
 

llvm/include/llvm/Object/ELFTypes.h

@@ -799,6 +799,7 @@ struct BBAddrMap {
   uint64_t Addr; // Function address
   // Struct representing the BBAddrMap information for one basic block.
   struct BBEntry {
+    uint32_t ID;     // Unique ID of this basic block.
     uint32_t Offset; // Offset of basic block relative to function start.
     uint32_t Size;   // Size of the basic block.
 
@@ -809,13 +810,13 @@ struct BBAddrMap {
     bool IsEHPad;        // If this is an exception handling block.
     bool CanFallThrough; // If this block can fall through to its next.
 
-    BBEntry(uint32_t Offset, uint32_t Size, uint32_t Metadata)
-        : Offset(Offset), Size(Size), HasReturn(Metadata & 1),
+    BBEntry(uint32_t ID, uint32_t Offset, uint32_t Size, uint32_t Metadata)
+        : ID(ID), Offset(Offset), Size(Size), HasReturn(Metadata & 1),
           HasTailCall(Metadata & (1 << 1)), IsEHPad(Metadata & (1 << 2)),
           CanFallThrough(Metadata & (1 << 3)){};
 
     bool operator==(const BBEntry &Other) const {
-      return Offset == Other.Offset && Size == Other.Size &&
+      return ID == Other.ID && Offset == Other.Offset && Size == Other.Size &&
              HasReturn == Other.HasReturn && HasTailCall == Other.HasTailCall &&
              IsEHPad == Other.IsEHPad && CanFallThrough == Other.CanFallThrough;
     }

llvm/include/llvm/ObjectYAML/ELFYAML.h

@@ -158,6 +158,7 @@ struct DynamicEntry {
 
 struct BBAddrMapEntry {
   struct BBEntry {
+    uint32_t ID;
     llvm::yaml::Hex64 AddressOffset;
     llvm::yaml::Hex64 Size;
     llvm::yaml::Hex64 Metadata;

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1347,7 +1347,8 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
   OutStreamer->pushSection();
   OutStreamer->switchSection(BBAddrMapSection);
   OutStreamer->AddComment("version");
-  OutStreamer->emitInt8(OutStreamer->getContext().getBBAddrMapVersion());
+  uint8_t BBAddrMapVersion = OutStreamer->getContext().getBBAddrMapVersion();
+  OutStreamer->emitInt8(BBAddrMapVersion);
   OutStreamer->AddComment("feature");
   OutStreamer->emitInt8(0);
   OutStreamer->AddComment("function address");
@@ -1359,12 +1360,19 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
   for (const MachineBasicBlock &MBB : MF) {
     const MCSymbol *MBBSymbol =
         MBB.isEntryBlock() ? FunctionSymbol : MBB.getSymbol();
+    // TODO: Remove this check when version 1 is deprecated.
+    if (BBAddrMapVersion > 1) {
+      OutStreamer->AddComment("BB id");
+      // Emit the BB ID for this basic block.
+      OutStreamer->emitULEB128IntValue(*MBB.getBBID());
+    }
     // Emit the basic block offset relative to the end of the previous block.
     // This is zero unless the block is padded due to alignment.
     emitLabelDifferenceAsULEB128(MBBSymbol, PrevMBBEndSymbol);
     // Emit the basic block size. When BBs have alignments, their size cannot
     // always be computed from their offsets.
     emitLabelDifferenceAsULEB128(MBB.getEndSymbol(), MBBSymbol);
+    // Emit the Metadata.
     OutStreamer->emitULEB128IntValue(getBBAddrMapMetadata(MBB));
     PrevMBBEndSymbol = MBB.getEndSymbol();
   }

llvm/lib/CodeGen/BasicBlockSections.cpp

@@ -71,8 +71,8 @@
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
-#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
 #include "llvm/CodeGen/BasicBlockSectionUtils.h"
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/Passes.h"
@@ -131,9 +131,9 @@ INITIALIZE_PASS(BasicBlockSections, "bbsections-prepare",
 
 // This function updates and optimizes the branching instructions of every basic
 // block in a given function to account for changes in the layout.
-static void updateBranches(
-    MachineFunction &MF,
-    const SmallVector<MachineBasicBlock *, 4> &PreLayoutFallThroughs) {
+static void
+updateBranches(MachineFunction &MF,
+               const SmallVector<MachineBasicBlock *> &PreLayoutFallThroughs) {
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   SmallVector<MachineOperand, 4> Cond;
   for (auto &MBB : MF) {
@@ -168,7 +168,7 @@ static void updateBranches(
 bool getBBClusterInfoForFunction(
     const MachineFunction &MF,
     BasicBlockSectionsProfileReader *BBSectionsProfileReader,
-    std::vector<Optional<BBClusterInfo>> &V) {
+    DenseMap<unsigned, BBClusterInfo> &V) {
 
   // Find the assoicated cluster information.
   std::pair<bool, SmallVector<BBClusterInfo, 4>> P =
@@ -183,13 +183,8 @@ bool getBBClusterInfoForFunction(
     return true;
   }
 
-  V.resize(MF.getNumBlockIDs());
-  for (auto bbClusterInfo : P.second) {
-    // Bail out if the cluster information contains invalid MBB numbers.
-    if (bbClusterInfo.MBBNumber >= MF.getNumBlockIDs())
-      return false;
-    V[bbClusterInfo.MBBNumber] = bbClusterInfo;
-  }
+  for (const BBClusterInfo &BBCI : P.second)
+    V[BBCI.BBID] = BBCI;
   return true;
 }
 
@@ -200,11 +195,12 @@ bool getBBClusterInfoForFunction(
 // clusters, they are moved into a single "Exception" section. Eventually,
 // clusters are ordered in increasing order of their IDs, with the "Exception"
 // and "Cold" succeeding all other clusters.
-// FuncBBClusterInfo represent the cluster information for basic blocks. If this
-// is empty, it means unique sections for all basic blocks in the function.
+// FuncBBClusterInfo represent the cluster information for basic blocks. It
+// maps from BBID of basic blocks to their cluster information. If this is
+// empty, it means unique sections for all basic blocks in the function.
 static void
 assignSections(MachineFunction &MF,
-               const std::vector<Optional<BBClusterInfo>> &FuncBBClusterInfo) {
+               const DenseMap<unsigned, BBClusterInfo> &FuncBBClusterInfo) {
   assert(MF.hasBBSections() && "BB Sections is not set for function.");
   // This variable stores the section ID of the cluster containing eh_pads (if
   // all eh_pads are one cluster). If more than one cluster contain eh_pads, we
@@ -219,15 +215,21 @@ assignSections(MachineFunction &MF,
     if (MF.getTarget().getBBSectionsType() == llvm::BasicBlockSection::All ||
         FuncBBClusterInfo.empty()) {
       // If unique sections are desired for all basic blocks of the function, we
-      // set every basic block's section ID equal to its number (basic block
-      // id). This further ensures that basic blocks are ordered canonically.
-      MBB.setSectionID({static_cast<unsigned int>(MBB.getNumber())});
-    } else if (FuncBBClusterInfo[MBB.getNumber()])
-      MBB.setSectionID(FuncBBClusterInfo[MBB.getNumber()]->ClusterID);
-    else {
-      // BB goes into the special cold section if it is not specified in the
-      // cluster info map.
-      MBB.setSectionID(MBBSectionID::ColdSectionID);
+      // set every basic block's section ID equal to its original position in
+      // the layout (which is equal to its number). This ensures that basic
+      // blocks are ordered canonically.
+      MBB.setSectionID(MBB.getNumber());
+    } else {
+      // TODO: Replace `getBBIDOrNumber` with `getBBID` once version 1 is
+      // deprecated.
+      auto I = FuncBBClusterInfo.find(MBB.getBBIDOrNumber());
+      if (I != FuncBBClusterInfo.end()) {
+        MBB.setSectionID(I->second.ClusterID);
+      } else {
+        // BB goes into the special cold section if it is not specified in the
+        // cluster info map.
+        MBB.setSectionID(MBBSectionID::ColdSectionID);
+      }
     }
 
     if (MBB.isEHPad() && EHPadsSectionID != MBB.getSectionID() &&
@@ -250,12 +252,14 @@ assignSections(MachineFunction &MF,
 
 void llvm::sortBasicBlocksAndUpdateBranches(
     MachineFunction &MF, MachineBasicBlockComparator MBBCmp) {
-  SmallVector<MachineBasicBlock *, 4> PreLayoutFallThroughs(
-      MF.getNumBlockIDs());
+  [[maybe_unused]] const MachineBasicBlock *EntryBlock = &MF.front();
+  SmallVector<MachineBasicBlock *> PreLayoutFallThroughs(MF.getNumBlockIDs());
   for (auto &MBB : MF)
     PreLayoutFallThroughs[MBB.getNumber()] = MBB.getFallThrough();
 
   MF.sort(MBBCmp);
+  assert(&MF.front() == EntryBlock &&
+         "Entry block should not be displaced by basic block sections");
 
   // Set IsBeginSection and IsEndSection according to the assigned section IDs.
   MF.assignBeginEndSections();
@@ -318,11 +322,14 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   if (BBSectionsType == BasicBlockSection::List &&
       hasInstrProfHashMismatch(MF))
     return true;
-
-  // Renumber blocks before sorting them for basic block sections.  This is
-  // useful during sorting, basic blocks in the same section will retain the
-  // default order.  This renumbering should also be done for basic block
-  // labels to match the profiles with the correct blocks.
+  // Renumber blocks before sorting them. This is useful during sorting,
+  // basic blocks in the same section will retain the default order.
+  // This renumbering should also be done for basic block labels to match the
+  // profiles with the correct blocks.
+  // For LLVM_BB_ADDR_MAP versions 2 and higher, this renumbering serves
+  // the different purpose of accessing the original layout positions and
+  // finding the original fallthroughs.
+  // TODO: Change the above comment accordingly when version 1 is deprecated.
   MF.RenumberBlocks();
 
   if (BBSectionsType == BasicBlockSection::Labels) {
@@ -332,7 +339,8 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
 
   BBSectionsProfileReader = &getAnalysis<BasicBlockSectionsProfileReader>();
 
-  std::vector<Optional<BBClusterInfo>> FuncBBClusterInfo;
+  // Map from BBID of blocks to their cluster information.
+  DenseMap<unsigned, BBClusterInfo> FuncBBClusterInfo;
   if (BBSectionsType == BasicBlockSection::List &&
       !getBBClusterInfoForFunction(MF, BBSectionsProfileReader,
                                    FuncBBClusterInfo))
@@ -372,8 +380,8 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
     // If the two basic block are in the same section, the order is decided by
     // their position within the section.
     if (XSectionID.Type == MBBSectionID::SectionType::Default)
-      return FuncBBClusterInfo[X.getNumber()]->PositionInCluster <
-             FuncBBClusterInfo[Y.getNumber()]->PositionInCluster;
+      return FuncBBClusterInfo.lookup(X.getBBIDOrNumber()).PositionInCluster <
+             FuncBBClusterInfo.lookup(Y.getBBIDOrNumber()).PositionInCluster;
     return X.getNumber() < Y.getNumber();
   };
 

llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp

@@ -93,23 +93,23 @@ static Error getBBClusterInfo(const MemoryBuffer *MBuf,
       if (FI == ProgramBBClusterInfo.end())
         return invalidProfileError(
             "Cluster list does not follow a function name specifier.");
-      SmallVector<StringRef, 4> BBIndexes;
-      S.split(BBIndexes, ' ');
+      SmallVector<StringRef, 4> BBIDs;
+      S.split(BBIDs, ' ');
       // Reset current cluster position.
       CurrentPosition = 0;
-      for (auto BBIndexStr : BBIndexes) {
-        unsigned long long BBIndex;
-        if (getAsUnsignedInteger(BBIndexStr, 10, BBIndex))
+      for (auto BBIDStr : BBIDs) {
+        unsigned long long BBID;
+        if (getAsUnsignedInteger(BBIDStr, 10, BBID))
           return invalidProfileError(Twine("Unsigned integer expected: '") +
-                                     BBIndexStr + "'.");
-        if (!FuncBBIDs.insert(BBIndex).second)
+                                     BBIDStr + "'.");
+        if (!FuncBBIDs.insert(BBID).second)
           return invalidProfileError(Twine("Duplicate basic block id found '") +
-                                     BBIndexStr + "'.");
-        if (!BBIndex && CurrentPosition)
+                                     BBIDStr + "'.");
+        if (BBID == 0 && CurrentPosition)
           return invalidProfileError("Entry BB (0) does not begin a cluster.");
 
-        FI->second.emplace_back(BBClusterInfo{
-            ((unsigned)BBIndex), CurrentCluster, CurrentPosition++});
+        FI->second.emplace_back(
+            BBClusterInfo{((unsigned)BBID), CurrentCluster, CurrentPosition++});
       }
       CurrentCluster++;
     } else { // This is a function name specifier.

llvm/lib/CodeGen/MIRParser/MILexer.cpp

@@ -272,11 +272,13 @@ static MIToken::TokenKind getIdentifierKind(StringRef Identifier) {
       .Case("pcsections", MIToken::kw_pcsections)
       .Case("cfi-type", MIToken::kw_cfi_type)
       .Case("bbsections", MIToken::kw_bbsections)
+      .Case("bb_id", MIToken::kw_bb_id)
       .Case("unknown-size", MIToken::kw_unknown_size)
       .Case("unknown-address", MIToken::kw_unknown_address)
       .Case("distinct", MIToken::kw_distinct)
       .Case("ir-block-address-taken", MIToken::kw_ir_block_address_taken)
-      .Case("machine-block-address-taken", MIToken::kw_machine_block_address_taken)
+      .Case("machine-block-address-taken",
+            MIToken::kw_machine_block_address_taken)
       .Default(MIToken::Identifier);
 }
 

llvm/lib/CodeGen/MIRParser/MILexer.h

@@ -128,6 +128,7 @@ struct MIToken {
     kw_pcsections,
     kw_cfi_type,
     kw_bbsections,
+    kw_bb_id,
     kw_unknown_size,
     kw_unknown_address,
     kw_ir_block_address_taken,

llvm/lib/CodeGen/MIRParser/MIParser.cpp

@@ -499,6 +499,7 @@ class MIParser {
   bool parseAlignment(uint64_t &Alignment);
   bool parseAddrspace(unsigned &Addrspace);
   bool parseSectionID(Optional<MBBSectionID> &SID);
+  bool parseBBID(Optional<unsigned> &BBID);
   bool parseOperandsOffset(MachineOperand &Op);
   bool parseIRValue(const Value *&V);
   bool parseMemoryOperandFlag(MachineMemOperand::Flags &Flags);
@@ -662,6 +663,18 @@ bool MIParser::parseSectionID(Optional<MBBSectionID> &SID) {
   return false;
 }
 
+// Parse Machine Basic Block ID.
+bool MIParser::parseBBID(Optional<unsigned> &BBID) {
+  assert(Token.is(MIToken::kw_bb_id));
+  lex();
+  unsigned Value = 0;
+  if (getUnsigned(Value))
+    return error("Unknown BB ID");
+  BBID = Value;
+  lex();
+  return false;
+}
+
 bool MIParser::parseBasicBlockDefinition(
     DenseMap<unsigned, MachineBasicBlock *> &MBBSlots) {
   assert(Token.is(MIToken::MachineBasicBlockLabel));
@@ -678,6 +691,7 @@ bool MIParser::parseBasicBlockDefinition(
   bool IsEHFuncletEntry = false;
   Optional<MBBSectionID> SectionID;
   uint64_t Alignment = 0;
+  Optional<unsigned> BBID;
   BasicBlock *BB = nullptr;
   if (consumeIfPresent(MIToken::lparen)) {
     do {
@@ -718,6 +732,10 @@ bool MIParser::parseBasicBlockDefinition(
         if (parseSectionID(SectionID))
           return true;
         break;
+      case MIToken::kw_bb_id:
+        if (parseBBID(BBID))
+          return true;
+        break;
       default:
         break;
       }
@@ -755,6 +773,13 @@ bool MIParser::parseBasicBlockDefinition(
     MBB->setSectionID(SectionID.value());
     MF.setBBSectionsType(BasicBlockSection::List);
   }
+  if (BBID.has_value()) {
+    // BBSectionsType is set to `List` if any basic blocks has `SectionID`.
+    // Here, we set it to `Labels` if it hasn't been set above.
+    if (!MF.hasBBSections())
+      MF.setBBSectionsType(BasicBlockSection::Labels);
+    MBB->setBBID(BBID.value());
+  }
   return false;
 }