Use BOLT to optimize librustc_driver

Author: Kobzol

src/bootstrap/compile.rs

@@ -887,6 +887,11 @@ impl Step for Rustc {
             cargo.arg("-p").arg(krate);
         }
 
+        if compiler.stage == 1 {
+            // Relocations are required for BOLT to work.
+            cargo.rustflag(&format!("-Clink-args=-Wl,-q"));
+        }
+
         let _guard = builder.msg_sysroot_tool(
             Kind::Build,
             compiler.stage,

src/tools/opt-dist/src/bolt.rs

@@ -1,7 +1,7 @@
 use anyhow::Context;
 
 use crate::exec::cmd;
-use crate::training::LlvmBoltProfile;
+use crate::training::BoltProfile;
 use camino::{Utf8Path, Utf8PathBuf};
 
 use crate::utils::io::copy_file;
@@ -40,7 +40,7 @@ pub fn with_bolt_instrumented<F: FnOnce() -> anyhow::Result<R>, R>(
 }
 
 /// Optimizes the file at `path` with BOLT in-place using the given `profile`.
-pub fn bolt_optimize(path: &Utf8Path, profile: &LlvmBoltProfile) -> anyhow::Result<()> {
+pub fn bolt_optimize(path: &Utf8Path, profile: &BoltProfile) -> 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.

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::{BoltProfile, LlvmPGOProfile, RustcPGOProfile};
 use camino::{Utf8Path, Utf8PathBuf};
 use std::collections::BTreeMap;
 use std::fs::File;
@@ -159,7 +159,7 @@ impl Bootstrap {
         self
     }
 
-    pub fn with_bolt_profile(mut self, profile: LlvmBoltProfile) -> Self {
+    pub fn with_bolt_profile(mut self, profile: BoltProfile) -> Self {
         self.cmd = self.cmd.arg("--reproducible-artifact").arg(profile.0.as_str());
         self
     }

src/tools/opt-dist/src/main.rs

@@ -12,7 +12,7 @@ use crate::environment::{Environment, EnvironmentBuilder};
 use crate::exec::{cmd, Bootstrap};
 use crate::tests::run_tests;
 use crate::timer::Timer;
-use crate::training::{gather_llvm_bolt_profiles, gather_llvm_profiles, gather_rustc_profiles};
+use crate::training::{gather_bolt_profiles, gather_llvm_profiles, gather_rustc_profiles};
 use crate::utils::io::{copy_directory, move_directory, reset_directory};
 use crate::utils::{
     clear_llvm_files, format_env_variables, print_binary_sizes, print_free_disk_space,
@@ -267,9 +267,13 @@ fn execute_pipeline(
                 ".so",
             )?;
 
+            log::info!("Optimizing {llvm_lib} with BOLT");
+
             // Instrument it and gather profiles
             let profile = with_bolt_instrumented(&llvm_lib, || {
-                stage.section("Gather profiles", |_| gather_llvm_bolt_profiles(env))
+                stage.section("Gather profiles", |_| {
+                    gather_bolt_profiles(env, "LLVM", llvm_benchmarks(env))
+                })
             })?;
             print_free_disk_space()?;
 
@@ -287,6 +291,35 @@ fn execute_pipeline(
         None
     };
 
+    let rustc_bolt_profile = if env.supports_bolt() {
+        // Stage 4: Build BOLT instrumented rustc
+        timer.section("Stage 4 (Rustc BOLT)", |stage| {
+            // Find the path to the `librustc_driver.so` file
+            let rustc_lib = io::find_file_in_dir(
+                &env.build_artifacts().join("stage2").join("lib"),
+                "librustc_driver",
+                ".so",
+            )?;
+
+            log::info!("Optimizing {rustc_lib} with BOLT");
+
+            // Instrument it and gather profiles
+            let profile = with_bolt_instrumented(&rustc_lib, || {
+                stage.section("Gather profiles", |_| {
+                    gather_bolt_profiles(env, "rustc", rustc_benchmarks(env))
+                })
+            })?;
+            print_free_disk_space()?;
+
+            // Now optimize the library with BOLT.
+            bolt_optimize(&rustc_lib, &profile).context("Could not optimize rustc with BOLT")?;
+
+            Ok(Some(profile))
+        })?
+    } else {
+        None
+    };
+
     let mut dist = Bootstrap::dist(env, &dist_args)
         .llvm_pgo_optimize(&llvm_pgo_profile)
         .rustc_pgo_optimize(&rustc_pgo_profile)
@@ -295,10 +328,13 @@ fn execute_pipeline(
     if let Some(llvm_bolt_profile) = llvm_bolt_profile {
         dist = dist.with_bolt_profile(llvm_bolt_profile);
     }
+    if let Some(rustc_bolt_profile) = rustc_bolt_profile {
+        dist = dist.with_bolt_profile(rustc_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))?;
+    timer.section("Stage 5 (final build)", |stage| dist.run(stage))?;
 
     // After dist has finished, run a subset of the test suite on the optimized artifacts to discover
     // possible regressions.

src/tools/opt-dist/src/training.rs

@@ -27,8 +27,6 @@ const RUSTC_PGO_CRATES: &[&str] = &[
     "bitmaps-3.1.0",
 ];
 
-const LLVM_BOLT_CRATES: &[&str] = LLVM_PGO_CRATES;
-
 fn init_compiler_benchmarks(
     env: &Environment,
     profiles: &[&str],
@@ -113,6 +111,14 @@ fn log_profile_stats(
     Ok(())
 }
 
+pub fn llvm_benchmarks(env: &dyn Environment) -> CmdBuilder {
+    init_compiler_benchmarks(env, &["Debug", "Opt"], &["Full"], LLVM_PGO_CRATES)
+}
+
+pub fn rustc_benchmarks(env: &dyn Environment) -> CmdBuilder {
+    init_compiler_benchmarks(env, &["Check", "Debug", "Opt"], &["All"], RUSTC_PGO_CRATES)
+}
+
 pub struct LlvmPGOProfile(pub Utf8PathBuf);
 
 pub fn gather_llvm_profiles(
@@ -122,9 +128,7 @@ pub fn gather_llvm_profiles(
     log::info!("Running benchmarks with PGO instrumented LLVM");
 
     with_log_group("Running benchmarks", || {
-        init_compiler_benchmarks(env, &["Debug", "Opt"], &["Full"], LLVM_PGO_CRATES)
-            .run()
-            .context("Cannot gather LLVM PGO profiles")
+        llvm_benchmarks(env).run().context("Cannot gather LLVM PGO profiles")
     })?;
 
     let merged_profile = env.artifact_dir().join("llvm-pgo.profdata");
@@ -157,7 +161,7 @@ pub fn gather_rustc_profiles(
     // Here we're profiling the `rustc` frontend, so we also include `Check`.
     // The benchmark set includes various stress tests that put the frontend under pressure.
     with_log_group("Running benchmarks", || {
-        init_compiler_benchmarks(env, &["Check", "Debug", "Opt"], &["All"], RUSTC_PGO_CRATES)
+        rustc_benchmarks(env)
             .env("LLVM_PROFILE_FILE", profile_template.as_str())
             .run()
             .context("Cannot gather rustc PGO profiles")
@@ -176,20 +180,22 @@ pub fn gather_rustc_profiles(
     Ok(RustcPGOProfile(merged_profile))
 }
 
-pub struct LlvmBoltProfile(pub Utf8PathBuf);
+pub struct BoltProfile(pub Utf8PathBuf);
 
-pub fn gather_llvm_bolt_profiles(env: &Environment) -> anyhow::Result<LlvmBoltProfile> {
-    log::info!("Running benchmarks with BOLT instrumented LLVM");
+pub fn gather_bolt_profiles(
+    env: &Environment,
+    name: &str,
+    benchmarks: CmdBuilder,
+) -> anyhow::Result<BoltProfile> {
+    log::info!("Running benchmarks with BOLT instrumented {name}");
 
     with_log_group("Running benchmarks", || {
-        init_compiler_benchmarks(env, &["Check", "Debug", "Opt"], &["Full"], LLVM_BOLT_CRATES)
-            .run()
-            .context("Cannot gather LLVM BOLT profiles")
+        benchmarks.run().with_context(|| "Cannot gather {name} BOLT profiles")
     })?;
 
-    let merged_profile = env.artifact_dir().join("llvm-bolt.profdata");
+    let merged_profile = env.artifact_dir().join(format!("{name}-bolt.profdata"));
     let profile_root = Utf8PathBuf::from("/tmp/prof.fdata");
-    log::info!("Merging LLVM BOLT profiles to {merged_profile}");
+    log::info!("Merging {name} BOLT profiles to {merged_profile}");
 
     let profiles: Vec<_> =
         glob::glob(&format!("{profile_root}*"))?.collect::<Result<Vec<_>, _>>()?;
@@ -204,7 +210,7 @@ pub fn gather_llvm_bolt_profiles(env: &Environment) -> anyhow::Result<LlvmBoltPr
             .context("Cannot merge BOLT profiles")
     })?;
 
-    log::info!("LLVM BOLT statistics");
+    log::info!("{name} BOLT statistics");
     log::info!(
         "{merged_profile}: {}",
         humansize::format_size(std::fs::metadata(merged_profile.as_std_path())?.len(), BINARY)
@@ -219,5 +225,14 @@ pub fn gather_llvm_bolt_profiles(env: &Environment) -> anyhow::Result<LlvmBoltPr
     log::info!("{profile_root}: {}", humansize::format_size(size, BINARY));
     log::info!("Profile file count: {}", profiles.len());
 
-    Ok(LlvmBoltProfile(merged_profile))
+    // Delete the gathered profiles
+    for profile in glob::glob(&format!("{profile_root}*"))?.into_iter() {
+        if let Ok(profile) = profile {
+            if let Err(error) = std::fs::remove_file(&profile) {
+                log::error!("Cannot delete BOLT profile {}: {error:?}", profile.display());
+            }
+        }
+    }
+
+    Ok(BoltProfile(merged_profile))
 }