[llvm] [BOLT] Add metadata pre-emit finalization interface (PR #79925)

[Maksim Panchenko via llvm-commits]

https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/79925

Some metadata needs to be updated/finalized before the binary context is emitted into the binary. Add the interface and use it for Linux ORC update invocation.

>From e9d05d686f8860f82e2e36999be7dd43e56f511b Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Tue, 21 Nov 2023 16:38:08 -0800
Subject: [PATCH] [BOLT] Add metadata pre-emit finalization interface

Some metadata needs to be updated/finalized before the binary context
is emitted into the binary. Add the interface and use it for Linux
ORC update invocation.
---
 bolt/include/bolt/Rewrite/MetadataManager.h  |  3 +++
 bolt/include/bolt/Rewrite/MetadataRewriter.h |  3 +++
 bolt/include/bolt/Rewrite/RewriteInstance.h  |  3 +++
 bolt/lib/Rewrite/LinuxKernelRewriter.cpp     | 10 +++++++---
 bolt/lib/Rewrite/MetadataManager.cpp         | 12 ++++++++++++
 bolt/lib/Rewrite/RewriteInstance.cpp         |  6 ++++++
 6 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/bolt/include/bolt/Rewrite/MetadataManager.h b/bolt/include/bolt/Rewrite/MetadataManager.h
index 1a4c394ff33bb..efbc74b4daba9 100644
--- a/bolt/include/bolt/Rewrite/MetadataManager.h
+++ b/bolt/include/bolt/Rewrite/MetadataManager.h
@@ -36,6 +36,9 @@ class MetadataManager {
   /// Execute metadata initializers after CFG was constructed for functions.
   void runInitializersPostCFG();
 
+  /// Run finalization step of rewriters before the binary is emitted.
+  void runFinalizersPreEmit();
+
   /// Run finalization step of rewriters after code has been emitted.
   void runFinalizersAfterEmit();
 };
diff --git a/bolt/include/bolt/Rewrite/MetadataRewriter.h b/bolt/include/bolt/Rewrite/MetadataRewriter.h
index 12908089db428..1e7e0381c1e98 100644
--- a/bolt/include/bolt/Rewrite/MetadataRewriter.h
+++ b/bolt/include/bolt/Rewrite/MetadataRewriter.h
@@ -52,6 +52,9 @@ class MetadataRewriter {
   /// Run the rewriter once the functions are in CFG state.
   virtual Error postCFGInitializer() { return Error::success(); }
 
+  /// Run the pass before the binary is emitted.
+  virtual Error preEmitFinalizer() { return Error::success(); }
+
   /// Finalize section contents based on the new context after the new code is
   /// emitted.
   virtual Error postEmitFinalizer() { return Error::success(); }
diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index 495a3f1d7b7b4..3074ae77c7ef1 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -183,6 +183,9 @@ class RewriteInstance {
   /// Process metadata in special sections after CFG is built for functions.
   void processMetadataPostCFG();
 
+  /// Make changes to metadata before the binary is emitted.
+  void finalizeMetadataPreEmit();
+
   /// Update debug and other auxiliary information in the file.
   void updateMetadata();
 
diff --git a/bolt/lib/Rewrite/LinuxKernelRewriter.cpp b/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
index 2f449da6206f2..0687c4836ced2 100644
--- a/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
+++ b/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
@@ -163,14 +163,18 @@ class LinuxKernelRewriter final : public MetadataRewriter {
     return Error::success();
   }
 
-  Error postEmitFinalizer() override {
-    updateLKMarkers();
-
+  Error preEmitFinalizer() override {
     if (Error E = rewriteORCTables())
       return E;
 
     return Error::success();
   }
+
+  Error postEmitFinalizer() override {
+    updateLKMarkers();
+
+    return Error::success();
+  }
 };
 
 Error LinuxKernelRewriter::markInstructions() {
diff --git a/bolt/lib/Rewrite/MetadataManager.cpp b/bolt/lib/Rewrite/MetadataManager.cpp
index 059b43f83328a..4ce44820d9eca 100644
--- a/bolt/lib/Rewrite/MetadataManager.cpp
+++ b/bolt/lib/Rewrite/MetadataManager.cpp
@@ -44,6 +44,18 @@ void MetadataManager::runInitializersPostCFG() {
   }
 }
 
+void MetadataManager::runFinalizersPreEmit() {
+  for (auto &Rewriter : Rewriters) {
+    LLVM_DEBUG(dbgs() << "BOLT-DEBUG: invoking " << Rewriter->getName()
+                      << " before emitting binary context\n");
+    if (Error E = Rewriter->preEmitFinalizer()) {
+      errs() << "BOLT-ERROR: while running " << Rewriter->getName()
+             << " before emit: " << toString(std::move(E)) << '\n';
+      exit(1);
+    }
+  }
+}
+
 void MetadataManager::runFinalizersAfterEmit() {
   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 f5a8a5b716874..d24bd18c08409 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -742,6 +742,8 @@ Error RewriteInstance::run() {
 
   runOptimizationPasses();
 
+  finalizeMetadataPreEmit();
+
   emitAndLink();
 
   updateMetadata();
@@ -3416,6 +3418,10 @@ void RewriteInstance::emitAndLink() {
   }
 }
 
+void RewriteInstance::finalizeMetadataPreEmit() {
+  MetadataManager.runFinalizersPreEmit();
+}
+
 void RewriteInstance::updateMetadata() {
   MetadataManager.runFinalizersAfterEmit();