Update SDT locations after bolt reordering

Summary: Update SDT locations in .note section to match the new location after bolt reorder the code.

(cherry picked from FBD15427779)

bolt/src/BinaryFunction.cpp

@@ -1702,11 +1702,19 @@ bool BinaryFunction::buildCFG() {
         updateOffset(LastInstrOffset);
     }
 
+    const auto InstrInputAddr = I->first + Address;
     bool IsSDTMarker =
-        MIB->isNoop(Instr) && BC.SDTMarkers.count(I->first + Address);
+        MIB->isNoop(Instr) && BC.SDTMarkers.count(InstrInputAddr);
 
-    if (IsSDTMarker)
+    if (IsSDTMarker) {
       HasSDTMarker = true;
+      DEBUG(dbgs() << "SDTMarker detected in the input at : "
+                   << utohexstr(InstrInputAddr) << "\n");
+
+      MIB->addAnnotation<uint64_t>(Instr, "SDTMarker", InstrInputAddr);
+      BC.SDTMarkers[InstrInputAddr].Label =
+          getOrCreateLocalLabel(InstrInputAddr);
+    }
 
     // Ignore nops except SDT markers. We use nops to derive alignment of the
     // next basic block. It will not always work, as some blocks are naturally
@@ -2685,6 +2693,17 @@ void BinaryFunction::emitBody(MCStreamer &Streamer, bool EmitColdPart,
         FirstInstr = false;
       }
 
+      // Emit SDT labels
+      if (BC.MIB->hasAnnotation(Instr, "SDTMarker")) {
+        auto OriginalAddress =
+            BC.MIB->tryGetAnnotationAs<uint64_t>(Instr, "SDTMarker").get();
+        auto *SDTLabel = BC.SDTMarkers[OriginalAddress].Label;
+        
+        //  A given symbol should only be emitted as a label once
+        if (SDTLabel->isUndefined())
+          Streamer.EmitLabel(SDTLabel);
+      }
+
       Streamer.EmitInstruction(Instr, *BC.STI);
       LastIsPrefix = BC.MIB->isPrefix(Instr);
     }

bolt/src/BinarySection.h

@@ -438,16 +438,22 @@ inline raw_ostream &operator<<(raw_ostream &OS, const BinarySection &Section) {
   return OS;
 }
 
-} // namespace bolt
-} // namespace llvm
-
 struct SDTMarkerInfo {
   uint64_t PC;
   uint64_t Base;
   uint64_t Semaphore;
-  llvm::StringRef Provider;
-  llvm::StringRef Name;
-  llvm::StringRef Args;
+  StringRef Provider;
+  StringRef Name;
+  StringRef Args;
+
+  /// The offset of PC within the note section
+  unsigned PCOffset;
+
+  /// A label that marks the location of the SDT nop instruction
+  MCSymbol *Label;
 };
 
+} // namespace bolt
+} // namespace llvm
+
 #endif

bolt/src/Passes/BinaryPasses.cpp

@@ -13,6 +13,7 @@
 #include "Passes/ReorderAlgorithm.h"
 #include "llvm/Support/Options.h"
 #include <numeric>
+#include <vector>
 
 #define DEBUG_TYPE "bolt-opts"
 
@@ -585,6 +586,8 @@ void FinalizeFunctions::runOnFunctions(BinaryContext &BC) {
 }
 
 void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
+  std::vector<std::pair<MCInst *, uint64_t>> SDTAnnotations;
+
   for (auto &It : BC.getBinaryFunctions()) {
     auto &BF = It.second;
     int64_t CurrentGnuArgsSize = 0;
@@ -610,13 +613,24 @@ void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
             II = std::next(InsertII);
           }
         }
+
+        if (BC.MIB->hasAnnotation(*II, "SDTMarker")) {
+          auto AnnotationValue =
+              BC.MIB->getAnnotationAs<uint64_t>(*II, "SDTMarker");
+          SDTAnnotations.push_back(std::make_pair(&(*II), AnnotationValue));
+        }
+
         BC.MIB->removeAllAnnotations(*II);
       }
     }
   }
