xtask/
pgo.rs

1//! PGO (Profile-Guided Optimization) utilities.
2
3use anyhow::Context;
4use std::env::consts::EXE_EXTENSION;
5use std::ffi::OsStr;
6use std::path::{Path, PathBuf};
7use xshell::{Cmd, Shell, cmd};
8
9use crate::flags::PgoTrainingCrate;
10
11/// Decorates `ra_build_cmd` to add PGO instrumentation, and then runs the PGO instrumented
12/// Rust Analyzer on itself to gather a PGO profile.
13pub(crate) fn gather_pgo_profile<'a>(
14    sh: &'a Shell,
15    ra_build_cmd: Cmd<'a>,
16    target: &str,
17    train_crate: PgoTrainingCrate,
18) -> anyhow::Result<PathBuf> {
19    let pgo_dir = std::path::absolute("rust-analyzer-pgo")?;
20    // Clear out any stale profiles
21    if pgo_dir.is_dir() {
22        std::fs::remove_dir_all(&pgo_dir)?;
23    }
24    std::fs::create_dir_all(&pgo_dir)?;
25
26    // Figure out a path to `llvm-profdata`
27    let target_libdir = cmd!(sh, "rustc --print=target-libdir")
28        .read()
29        .context("cannot resolve target-libdir from rustc")?;
30    let target_bindir = PathBuf::from(target_libdir).parent().unwrap().join("bin");
31    let llvm_profdata = target_bindir.join("llvm-profdata").with_extension(EXE_EXTENSION);
32
33    // Build RA with PGO instrumentation
34    let cmd_gather =
35        ra_build_cmd.env("RUSTFLAGS", format!("-Cprofile-generate={}", pgo_dir.to_str().unwrap()));
36    cmd_gather.run().context("cannot build rust-analyzer with PGO instrumentation")?;
37
38    let (train_path, label) = match &train_crate {
39        PgoTrainingCrate::RustAnalyzer => (PathBuf::from("."), "itself"),
40        PgoTrainingCrate::GitHub(repo) => {
41            (download_crate_for_training(sh, &pgo_dir, repo)?, repo.as_str())
42        }
43    };
44
45    // Run RA either on itself or on a downloaded crate
46    eprintln!("Training RA on {label}...");
47    cmd!(
48        sh,
49        "target/{target}/release/rust-analyzer analysis-stats -q --run-all-ide-things {train_path}"
50    )
51    .run()
52    .context("cannot generate PGO profiles")?;
53
54    // Merge profiles into a single file
55    let merged_profile = pgo_dir.join("merged.profdata");
56    let profile_files = std::fs::read_dir(pgo_dir)?
57        .filter_map(|entry| {
58            let entry = entry.ok()?;
59            if entry.path().extension() == Some(OsStr::new("profraw")) {
60                Some(entry.path().to_str().unwrap().to_owned())
61            } else {
62                None
63            }
64        })
65        .collect::<Vec<_>>();
66
67    if profile_files.is_empty() {
68        anyhow::bail!(
69            "rust-analyzer analysis-stats produced no pgo files. This is a bug in rust-analyzer; please file an issue."
70        );
71    }
72
73    cmd!(sh, "{llvm_profdata} merge {profile_files...} -o {merged_profile}").run().context(
74        "cannot merge PGO profiles. Do you have the rustup `llvm-tools` component installed?",
75    )?;
76
77    Ok(merged_profile)
78}
79
80/// Downloads a crate from GitHub, stores it into `pgo_dir` and returns a path to it.
81fn download_crate_for_training(sh: &Shell, pgo_dir: &Path, repo: &str) -> anyhow::Result<PathBuf> {
82    let mut it = repo.splitn(2, '@');
83    let repo = it.next().unwrap();
84    let revision = it.next();
85
86    // FIXME: switch to `--revision` here around 2035 or so
87    let revision =
88        if let Some(revision) = revision { &["--branch", revision] as &[&str] } else { &[] };
89
90    let normalized_path = repo.replace("/", "-");
91    let target_path = pgo_dir.join(normalized_path);
92    cmd!(sh, "git clone --depth 1 https://github.com/{repo} {revision...} {target_path}")
93        .run()
94        .with_context(|| "cannot download PGO training crate from {repo}")?;
95
96    Ok(target_path)
97}
98
99/// Helper function to create a build command for rust-analyzer
100pub(crate) fn build_command<'a>(
101    sh: &'a Shell,
102    command: &str,
103    target_name: &str,
104    features: &[&str],
105) -> Cmd<'a> {
106    cmd!(
107        sh,
108        "cargo {command} --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --target {target_name} {features...} --release"
109    )
110}
111
112pub(crate) fn apply_pgo_to_cmd<'a>(cmd: Cmd<'a>, profile_path: &Path) -> Cmd<'a> {
113    cmd.env("RUSTFLAGS", format!("-Cprofile-use={}", profile_path.to_str().unwrap()))
114}
xtask/pgo.rs

xtask/
pgo.rs
