Merge pull request #19685 from fbernier/pgo-install

Add PGO support to install
This commit is contained in:
Lukas Wirth 2025-04-29 09:46:09 +00:00 committed by GitHub
commit 9f5b8844bd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 175 additions and 132 deletions

View File

@ -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<PathBuf> {
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<PathBuf> {
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),

View File

@ -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<Self, Self::Err> {
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<String>,
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<PgoTrainingCrate>,
}
#[derive(Debug)]
@ -143,25 +163,6 @@ pub struct RustcPush {
pub branch: Option<String>,
}
#[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<Self, Self::Err> {
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<ProcMacroServerOpt> {
if !self.proc_macro_server {

View File

@ -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<PgoTrainingCrate>,
}
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(())
}

View File

@ -22,6 +22,7 @@ mod codegen;
mod dist;
mod install;
mod metrics;
mod pgo;
mod publish;
mod release;
mod tidy;

105
xtask/src/pgo.rs Normal file
View File

@ -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<PathBuf> {
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<PathBuf> {
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()))
}

View File

@ -1,5 +1,7 @@
use std::path::{Path, PathBuf};
use xshell::{Shell, cmd};
pub(crate) fn list_rust_files(dir: &Path) -> Vec<PathBuf> {
let mut res = list_files(dir);
res.retain(|it| {
@ -29,3 +31,13 @@ pub(crate) fn list_files(dir: &Path) -> Vec<PathBuf> {
}
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),
},
}
}