From a16925d5f37c24c0839903e3441ca55f7be707f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Ber=C3=A1nek?= Date: Tue, 11 Jul 2023 17:26:17 +0200 Subject: [PATCH] Implement BOLT optimization in the `opt-dist` tool --- Cargo.lock | 1 + src/tools/opt-dist/Cargo.toml | 1 + src/tools/opt-dist/src/bolt.rs | 98 ++++++++++++++++++++++++++++++ src/tools/opt-dist/src/exec.rs | 13 ++-- src/tools/opt-dist/src/main.rs | 43 ++++++++----- src/tools/opt-dist/src/utils/io.rs | 32 ++++++++++ 6 files changed, 164 insertions(+), 24 deletions(-) create mode 100644 src/tools/opt-dist/src/bolt.rs 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 anyhow::Result, R>( + path: &Utf8Path, + func: F, +) -> anyhow::Result { + // 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 { + 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(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, D: AsRef>(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::, _>>()?) } + +/// 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 { + let mut files = glob::glob(&format!("{directory}/{prefix}*{suffix}"))? + .into_iter() + .collect::, _>>()?; + 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}")), + } +}