[MC] Relax fragments eagerly

MaskRay · MaskRay · commit 9d0754ada5db · 2024-06-09T23:05:05.000-07:00
Lazy relaxation caused hash table lookups (`getFragmentOffset`) and complex use/compute interdependencies. Some expressions involding forward declared symbols (e.g. `subsection-if.s`) cannot be computed. Recursion detection requires complex `IsBeingLaidOut` (https://reviews.llvm.org/D79570). D76114's `invalidateFragmentsFrom` makes lazy relaxation even less useful. Switch to eager relaxation to greatly simplify code and resolve these issues. This change also removes a `getPrevNode` use, which makes it more feasible to replace the fragment representation, which might yield a large peak RSS win. Minor downsides: The number of section relaxations may increase (offset by avoiding the hash table lookup). For relax-recompute-align.s, the computed layout is not optimal.
diff --git a/llvm/include/llvm/MC/MCAsmLayout.h b/llvm/include/llvm/MC/MCAsmLayout.h
@@ -31,37 +31,21 @@ class MCAsmLayout {
   /// List of sections in layout order.
   llvm::SmallVector<MCSection *, 16> SectionOrder;
 
-  /// The last fragment which was laid out, or 0 if nothing has been laid
-  /// out. Fragments are always laid out in order, so all fragments with a
-  /// lower ordinal will be valid.
-  mutable DenseMap<const MCSection *, MCFragment *> LastValidFragment;
-
-  /// Make sure that the layout for the given fragment is valid, lazily
-  /// computing it if necessary.
+  /// Compute the layout for the section if necessary.
   void ensureValid(const MCFragment *F) const;
 
-  /// Is the layout for this fragment valid?
-  bool isFragmentValid(const MCFragment *F) const;
-
 public:
   MCAsmLayout(MCAssembler &Assembler);
 
   /// Get the assembler object this is a layout for.
   MCAssembler &getAssembler() const { return Assembler; }
 
-  /// \returns whether the offset of fragment \p F can be obtained via
-  /// getFragmentOffset.
-  bool canGetFragmentOffset(const MCFragment *F) const;
-
   /// Invalidate the fragments starting with F because it has been
   /// resized. The fragment's size should have already been updated, but
   /// its bundle padding will be recomputed.
   void invalidateFragmentsFrom(MCFragment *F);
 
-  /// Perform layout for a single fragment, assuming that the previous
-  /// fragment has already been laid out correctly, and the parent section has
-  /// been initialized.
-  void layoutFragment(MCFragment *Fragment);
+  void layoutBundle(MCFragment *F);
 
   /// \name Section Access (in layout order)
   /// @{
diff --git a/llvm/include/llvm/MC/MCAssembler.h b/llvm/include/llvm/MC/MCAssembler.h
@@ -203,10 +203,6 @@ class MCAssembler {
   /// were adjusted.
   bool layoutOnce(MCAsmLayout &Layout);
 
-  /// Perform one layout iteration of the given section and return true
-  /// if any offsets were adjusted.
-  bool layoutSectionOnce(MCAsmLayout &Layout, MCSection &Sec);
-
   /// Perform relaxation on a single fragment - returns true if the fragment
   /// changes as a result of relaxation.
   bool relaxFragment(MCAsmLayout &Layout, MCFragment &F);
diff --git a/llvm/include/llvm/MC/MCFragment.h b/llvm/include/llvm/MC/MCFragment.h
@@ -70,9 +70,6 @@ class MCFragment : public ilist_node_with_parent<MCFragment, MCSection> {
 
   FragmentType Kind;
 
-  /// Whether fragment is being laid out.
-  bool IsBeingLaidOut;
-
 protected:
   bool HasInstructions;
   bool LinkerRelaxable = false;
diff --git a/llvm/include/llvm/MC/MCSection.h b/llvm/include/llvm/MC/MCSection.h
@@ -89,6 +89,8 @@ class MCSection {
   /// Whether this section has had instructions emitted into it.
   bool HasInstructions : 1;
 
+  bool HasLayout : 1;
+
   bool IsRegistered : 1;
 
   MCDummyFragment DummyFragment;
@@ -166,6 +168,9 @@ 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; }
 
diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
@@ -66,7 +66,6 @@ STATISTIC(EmittedFillFragments,
 STATISTIC(EmittedNopsFragments, "Number of emitted assembler fragments - nops");
 STATISTIC(EmittedOrgFragments, "Number of emitted assembler fragments - org");
 STATISTIC(evaluateFixup, "Number of evaluated fixups");
-STATISTIC(FragmentLayouts, "Number of fragment layouts");
 STATISTIC(ObjectBytes, "Number of emitted object file bytes");
 STATISTIC(RelaxationSteps, "Number of assembler layout and relaxation steps");
 STATISTIC(RelaxedInstructions, "Number of relaxed instructions");
@@ -404,29 +403,7 @@ uint64_t MCAssembler::computeFragmentSize(const MCAsmLayout &Layout,
   llvm_unreachable("invalid fragment kind");
 }
 
-void MCAsmLayout::layoutFragment(MCFragment *F) {
-  MCFragment *Prev = F->getPrevNode();
-
-  // We should never try to recompute something which is valid.
-  assert(!isFragmentValid(F) && "Attempt to recompute a valid fragment!");
-  // We should never try to compute the fragment layout if its predecessor
-  // isn't valid.
-  assert((!Prev || isFragmentValid(Prev)) &&
-         "Attempt to compute fragment before its predecessor!");
-
-  assert(!F->IsBeingLaidOut && "Already being laid out!");
-  F->IsBeingLaidOut = true;
-
-  ++stats::FragmentLayouts;
-
-  // Compute fragment offset and size.
-  if (Prev)
-    F->Offset = Prev->Offset + getAssembler().computeFragmentSize(*this, *Prev);
-  else
-    F->Offset = 0;
-  F->IsBeingLaidOut = false;
-  LastValidFragment[F->getParent()] = F;
-
+void MCAsmLayout::layoutBundle(MCFragment *F) {
   // If bundling is enabled and this fragment has instructions in it, it has to
   // obey the bundling restrictions. With padding, we'll have:
   //
@@ -454,21 +431,40 @@ void MCAsmLayout::layoutFragment(MCFragment *F) {
   // within-fragment padding (which would produce less padding when N is less
   // than the bundle size), but for now we don't.
   //
-  if (Assembler.isBundlingEnabled() && F->hasInstructions()) {
-    assert(isa<MCEncodedFragment>(F) &&
-           "Only MCEncodedFragment implementations have instructions");
-    MCEncodedFragment *EF = cast<MCEncodedFragment>(F);
-    uint64_t FSize = Assembler.computeFragmentSize(*this, *EF);
-
-    if (!Assembler.getRelaxAll() && FSize > Assembler.getBundleAlignSize())
-      report_fatal_error("Fragment can't be larger than a bundle size");
-
-    uint64_t RequiredBundlePadding =
-        computeBundlePadding(Assembler, EF, EF->Offset, FSize);
-    if (RequiredBundlePadding > UINT8_MAX)
-      report_fatal_error("Padding cannot exceed 255 bytes");
-    EF->setBundlePadding(static_cast<uint8_t>(RequiredBundlePadding));
-    EF->Offset += RequiredBundlePadding;
+  assert(isa<MCEncodedFragment>(F) &&
+         "Only MCEncodedFragment implementations have instructions");
+  MCEncodedFragment *EF = cast<MCEncodedFragment>(F);
+  uint64_t FSize = Assembler.computeFragmentSize(*this, *EF);
+
+  if (!Assembler.getRelaxAll() && FSize > Assembler.getBundleAlignSize())
+    report_fatal_error("Fragment can't be larger than a bundle size");
+
+  uint64_t RequiredBundlePadding =
+      computeBundlePadding(Assembler, EF, EF->Offset, FSize);
+  if (RequiredBundlePadding > UINT8_MAX)
+    report_fatal_error("Padding cannot exceed 255 bytes");
+  EF->setBundlePadding(static_cast<uint8_t>(RequiredBundlePadding));
+  EF->Offset += RequiredBundlePadding;
+}
+
+uint64_t MCAsmLayout::getFragmentOffset(const MCFragment *F) const {
+  ensureValid(F);
+  return F->Offset;
+}
+
+void MCAsmLayout::ensureValid(const MCFragment *Frag) const {
+  MCSection &Sec = *Frag->getParent();
+  if (Sec.hasLayout())
+    return;
+  Sec.setHasLayout(true);
+  uint64_t Offset = 0;
+  for (MCFragment &F : Sec) {
+    F.Offset = Offset;
+    if (Assembler.isBundlingEnabled() && F.hasInstructions()) {
+      const_cast<MCAsmLayout *>(this)->layoutBundle(&F);
+      Offset = F.Offset;
+    }
+    Offset += getAssembler().computeFragmentSize(*this, F);
   }
 }
 
@@ -848,7 +844,7 @@ void MCAssembler::layout(MCAsmLayout &Layout) {
     // another. If any fragment has changed size, we have to re-layout (and
     // as a result possibly further relax) all.
     for (MCSection &Sec : *this)
-      Layout.invalidateFragmentsFrom(&*Sec.begin());
+      Sec.setHasLayout(false);
   }
 
   DEBUG_WITH_TYPE("mc-dump", {
@@ -1109,7 +1105,6 @@ bool MCAssembler::relaxBoundaryAlign(MCAsmLayout &Layout,
   if (NewSize == BF.getSize())
     return false;
   BF.setSize(NewSize);
-  Layout.invalidateFragmentsFrom(&BF);
   return true;
 }
 
@@ -1219,47 +1214,19 @@ bool MCAssembler::relaxFragment(MCAsmLayout &Layout, MCFragment &F) {
   }
 }
 
-bool MCAssembler::layoutSectionOnce(MCAsmLayout &Layout, MCSection &Sec) {
-  // Holds the first fragment which needed relaxing during this layout. It will
-  // remain NULL if none were relaxed.
-  // When a fragment is relaxed, all the fragments following it should get
-  // invalidated because their offset is going to change.
-  MCFragment *FirstRelaxedFragment = nullptr;
-
-  // Attempt to relax all the fragments in the section.
-  for (MCFragment &Frag : Sec) {
-    // Check if this is a fragment that needs relaxation.
-    bool RelaxedFrag = relaxFragment(Layout, Frag);
-    if (RelaxedFrag && !FirstRelaxedFragment)
-      FirstRelaxedFragment = &Frag;
-  }
-  if (FirstRelaxedFragment) {
-    Layout.invalidateFragmentsFrom(FirstRelaxedFragment);
-    return true;
-  }
-  return false;
-}
-
 bool MCAssembler::layoutOnce(MCAsmLayout &Layout) {
   ++stats::RelaxationSteps;
 
-  bool WasRelaxed = false;
-  for (MCSection &Sec : *this) {
-    while (layoutSectionOnce(Layout, Sec))
-      WasRelaxed = true;
-  }
-
-  return WasRelaxed;
+  bool Changed = false;
+  for (MCSection &Sec : *this)
+    for (MCFragment &Frag : Sec)
+      if (relaxFragment(Layout, Frag))
+        Changed = true;
+  return Changed;
 }
 
 void MCAssembler::finishLayout(MCAsmLayout &Layout) {
   assert(getBackendPtr() && "Expected assembler backend");
-  // The layout is done. Mark every fragment as valid.
-  for (unsigned int i = 0, n = Layout.getSectionOrder().size(); i != n; ++i) {
-    MCSection &Section = *Layout.getSectionOrder()[i];
-    Layout.getFragmentOffset(&*Section.getFragmentList().rbegin());
-    computeFragmentSize(Layout, *Section.getFragmentList().rbegin());
-  }
   getBackend().finishLayout(*this, Layout);
 }
 
diff --git a/llvm/lib/MC/MCExpr.cpp b/llvm/lib/MC/MCExpr.cpp
@@ -645,10 +645,6 @@ static void AttemptToFoldSymbolOffsetDifference(
       Addend += SA.getOffset() - SB.getOffset();
       return FinalizeFolding();
     }
-    // One of the symbol involved is part of a fragment being laid out. Quit now
-    // to avoid a self loop.
-    if (!Layout->canGetFragmentOffset(FA) || !Layout->canGetFragmentOffset(FB))
-      return;
 
     // Eagerly evaluate when layout is finalized.
     Addend += Layout->getSymbolOffset(A->getSymbol()) -
diff --git a/llvm/lib/MC/MCFragment.cpp b/llvm/lib/MC/MCFragment.cpp
@@ -39,64 +39,8 @@ MCAsmLayout::MCAsmLayout(MCAssembler &Asm) : Assembler(Asm) {
       SectionOrder.push_back(&Sec);
 }
 
-bool MCAsmLayout::isFragmentValid(const MCFragment *F) const {
-  const MCSection *Sec = F->getParent();
-  const MCFragment *LastValid = LastValidFragment.lookup(Sec);
-  if (!LastValid)
-    return false;
-  assert(LastValid->getParent() == Sec);
-  return F->getLayoutOrder() <= LastValid->getLayoutOrder();
-}
-
-bool MCAsmLayout::canGetFragmentOffset(const MCFragment *F) const {
-  MCSection *Sec = F->getParent();
-  MCSection::iterator I;
-  if (MCFragment *LastValid = LastValidFragment[Sec]) {
-    // Fragment already valid, offset is available.
-    if (F->getLayoutOrder() <= LastValid->getLayoutOrder())
-      return true;
-    I = ++MCSection::iterator(LastValid);
-  } else
-    I = Sec->begin();
-
-  // A fragment ordered before F is currently being laid out.
-  const MCFragment *FirstInvalidFragment = &*I;
-  if (FirstInvalidFragment->IsBeingLaidOut)
-    return false;
-
-  return true;
-}
-
 void MCAsmLayout::invalidateFragmentsFrom(MCFragment *F) {
-  // If this fragment wasn't already valid, we don't need to do anything.
-  if (!isFragmentValid(F))
-    return;
-
-  // Otherwise, reset the last valid fragment to the previous fragment
-  // (if this is the first fragment, it will be NULL).
-  LastValidFragment[F->getParent()] = F->getPrevNode();
-}
-
-void MCAsmLayout::ensureValid(const MCFragment *F) const {
-  MCSection *Sec = F->getParent();
-  MCSection::iterator I;
-  if (MCFragment *Cur = LastValidFragment[Sec])
-    I = ++MCSection::iterator(Cur);
-  else
-    I = Sec->begin();
-
-  // Advance the layout position until the fragment is valid.
-  while (!isFragmentValid(F)) {
-    assert(I != Sec->end() && "Layout bookkeeping error");
-    const_cast<MCAsmLayout *>(this)->layoutFragment(&*I);
-    ++I;
-  }
-}
-
-uint64_t MCAsmLayout::getFragmentOffset(const MCFragment *F) const {
-  ensureValid(F);
-  assert(F->Offset != ~UINT64_C(0) && "Address not set!");
-  return F->Offset;
+  F->getParent()->setHasLayout(false);
 }
 
 // Simple getSymbolOffset helper for the non-variable case.
@@ -258,7 +202,7 @@ void ilist_alloc_traits<MCFragment>::deleteNode(MCFragment *V) { V->destroy(); }
 MCFragment::MCFragment(FragmentType Kind, bool HasInstructions,
                        MCSection *Parent)
     : Parent(Parent), Atom(nullptr), Offset(~UINT64_C(0)), LayoutOrder(0),
-      Kind(Kind), IsBeingLaidOut(false), HasInstructions(HasInstructions) {
+      Kind(Kind), HasInstructions(HasInstructions) {
   if (Parent && !isa<MCDummyFragment>(*this))
     Parent->addFragment(*this);
 }
diff --git a/llvm/lib/MC/MCSection.cpp b/llvm/lib/MC/MCSection.cpp
@@ -23,8 +23,8 @@ using namespace llvm;
 MCSection::MCSection(SectionVariant V, StringRef Name, SectionKind K,
                      MCSymbol *Begin)
     : Begin(Begin), BundleGroupBeforeFirstInst(false), HasInstructions(false),
-      IsRegistered(false), DummyFragment(this), Name(Name), Variant(V),
-      Kind(K) {}
+      HasLayout(false), IsRegistered(false), DummyFragment(this), Name(Name),
+      Variant(V), Kind(K) {}
 
 MCSymbol *MCSection::getEndSymbol(MCContext &Ctx) {
   if (!End)
diff --git a/llvm/test/MC/ELF/layout-interdependency.s b/llvm/test/MC/ELF/layout-interdependency.s
@@ -1,12 +1,10 @@
-# RUN: not llvm-mc --filetype=obj %s -o /dev/null 2>&1 | FileCheck %s
-# REQUIRES: object-emission
-# UNSUPPORTED: target={{.*}}-zos{{.*}}
+# RUN: not llvm-mc -filetype=obj -triple=x86_64 %s -o /dev/null 2>&1 | FileCheck %s
 
 fct_end:
 
-# CHECK: layout-interdependency.s:[[#@LINE+1]]:7: error: expected assembly-time absolute expression
+# CHECK: layout-interdependency.s:[[#@LINE+1]]:7: error: invalid number of bytes
 .fill (data_start - fct_end), 1, 42
-# CHECK: layout-interdependency.s:[[#@LINE+1]]:7: error: expected assembly-time absolute expression
+# CHECK: layout-interdependency.s:[[#@LINE+1]]:7: error: invalid number of bytes
 .fill (fct_end - data_start), 1, 42
 
 data_start:
diff --git a/llvm/test/MC/ELF/relax-recompute-align.s b/llvm/test/MC/ELF/relax-recompute-align.s
@@ -1,14 +1,14 @@
 // RUN: llvm-mc -filetype=obj -triple i386 %s -o - | llvm-objdump -d --no-show-raw-insn - | FileCheck %s
 
-// This is a case where llvm-mc computes a better layout than Darwin 'as'. This
+/// This is a case where the computed layout is not optimal. The
 // issue is that after the first jmp slides, the .align size must be
 // recomputed -- otherwise the second jump will appear to be out-of-range for a
 // 1-byte jump.
 
 // CHECK:            int3
-// CHECK-NEXT:  ce:  int3
-// CHECK:       d0:  pushal
-// CHECK:      130:  jl      0xd0
+// CHECK-NEXT:  d2:  int3
+// CHECK:       e0:  pushal
+// CHECK:      140:  jl      0xe0
 
 L0:
         .space 0x8a, 0x90
diff --git a/llvm/test/MC/ELF/subsection-if.s b/llvm/test/MC/ELF/subsection-if.s
@@ -1,8 +1,10 @@
 # RUN: llvm-mc -filetype=obj -triple=x86_64 %s -o %t
 # RUN: llvm-readelf -x .text %t | FileCheck %s
-# RUN: not llvm-mc -filetype=obj -triple=x86_64 --defsym ERR=1 %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR
+# RUN: llvm-mc -filetype=obj -triple=x86_64 --defsym ERR=1 %s -o %t1
+# RUN: llvm-readelf -x .text %t1 | FileCheck %s --check-prefix=CHECK1
 
 # CHECK: 0x00000000 9090
+# CHECK1: 0x00000000 90909090 90
 
 .subsection 1
 661:
