Implement BOLT optimization in the `opt-dist` tool

6 files changed, 164 insertions, 24 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 45959c039e7..53e1e4d7567 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2495,6 +2495,7 @@ dependencies = [
  "serde_json",
  "sysinfo",
  "tar",
+ "tempfile",
  "xz",
  "zip",
 ]
diff --git a/src/tools/opt-dist/Cargo.toml b/src/tools/opt-dist/Cargo.toml
index 5a1794d3336..3f7dba81c3a 100644
--- a/src/tools/opt-dist/Cargo.toml
+++ b/src/tools/opt-dist/Cargo.toml
@@ -20,3 +20,4 @@ xz = "0.1"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 glob = "0.3"
+tempfile = "3.5"
diff --git a/src/tools/opt-dist/src/bolt.rs b/src/tools/opt-dist/src/bolt.rs
new file mode 100644
index 00000000000..161ba4f3291
--- /dev/null
+++ b/src/tools/opt-dist/src/bolt.rs
@@ -0,0 +1,98 @@
+use anyhow::Context;
+
+use crate::exec::cmd;
+use crate::training::LlvmBoltProfile;
+use camino::{Utf8Path, Utf8PathBuf};
+
+use crate::utils::io::copy_file;
+
+/// Instruments an artifact at the given `path` (in-place) with BOLT and then calls `func`.
+/// After this function finishes, the original file will be restored.
+pub fn with_bolt_instrumented<F: FnOnce() -> anyhow::Result<R>, R>(
+    path: &Utf8Path,
+    func: F,
+) -> anyhow::Result<R> {
+    // Back up the original file.
+    // It will be restored to its original state when this function exits.
+    // By copying it, we break any existing hard links, so that they are not affected by the
+    // instrumentation.
+    let _backup_file = BackedUpFile::new(path)?;
+
+    let instrumented_path = tempfile::NamedTempFile::new()?.into_temp_path();
+
+    // Instrument the original file with BOLT, saving the result into `instrumented_path`
+    cmd(&["llvm-bolt"])
+        .arg("-instrument")
+        .arg(path)
+        // Make sure that each process will write its profiles into a separate file
+        .arg("--instrumentation-file-append-pid")
+        .arg("-o")
+        .arg(instrumented_path.display())
+        .run()
+        .with_context(|| anyhow::anyhow!("Could not instrument {path} using BOLT"))?;
+
+    // Copy the instrumented artifact over the original one
+    copy_file(&instrumented_path, path)?;
+
+    // Run the function that will make use of the instrumented artifact.
+    // The original file will be restored when `_backup_file` is dropped.
+    func()
+}
+
+/// Optimizes the file at `path` with BOLT in-place using the given `profile`.
+pub fn bolt_optimize(path: &Utf8Path, profile: LlvmBoltProfile) -> anyhow::Result<()> {
+    // Copy the artifact to a new location, so that we do not use the same input and output file.
+    // BOLT cannot handle optimizing when the input and output is the same file, because it performs
+    // in-place patching.
+    let temp_path = tempfile::NamedTempFile::new()?.into_temp_path();
+    copy_file(path, &temp_path)?;
+
+    cmd(&["llvm-bolt"])
+        .arg(temp_path.display())
+        .arg("-data")
+        .arg(&profile.0)
+        .arg("-o")
+        .arg(path)
+        // Reorder basic blocks within functions
+        .arg("-reorder-blocks=ext-tsp")
+        // Reorder functions within the binary
+        .arg("-reorder-functions=hfsort+")
+        // Split function code into hot and code regions
+        .arg("-split-functions")
+        // Split as many basic blocks as possible
+        .arg("-split-all-cold")
+        // Move jump tables to a separate section
+        .arg("-jump-tables=move")
+        // Fold functions with identical code
+        .arg("-icf=1")
+        // Try to reuse old text segments to reduce binary size
+        .arg("--use-old-text")
+        // Update DWARF debug info in the final binary
+        .arg("-update-debug-sections")
+        // Print optimization statistics
+        .arg("-dyno-stats")
+        .run()
+        .with_context(|| anyhow::anyhow!("Could not optimize {path} with BOLT"))?;
+
+    Ok(())
+}
+
+/// Copies a file to a temporary location and restores it (copies it back) when it is dropped.
+pub struct BackedUpFile {
+    original: Utf8PathBuf,
+    backup: tempfile::TempPath,
+}
+
+impl BackedUpFile {
+    pub fn new(file: &Utf8Path) -> anyhow::Result<Self> {
+        let temp_path = tempfile::NamedTempFile::new()?.into_temp_path();
+        copy_file(file, &temp_path)?;
+        Ok(Self { backup: temp_path, original: file.to_path_buf() })
+    }
+}
+
+impl Drop for BackedUpFile {
+    fn drop(&mut self) {
+        copy_file(&self.backup, &self.original).expect("Cannot restore backed up file");
+    }
+}
diff --git a/src/tools/opt-dist/src/exec.rs b/src/tools/opt-dist/src/exec.rs
index 3777c7c9718..100b3ada7f2 100644
--- a/src/tools/opt-dist/src/exec.rs
+++ b/src/tools/opt-dist/src/exec.rs
@@ -1,7 +1,7 @@
 use crate::environment::Environment;
 use crate::metrics::{load_metrics, record_metrics};
 use crate::timer::TimerSection;
