[MC] Compute fragment offsets eagerly

This builds on top of commit 9d0754a
("[MC] Relax fragments eagerly") and relaxes fragments eagerly to
eliminate MCSection::HasLayout and `getFragmentOffset` overhead.

Note: The removed `#ifndef NDEBUG` code (disabled by default) in
X86AsmBackend::finishLayout was problematic, as (a) !NDEBUG and NDEBUG
builds evaluated fragment offsets at different times before this patch
(b) one iteration might not be sufficient to converge. There might be
some edge cases that it did not handle. Anyhow, this patch probably
makes it work for more cases.

Author: MaskRay

llvm/include/llvm/MC/MCAsmBackend.h

@@ -217,8 +217,9 @@ class MCAsmBackend {
   virtual bool writeNopData(raw_ostream &OS, uint64_t Count,
                             const MCSubtargetInfo *STI) const = 0;
 
-  /// Give backend an opportunity to finish layout after relaxation
-  virtual void finishLayout(MCAssembler const &Asm) const {}
+  // Return true if fragment offsets have been adjusted and an extra layout
+  // iteration is needed.
+  virtual bool finishLayout(const MCAssembler &Asm) const { return false; }
 
   /// Handle any target-specific assembler flags. By default, do nothing.
   virtual void handleAssemblerFlag(MCAssemblerFlag Flag) {}

llvm/include/llvm/MC/MCAssembler.h

@@ -126,6 +126,7 @@ class MCAssembler {
   /// Check whether the given fragment needs relaxation.
   bool fragmentNeedsRelaxation(const MCRelaxableFragment *IF) const;
 
+  void layoutSection(MCSection &Sec);
   /// Perform one layout iteration and return true if any offsets
   /// were adjusted.
   bool layoutOnce();
@@ -172,10 +173,9 @@ class MCAssembler {
   uint64_t computeFragmentSize(const MCFragment &F) const;
 
   void layoutBundle(MCFragment *Prev, MCFragment *F) const;
-  void ensureValid(MCSection &Sec) const;
 
   // Get the offset of the given fragment inside its containing section.
-  uint64_t getFragmentOffset(const MCFragment &F) const;
+  uint64_t getFragmentOffset(const MCFragment &F) const { return F.Offset; }
 
   uint64_t getSectionAddressSize(const MCSection &Sec) const;
   uint64_t getSectionFileSize(const MCSection &Sec) const;

llvm/include/llvm/MC/MCSection.h

@@ -99,8 +99,6 @@ class MCSection {
   /// Whether this section has had instructions emitted into it.
   bool HasInstructions : 1;
 
-  bool HasLayout : 1;
-
   bool IsRegistered : 1;
 
   bool IsText : 1;
@@ -169,9 +167,6 @@ class MCSection {
   bool hasInstructions() const { return HasInstructions; }
   void setHasInstructions(bool Value) { HasInstructions = Value; }
 
-  bool hasLayout() const { return HasLayout; }
-  void setHasLayout(bool Value) { HasLayout = Value; }
-
   bool isRegistered() const { return IsRegistered; }
   void setIsRegistered(bool Value) { IsRegistered = Value; }
 

llvm/lib/MC/MCAssembler.cpp

@@ -438,28 +438,6 @@ void MCAssembler::layoutBundle(MCFragment *Prev, MCFragment *F) const {
       DF->Offset = EF->Offset;
 }
 
-void MCAssembler::ensureValid(MCSection &Sec) const {
-  if (Sec.hasLayout())
-    return;
-  Sec.setHasLayout(true);
-  MCFragment *Prev = nullptr;
-  uint64_t Offset = 0;
-  for (MCFragment &F : Sec) {
-    F.Offset = Offset;
-    if (isBundlingEnabled() && F.hasInstructions()) {
-      layoutBundle(Prev, &F);
-      Offset = F.Offset;
-    }
-    Offset += computeFragmentSize(F);
-    Prev = &F;
-  }
-}
-
-uint64_t MCAssembler::getFragmentOffset(const MCFragment &F) const {
-  ensureValid(*F.getParent());
-  return F.Offset;
-}
-
 // Simple getSymbolOffset helper for the non-variable case.
 static bool getLabelOffset(const MCAssembler &Asm, const MCSymbol &S,
                            bool ReportError, uint64_t &Val) {
@@ -944,22 +922,20 @@ void MCAssembler::layout() {
 
   // Layout until everything fits.
   this->HasLayout = true;
+  for (MCSection &Sec : *this)
+    layoutSection(Sec);
   while (layoutOnce()) {
-    if (getContext().hadError())
-      return;
-    // Size of fragments in one section can depend on the size of fragments in
-    // another. If any fragment has changed size, we have to re-layout (and
-    // as a result possibly further relax) all.
-    for (MCSection &Sec : *this)
-      Sec.setHasLayout(false);
   }
 
   DEBUG_WITH_TYPE("mc-dump", {
       errs() << "assembler backend - post-relaxation\n--\n";
       dump(); });
 
-  // Finalize the layout, including fragment lowering.
-  getBackend().finishLayout(*this);
+  // Some targets might want to adjust fragment offsets. If so, perform another
+  // layout loop.
+  if (getBackend().finishLayout(*this))
+    for (MCSection &Sec : *this)
+      layoutSection(Sec);
 
   DEBUG_WITH_TYPE("mc-dump", {
       errs() << "assembler backend - final-layout\n--\n";
@@ -1312,15 +1288,42 @@ bool MCAssembler::relaxFragment(MCFragment &F) {
   }
 }
 
+void MCAssembler::layoutSection(MCSection &Sec) {
+  MCFragment *Prev = nullptr;
+  uint64_t Offset = 0;
+  for (MCFragment &F : Sec) {
+    F.Offset = Offset;
+    if (LLVM_UNLIKELY(isBundlingEnabled())) {
+      if (F.hasInstructions()) {
+        layoutBundle(Prev, &F);
+        Offset = F.Offset;
+      }
+      Prev = &F;
+    }
+    Offset += computeFragmentSize(F);
+  }
+}
+
 bool MCAssembler::layoutOnce() {
   ++stats::RelaxationSteps;
 
-  bool Changed = false;
-  for (MCSection &Sec : *this)
-    for (MCFragment &Frag : Sec)
-      if (relaxFragment(Frag))
-        Changed = true;
-  return Changed;
+  // Size of fragments in one section can depend on the size of fragments in
+  // another. If any fragment has changed size, we have to re-layout (and
+  // as a result possibly further relax) all.
+  bool ChangedAny = false;
+  for (MCSection &Sec : *this) {
+    for (;;) {
+      bool Changed = false;
+      for (MCFragment &F : Sec)
+        if (relaxFragment(F))
+          Changed = true;
+      ChangedAny |= Changed;
+      if (!Changed)
+        break;
+      layoutSection(Sec);
+    }
+  }
+  return ChangedAny;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)

llvm/lib/MC/MCSection.cpp

@@ -23,8 +23,8 @@ using namespace llvm;
 MCSection::MCSection(SectionVariant V, StringRef Name, bool IsText,
                      bool IsVirtual, MCSymbol *Begin)
     : Begin(Begin), BundleGroupBeforeFirstInst(false), HasInstructions(false),
-      HasLayout(false), IsRegistered(false), IsText(IsText),
-      IsVirtual(IsVirtual), Name(Name), Variant(V) {
+      IsRegistered(false), IsText(IsText), IsVirtual(IsVirtual), Name(Name),
+      Variant(V) {
   DummyFragment.setParent(this);
   // The initial subsection number is 0. Create a fragment list.
   CurFragList = &Subsections.emplace_back(0u, FragList{}).second;

llvm/lib/Target/Hexagon/MCTargetDesc/HexagonAsmBackend.cpp

@@ -702,7 +702,7 @@ class HexagonAsmBackend : public MCAsmBackend {
     return true;
   }
 
-  void finishLayout(MCAssembler const &Asm) const override {
+  bool finishLayout(const MCAssembler &Asm) const override {
     SmallVector<MCFragment *> Frags;
     for (MCSection &Sec : Asm) {
       Frags.clear();
@@ -747,7 +747,6 @@ class HexagonAsmBackend : public MCAsmBackend {
               //assert(!Error);
               (void)Error;
               ReplaceInstruction(Asm.getEmitter(), RF, Inst);
-              Sec.setHasLayout(false);
               Size = 0; // Only look back one instruction
               break;
             }
@@ -757,6 +756,7 @@ class HexagonAsmBackend : public MCAsmBackend {
         }
       }
     }
+    return true;
   }
 }; // class HexagonAsmBackend
 

llvm/lib/Target/X86/MCTargetDesc/X86AsmBackend.cpp

@@ -201,7 +201,7 @@ class X86AsmBackend : public MCAsmBackend {
   bool padInstructionEncoding(MCRelaxableFragment &RF, MCCodeEmitter &Emitter,
                               unsigned &RemainingSize) const;
 
-  void finishLayout(const MCAssembler &Asm) const override;
+  bool finishLayout(const MCAssembler &Asm) const override;
 
   unsigned getMaximumNopSize(const MCSubtargetInfo &STI) const override;
 
@@ -856,15 +856,15 @@ bool X86AsmBackend::padInstructionEncoding(MCRelaxableFragment &RF,
   return Changed;
 }
 
-void X86AsmBackend::finishLayout(MCAssembler const &Asm) const {
+bool X86AsmBackend::finishLayout(const MCAssembler &Asm) const {
   // See if we can further relax some instructions to cut down on the number of
   // nop bytes required for code alignment.  The actual win is in reducing
   // instruction count, not number of bytes.  Modern X86-64 can easily end up
   // decode limited.  It is often better to reduce the number of instructions
   // (i.e. eliminate nops) even at the cost of increasing the size and
   // complexity of others.
   if (!X86PadForAlign && !X86PadForBranchAlign)
-    return;
+    return false;
 
   // The processed regions are delimitered by LabeledFragments. -g may have more
   // MCSymbols and therefore different relaxation results. X86PadForAlign is
@@ -911,9 +911,6 @@ void X86AsmBackend::finishLayout(MCAssembler const &Asm) const {
         continue;
       }
 
-#ifndef NDEBUG
-      const uint64_t OrigOffset = Asm.getFragmentOffset(F);
-#endif
       const uint64_t OrigSize = Asm.computeFragmentSize(F);
 
       // To keep the effects local, prefer to relax instructions closest to
@@ -926,8 +923,7 @@ void X86AsmBackend::finishLayout(MCAssembler const &Asm) const {
         // Give the backend a chance to play any tricks it wishes to increase
         // the encoding size of the given instruction.  Target independent code
         // will try further relaxation, but target's may play further tricks.
-        if (padInstructionEncoding(RF, Asm.getEmitter(), RemainingSize))
-          Sec.setHasLayout(false);
+        padInstructionEncoding(RF, Asm.getEmitter(), RemainingSize);
 
         // If we have an instruction which hasn't been fully relaxed, we can't
         // skip past it and insert bytes before it.  Changing its starting
@@ -944,14 +940,6 @@ void X86AsmBackend::finishLayout(MCAssembler const &Asm) const {
       if (F.getKind() == MCFragment::FT_BoundaryAlign)
         cast<MCBoundaryAlignFragment>(F).setSize(RemainingSize);
 
-#ifndef NDEBUG
-      const uint64_t FinalOffset = Asm.getFragmentOffset(F);
-      const uint64_t FinalSize = Asm.computeFragmentSize(F);
-      assert(OrigOffset + OrigSize == FinalOffset + FinalSize &&
-             "can't move start of next fragment!");
-      assert(FinalSize == RemainingSize && "inconsistent size computation?");
-#endif
-
       // If we're looking at a boundary align, make sure we don't try to pad
       // its target instructions for some following directive.  Doing so would
       // break the alignment of the current boundary align.
@@ -965,11 +953,7 @@ void X86AsmBackend::finishLayout(MCAssembler const &Asm) const {
     }
   }
 
-  // The layout is done. Mark every fragment as valid.
-  for (MCSection &Section : Asm) {
-    Asm.getFragmentOffset(*Section.curFragList()->Tail);
-    Asm.computeFragmentSize(*Section.curFragList()->Tail);
-  }
+  return true;
 }
 
 unsigned X86AsmBackend::getMaximumNopSize(const MCSubtargetInfo &STI) const {