Allow passing a specific crate for PGO training

Author: Kobzol

.github/workflows/release.yaml

@@ -88,11 +88,11 @@ jobs:
 
       - name: Dist (plain)
         if: ${{ !matrix.zig_target }}
-        run: cargo xtask dist --client-patch-version ${{ github.run_number }} ${{ matrix.pgo && '--pgo' || ''}}
+        run: cargo xtask dist --client-patch-version ${{ github.run_number }} ${{ matrix.pgo && format('--pgo {0}', matrix.pgo) || ''}}
 
       - name: Dist (using zigbuild)
         if: ${{ matrix.zig_target }}
-        run: RA_TARGET=${{ matrix.zig_target}} cargo xtask dist --client-patch-version ${{ github.run_number }} --zig ${{ matrix.pgo && '--pgo' || ''}}
+        run: RA_TARGET=${{ matrix.zig_target}} cargo xtask dist --client-patch-version ${{ github.run_number }} --zig ${{ matrix.pgo && format('--pgo {0}', matrix.pgo) || ''}}
 
       - run: npm ci
         working-directory: editors/code

xtask/src/dist.rs

@@ -12,6 +12,7 @@ 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},
@@ -93,7 +94,7 @@ fn dist_server(
     target: &Target,
     allocator: Malloc,
     zig: bool,
-    pgo: bool,
+    pgo: Option<PgoTrainingCrate>,
 ) -> anyhow::Result<()> {
     let _e = sh.push_env("CFG_RELEASE", release);
     let _e = sh.push_env("CARGO_PROFILE_RELEASE_LTO", "thin");
@@ -111,11 +112,12 @@ fn dist_server(
     let features = allocator.to_features();
     let command = if linux_target && zig { "zigbuild" } else { "build" };
 
-    let pgo_profile = if pgo {
+    let pgo_profile = if let Some(train_crate) = pgo {
         Some(gather_pgo_profile(
             sh,
             build_command(sh, command, &target_name, features),
             &target_name,
+            train_crate,
         )?)
     } else {
         None
@@ -155,8 +157,9 @@ 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("ra-pgo-profiles")?;
+    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)?;
@@ -175,12 +178,21 @@ fn gather_pgo_profile<'a>(
         ra_build_cmd.env("RUSTFLAGS", format!("-Cprofile-generate={}", pgo_dir.to_str().unwrap()));
     cmd_gather.run().context("cannot build rust-analyzer with PGO instrumentation")?;
 
-    // Run RA on itself to gather profiles
-    let train_crate = ".";
+    let (train_path, label) = match &train_crate {
+        PgoTrainingCrate::RustAnalyzer => (PathBuf::from("."), "itself"),
+        PgoTrainingCrate::GitHub(url) => {
+            (download_crate_for_training(sh, &pgo_dir, url)?, url.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 {train_crate} --run-all-ide-things"
+        "target/{target}/release/rust-analyzer analysis-stats --run-all-ide-things {train_path}"
     )
+    // analysis-stats produces an enormous amount of output on stdout
+    .ignore_stdout()
     .run()
     .context("cannot generate PGO profiles")?;
 
@@ -201,6 +213,17 @@ fn gather_pgo_profile<'a>(
     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, url: &str) -> anyhow::Result<PathBuf> {
+    let normalized_path = url.replace("/", "-");
+    let target_path = pgo_dir.join(normalized_path);
+    cmd!(sh, "git clone --depth 1 https://github.com/{url} {target_path}")
+        .run()
+        .with_context(|| "cannot download PGO training crate from {url}")?;
+
+    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)?);

xtask/src/flags.rs

@@ -60,7 +60,7 @@ xflags::xflags! {
             /// Use cargo-zigbuild
             optional --zig
             /// Apply PGO optimizations
-            optional --pgo
+            optional --pgo pgo: PgoTrainingCrate
         }
         /// Read a changelog AsciiDoc file and update the GitHub Releases entry in Markdown.
         cmd publish-release-notes {
@@ -143,13 +143,32 @@ 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,
     pub jemalloc: bool,
     pub client_patch_version: Option<String>,
     pub zig: bool,
-    pub pgo: bool,
+    pub pgo: Option<PgoTrainingCrate>,
 }
 
 #[derive(Debug)]