[MC] Add MCFragment allocation helpers

[MaskRay]

allocFragment might be changed to a placement new when the allocation
strategy changes.

allocInitialFragment is to deduplicate the following pattern

  auto *F = new MCDataFragment();
  Result->addFragment(*F);
  F->setParent(Result);

[llvmbot]

@llvm/pr-subscribers-backend-x86

@llvm/pr-subscribers-mc

Author: Fangrui Song (MaskRay)

Changes

to make it easy to change the allocation strategy.

---

Full diff: https://github.com/llvm/llvm-project/pull/95197.diff

6 Files Affected:

• (modified) llvm/include/llvm/MC/MCContext.h (+7)
• (modified) llvm/include/llvm/MC/MCObjectStreamer.h (+6)
• (modified) llvm/lib/MC/MCContext.cpp (+14-20)
• (modified) llvm/lib/MC/MCELFStreamer.cpp (+2-4)
• (modified) llvm/lib/MC/MCMachOStreamer.cpp (+1-1)
• (modified) llvm/lib/MC/MCObjectStreamer.cpp (+12-16)

diff --git a/llvm/include/llvm/MC/MCContext.h b/llvm/include/llvm/MC/MCContext.h
index 72eae85467dc9..f5938475510bf 100644
--- a/llvm/include/llvm/MC/MCContext.h
+++ b/llvm/include/llvm/MC/MCContext.h
@@ -44,6 +44,7 @@ namespace llvm {
 
 class CodeViewContext;
 class MCAsmInfo;
+class MCDataFragment;
 class MCInst;
 class MCLabel;
 class MCObjectFileInfo;
@@ -345,6 +346,8 @@ class MCContext {
   void reportCommon(SMLoc Loc,
                     std::function<void(SMDiagnostic &, const SourceMgr *)>);
 
+  MCDataFragment *allocInitialFragment(MCSection &Sec);
+
   MCSymbol *createSymbolImpl(const StringMapEntry<bool> *Name,
                              bool IsTemporary);
   MCSymbol *createSymbol(StringRef Name, bool AlwaysAddSuffix,
@@ -438,6 +441,10 @@ class MCContext {
   /// Create and return a new MC instruction.
   MCInst *createMCInst();
 
+  template <typename F, typename... Args> F *allocFragment(Args &&...args) {
+    return new F(std::forward<Args>(args)...);
+  }
+
   /// \name Symbol Management
   /// @{
 
diff --git a/llvm/include/llvm/MC/MCObjectStreamer.h b/llvm/include/llvm/MC/MCObjectStreamer.h
index c0a337f5ea45e..555c4e1d3f1a1 100644
--- a/llvm/include/llvm/MC/MCObjectStreamer.h
+++ b/llvm/include/llvm/MC/MCObjectStreamer.h
@@ -90,6 +90,12 @@ class MCObjectStreamer : public MCStreamer {
 
   MCFragment *getCurrentFragment() const;
 
+  template <typename F, typename... Args> F *allocAndAdd(Args &&...args) {
+    F *Frag = new F(std::forward<Args>(args)...);
+    insert(Frag);
+    return Frag;
+  }
+
   void insert(MCFragment *F) {
     flushPendingLabels(F);
     MCSection *CurSection = getCurrentSectionOnly();
diff --git a/llvm/lib/MC/MCContext.cpp b/llvm/lib/MC/MCContext.cpp
index 1590054717960..82ac5f04c9076 100644
--- a/llvm/lib/MC/MCContext.cpp
+++ b/llvm/lib/MC/MCContext.cpp
@@ -195,6 +195,15 @@ MCInst *MCContext::createMCInst() {
   return new (MCInstAllocator.Allocate()) MCInst;
 }
 
+// Allocate the initial MCDataFragment for the begin symbol.
+MCDataFragment *MCContext::allocInitialFragment(MCSection &Sec) {
+  assert(!Sec.curFragList()->Head);
+  auto *F = allocFragment<MCDataFragment>();
+  F->setParent(&Sec);
+  Sec.addFragment(*F);
+  return F;
+}
+
 //===----------------------------------------------------------------------===//
 // Symbol Manipulation
 //===----------------------------------------------------------------------===//
@@ -497,11 +506,8 @@ MCSectionELF *MCContext::createELFSectionImpl(StringRef Section, unsigned Type,
       MCSectionELF(Section, Type, Flags, K, EntrySize, Group, Comdat, UniqueID,
                    R, LinkedToSym);
 
-  auto *F = new MCDataFragment();
-  Ret->addFragment(*F);
-  F->setParent(Ret);
+  auto *F = allocInitialFragment(*Ret);
   R->setFragment(F);
-
   return Ret;
 }
 
@@ -797,11 +803,8 @@ MCSectionWasm *MCContext::getWasmSection(const Twine &Section, SectionKind Kind,
       MCSectionWasm(CachedName, Kind, Flags, GroupSym, UniqueID, Begin);
   Entry.second = Result;
 
-  auto *F = new MCDataFragment();
-  Result->addFragment(*F);
-  F->setParent(Result);
+  auto *F = allocInitialFragment(*Result);
   Begin->setFragment(F);
-
   return Result;
 }
 
@@ -863,10 +866,7 @@ MCSectionXCOFF *MCContext::getXCOFFSection(
 
   Entry.second = Result;
 
-  auto *F = new MCDataFragment();
-  Result->addFragment(*F);
-  F->setParent(Result);
-
+  auto *F = allocInitialFragment(*Result);
   if (Begin)
     Begin->setFragment(F);
 
@@ -886,10 +886,7 @@ MCSectionSPIRV *MCContext::getSPIRVSection() {
   MCSectionSPIRV *Result = new (SPIRVAllocator.Allocate())
       MCSectionSPIRV(SectionKind::getText(), Begin);
 
-  auto *F = new MCDataFragment();
-  Result->addFragment(*F);
-  F->setParent(Result);
-
+  allocInitialFragment(*Result);
   return Result;
 }
 
@@ -909,10 +906,7 @@ MCSectionDXContainer *MCContext::getDXContainerSection(StringRef Section,
       new (DXCAllocator.Allocate()) MCSectionDXContainer(Name, K, nullptr);
 
   // The first fragment will store the header
-  auto *F = new MCDataFragment();
-  MapIt->second->addFragment(*F);
-  F->setParent(MapIt->second);
-
+  allocInitialFragment(*MapIt->second);
   return MapIt->second;
 }
 
diff --git a/llvm/lib/MC/MCELFStreamer.cpp b/llvm/lib/MC/MCELFStreamer.cpp
index 23e926c3a9d14..b2bef3934aa9a 100644
--- a/llvm/lib/MC/MCELFStreamer.cpp
+++ b/llvm/lib/MC/MCELFStreamer.cpp
@@ -596,14 +596,12 @@ void MCELFStreamer::emitInstToData(const MCInst &Inst,
       // Optimize memory usage by emitting the instruction to a
       // MCCompactEncodedInstFragment when not in a bundle-locked group and
       // there are no fixups registered.
-      MCCompactEncodedInstFragment *CEIF = new MCCompactEncodedInstFragment();
-      insert(CEIF);
+      auto *CEIF = allocAndAdd<MCCompactEncodedInstFragment>();
       CEIF->getContents().append(Code.begin(), Code.end());
       CEIF->setHasInstructions(STI);
       return;
     } else {
-      DF = new MCDataFragment();
-      insert(DF);
+      DF = allocAndAdd<MCDataFragment>();
     }
     if (Sec.getBundleLockState() == MCSection::BundleLockedAlignToEnd) {
       // If this fragment is for a group marked "align_to_end", set a flag
diff --git a/llvm/lib/MC/MCMachOStreamer.cpp b/llvm/lib/MC/MCMachOStreamer.cpp
index 10f9988b9d16a..7c3f3fa84bd1a 100644
--- a/llvm/lib/MC/MCMachOStreamer.cpp
+++ b/llvm/lib/MC/MCMachOStreamer.cpp
@@ -199,7 +199,7 @@ void MCMachOStreamer::emitLabel(MCSymbol *Symbol, SMLoc Loc) {
   // We have to create a new fragment if this is an atom defining symbol,
   // fragments cannot span atoms.
   if (getAssembler().isSymbolLinkerVisible(*Symbol))
-    insert(new MCDataFragment());
+    allocAndAdd<MCDataFragment>();
 
   MCObjectStreamer::emitLabel(Symbol, Loc);
 
diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index bf1ce76cdc14b..37f90e8bd10bf 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -224,10 +224,8 @@ static bool canReuseDataFragment(const MCDataFragment &F,
 MCDataFragment *
 MCObjectStreamer::getOrCreateDataFragment(const MCSubtargetInfo *STI) {
   MCDataFragment *F = dyn_cast_or_null<MCDataFragment>(getCurrentFragment());
-  if (!F || !canReuseDataFragment(*F, *Assembler, STI)) {
-    F = new MCDataFragment();
-    insert(F);
-  }
+  if (!F || !canReuseDataFragment(*F, *Assembler, STI))
+    F = allocAndAdd<MCDataFragment>();
   return F;
 }
 
@@ -343,7 +341,7 @@ void MCObjectStreamer::emitULEB128Value(const MCExpr *Value) {
     emitULEB128IntValue(IntValue);
     return;
   }
-  insert(new MCLEBFragment(*Value, false));
+  allocAndAdd<MCLEBFragment>(*Value, false);
 }
 
 void MCObjectStreamer::emitSLEB128Value(const MCExpr *Value) {
@@ -352,7 +350,7 @@ void MCObjectStreamer::emitSLEB128Value(const MCExpr *Value) {
     emitSLEB128IntValue(IntValue);
     return;
   }
-  insert(new MCLEBFragment(*Value, true));
+  allocAndAdd<MCLEBFragment>(*Value, true);
 }
 
 void MCObjectStreamer::emitWeakReference(MCSymbol *Alias,
@@ -470,9 +468,7 @@ void MCObjectStreamer::emitInstToFragment(const MCInst &Inst,
 
   // Always create a new, separate fragment here, because its size can change
   // during relaxation.
-  MCRelaxableFragment *IF = new MCRelaxableFragment(Inst, STI);
-  insert(IF);
-
+  auto *IF = allocAndAdd<MCRelaxableFragment>(Inst, STI);
   SmallString<128> Code;
   getAssembler().getEmitter().encodeInstruction(Inst, Code, IF->getFixups(),
                                                 STI);
@@ -544,7 +540,7 @@ void MCObjectStreamer::emitDwarfAdvanceLineAddr(int64_t LineDelta,
     return;
   }
   const MCExpr *AddrDelta = buildSymbolDiff(*this, Label, LastLabel, SMLoc());
-  insert(new MCDwarfLineAddrFragment(LineDelta, *AddrDelta));
+  allocAndAdd<MCDwarfLineAddrFragment>(LineDelta, *AddrDelta);
 }
 
 void MCObjectStreamer::emitDwarfLineEndEntry(MCSection *Section,
@@ -569,7 +565,7 @@ void MCObjectStreamer::emitDwarfAdvanceFrameAddr(const MCSymbol *LastLabel,
                                                  const MCSymbol *Label,
                                                  SMLoc Loc) {
   const MCExpr *AddrDelta = buildSymbolDiff(*this, Label, LastLabel, Loc);
-  insert(new MCDwarfCallFrameFragment(*AddrDelta, nullptr));
+  allocAndAdd<MCDwarfCallFrameFragment>(*AddrDelta, nullptr);
 }
 
 void MCObjectStreamer::emitCVLocDirective(unsigned FunctionId, unsigned FileNo,
@@ -640,7 +636,7 @@ void MCObjectStreamer::emitValueToAlignment(Align Alignment, int64_t Value,
                                             unsigned MaxBytesToEmit) {
   if (MaxBytesToEmit == 0)
     MaxBytesToEmit = Alignment.value();
-  insert(new MCAlignFragment(Alignment, Value, ValueSize, MaxBytesToEmit));
+  allocAndAdd<MCAlignFragment>(Alignment, Value, ValueSize, MaxBytesToEmit);
 
   // Update the maximum alignment on the current section if necessary.
   MCSection *CurSec = getCurrentSectionOnly();
@@ -657,7 +653,7 @@ void MCObjectStreamer::emitCodeAlignment(Align Alignment,
 void MCObjectStreamer::emitValueToOffset(const MCExpr *Offset,
                                          unsigned char Value,
                                          SMLoc Loc) {
-  insert(new MCOrgFragment(*Offset, Value, Loc));
+  allocAndAdd<MCOrgFragment>(*Offset, Value, Loc);
 }
 
 // Associate DTPRel32 fixup with data and resize data area
@@ -844,7 +840,7 @@ void MCObjectStreamer::emitFill(const MCExpr &NumBytes, uint64_t FillValue,
   flushPendingLabels(DF, DF->getContents().size());
 
   assert(getCurrentSectionOnly() && "need a section");
-  insert(new MCFillFragment(FillValue, 1, NumBytes, Loc));
+  allocAndAdd<MCFillFragment>(FillValue, 1, NumBytes, Loc);
 }
 
 void MCObjectStreamer::emitFill(const MCExpr &NumValues, int64_t Size,
@@ -874,7 +870,7 @@ void MCObjectStreamer::emitFill(const MCExpr &NumValues, int64_t Size,
   flushPendingLabels(DF, DF->getContents().size());
 
   assert(getCurrentSectionOnly() && "need a section");
-  insert(new MCFillFragment(Expr, Size, NumValues, Loc));
+  allocAndAdd<MCFillFragment>(Expr, Size, NumValues, Loc);
 }
 
 void MCObjectStreamer::emitNops(int64_t NumBytes, int64_t ControlledNopLength,
@@ -885,7 +881,7 @@ void MCObjectStreamer::emitNops(int64_t NumBytes, int64_t ControlledNopLength,
 
   assert(getCurrentSectionOnly() && "need a section");
 
-  insert(new MCNopsFragment(NumBytes, ControlledNopLength, Loc, STI));
+  allocAndAdd<MCNopsFragment>(NumBytes, ControlledNopLength, Loc, STI);
 }
 
 void MCObjectStreamer::emitFileDirective(StringRef Filename) {

[aengelke]

Ok, but what about the other places with new MCXXXFragment? MCPseudoProbe, MCELFStreamer and MCMachOStreamer have a few insert(new Fragment) calls, which could be trivially replaced by allocAndAdd.

[aengelke]

Is the plan to remove the Parent parameter from MCFragment? Otherwise, shouldn't allocFragment<MCDataFragment>(&Sec); work?

[MaskRay]

Yes, I plan to remove the Parent parameter.

[aengelke]

Use allocFragment from context?

[MaskRay]

MCObjectStreamer.h cannot include MCContext.h, so we cannot use allocFragment.
I believe in the end we will have different allocation functions. The primary one from MCObjectStreamer.h will be made to a bump allocator., but operator new might remain.

[aengelke]

Hmm. What was the reason to put allocation inside the streamer? My thinking was to keep allocation in MCContext (maybe without merging allocation and adding), because it already manages all other resources, and then all fragments can go to the same allocator. Mixing bump ptr allocator and operator new makes it more difficult to decide how to free a fragment.

[MaskRay]

Agree that allocAndAdd is not useful. Removed.

[aengelke]

Ok. What about new Fragment() in X86/MCTargetDesc/X86AsmBackend.cpp and MCPseudoProbe.cpp? Then all new Fragment() would be gone, I think.

[MaskRay]

Ok. What about new Fragment() in X86/MCTargetDesc/X86AsmBackend.cpp and MCPseudoProbe.cpp? Then all new Fragment() would be gone, I think.

Update X86/MCTargetDesc/X86AsmBackend.cpp and MCPseudoProbe.cpp and MCCodeView.cpp.

Perhaps MCSection::flushPendingLabels new MCDataFragment is the only remaining one.
If we can remove aligned bundling, flushPendingLabels can go away
https://discourse.llvm.org/t/mc-removing-aligned-bundling-support/79518/5

[aengelke]

LGTM. Re flushPendingLabels -- maybe for now pass a context as parameter so that the new() is gone there, too? We then could switch to MCContext's allocator.

[MaskRay]

LGTM. Re flushPendingLabels -- maybe for now pass a context as parameter so that the new() is gone there, too? We then could switch to MCContext's allocator.

Thanks. Let me change flushPendingLabels separately...

I guess a bump allocator helps glibc malloc but less for mimalloc.