Linux kernel marker to update special sections

Summary: This diff adds SDT marker like LK marker to update special lk sections

(cherry picked from FBD22932157)

bolt/src/BinaryContext.h

@@ -579,6 +579,10 @@ public:
   /// Map SDT locations to SDT markers info
   std::unordered_map<uint64_t, SDTMarkerInfo> SDTMarkers;
 
+  /// Map linux kernel program locations/instructions to their pointers in
+  /// special linux kernel sections
+  std::unordered_map<uint64_t, std::vector<LKInstructionMarkerInfo>> LKMarkers;
+
   BinaryContext(std::unique_ptr<MCContext> Ctx,
                 std::unique_ptr<DWARFContext> DwCtx,
                 std::unique_ptr<Triple> TheTriple,

bolt/src/BinaryFunction.cpp

@@ -1979,9 +1979,10 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
     const auto InstrInputAddr = I->first + Address;
     bool IsSDTMarker =
         MIB->isNoop(Instr) && BC.SDTMarkers.count(InstrInputAddr);
-    if (IsSDTMarker) {
+    bool IsLKMarker = BC.LKMarkers.count(InstrInputAddr);
+    if (IsSDTMarker || IsLKMarker) {
       HasSDTMarker = true;
-      DEBUG(dbgs() << "SDTMarker detected in the input at : "
+      DEBUG(dbgs() << "SDTMarker or LKMarker detected in the input at : "
                    << utohexstr(InstrInputAddr) << "\n");
       if (!MIB->hasAnnotation(Instr, "Offset")) {
         MIB->addAnnotation(Instr, "Offset", static_cast<uint32_t>(Offset),
@@ -1992,7 +1993,7 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
     // 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
     // aligned, but it's just part of heuristic for block alignment.
-    if (MIB->isNoop(Instr) && !PreserveNops && !IsSDTMarker) {
+    if (MIB->isNoop(Instr) && !PreserveNops && !IsSDTMarker && !IsLKMarker) {
       IsLastInstrNop = true;
       continue;
     }

bolt/src/BinarySection.h

@@ -515,6 +515,15 @@ struct SDTMarkerInfo {
   unsigned PCOffset;
 };
 
+/// Linux Kernel special sections point to a specific instruction in many cases.
+/// Unlike SDTMarkerInfo, these markers can come from different sections.
+struct LKInstructionMarkerInfo {
+  uint64_t SectionOffset;
+  int32_t PCRelativeOffset;
+  bool IsPCRelative;
+  StringRef SectionName;
+};
+
 } // namespace bolt
 } // namespace llvm
 

bolt/src/RewriteInstance.cpp

@@ -2820,6 +2820,7 @@ void RewriteInstance::emitAndLink() {
 
 void RewriteInstance::updateMetadata() {
   updateSDTMarkers();
+  updateLKMarkers();
 
   if (opts::UpdateDebugSections) {
     NamedRegionTimer T("updateDebugInfo", "update debug info", TimerGroupName,
@@ -2850,6 +2851,74 @@ void RewriteInstance::updateSDTMarkers() {
   }
 }
 
+void RewriteInstance::updateLKMarkers() {
+  if (BC->LKMarkers.size() == 0) {
+    return;
+  }
+
+  NamedRegionTimer T("updateLKMarkers", "update LK markers", TimerGroupName,
+                     TimerGroupDesc, opts::TimeRewrite);
+
+  std::unordered_map<std::string, uint64_t> PatchCounts;
+  for (auto &LKMarkerInfoKV : BC->LKMarkers) {
+    const auto OriginalAddress = LKMarkerInfoKV.first;
+    const auto *F =
+        BC->getBinaryFunctionContainingAddress(OriginalAddress, false, true);
+    if (!F) {
+      continue;
+    }
+    uint64_t NewAddress = F->translateInputToOutputAddress(OriginalAddress);
+    if (NewAddress == 0) {
+      continue;
+    }
+    // rather than making address range of BBL 64_bit use base for LK BBLs
+    if (OriginalAddress >= 0xffffffff00000000 && NewAddress < 0xffffffff) {
+      NewAddress = NewAddress + 0xffffffff00000000;
+    }
+    if (OriginalAddress == NewAddress) {
+      continue;
+    }
+    uint64_t NumEntries = LKMarkerInfoKV.second.size();
+    if (NumEntries > 1) {
+      DEBUG(dbgs() << "Original linux kernel address 0x"
+                   << Twine::utohexstr(OriginalAddress) << " belongs to "
+                   << NumEntries << " marker entries.\n";);
+    }
+    for (auto &LKMarkerInfo : LKMarkerInfoKV.second) {
+      const auto SectionName = LKMarkerInfo.SectionName;
+      SimpleBinaryPatcher *LKPatcher;
+      if (SectionPatchers.find(SectionName) != SectionPatchers.end()) {
+        LKPatcher = static_cast<SimpleBinaryPatcher *>(
+            SectionPatchers[SectionName].get());
+        PatchCounts[SectionName]++;
+      } else {
+        DEBUG(dbgs() << "Starting the patch for section, " << SectionName
+                     << '\n');
+        PatchCounts[SectionName] = 1;
+        SectionPatchers[SectionName] = llvm::make_unique<SimpleBinaryPatcher>();
+        LKPatcher = static_cast<SimpleBinaryPatcher *>(
+            SectionPatchers[SectionName].get());
+      }
+      DEBUG(dbgs() << "LK patching from address 0x"
+                   << Twine::utohexstr(OriginalAddress) << ','
+                   << " to address 0x" << Twine::utohexstr(NewAddress) << '\n');
+      if (LKMarkerInfo.IsPCRelative) {
+        LKPatcher->addLE32Patch(LKMarkerInfo.SectionOffset,
+                                NewAddress - OriginalAddress +
+                                    LKMarkerInfo.PCRelativeOffset);
+      } else {
+        LKPatcher->addLE64Patch(LKMarkerInfo.SectionOffset, NewAddress);
+      }
+    }
+  }
+  outs() << "BOLT-INFO: patching linux kernel sections. Total patches per "
+            "section are as follows:\n";
+  for (const auto &KV : PatchCounts) {
+    outs() << "  Section: " << KV.first << ", patch-counts: " << KV.second
+           << '\n';
+  }
+}
+
 void RewriteInstance::mapFileSections(orc::VModuleKey Key) {
   mapCodeSections(Key);
   mapDataSections(Key);

bolt/src/RewriteInstance.h

@@ -153,6 +153,9 @@ private:
   /// Update SDTMarkers' locations for the output binary.
   void updateSDTMarkers();
 
+  /// Update LKMarkers' locations for the output binary.
+  void updateLKMarkers();
+
   /// Return the list of code sections in the output order.
   std::vector<BinarySection *> getCodeSections();