Create BOLT build steps to avoid running BOLT multiple times on the same file

2 files changed, 127 insertions, 19 deletions

diff --git a/src/bootstrap/bolt.rs b/src/bootstrap/bolt.rs
index d424852b845..973dc4f602b 100644
--- a/src/bootstrap/bolt.rs
+++ b/src/bootstrap/bolt.rs
@@ -1,47 +1,40 @@
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::process::Command;
 
-/// Uses the `llvm-bolt` binary to instrument the binary/library at the given `path` with BOLT.
+/// Uses the `llvm-bolt` binary to instrument the artifact at the given `path` with BOLT.
 /// When the instrumented artifact is executed, it will generate BOLT profiles into
 /// `/tmp/prof.fdata.<pid>.fdata`.
-/// Returns a path to the instrumented artifact, created in a temporary directory.
-pub fn instrument_with_bolt(path: &Path) -> PathBuf {
-    let dir = std::env::temp_dir();
-    let instrumented_path = dir.join(path.file_name().unwrap());
-
+/// Creates the instrumented artifact at `output_path`.
+pub fn instrument_with_bolt(path: &Path, output_path: &Path) {
     let status = Command::new("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)
+        .arg(output_path)
         .status()
         .expect("Could not instrument artifact using BOLT");
 
     if !status.success() {
         panic!("Could not instrument {} with BOLT, exit code {:?}", path.display(), status.code());
     }
-    instrumented_path
 }
 
-/// Uses the `llvm-bolt` binary to optimize the binary/library at the given `path` with BOLT,
+/// Uses the `llvm-bolt` binary to optimize the artifact at the given `path` with BOLT,
 /// using merged profiles from `profile_path`.
 ///
 /// The recorded profiles have to be merged using the `merge-fdata` tool from LLVM and the merged
 /// profile path should be then passed to this function.
 ///
-/// Returns a path to the optimized artifact, created in a temporary directory.
-pub fn optimize_with_bolt(path: &Path, profile_path: &Path) -> PathBuf {
-    let dir = std::env::temp_dir();
-    let optimized_path = dir.join(path.file_name().unwrap());
-
+/// Creates the optimized artifact at `output_path`.
+pub fn optimize_with_bolt(path: &Path, profile_path: &Path, output_path: &Path) {
     let status = Command::new("llvm-bolt")
         .arg(&path)
         .arg("-data")
         .arg(&profile_path)
         .arg("-o")
-        .arg(&optimized_path)
+        .arg(output_path)
         // Reorder basic blocks within functions
         .arg("-reorder-blocks=ext-tsp")
         // Reorder functions within the binary
@@ -66,5 +59,4 @@ pub fn optimize_with_bolt(path: &Path, profile_path: &Path) -> PathBuf {
     if !status.success() {
         panic!("Could not optimize {} with BOLT, exit code {:?}", path.display(), status.code());
     }
-    optimized_path
 }
diff --git a/src/bootstrap/dist.rs b/src/bootstrap/dist.rs
index f6f6f4cc8ea..1c0957bc35e 100644
--- a/src/bootstrap/dist.rs
+++ b/src/bootstrap/dist.rs
@@ -18,6 +18,7 @@ use std::process::Command;
 
 use object::read::archive::ArchiveFile;
 use object::BinaryFormat;
+use sha2::Digest;
 
 use crate::bolt::{instrument_with_bolt, optimize_with_bolt};
 use crate::builder::{Builder, Kind, RunConfig, ShouldRun, Step};
@@ -1915,9 +1916,9 @@ fn install_llvm_file(builder: &Builder<'_>, source: &Path, destination: &Path) {
     // We perform the instrumentation/optimization here, on the fly, just before they are being
     // packaged into some destination directory.
     let postprocessed = if builder.config.llvm_bolt_profile_generate {
-        instrument_with_bolt(source)
+        builder.ensure(BoltInstrument::new(source.to_path_buf()))
     } else if let Some(path) = &builder.config.llvm_bolt_profile_use {
-        optimize_with_bolt(source, &Path::new(&path))
+        builder.ensure(BoltOptimize::new(source.to_path_buf(), path.into()))
     } else {
         source.to_path_buf()
     };
@@ -2007,6 +2008,121 @@ pub fn maybe_install_llvm_runtime(builder: &Builder<'_>, target: TargetSelection
     }
 }
 
+/// Creates an output path to a BOLT-manipulated artifact for the given `file`.
+/// The hash of the file is used to make sure that we don't mix BOLT artifacts amongst different
+/// files with the same name.
+///
+/// We need to keep the file-name the same though, to make sure that copying the manipulated file
+/// to a directory will not change the final file path.
+fn create_bolt_output_path(builder: &Builder<'_>, file: &Path, hash: &str) -> PathBuf {
+    let directory = builder.out.join("bolt").join(hash);
+    t!(fs::create_dir_all(&directory));
+    directory.join(file.file_name().unwrap())
+}
+
+/// Instrument the provided file with BOLT.
+/// Returns a path to the instrumented artifact.
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub struct BoltInstrument {
+    file: PathBuf,
+    hash: String,
+}
+
+impl BoltInstrument {
+    fn new(file: PathBuf) -> Self {
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(t!(fs::read(&file)));
+        let hash = hex::encode(hasher.finalize().as_slice());
+
+        Self { file, hash }
+    }
+}
+
+impl Step for BoltInstrument {
+    type Output = PathBuf;
+
+    const ONLY_HOSTS: bool = false;
+    const DEFAULT: bool = false;
+
+    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
+        run.never()
+    }
+
+    fn make_run(_run: RunConfig<'_>) {}
+
+    fn run(self, builder: &Builder<'_>) -> PathBuf {
+        if builder.build.config.dry_run() {
+            return self.file.clone();
+        }
+
+        if builder.build.config.llvm_from_ci {
+            println!("warning: trying to use BOLT with LLVM from CI, this will probably not work");
+        }
+
+        println!("Instrumenting {} with BOLT", self.file.display());
+
+        let output_path = create_bolt_output_path(builder, &self.file, &self.hash);
+        if !output_path.is_file() {
+            instrument_with_bolt(&self.file, &output_path);
+        }
+        output_path
+    }
+}
+
+/// Optimize the provided file with BOLT.
+/// Returns a path to the optimized artifact.
+///
+/// The hash is stored in the step to make sure that we don't optimize the same file
+/// twice (even under  different file paths).
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub struct BoltOptimize {
+    file: PathBuf,
+    profile: PathBuf,
+    hash: String,
+}
+
+impl BoltOptimize {
+    fn new(file: PathBuf, profile: PathBuf) -> Self {
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(t!(fs::read(&file)));
+        hasher.update(t!(fs::read(&profile)));
+        let hash = hex::encode(hasher.finalize().as_slice());
+
+        Self { file, profile, hash }
+    }
+}
+
+impl Step for BoltOptimize {
+    type Output = PathBuf;
+
+    const ONLY_HOSTS: bool = false;
+    const DEFAULT: bool = false;
+
+    fn should_run(run: ShouldRun<'_>) -> ShouldRun<'_> {
+        run.never()
+    }
+
+    fn make_run(_run: RunConfig<'_>) {}
+
+    fn run(self, builder: &Builder<'_>) -> PathBuf {
+        if builder.build.config.dry_run() {
+            return self.file.clone();
+        }
+
+        if builder.build.config.llvm_from_ci {
+            println!("warning: trying to use BOLT with LLVM from CI, this will probably not work");
+        }
+
+        println!("Optimizing {} with BOLT", self.file.display());
+
+        let output_path = create_bolt_output_path(builder, &self.file, &self.hash);
+        if !output_path.is_file() {
+            optimize_with_bolt(&self.file, &self.profile, &output_path);
+        }
+        output_path
+    }
+}
+
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 pub struct LlvmTools {
     pub target: TargetSelection,