[BOLT] Introduce MetadataRewriter interface

[maksfb]

Summary

Introduce the MetadataRewriter interface to handle updates for various
types of auxiliary data stored in a binary file.

To implement metadata processing using this new interface, all metadata
rewriters should derive from the RewriterBase class and implement
one or more of the following methods, depending on the timing of metadata
read and write operations:

• preCFGInitializer()
• postCFGInitializer() // TBD
• preEmitFinalizer() // TBD
• postEmitFinalizer()

By adopting this approach, we aim to simplify the RewriteInstance class
and improve its scalability to accommodate new extensions of file formats,
including various metadata types of the Linux Kernel.

[Amir]

Looks good and nicely factored.

bolt/include/bolt/Rewrite/RewriteManager.h

48 ↗

(On Diff #535567)

Why a global instance and not a RewriteInstance member variable?

[jobnoorman]

bolt/include/bolt/Rewrite/RewriteManager.h

48 ↗

(On Diff #535567)

I would also recommend to not make this a singleton. It is (in theory) possible to have multiple RewriteInstance objects but having a global RewriteManager would make this (more) difficult.

bolt/include/bolt/Rewrite/RewriterBase.h

23 ↗

(On Diff #535567)

Is this only supposed to be used for rewriting metadata? If so, maybe this needs a more specific name?

For example, at some point I was considering to rewrite instructions pre-CFG to support RISC-V linker relaxation. Would this interface have been suitable for that? (Let's ignore for now whether that's actually a good idea to do :))

[Amir]

bolt/include/bolt/Rewrite/RewriteManager.h

48 ↗

(On Diff #535567)

Yes, in fact, llvm-boltdiff constructs two RI objects.

[maksfb]

Get rid of the singleton and some cleanup.

[maksfb]

bolt/include/bolt/Rewrite/RewriterBase.h

23 ↗

(On Diff #535567)

How about MetadataRewriterBase? I'm thinking the instances/implementations should be generic enough, but not modify the instructions directly. I.e. they can annotate instructions and introduce pseudo instructions to the stream, but otherwise keep the disassembly intact.

The pre-CFG modification you are talking about is better suited for a BinaryPass that modifies BinaryFunction code. We just have to teach the BinaryPassManager to operate on pre-CFG state.

[jobnoorman]

bolt/include/bolt/Rewrite/RewriterBase.h

23 ↗

(On Diff #535567)

How about MetadataRewriterBase?

I'd say that's better. It would make me think twice about using it to rewrite instructions :)

[maksfb]

Rename the new interface to MetadataRewriter.

bolt/include/bolt/Rewrite/MetadataManager.h

+//===- bolt/Rewrite/MetadataManager.h ---------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_REWRITE_METADATA_MANAGER_H
+#define BOLT_REWRITE_METADATA_MANAGER_H
+
+#include "bolt/Rewrite/MetadataRewriter.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Error.h"
+
+namespace llvm {
+namespace bolt {
+
+class BinaryContext;
+
+/// This class manages a collection of metadata handlers/rewriters.
+/// It is responsible for registering new rewriters and invoking them at
+/// certain stages of the binary processing pipeline.
+class MetadataManager {
+  using RewritersListType = SmallVector<std::unique_ptr<MetadataRewriter>, 1>;
+  RewritersListType Rewriters;
+
+public:
+  /// Register a new \p Rewriter.
+  void registerRewriter(std::unique_ptr<MetadataRewriter> Rewriter);
+
+  /// Execute initialization of rewriters while functions are disassembled, but
+  /// CFG is not yet built.
+  void runInitializersPreCFG();
+
+  /// Run finalization step of rewriters after code has been emitted.
+  void runFinalizersAfterEmit();
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif // BOLT_REWRITE_METADATA_MANAGER_H

bolt/include/bolt/Rewrite/MetadataRewriter.h

+//===- bolt/Rewrite/MetadataRewriter.h --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Interface for reading and updating metadata in a file.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_REWRITE_METADATA_REWRITER_H
+#define BOLT_REWRITE_METADATA_REWRITER_H
+
+#include "bolt/Core/BinaryContext.h"
+#include "llvm/Support/Error.h"
+
+namespace llvm {
+namespace bolt {
+
+/// Base class for handling file sections with metadata. In this context,
+/// metadata encompasses a wide range of data that references code and other
+/// data. Such metadata may or may not have an impact on program execution.
+/// Examples include: debug information, unwind information, exception handling
+/// tables, etc.
+//
+/// The metadata can occupy a section (e.g. .note.stapsdt), span a number of
+/// sections (e.g.,  DWARF debug info), or exist as subsection of another
+/// section in the binary (e.g., static-key jump tables embedded in .rodata
+/// section in the Linux Kernel).
+class MetadataRewriter {
+  /// The name of the data type handled by an instance of this class.
+  StringRef Name;
+
+protected:
+  /// Provides access to the binary context.
+  BinaryContext &BC;
+
+  MetadataRewriter(StringRef Name, BinaryContext &BC) : Name(Name), BC(BC) {}
+
+public:
+  virtual ~MetadataRewriter() = default;
+
+  /// Return name for the rewriter.
+  StringRef getName() const { return Name; }
+
+  /// Interface for modifying/annotating functions in the binary based on the
+  /// contents of the section. Functions are in pre-cfg state.
+  virtual Error preCFGInitializer() { return Error::success(); }
+
+  /// Finalize section contents based on the new context after the new code is
+  /// emitted.
+  virtual Error postEmitFinalizer() { return Error::success(); }
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif // BOLT_REWRITE_METADATA_REWRITER_H

bolt/include/bolt/Rewrite/MetadataRewriters.h

+//===- bolt/Rewrite/MetadataRewriters.h -------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_REWRITE_METADATA_REWRITERS_H
+#define BOLT_REWRITE_METADATA_REWRITERS_H
+
+#include <memory>
+
+namespace llvm {
+namespace bolt {
+
+class MetadataRewriter;
+
+/// The list of rewriter build functions (e.g. createNewMetadataRewriter()) that
+/// return std::unique_ptr<MetadataRewriter>.
+
+} // namespace bolt
+} // namespace llvm
+
+#endif // BOLT_REWRITE_METADATA_REWRITERS_H

bolt/include/bolt/Rewrite/RewriteInstance.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BOLT_REWRITE_REWRITE_INSTANCE_H
 #define BOLT_REWRITE_REWRITE_INSTANCE_H
 
 #include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/Linker.h"
+#include "bolt/Rewrite/MetadataManager.h"
 #include "bolt/Utils/NameResolver.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/MC/StringTableBuilder.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Error.h"
 #include <map>
 #include <set>
const ProfileReaderBase *getProfileReader() const {
     return ProfileReader.get();
   }
 
 private:
   /// Populate array of binary functions and other objects of interest
   /// from meta data in the file.
   void discoverFileObjects();
 
+  /// Create and initialize metadata rewriters for this instance.
+  void initializeMetadataManager();
+
   /// Process fragments, locate parent functions.
   void registerFragments();
 
   /// Read info from special sections. E.g. eh_frame and .gcc_except_table
   /// for exception and stack unwinding information.
   Error readSpecialSections();
 
   /// Adjust supplied command-line options based on input data.
private:
   /// Write code and data into an intermediary object file, map virtual to real
   /// addresses and link the object file, resolving all relocations and
   /// performing final relaxation.
   void emitAndLink();
 
   /// Link additional runtime code to support instrumentation.
   void linkRuntime();
 
+  /// Process metadata in special sections before CFG is built for functions.
+  void processMetadataPreCFG();
+
   /// Update debug and other auxiliary information in the file.
   void updateMetadata();
 
   /// Update SDTMarkers' locations for the output binary.
   void updateSDTMarkers();
 
   /// Update LKMarkers' locations for the output binary.
   void updateLKMarkers();
public:
   static bool isKSymtabSection(StringRef SectionName);
 
   /// Adds Debug section to overwrite.
   static void addToDebugSectionsToOverwrite(const char *Section) {
     DebugSectionsToOverwrite.emplace_back(Section);
   }
 
 private:
+  /// Manage a pipeline of metadata handlers.
+  MetadataManager MetadataManager;
+
   /// Get the contents of the LSDA section for this binary.
   ArrayRef<uint8_t> getLSDAData();
 
   /// Get the mapped address of the LSDA section for this binary.
   uint64_t getLSDAAddress();
 
   static const char TimerGroupName[];
 

bolt/lib/Rewrite/CMakeLists.txt

 
 add_llvm_library(LLVMBOLTRewrite
   BinaryPassManager.cpp
   BoltDiff.cpp
   DWARFRewriter.cpp
   ExecutableFileMemoryManager.cpp
   JITLinkLinker.cpp
   MachORewriteInstance.cpp
+  MetadataManager.cpp
   RewriteInstance.cpp
 
   DISABLE_LLVM_LINK_LLVM_DYLIB
 
   LINK_LIBS
   ${LLVM_PTHREAD_LIB}
   )
 

bolt/lib/Rewrite/MetadataManager.cpp

+//===- bolt/Rewrite/MetadataManager.cpp -----------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Rewrite/MetadataManager.h"
+#include "llvm/Support/Debug.h"
+
+#undef DEBUG_TYPE
+#define DEBUG_TYPE "bolt-metadata"
+
+using namespace llvm;
+using namespace bolt;
+
+void MetadataManager::registerRewriter(
+    std::unique_ptr<MetadataRewriter> Rewriter) {
+  Rewriters.emplace_back(std::move(Rewriter));
+}
+
+void MetadataManager::runInitializersPreCFG() {
+  for (auto &Rewriter : Rewriters) {
+    LLVM_DEBUG(dbgs() << "BOLT-DEBUG: invoking " << Rewriter->getName()
+                      << " before CFG construction\n");
+    if (Error E = Rewriter->preCFGInitializer()) {
+      errs() << "BOLT-ERROR: while running " << Rewriter->getName()
+             << " in pre-CFG state: " << toString(std::move(E)) << '\n';
+      exit(1);
+    }
+  }
+}
+
+void MetadataManager::runFinalizersAfterEmit() {
+  for (auto &Rewriter : Rewriters) {
+    LLVM_DEBUG(dbgs() << "BOLT-DEBUG: invoking " << Rewriter->getName()
+                      << " after emit\n");
+    if (Error E = Rewriter->postEmitFinalizer()) {
+      errs() << "BOLT-ERROR: while running " << Rewriter->getName()
+             << " after emit: " << toString(std::move(E)) << '\n';
+      exit(1);
+    }
+  }
+}

bolt/lib/Rewrite/RewriteInstance.cpp

 #include "bolt/Profile/DataAggregator.h"
 #include "bolt/Profile/DataReader.h"
 #include "bolt/Profile/YAMLProfileReader.h"
 #include "bolt/Profile/YAMLProfileWriter.h"
 #include "bolt/Rewrite/BinaryPassManager.h"
 #include "bolt/Rewrite/DWARFRewriter.h"
 #include "bolt/Rewrite/ExecutableFileMemoryManager.h"
 #include "bolt/Rewrite/JITLinkLinker.h"
+#include "bolt/Rewrite/MetadataRewriters.h"
 #include "bolt/RuntimeLibs/HugifyRuntimeLibrary.h"
 #include "bolt/RuntimeLibs/InstrumentationRuntimeLibrary.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "bolt/Utils/Utils.h"
 #include "llvm/ADT/AddressRanges.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/DebugInfo/DWARF/DWARFDebugFrame.h"
Error RewriteInstance::run() {
   }
 
   selectFunctionsToProcess();
 
   readDebugInfo();
 
   disassembleFunctions();
 
-  processProfileDataPreCFG();
+  processMetadataPreCFG();
 
   buildFunctionsCFG();
 
   processProfileData();
 
   postProcessFunctions();
 
   if (opts::DiffOnly)
if (!BC->hasSymbolsWithFileName() && ProfileReader->hasLocalsWithFileName()) {
     errs() << "BOLT-ERROR: input binary does not have local file symbols "
               "but profile data includes function names with embedded file "
               "names. It appears that the input binary was stripped while a "
               "profiled binary was not\n";
     exit(1);
   }
 }
 
+void RewriteInstance::initializeMetadataManager() {
+  // TODO
+}
+
+void RewriteInstance::processMetadataPreCFG() {
+  initializeMetadataManager();
+
+  MetadataManager.runInitializersPreCFG();
+
+  processProfileDataPreCFG();
+}
+
 void RewriteInstance::processProfileDataPreCFG() {
   if (!ProfileReader)
     return;
 
   NamedRegionTimer T("processprofile-precfg", "process profile data pre-CFG",
                      TimerGroupName, TimerGroupDesc, opts::TimeRewrite);
 
   if (Error E = ProfileReader->readProfilePreCFG(*BC.get()))
if (opts::KeepTmp) {
     TempOut->keep();
     outs() << "BOLT-INFO: intermediary output object file saved for debugging "
               "purposes: "
            << OutObjectPath << "\n";
   }
 }
 
 void RewriteInstance::updateMetadata() {
+  MetadataManager.runFinalizersAfterEmit();
+
+  // TODO: use MetadataManager for updates.
   updateSDTMarkers();
   updateLKMarkers();
   parsePseudoProbe();
   updatePseudoProbes();
 
   if (opts::UpdateDebugSections) {
     NamedRegionTimer T("updateDebugInfo", "update debug info", TimerGroupName,
                        TimerGroupDesc, opts::TimeRewrite);