[BOLT][NFC] Pre-disasm metadata rewriters

[aaupov]

We need a hook for metadata rewriters after functions are identified
but before functions are disassembled. Currently we only have
section rewriters (after storage is discovered but before functions
are identified) and pre-CFG rewriters (after functions are disassembled).

The reason for that particular location is that BOLT finds and uses
jump tables during disassembly. If we pre-parse jump tables through
#132114, we need them before disassembly.

Additionally, debug info pre-parsing can be moved into that hook.

Depends on #132110.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

We need a hook for metadata rewriters after functions are identified
but before functions are disassembled. Currently we only have
section rewriters (after storage is discovered but before functions
are identified) and pre-CFG rewriters (after functions are disassembled).

The reason for that particular location is that BOLT finds and uses
jump tables during disassembly. If we pre-parse jump tables through
#132114, we need them before disassembly.

Additionally, debug info pre-parsing can be moved into that hook.

Depends on #132110.

---

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

5 Files Affected:

• (modified) bolt/include/bolt/Rewrite/MetadataManager.h (+4)
• (modified) bolt/include/bolt/Rewrite/MetadataRewriter.h (+4)
• (modified) bolt/include/bolt/Rewrite/RewriteInstance.h (+3)
• (modified) bolt/lib/Rewrite/MetadataManager.cpp (+12)
• (modified) bolt/lib/Rewrite/RewriteInstance.cpp (+9-1)

diff --git a/bolt/include/bolt/Rewrite/MetadataManager.h b/bolt/include/bolt/Rewrite/MetadataManager.h
index 6001b70f625e2..cc6e3f98f5e82 100644
--- a/bolt/include/bolt/Rewrite/MetadataManager.h
+++ b/bolt/include/bolt/Rewrite/MetadataManager.h
@@ -31,6 +31,10 @@ class MetadataManager {
   /// Run initializers after sections are discovered.
   void runSectionInitializers();
 
+  /// Execute metadata initializers when functions are discovered but not yet
+  /// disassembled.
+  void runInitializersPreDisasm();
+
   /// Execute initialization of rewriters while functions are disassembled, but
   /// CFG is not yet built.
   void runInitializersPreCFG();
diff --git a/bolt/include/bolt/Rewrite/MetadataRewriter.h b/bolt/include/bolt/Rewrite/MetadataRewriter.h
index 6ff8f0af7a8e6..d39500c83814c 100644
--- a/bolt/include/bolt/Rewrite/MetadataRewriter.h
+++ b/bolt/include/bolt/Rewrite/MetadataRewriter.h
@@ -49,6 +49,10 @@ class MetadataRewriter {
   /// but before functions are discovered.
   virtual Error sectionInitializer() { return Error::success(); }
 
+  /// Run initialization after the functions are identified but not yet
+  /// disassembled.
+  virtual Error preDisasmInitializer() { return Error::success(); }
+
   /// 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(); }
diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index fdd65bbd535f7..a8e627711db9e 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -181,6 +181,9 @@ class RewriteInstance {
   /// Process metadata in sections before functions are discovered.
   void processSectionMetadata();
 
+  /// Process metadata in special sections before functions are disassembled.
+  void processMetadataPreDisasm();
+
   /// Process metadata in special sections before CFG is built for functions.
   void processMetadataPreCFG();
 
diff --git a/bolt/lib/Rewrite/MetadataManager.cpp b/bolt/lib/Rewrite/MetadataManager.cpp
index 713d2e47b6efa..8114e156f5a96 100644
--- a/bolt/lib/Rewrite/MetadataManager.cpp
+++ b/bolt/lib/Rewrite/MetadataManager.cpp
@@ -32,6 +32,18 @@ void MetadataManager::runSectionInitializers() {
   }
 }
 
+void MetadataManager::runInitializersPreDisasm() {
+  for (auto &Rewriter : Rewriters) {
+    LLVM_DEBUG(dbgs() << "BOLT-DEBUG: invoking " << Rewriter->getName()
+                      << " after reading sections\n");
+    if (Error E = Rewriter->preDisasmInitializer()) {
+      errs() << "BOLT-ERROR: while running " << Rewriter->getName()
+             << " in pre-disasm state: " << toString(std::move(E)) << '\n';
+      exit(1);
+    }
+  }
+}
+
 void MetadataManager::runInitializersPreCFG() {
   for (auto &Rewriter : Rewriters) {
     LLVM_DEBUG(dbgs() << "BOLT-DEBUG: invoking " << Rewriter->getName()
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index a97762063eb1e..32f2cfee60053 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -694,7 +694,7 @@ Error RewriteInstance::run() {
 
   selectFunctionsToProcess();
 
-  readDebugInfo();
+  processMetadataPreDisasm();
 
   disassembleFunctions();
 
@@ -3237,6 +3237,14 @@ void RewriteInstance::processSectionMetadata() {
   MetadataManager.runSectionInitializers();
 }
 
+void RewriteInstance::processMetadataPreDisasm() {
+  NamedRegionTimer T("processmetadata-predisasm", "process metadata pre-disasm",
+                     TimerGroupName, TimerGroupDesc, opts::TimeRewrite);
+  MetadataManager.runInitializersPreDisasm();
+
+  readDebugInfo();
+}
+
 void RewriteInstance::processMetadataPreCFG() {
   NamedRegionTimer T("processmetadata-precfg", "process metadata pre-CFG",
                      TimerGroupName, TimerGroupDesc, opts::TimeRewrite);