Extend BasicBlock sections to allow specifying clusters of basic blocks

in the same section.

This allows specifying BasicBlock clusters like the following example:
!foo
!!0 1 2
!!4
This places basic blocks 0, 1, and 2 in one section in this order, and
places basic block #4 in a single section of its own.

Author: rlavaee

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -141,6 +141,10 @@ class AsmPrinter : public MachineFunctionPass {
   MCSymbol *CurrentFnEnd = nullptr;
   MCSymbol *CurExceptionSym = nullptr;
 
+  // The symbol used to represent the start of the current BB section of the
+  // function. This is used to calculate the size of the BB section.
+  MCSymbol *CurrentSectionBeginSym = nullptr;
+
   // The garbage collection metadata printer table.
   void *GCMetadataPrinters = nullptr; // Really a DenseMap.
 

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -46,17 +46,35 @@ class raw_ostream;
 class TargetRegisterClass;
 class TargetRegisterInfo;
 
-enum MachineBasicBlockSection : unsigned {
-  ///  This is also the order of sections in a function.  Basic blocks that are
-  ///  part of the original function section (entry block) come first, followed
-  ///  by exception handling basic blocks, cold basic blocks and finally basic
-  //   blocks that need unique sections.
-  MBBS_Entry,
-  MBBS_Exception,
-  MBBS_Cold,
-  MBBS_Unique,
-  ///  None implies no sections for any basic block, the default.
-  MBBS_None,
+// This structure uniquely identifies a basic block section.
+// Possible values are
+//  {Type: Default, Number: (unsigned)} (These are regular section IDs)
+//  {Type: Exception, Number: 0}  (ExceptionSectionID)
+//  {Type: Cold, Number: 0}  (ColdSectionID)
+struct MBBSectionID {
+  enum SectionType {
+    Default = 0, // Regular section (these sections are distinguished by the
+                 // Number field).
+    Exception,   // Special section type for exception handling blocks
+    Cold,        // Special section type for cold blocks
+  } Type;
+  unsigned Number;
+
+  MBBSectionID(unsigned N) : Type(Default), Number(N) {}
+
+  // Special unique sections for cold and exception blocks.
+  const static MBBSectionID ColdSectionID;
+  const static MBBSectionID ExceptionSectionID;
+
+  bool operator==(const MBBSectionID &Other) const {
+    return Type == Other.Type && Number == Other.Number;
+  }
+
+  bool operator!=(const MBBSectionID &Other) const { return !(*this == Other); }
+
+private:
+  // This is only used to construct the special cold and exception sections.
+  MBBSectionID(SectionType T) : Type(T), Number(0) {}
 };
 
 template <> struct ilist_traits<MachineInstr> {
@@ -143,8 +161,14 @@ class MachineBasicBlock
   /// Indicate that this basic block is the entry block of a cleanup funclet.
   bool IsCleanupFuncletEntry = false;
 
-  /// Stores the Section type of the basic block with basic block sections.
-  MachineBasicBlockSection SectionType = MBBS_None;
+  /// With basic block sections, this stores the Section ID of the basic block.
+  MBBSectionID SectionID{0};
+
+  // Indicate that this basic block begins a section.
+  bool IsBeginSection = false;
+
+  // Indicate that this basic block ends a section.
+  bool IsEndSection = false;
 
   /// Default target of the callbr of a basic block.
   bool InlineAsmBrDefaultTarget = false;
@@ -435,16 +459,20 @@ class MachineBasicBlock
   void setIsCleanupFuncletEntry(bool V = true) { IsCleanupFuncletEntry = V; }
 
   /// Returns true if this block begins any section.
-  bool isBeginSection() const;
+  bool isBeginSection() const { return IsBeginSection; }
 
   /// Returns true if this block ends any section.
-  bool isEndSection() const;
+  bool isEndSection() const { return IsEndSection; }
 
-  /// Returns the type of section this basic block belongs to.
-  MachineBasicBlockSection getSectionType() const { return SectionType; }
+  void setIsBeginSection(bool V = true) { IsBeginSection = V; }
 
-  /// Indicate that the basic block belongs to a Section Type.
-  void setSectionType(MachineBasicBlockSection V) { SectionType = V; }
+  void setIsEndSection(bool V = true) { IsEndSection = V; }
+
+  /// Returns the section ID of this basic block.
+  MBBSectionID getSectionID() const { return SectionID; }
+
+  /// Sets the section ID for this basic block.
+  void setSectionID(MBBSectionID V) { SectionID = V; }
 
   /// Returns true if this is the indirect dest of an INLINEASM_BR.
   bool isInlineAsmBrIndirectTarget(const MachineBasicBlock *Tgt) const {
@@ -485,10 +513,9 @@ class MachineBasicBlock
   void moveAfter(MachineBasicBlock *NewBefore);
 
   /// Returns true if this and MBB belong to the same section.
-  bool sameSection(const MachineBasicBlock *MBB) const;
-
-  /// Returns the basic block that ends the section which contains this one.
-  const MachineBasicBlock *getSectionEndMBB() const;
+  bool sameSection(const MachineBasicBlock *MBB) const {
+    return getSectionID() == MBB->getSectionID();
+  }
 
   /// Update the terminator instructions in block to account for changes to the
   /// layout. If the block previously used a fallthrough, it may now need a
@@ -876,12 +903,6 @@ class MachineBasicBlock
   /// Return the MCSymbol for this basic block.
   MCSymbol *getSymbol() const;
 
-  /// Sets the MCSymbol corresponding to the end of this basic block.
-  void setEndMCSymbol(MCSymbol *Sym) { EndMCSymbol = Sym; }
-
-  /// Returns the MCSymbol corresponding to the end of this basic block.
-  MCSymbol *getEndMCSymbol() const { return EndMCSymbol; }
-
   Optional<uint64_t> getIrrLoopHeaderWeight() const {
     return IrrLoopHeaderWeight;
   }

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -346,11 +346,6 @@ class MachineFunction {
   /// Section Type for basic blocks, only relevant with basic block sections.
   BasicBlockSection BBSectionsType = BasicBlockSection::None;
 
-  /// With Basic Block Sections, this stores the bb ranges of cold and
-  /// exception sections.
-  std::pair<int, int> ColdSectionRange = {-1, -1};
-  std::pair<int, int> ExceptionSectionRange = {-1, -1};
-
   /// List of C++ TypeInfo used.
   std::vector<const GlobalValue *> TypeInfos;
 
@@ -508,22 +503,13 @@ class MachineFunction {
 
   void setBBSectionsType(BasicBlockSection V) { BBSectionsType = V; }
 
-  void setSectionRange();
-
-  /// Returns true if this basic block number starts a cold or exception
-  /// section.
-  bool isSectionStartMBB(int N) const {
-    return (N == ColdSectionRange.first || N == ExceptionSectionRange.first);
-  }
-
-  /// Returns true if this basic block ends a cold or exception section.
-  bool isSectionEndMBB(int N) const {
-    return (N == ColdSectionRange.second || N == ExceptionSectionRange.second);
-  }
-
   /// Creates basic block Labels for this function.
   void createBBLabels();
 
+  /// Assign IsBeginSection IsEndSection fields for basic blocks in this
+  /// function.
+  void assignBeginEndSections();
+
   /// getTarget - Return the target machine this machine code is compiled with
   const LLVMTargetMachine &getTarget() const { return Target; }
 

llvm/include/llvm/CodeGen/TargetLoweringObjectFileImpl.h

@@ -69,10 +69,6 @@ class TargetLoweringObjectFileELF : public TargetLoweringObjectFile {
                                  const MachineBasicBlock &MBB,
                                  const TargetMachine &TM) const override;
 
-  MCSection *getNamedSectionForMachineBasicBlock(
-      const Function &F, const MachineBasicBlock &MBB, const TargetMachine &TM,
-      const char *Suffix) const override;
-
   bool shouldPutJumpTableInFunctionSection(bool UsesLabelDifference,
                                            const Function &F) const override;
 

llvm/include/llvm/Target/TargetLoweringObjectFile.h

@@ -92,10 +92,6 @@ class TargetLoweringObjectFile : public MCObjectFileInfo {
                                  const MachineBasicBlock &MBB,
                                  const TargetMachine &TM) const;
 
-  virtual MCSection *getNamedSectionForMachineBasicBlock(
-      const Function &F, const MachineBasicBlock &MBB, const TargetMachine &TM,
-      const char *Suffix) const;
-
   /// Classify the specified global variable into a set of target independent
   /// categories embodied in SectionKind.
   static SectionKind getKindForGlobal(const GlobalObject *GO,

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1097,14 +1097,6 @@ void AsmPrinter::emitFunctionBody() {
   // Print out code for the function.
   bool HasAnyRealCode = false;
   int NumInstsInFunction = 0;
-  bool emitBBSections = MF->hasBBSections();
-  MachineBasicBlock *EndOfRegularSectionMBB = nullptr;
-  if (emitBBSections) {
-    EndOfRegularSectionMBB =
-        const_cast<MachineBasicBlock *>(MF->front().getSectionEndMBB());
-    assert(EndOfRegularSectionMBB->isEndSection() &&
-           "The MBB at the end of the regular section must end a section");
-  }
 
   for (auto &MBB : *MF) {
     // Print a label for the basic block.
@@ -1185,17 +1177,41 @@ void AsmPrinter::emitFunctionBody() {
         }
       }
     }
-    if (&MBB != EndOfRegularSectionMBB &&
-        (MF->hasBBLabels() || MBB.isEndSection())) {
-      // Emit size directive for the size of this basic block.  Create a symbol
-      // for the end of the basic block.
-      MCSymbol *CurrentBBEnd = OutContext.createTempSymbol();
+
+    // We need a temporary symbol for the end of this basic block, if either we
+    // have BBLabels enabled and we want to emit size directive for the BBs, or
+    // if this basic blocks marks the end of a section (except the section
+    // containing the entry basic block as the end symbol for that section is
+    // CurrentFnEnd).
+    MCSymbol *CurrentBBEnd = nullptr;
+    if ((MAI->hasDotTypeDotSizeDirective() && MF->hasBBLabels()) ||
+        (MBB.isEndSection() && !MBB.sameSection(&MF->front()))) {
+      CurrentBBEnd = OutContext.createTempSymbol();
+      OutStreamer->emitLabel(CurrentBBEnd);
+    }
+
+    // Helper for emitting the size directive associated with a basic block
+    // symbol.
+    auto emitELFSizeDirective = [&](MCSymbol *SymForSize) {
+      assert(CurrentBBEnd && "Basicblock end symbol not set!");
       const MCExpr *SizeExp = MCBinaryExpr::createSub(
           MCSymbolRefExpr::create(CurrentBBEnd, OutContext),
-          MCSymbolRefExpr::create(MBB.getSymbol(), OutContext), OutContext);
-      OutStreamer->emitLabel(CurrentBBEnd);
-      MBB.setEndMCSymbol(CurrentBBEnd);
-      OutStreamer->emitELFSize(MBB.getSymbol(), SizeExp);
+          MCSymbolRefExpr::create(SymForSize, OutContext), OutContext);
+      OutStreamer->emitELFSize(SymForSize, SizeExp);
+    };
+
+    // Emit size directive for the size of each basic block, if BBLabels is
+    // enabled.
+    if (MAI->hasDotTypeDotSizeDirective() && MF->hasBBLabels())
+      emitELFSizeDirective(MBB.getSymbol());
+
+    // Emit size directive for the size of each basic block section once we
+    // get to the end of that section.
+    if (MBB.isEndSection()) {
+      if (!MBB.sameSection(&MF->front())) {
+        if (MAI->hasDotTypeDotSizeDirective())
+          emitELFSizeDirective(CurrentSectionBeginSym);
+      }
     }
     emitBasicBlockEnd(MBB);
   }
@@ -1230,9 +1246,8 @@ void AsmPrinter::emitFunctionBody() {
     }
   }
 
-  // Switch to the original section if basic block sections was used.
-  if (emitBBSections)
-    OutStreamer->SwitchSection(MF->getSection());
+  // Switch to the original section in case basic block sections was used.
+  OutStreamer->SwitchSection(MF->getSection());
 
   const Function &F = MF->getFunction();
   for (const auto &BB : F) {
@@ -1249,7 +1264,7 @@ void AsmPrinter::emitFunctionBody() {
   emitFunctionBodyEnd();
 
   if (needFuncLabelsForEHOrDebugInfo(*MF, MMI) ||
-      MAI->hasDotTypeDotSizeDirective() || emitBBSections) {
+      MAI->hasDotTypeDotSizeDirective()) {
     // Create a symbol for the end of function.
     CurrentFnEnd = createTempSymbol("func_end");
     OutStreamer->emitLabel(CurrentFnEnd);
@@ -1272,8 +1287,6 @@ void AsmPrinter::emitFunctionBody() {
     HI.Handler->markFunctionEnd();
   }
 
-  if (emitBBSections)
-    EndOfRegularSectionMBB->setEndMCSymbol(CurrentFnEnd);
 
   // Print out jump tables referenced by the function.
   emitJumpTableInfo();
@@ -1753,6 +1766,7 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
 
   CurrentFnSymForSize = CurrentFnSym;
   CurrentFnBegin = nullptr;
+  CurrentSectionBeginSym = nullptr;
   CurExceptionSym = nullptr;
   bool NeedsLocalForSize = MAI->needsLocalForSize();
   if (F.hasFnAttribute("patchable-function-entry") ||
@@ -2981,7 +2995,6 @@ static void emitBasicBlockLoopComments(const MachineBasicBlock &MBB,
 /// MachineBasicBlock, an alignment (if present) and a comment describing
 /// it if appropriate.
 void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
-  bool BBSections = MF->hasBBSections();
   // End the previous funclet and start a new one.
   if (MBB.isEHFuncletEntry()) {
     for (const HandlerInfo &HI : Handlers) {
@@ -2991,11 +3004,9 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
   }
 
   // Emit an alignment directive for this block, if needed.
-  if (MBB.pred_empty() || !BBSections) {
-    const Align Alignment = MBB.getAlignment();
-    if (Alignment != Align(1))
-      emitAlignment(Alignment);
-  }
+  const Align Alignment = MBB.getAlignment();
+  if (Alignment != Align(1))
+    emitAlignment(Alignment);
 
   // If the block has its address taken, emit any labels that were used to
   // reference the block.  It is possible that there is more than one label
@@ -3027,9 +3038,8 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
     emitBasicBlockLoopComments(MBB, MLI, *this);
   }
 
-  bool emitBBLabels = BBSections || MF->hasBBLabels();
   if (MBB.pred_empty() ||
-      (!emitBBLabels && isBlockOnlyReachableByFallthrough(&MBB) &&
+      (!MF->hasBBLabels() && isBlockOnlyReachableByFallthrough(&MBB) &&
        !MBB.isEHFuncletEntry() && !MBB.hasLabelMustBeEmitted())) {
     if (isVerbose()) {
       // NOTE: Want this comment at start of line, don't emit with AddComment.
@@ -3040,23 +3050,12 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
     if (isVerbose() && MBB.hasLabelMustBeEmitted()) {
       OutStreamer->AddComment("Label of block must be emitted");
     }
-    // With -fbasicblock-sections, a basic block can start a new section.
-    if (MBB.getSectionType() == MachineBasicBlockSection::MBBS_Exception) {
-      // Create the exception section for this function.
-      OutStreamer->SwitchSection(
-          getObjFileLowering().getNamedSectionForMachineBasicBlock(
-              MF->getFunction(), MBB, TM, ".eh"));
-    } else if (MBB.getSectionType() == MachineBasicBlockSection::MBBS_Cold) {
-      // Create the cold section here.
-      OutStreamer->SwitchSection(
-          getObjFileLowering().getNamedSectionForMachineBasicBlock(
-              MF->getFunction(), MBB, TM, ".unlikely"));
-    } else if (MBB.isBeginSection() && MBB.isEndSection()) {
+    // Switch to a new section if this basic block must begin a section.
+    if (MBB.isBeginSection()) {
       OutStreamer->SwitchSection(
           getObjFileLowering().getSectionForMachineBasicBlock(MF->getFunction(),
                                                               MBB, TM));
-    } else if (BBSections) {
-      OutStreamer->SwitchSection(MF->getSection());
+      CurrentSectionBeginSym = MBB.getSymbol();
     }
     OutStreamer->emitLabel(MBB.getSymbol());
   }
@@ -3090,7 +3089,7 @@ void AsmPrinter::emitVisibility(MCSymbol *Sym, unsigned Visibility,
 /// the predecessor and this block is a fall-through.
 bool AsmPrinter::
 isBlockOnlyReachableByFallthrough(const MachineBasicBlock *MBB) const {
-  // With BasicBlock Sections, no block is a fall through.
+  // With BasicBlock Sections, beginning of the section is not a fallthrough.
   if (MBB->isBeginSection())
     return false;