[llvm] [MC][ELF] Emit instructions directly into fragment (PR #94950)

[Alexis Engelke via llvm-commits]

https://github.com/aengelke updated https://github.com/llvm/llvm-project/pull/94950

>From 65b3b4687b4f5da25b36d14efb9d222339cbda83 Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Sat, 8 Jun 2024 13:56:31 +0200
Subject: [PATCH] [MC][ELF] Emit instructions directly into fragment

Avoid needless copying of instructions and fixups.
---
 llvm/lib/MC/MCELFStreamer.cpp                 | 48 ++++++++++++++-----
 llvm/lib/MC/MCObjectStreamer.cpp              |  6 +--
 .../MCTargetDesc/SystemZMCCodeEmitter.cpp     |  1 -
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/MC/MCELFStreamer.cpp b/llvm/lib/MC/MCELFStreamer.cpp
index e6e6b7d19dff4..4741b0c3a9c4e 100644
--- a/llvm/lib/MC/MCELFStreamer.cpp
+++ b/llvm/lib/MC/MCELFStreamer.cpp
@@ -511,12 +511,6 @@ static void CheckBundleSubtargets(const MCSubtargetInfo *OldSTI,
 void MCELFStreamer::emitInstToData(const MCInst &Inst,
                                    const MCSubtargetInfo &STI) {
   MCAssembler &Assembler = getAssembler();
-  SmallVector<MCFixup, 4> Fixups;
-  SmallString<256> Code;
-  Assembler.getEmitter().encodeInstruction(Inst, Code, Fixups, STI);
-
-  for (auto &Fixup : Fixups)
-    fixSymbolsInTLSFixups(Fixup.getValue());
 
   // There are several possibilities here:
   //
@@ -535,7 +529,16 @@ void MCELFStreamer::emitInstToData(const MCInst &Inst,
   //   the group, though.
   MCDataFragment *DF;
 
+  // When bundling is enabled, we can't just append to the data fragment, as it
+  // might need to be a MCCompactEncodedInstFragment for zero fixups.
   if (Assembler.isBundlingEnabled()) {
+    SmallVector<MCFixup, 4> Fixups;
+    SmallString<256> Code;
+    Assembler.getEmitter().encodeInstruction(Inst, Code, Fixups, STI);
+
+    for (auto &Fixup : Fixups)
+      fixSymbolsInTLSFixups(Fixup.getValue());
+
     MCSection &Sec = *getCurrentSectionOnly();
     if (isBundleLocked() && !Sec.isBundleGroupBeforeFirstInst()) {
       // If we are bundle-locked, we re-use the current fragment.
@@ -546,6 +549,9 @@ 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.
+      //
+      // Apparently, this is not just a performance optimization? Using an
+      // MCDataFragment at this point causes test failures.
       MCCompactEncodedInstFragment *CEIF =
           getContext().allocFragment<MCCompactEncodedInstFragment>();
       insert(CEIF);
@@ -567,21 +573,39 @@ void MCELFStreamer::emitInstToData(const MCInst &Inst,
     // We're now emitting an instruction in a bundle group, so this flag has
     // to be turned off.
     Sec.setBundleGroupBeforeFirstInst(false);
-  } else {
-    DF = getOrCreateDataFragment(&STI);
+
+    for (auto &Fixup : Fixups) {
+      Fixup.setOffset(Fixup.getOffset() + DF->getContents().size());
+      DF->getFixups().push_back(Fixup);
+    }
+
+    DF->setHasInstructions(STI);
+    if (!Fixups.empty() && Fixups.back().getTargetKind() ==
+                               getAssembler().getBackend().RelaxFixupKind)
+      DF->setLinkerRelaxable();
+
+    DF->getContents().append(Code.begin(), Code.end());
+    return;
   }
 
-  // Add the fixups and data.
+  DF = getOrCreateDataFragment(&STI);
+
+  // Emit instruction directly into data fragment.
+  size_t FixupStartIndex = DF->getFixups().size();
+  size_t CodeOffset = DF->getContents().size();
+  Assembler.getEmitter().encodeInstruction(Inst, DF->getContents(),
+                                           DF->getFixups(), STI);
+
+  auto Fixups = MutableArrayRef(DF->getFixups()).slice(FixupStartIndex);
   for (auto &Fixup : Fixups) {
-    Fixup.setOffset(Fixup.getOffset() + DF->getContents().size());
-    DF->getFixups().push_back(Fixup);
+    Fixup.setOffset(Fixup.getOffset() + CodeOffset);
+    fixSymbolsInTLSFixups(Fixup.getValue());
   }
 
   DF->setHasInstructions(STI);
   if (!Fixups.empty() && Fixups.back().getTargetKind() ==
                              getAssembler().getBackend().RelaxFixupKind)
     DF->setLinkerRelaxable();
-  DF->getContents().append(Code.begin(), Code.end());
 }
 
 void MCELFStreamer::emitBundleAlignMode(Align Alignment) {
diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index afe5da6e5fbb9..a72e34fe6fd33 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -393,10 +393,8 @@ void MCObjectStreamer::emitInstToFragment(const MCInst &Inst,
       getContext().allocFragment<MCRelaxableFragment>(Inst, STI);
   insert(IF);
 
-  SmallString<128> Code;
-  getAssembler().getEmitter().encodeInstruction(Inst, Code, IF->getFixups(),
-                                                STI);
-  IF->getContents().append(Code.begin(), Code.end());
+  getAssembler().getEmitter().encodeInstruction(Inst, IF->getContents(),
+                                                IF->getFixups(), STI);
 }
 
 #ifndef NDEBUG
diff --git a/llvm/lib/Target/SystemZ/MCTargetDesc/SystemZMCCodeEmitter.cpp b/llvm/lib/Target/SystemZ/MCTargetDesc/SystemZMCCodeEmitter.cpp
index a6285a2ccf9d1..b161eed95d6e2 100644
--- a/llvm/lib/Target/SystemZ/MCTargetDesc/SystemZMCCodeEmitter.cpp
+++ b/llvm/lib/Target/SystemZ/MCTargetDesc/SystemZMCCodeEmitter.cpp
@@ -172,7 +172,6 @@ uint64_t SystemZMCCodeEmitter::getImmOpValue(const MCInst &MI, unsigned OpNum,
     uint32_t BitOffset = MIBitSize - RawBitOffset - OpBitSize;
     Fixups.push_back(MCFixup::create(BitOffset >> 3, MO.getExpr(),
                                      (MCFixupKind)Kind, MI.getLoc()));
-    assert(Fixups.size() <= 2 && "More than two memory operands in MI?");
     return 0;
   }
   llvm_unreachable("Unexpected operand type!");