-
+  
   // Release all memory taken by annotations.
   BC.MIB->freeAnnotations();
+
+  // Reinsert SDT annotations since we need them during code emition.
+  for (auto &Item : SDTAnnotations)
+    BC.MIB->addAnnotation<uint64_t>(*Item.first, "SDTMarker", Item.second);
 }
 
 namespace {

bolt/src/RewriteInstance.cpp

@@ -911,13 +911,13 @@ void RewriteInstance::parseSDTNotes() {
 
     // Parse description
     SDTMarkerInfo Marker;
+    Marker.PCOffset = Offset;
     Marker.PC = DE.getU64(&Offset);
     Marker.Base = DE.getU64(&Offset);
     Marker.Semaphore = DE.getU64(&Offset);
     Marker.Provider = DE.getCStr(&Offset);
     Marker.Name = DE.getCStr(&Offset);
     Marker.Args = DE.getCStr(&Offset);
-
     Offset = alignTo(Offset, 4);
     BC->SDTMarkers[Marker.PC] = Marker;
   }
@@ -2924,6 +2924,9 @@ void RewriteInstance::emitSections() {
 
   cantFail(OLT->emitAndFinalize(K));
 
+  // Release all memory taken by annotations.
+  BC->MIB->freeAnnotations();
+
   // Once the code is emitted, we can rename function sections to actual
   // output sections and de-register sections used for emission.
   if (!BC->HasRelocations) {
@@ -3279,6 +3282,9 @@ void RewriteInstance::mapDataSections(orc::VModuleKey Key) {
 }
 
 void RewriteInstance::updateOutputValues(const MCAsmLayout &Layout) {
+  SectionPatchers[".note.stapsdt"] = llvm::make_unique<SimpleBinaryPatcher>();
+  auto *SDTNotePatcher = static_cast<SimpleBinaryPatcher *>(
+      SectionPatchers[".note.stapsdt"].get());
 
   auto updateOutputValue = [&](BinaryFunction &Function) {
     if (!Function.isEmitted()) {
@@ -3326,6 +3332,33 @@ void RewriteInstance::updateOutputValues(const MCAsmLayout &Layout) {
           Layout.getSymbolOffset(*Function.getFunctionEndLabel()));
     }
 
+    // Create patches that update .note.stapsdt section to reflect the new
+    // locations of the SDT markers
+    if (Function.hasSDTMarker()) {
+      for (auto BBI = Function.layout_begin(), BBE = Function.layout_end();
+           BBI != BBE; ++BBI) {
+        auto *BB = *BBI;
+        const auto BBBaseAddress = BB->isCold() ? ColdBaseAddress : BaseAddress;
+
+        for (auto &Instr : *BB) {
+          if (BC->MIB->hasAnnotation(Instr, "SDTMarker")) {
+            const auto OriginalSDTAddress =
+                BC->MIB->getAnnotationAs<uint64_t>(Instr, "SDTMarker");
+            const auto NewSDTAddress =
+                Layout.getSymbolOffset(
+                    *BC->SDTMarkers[OriginalSDTAddress].Label) +
+                BBBaseAddress;
+            DEBUG(dbgs() << "SDTMarker at :" << utohexstr(OriginalSDTAddress)
+                         << "moved to :" << utohexstr(NewSDTAddress) << "\n");
+
+            SDTNotePatcher->addLE64Patch(
+                BC->SDTMarkers[OriginalSDTAddress].PCOffset, NewSDTAddress);
+            BC->MIB->removeAnnotation(Instr, "SDTMarker");
+          }
+        }
+      }
+    }
+
     // Update basic block output ranges only for the debug info or if we have
     // secondary entry points in the symbol table to update
     if (!opts::UpdateDebugSections && !Function.isMultiEntry())
@@ -4802,7 +4835,7 @@ bool RewriteInstance::willOverwriteSection(StringRef SectionName) {
     if (SectionName == OverwriteName)
       return true;
   }
-
+ 
   auto Section = BC->getUniqueSectionByName(SectionName);
   return Section && Section->isAllocatable() && Section->isFinalized();
 }