From e2c48d3c457111c19476f4ec54f6cb8d30eaa93a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Bernier?= Date: Fri, 18 Apr 2025 00:27:08 -0400 Subject: [PATCH] Add PGO support to install Enable PGO optimizations when installing rust-analyzer with the --pgo flag. This mirrors functionality already available in dist command, allowing developers to create optimized local builds. Example: cargo xtask install --server --pgo clap-rs/clap PGO code has been extracted to a dedicated module for reuse. --- xtask/src/dist.rs | 109 +++---------------------------------------- xtask/src/flags.rs | 52 +++++++++++---------- xtask/src/install.rs | 28 ++++++++--- xtask/src/main.rs | 1 + xtask/src/pgo.rs | 105 +++++++++++++++++++++++++++++++++++++++++ xtask/src/util.rs | 12 +++++ 6 files changed, 175 insertions(+), 132 deletions(-) create mode 100644 xtask/src/pgo.rs diff --git a/xtask/src/dist.rs b/xtask/src/dist.rs index b3d6f06b07..dbfecdbe11 100644 --- a/xtask/src/dist.rs +++ b/xtask/src/dist.rs @@ -1,9 +1,6 @@ use anyhow::Context; use flate2::{Compression, write::GzEncoder}; -use std::env::consts::EXE_EXTENSION; -use std::ffi::OsStr; use std::{ - env, fs::File, io::{self, BufWriter}, path::{Path, PathBuf}, @@ -12,11 +9,11 @@ use time::OffsetDateTime; use xshell::{Cmd, Shell, cmd}; use zip::{DateTime, ZipWriter, write::SimpleFileOptions}; -use crate::flags::PgoTrainingCrate; use crate::{ date_iso, - flags::{self, Malloc}, + flags::{self, Malloc, PgoTrainingCrate}, project_root, + util::detect_target, }; const VERSION_STABLE: &str = "0.3"; @@ -28,7 +25,7 @@ impl flags::Dist { let stable = sh.var("GITHUB_REF").unwrap_or_default().as_str() == "refs/heads/release"; let project_root = project_root(); - let target = Target::get(&project_root); + let target = Target::get(&project_root, sh); let allocator = self.allocator(); let dist = project_root.join("dist"); sh.remove_path(&dist)?; @@ -113,9 +110,9 @@ fn dist_server( let command = if linux_target && zig { "zigbuild" } else { "build" }; let pgo_profile = if let Some(train_crate) = pgo { - Some(gather_pgo_profile( + Some(crate::pgo::gather_pgo_profile( sh, - build_command(sh, command, &target_name, features), + crate::pgo::build_command(sh, command, &target_name, features), &target_name, train_crate, )?) @@ -151,85 +148,6 @@ fn build_command<'a>( ) } -/// Decorates `ra_build_cmd` to add PGO instrumentation, and then runs the PGO instrumented -/// Rust Analyzer on itself to gather a PGO profile. -fn gather_pgo_profile<'a>( - sh: &'a Shell, - ra_build_cmd: Cmd<'a>, - target: &str, - train_crate: PgoTrainingCrate, -) -> anyhow::Result { - let pgo_dir = std::path::absolute("rust-analyzer-pgo")?; - // Clear out any stale profiles - if pgo_dir.is_dir() { - std::fs::remove_dir_all(&pgo_dir)?; - } - std::fs::create_dir_all(&pgo_dir)?; - - // Figure out a path to `llvm-profdata` - let target_libdir = cmd!(sh, "rustc --print=target-libdir") - .read() - .context("cannot resolve target-libdir from rustc")?; - let target_bindir = PathBuf::from(target_libdir).parent().unwrap().join("bin"); - let llvm_profdata = target_bindir.join("llvm-profdata").with_extension(EXE_EXTENSION); - - // Build RA with PGO instrumentation - let cmd_gather = - ra_build_cmd.env("RUSTFLAGS", format!("-Cprofile-generate={}", pgo_dir.to_str().unwrap())); - cmd_gather.run().context("cannot build rust-analyzer with PGO instrumentation")?; - - let (train_path, label) = match &train_crate { - PgoTrainingCrate::RustAnalyzer => (PathBuf::from("."), "itself"), - PgoTrainingCrate::GitHub(repo) => { - (download_crate_for_training(sh, &pgo_dir, repo)?, repo.as_str()) - } - }; - - // Run RA either on itself or on a downloaded crate - eprintln!("Training RA on {label}..."); - cmd!( - sh, - "target/{target}/release/rust-analyzer analysis-stats -q --run-all-ide-things {train_path}" - ) - .run() - .context("cannot generate PGO profiles")?; - - // Merge profiles into a single file - let merged_profile = pgo_dir.join("merged.profdata"); - let profile_files = std::fs::read_dir(pgo_dir)?.filter_map(|entry| { - let entry = entry.ok()?; - if entry.path().extension() == Some(OsStr::new("profraw")) { - Some(entry.path().to_str().unwrap().to_owned()) - } else { - None - } - }); - cmd!(sh, "{llvm_profdata} merge {profile_files...} -o {merged_profile}").run().context( - "cannot merge PGO profiles. Do you have the rustup `llvm-tools` component installed?", - )?; - - Ok(merged_profile) -} - -/// Downloads a crate from GitHub, stores it into `pgo_dir` and returns a path to it. -fn download_crate_for_training(sh: &Shell, pgo_dir: &Path, repo: &str) -> anyhow::Result { - let mut it = repo.splitn(2, '@'); - let repo = it.next().unwrap(); - let revision = it.next(); - - // FIXME: switch to `--revision` here around 2035 or so - let revision = - if let Some(revision) = revision { &["--branch", revision] as &[&str] } else { &[] }; - - let normalized_path = repo.replace("/", "-"); - let target_path = pgo_dir.join(normalized_path); - cmd!(sh, "git clone --depth 1 https://github.com/{repo} {revision...} {target_path}") - .run() - .with_context(|| "cannot download PGO training crate from {repo}")?; - - Ok(target_path) -} - fn gzip(src_path: &Path, dest_path: &Path) -> anyhow::Result<()> { let mut encoder = GzEncoder::new(File::create(dest_path)?, Compression::best()); let mut input = io::BufReader::new(File::open(src_path)?); @@ -283,21 +201,8 @@ struct Target { } impl Target { - fn get(project_root: &Path) -> Self { - let name = match env::var("RA_TARGET") { - Ok(target) => target, - _ => { - if cfg!(target_os = "linux") { - "x86_64-unknown-linux-gnu".to_owned() - } else if cfg!(target_os = "windows") { - "x86_64-pc-windows-msvc".to_owned() - } else if cfg!(target_os = "macos") { - "x86_64-apple-darwin".to_owned() - } else { - panic!("Unsupported OS, maybe try setting RA_TARGET") - } - } - }; + fn get(project_root: &Path, sh: &Shell) -> Self { + let name = detect_target(sh); let (name, libc_suffix) = match name.split_once('.') { Some((l, r)) => (l.to_owned(), Some(r.to_owned())), None => (name, None), diff --git a/xtask/src/flags.rs b/xtask/src/flags.rs index 700806d178..2fd471b35c 100644 --- a/xtask/src/flags.rs +++ b/xtask/src/flags.rs @@ -4,6 +4,25 @@ use std::{fmt, str::FromStr}; use crate::install::{ClientOpt, ProcMacroServerOpt, ServerOpt}; +#[derive(Debug, Clone)] +pub enum PgoTrainingCrate { + // Use RA's own sources for PGO training + RustAnalyzer, + // Download a Rust crate from `https://github.com/{0}` and use it for PGO training. + GitHub(String), +} + +impl FromStr for PgoTrainingCrate { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "rust-analyzer" => Ok(Self::RustAnalyzer), + url => Ok(Self::GitHub(url.to_owned())), + } + } +} + xflags::xflags! { src "./src/flags.rs" @@ -29,6 +48,9 @@ xflags::xflags! { /// build in release with debug info set to 2. optional --dev-rel + + /// Apply PGO optimizations + optional --pgo pgo: PgoTrainingCrate } cmd fuzz-tests {} @@ -109,18 +131,16 @@ pub enum XtaskCmd { Tidy(Tidy), } -#[derive(Debug)] -pub struct Tidy {} - #[derive(Debug)] pub struct Install { pub client: bool, pub code_bin: Option, pub server: bool, - pub proc_macro_server: bool, pub mimalloc: bool, pub jemalloc: bool, + pub proc_macro_server: bool, pub dev_rel: bool, + pub pgo: Option, } #[derive(Debug)] @@ -143,25 +163,6 @@ pub struct RustcPush { pub branch: Option, } -#[derive(Debug)] -pub enum PgoTrainingCrate { - // Use RA's own sources for PGO training - RustAnalyzer, - // Download a Rust crate from `https://github.com/{0}` and use it for PGO training. - GitHub(String), -} - -impl FromStr for PgoTrainingCrate { - type Err = String; - - fn from_str(s: &str) -> Result { - match s { - "rust-analyzer" => Ok(Self::RustAnalyzer), - url => Ok(Self::GitHub(url.to_owned())), - } - } -} - #[derive(Debug)] pub struct Dist { pub mimalloc: bool, @@ -195,6 +196,9 @@ pub struct Codegen { pub check: bool, } +#[derive(Debug)] +pub struct Tidy; + impl Xtask { #[allow(dead_code)] pub fn from_env_or_exit() -> Self { @@ -324,7 +328,7 @@ impl Install { } else { Malloc::System }; - Some(ServerOpt { malloc, dev_rel: self.dev_rel }) + Some(ServerOpt { malloc, dev_rel: self.dev_rel, pgo: self.pgo.clone() }) } pub(crate) fn proc_macro_server(&self) -> Option { if !self.proc_macro_server { diff --git a/xtask/src/install.rs b/xtask/src/install.rs index 4e2093f069..f0cc445dfa 100644 --- a/xtask/src/install.rs +++ b/xtask/src/install.rs @@ -5,7 +5,10 @@ use std::{env, path::PathBuf, str}; use anyhow::{Context, bail, format_err}; use xshell::{Shell, cmd}; -use crate::flags::{self, Malloc}; +use crate::{ + flags::{self, Malloc, PgoTrainingCrate}, + util::detect_target, +}; impl flags::Install { pub(crate) fn run(self, sh: &Shell) -> anyhow::Result<()> { @@ -35,6 +38,7 @@ const VS_CODES: &[&str] = &["code", "code-exploration", "code-insiders", "codium pub(crate) struct ServerOpt { pub(crate) malloc: Malloc, pub(crate) dev_rel: bool, + pub(crate) pgo: Option, } pub(crate) struct ProcMacroServerOpt { @@ -135,21 +139,33 @@ fn install_server(sh: &Shell, opts: ServerOpt) -> anyhow::Result<()> { let features = opts.malloc.to_features(); let profile = if opts.dev_rel { "dev-rel" } else { "release" }; - let cmd = cmd!( + let mut install_cmd = cmd!( sh, "cargo install --path crates/rust-analyzer --profile={profile} --locked --force --features force-always-assert {features...}" ); - cmd.run()?; + + if let Some(train_crate) = opts.pgo { + let build_cmd = cmd!( + sh, + "cargo build --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --profile={profile} --locked --features force-always-assert {features...}" + ); + + let target = detect_target(sh); + let profile = crate::pgo::gather_pgo_profile(sh, build_cmd, &target, train_crate)?; + install_cmd = crate::pgo::apply_pgo_to_cmd(install_cmd, &profile); + } + + install_cmd.run()?; Ok(()) } fn install_proc_macro_server(sh: &Shell, opts: ProcMacroServerOpt) -> anyhow::Result<()> { let profile = if opts.dev_rel { "dev-rel" } else { "release" }; - let cmd = cmd!( + cmd!( sh, "cargo +nightly install --path crates/proc-macro-srv-cli --profile={profile} --locked --force --features sysroot-abi" - ); - cmd.run()?; + ).run()?; + Ok(()) } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 52ea896c73..aaa8d0e1d4 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -22,6 +22,7 @@ mod codegen; mod dist; mod install; mod metrics; +mod pgo; mod publish; mod release; mod tidy; diff --git a/xtask/src/pgo.rs b/xtask/src/pgo.rs new file mode 100644 index 0000000000..7f7b3311d9 --- /dev/null +++ b/xtask/src/pgo.rs @@ -0,0 +1,105 @@ +//! PGO (Profile-Guided Optimization) utilities. + +use anyhow::Context; +use std::env::consts::EXE_EXTENSION; +use std::ffi::OsStr; +use std::path::{Path, PathBuf}; +use xshell::{Cmd, Shell, cmd}; + +use crate::flags::PgoTrainingCrate; + +/// Decorates `ra_build_cmd` to add PGO instrumentation, and then runs the PGO instrumented +/// Rust Analyzer on itself to gather a PGO profile. +pub(crate) fn gather_pgo_profile<'a>( + sh: &'a Shell, + ra_build_cmd: Cmd<'a>, + target: &str, + train_crate: PgoTrainingCrate, +) -> anyhow::Result { + let pgo_dir = std::path::absolute("rust-analyzer-pgo")?; + // Clear out any stale profiles + if pgo_dir.is_dir() { + std::fs::remove_dir_all(&pgo_dir)?; + } + std::fs::create_dir_all(&pgo_dir)?; + + // Figure out a path to `llvm-profdata` + let target_libdir = cmd!(sh, "rustc --print=target-libdir") + .read() + .context("cannot resolve target-libdir from rustc")?; + let target_bindir = PathBuf::from(target_libdir).parent().unwrap().join("bin"); + let llvm_profdata = target_bindir.join("llvm-profdata").with_extension(EXE_EXTENSION); + + // Build RA with PGO instrumentation + let cmd_gather = + ra_build_cmd.env("RUSTFLAGS", format!("-Cprofile-generate={}", pgo_dir.to_str().unwrap())); + cmd_gather.run().context("cannot build rust-analyzer with PGO instrumentation")?; + + let (train_path, label) = match &train_crate { + PgoTrainingCrate::RustAnalyzer => (PathBuf::from("."), "itself"), + PgoTrainingCrate::GitHub(repo) => { + (download_crate_for_training(sh, &pgo_dir, repo)?, repo.as_str()) + } + }; + + // Run RA either on itself or on a downloaded crate + eprintln!("Training RA on {label}..."); + cmd!( + sh, + "target/{target}/release/rust-analyzer analysis-stats -q --run-all-ide-things {train_path}" + ) + .run() + .context("cannot generate PGO profiles")?; + + // Merge profiles into a single file + let merged_profile = pgo_dir.join("merged.profdata"); + let profile_files = std::fs::read_dir(pgo_dir)?.filter_map(|entry| { + let entry = entry.ok()?; + if entry.path().extension() == Some(OsStr::new("profraw")) { + Some(entry.path().to_str().unwrap().to_owned()) + } else { + None + } + }); + cmd!(sh, "{llvm_profdata} merge {profile_files...} -o {merged_profile}").run().context( + "cannot merge PGO profiles. Do you have the rustup `llvm-tools` component installed?", + )?; + + Ok(merged_profile) +} + +/// Downloads a crate from GitHub, stores it into `pgo_dir` and returns a path to it. +fn download_crate_for_training(sh: &Shell, pgo_dir: &Path, repo: &str) -> anyhow::Result { + let mut it = repo.splitn(2, '@'); + let repo = it.next().unwrap(); + let revision = it.next(); + + // FIXME: switch to `--revision` here around 2035 or so + let revision = + if let Some(revision) = revision { &["--branch", revision] as &[&str] } else { &[] }; + + let normalized_path = repo.replace("/", "-"); + let target_path = pgo_dir.join(normalized_path); + cmd!(sh, "git clone --depth 1 https://github.com/{repo} {revision...} {target_path}") + .run() + .with_context(|| "cannot download PGO training crate from {repo}")?; + + Ok(target_path) +} + +/// Helper function to create a build command for rust-analyzer +pub(crate) fn build_command<'a>( + sh: &'a Shell, + command: &str, + target_name: &str, + features: &[&str], +) -> Cmd<'a> { + cmd!( + sh, + "cargo {command} --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --target {target_name} {features...} --release" + ) +} + +pub(crate) fn apply_pgo_to_cmd<'a>(cmd: Cmd<'a>, profile_path: &Path) -> Cmd<'a> { + cmd.env("RUSTFLAGS", format!("-Cprofile-use={}", profile_path.to_str().unwrap())) +} diff --git a/xtask/src/util.rs b/xtask/src/util.rs index 39f52938c8..a740ad6afd 100644 --- a/xtask/src/util.rs +++ b/xtask/src/util.rs @@ -1,5 +1,7 @@ use std::path::{Path, PathBuf}; +use xshell::{Shell, cmd}; + pub(crate) fn list_rust_files(dir: &Path) -> Vec { let mut res = list_files(dir); res.retain(|it| { @@ -29,3 +31,13 @@ pub(crate) fn list_files(dir: &Path) -> Vec { } res } + +pub(crate) fn detect_target(sh: &Shell) -> String { + match std::env::var("RA_TARGET") { + Ok(target) => target, + _ => match cmd!(sh, "rustc --print=host-tuple").read() { + Ok(target) => target, + Err(e) => panic!("Failed to detect target: {}\nPlease set RA_TARGET explicitly", e), + }, + } +}