-use crate::training::{LlvmBoltProfile, LlvmPGOProfile, RustcPGOProfile};
+use crate::training::{LlvmPGOProfile, RustcPGOProfile};
 use camino::{Utf8Path, Utf8PathBuf};
 use std::collections::BTreeMap;
 use std::fs::File;
@@ -16,7 +16,7 @@ pub struct CmdBuilder {
 }
 
 impl CmdBuilder {
-    pub fn arg(mut self, arg: &str) -> Self {
+    pub fn arg<S: ToString>(mut self, arg: S) -> Self {
         self.args.push(arg.to_string());
         self
     }
@@ -154,13 +154,8 @@ impl Bootstrap {
         self
     }
 
-    pub fn llvm_bolt_instrument(mut self) -> Self {
-        self.cmd = self.cmd.arg("--llvm-bolt-profile-generate");
-        self
-    }
-
-    pub fn llvm_bolt_optimize(mut self, profile: &LlvmBoltProfile) -> Self {
-        self.cmd = self.cmd.arg("--llvm-bolt-profile-use").arg(profile.0.as_str());
+    pub fn with_llvm_bolt_ldflags(mut self) -> Self {
+        self.cmd = self.cmd.arg("--set").arg("llvm.ldflags=-Wl,-q");
         self
     }
 
diff --git a/src/tools/opt-dist/src/main.rs b/src/tools/opt-dist/src/main.rs
index 08f5d61000d..256317a3bcc 100644
--- a/src/tools/opt-dist/src/main.rs
+++ b/src/tools/opt-dist/src/main.rs
@@ -1,5 +1,7 @@
+use crate::bolt::{bolt_optimize, with_bolt_instrumented};
 use anyhow::Context;
 use log::LevelFilter;
+use utils::io;
 
 use crate::environment::{create_environment, Environment};
 use crate::exec::Bootstrap;
@@ -12,6 +14,7 @@ use crate::utils::{
     with_log_group,
 };
 
+mod bolt;
 mod environment;
 mod exec;
 mod metrics;
@@ -92,41 +95,51 @@ fn execute_pipeline(
         Ok(profile)
     })?;
 
-    let llvm_bolt_profile = if env.supports_bolt() {
+    if env.supports_bolt() {
         // Stage 3: Build BOLT instrumented LLVM
         // We build a PGO optimized LLVM in this step, then instrument it with BOLT and gather BOLT profiles.
         // Note that we don't remove LLVM artifacts after this step, so that they are reused in the final dist build.
         // BOLT instrumentation is performed "on-the-fly" when the LLVM library is copied to the sysroot of rustc,
         // therefore the LLVM artifacts on disk are not "tainted" with BOLT instrumentation and they can be reused.
         timer.section("Stage 3 (LLVM BOLT)", |stage| {
-            stage.section("Build BOLT instrumented LLVM", |stage| {
+            stage.section("Build PGO optimized LLVM", |stage| {
                 Bootstrap::build(env)
-                    .llvm_bolt_instrument()
+                    .with_llvm_bolt_ldflags()
                     .llvm_pgo_optimize(&llvm_pgo_profile)
                     .avoid_rustc_rebuild()
                     .run(stage)
             })?;
 
-            let profile = stage.section("Gather profiles", |_| gather_llvm_bolt_profiles(env))?;
+            // Find the path to the `libLLVM.so` file
+            let llvm_lib = io::find_file_in_dir(
+                &env.build_artifacts().join("stage2").join("lib"),
+                "libLLVM",
+                ".so",
+            )?;
+
+            // Instrument it and gather profiles
+            let profile = with_bolt_instrumented(&llvm_lib, || {
+                stage.section("Gather profiles", |_| gather_llvm_bolt_profiles(env))
+            })?;
             print_free_disk_space()?;
 
-            // LLVM is not being cleared here, we want to reuse the previous PGO-optimized build
+            // Now optimize the library with BOLT. The `libLLVM-XXX.so` library is actually hard-linked
+            // from several places, and this specific path (`llvm_lib`) will *not* be packaged into
+            // the final dist build. However, when BOLT optimizes an artifact, it does so *in-place*,
+            // therefore it will actually optimize all the hard links, which means that the final
+            // packaged `libLLVM.so` file *will* be BOLT optimized.
+            bolt_optimize(&llvm_lib, profile).context("Could not optimize LLVM with BOLT")?;
 
-            Ok(Some(profile))
-        })?
-    } else {
-        None
-    };
+            // LLVM is not being cleared here, we want to use the BOLT-optimized LLVM
+            Ok(())
+        })?;
+    }
 
-    let mut dist = Bootstrap::dist(env, &dist_args)
+    let dist = Bootstrap::dist(env, &dist_args)
         .llvm_pgo_optimize(&llvm_pgo_profile)
         .rustc_pgo_optimize(&rustc_pgo_profile)
         .avoid_rustc_rebuild();
 
-    if let Some(llvm_bolt_profile) = llvm_bolt_profile {
-        dist = dist.llvm_bolt_optimize(&llvm_bolt_profile);
-    }
-
     // Final stage: Assemble the dist artifacts
     // The previous PGO optimized rustc build and PGO optimized LLVM builds should be reused.
     timer.section("Stage 4 (final build)", |stage| dist.run(stage))?;
diff --git a/src/tools/opt-dist/src/utils/io.rs b/src/tools/opt-dist/src/utils/io.rs
index aab078067af..f333ed481b0 100644
--- a/src/tools/opt-dist/src/utils/io.rs
+++ b/src/tools/opt-dist/src/utils/io.rs
@@ -2,6 +2,7 @@ use anyhow::Context;
 use camino::{Utf8Path, Utf8PathBuf};
 use fs_extra::dir::CopyOptions;
 use std::fs::File;
+use std::path::Path;
 
 /// Delete and re-create the directory.
 pub fn reset_directory(path: &Utf8Path) -> anyhow::Result<()> {
@@ -17,6 +18,12 @@ pub fn copy_directory(src: &Utf8Path, dst: &Utf8Path) -> anyhow::Result<()> {
     Ok(())
 }
 
+pub fn copy_file<S: AsRef<Path>, D: AsRef<Path>>(src: S, dst: D) -> anyhow::Result<()> {
+    log::info!("Copying file {} to {}", src.as_ref().display(), dst.as_ref().display());
+    std::fs::copy(src.as_ref(), dst.as_ref())?;
+    Ok(())
+}
+
 #[allow(unused)]
 pub fn move_directory(src: &Utf8Path, dst: &Utf8Path) -> anyhow::Result<()> {
     log::info!("Moving directory {src} to {dst}");
@@ -60,3 +67,28 @@ pub fn get_files_from_dir(
         .map(|p| p.map(|p| Utf8PathBuf::from_path_buf(p).unwrap()))
         .collect::<Result<Vec<_>, _>>()?)
 }
+
+/// Finds a single file in the specified `directory` with the given `prefix` and `suffix`.
+pub fn find_file_in_dir(
+    directory: &Utf8Path,
+    prefix: &str,
+    suffix: &str,
+) -> anyhow::Result<Utf8PathBuf> {
+    let mut files = glob::glob(&format!("{directory}/{prefix}*{suffix}"))?
+        .into_iter()
+        .collect::<Result<Vec<_>, _>>()?;
+    match files.pop() {
+        Some(file) => {
+            if !files.is_empty() {
+                files.push(file);
+                Err(anyhow::anyhow!(
+                    "More than one file with prefix {prefix} found in {directory}: {:?}",
+                    files
+                ))
+            } else {
+                Ok(Utf8PathBuf::from_path_buf(file).unwrap())
+            }
+        }
+        None => Err(anyhow::anyhow!("No file with prefix {prefix} found in {directory}")),
+    }
+}