Use PGO for rustdoc

bootstrap.example.toml

@@ -985,6 +985,11 @@
 # Instrument the Rust compiler, so that when executed, it will gather profiles
 # to this path.
 #pgo.rustc.generate = "/tmp/profiles/foo.profraw"
+# Use the following profile to PGO optimize rustdoc.
+#pgo.rustdoc.use = "/tmp/profiles/foo.profraw"
+# Instrument rustdoc, so that when executed, it will gather profiles
+# to this path.
+#pgo.rustdoc.generate = "/tmp/profiles/foo.profraw"
 # Use the following profile to PGO optimize LLVM.
 #pgo.llvm.use = "/tmp/profiles/foo.profraw"
 # Instrument LLVM, so that when executed, it will gather profiles

src/bootstrap/src/core/build_steps/compile.rs

@@ -24,7 +24,7 @@ use crate::core::build_steps::tool::{RustcPrivateCompilers, SourceType, copy_lld
 use crate::core::build_steps::{dist, llvm};
 use crate::core::builder;
 use crate::core::builder::{
-    Builder, Cargo, Kind, RunConfig, ShouldRun, Step, StepMetadata, crate_description,
+    Builder, Cargo, Kind, RunConfig, ShouldRun, Step, StepMetadata, apply_pgo, crate_description,
 };
 use crate::core::config::toml::target::DefaultLinuxLinkerOverride;
 use crate::core::config::{
@@ -1288,37 +1288,7 @@ pub fn rustc_cargo(
         cargo.rustflag("-Clink-args=-Wl,--icf=all");
     }
 
-    let is_collecting = if let Some(path) = &builder.config.rust_pgo.generate_profile {
-        if build_compiler.stage == 1 {
-            cargo
-                .rustflag(&format!("-Cprofile-generate={}", path.to_str().expect("non-UTF8 path")));
-            // Apparently necessary to avoid overflowing the counters during
-            // a Cargo build profile
-            cargo.rustflag("-Cllvm-args=-vp-counters-per-site=4");
-            true
-        } else {
-            false
-        }
-    } else if let Some(path) = &builder.config.rust_pgo.use_profile {
-        if build_compiler.stage == 1 {
-            cargo.rustflag(&format!("-Cprofile-use={}", path.to_str().expect("non-UTF8 path")));
-            if builder.is_verbose() {
-                cargo.rustflag("-Cllvm-args=-pgo-warn-missing-function");
-            }
-            true
-        } else {
-            false
-        }
-    } else {
-        false
-    };
-    if is_collecting {
-        // Ensure paths to Rust sources are relative, not absolute.
-        cargo.rustflag(&format!(
-            "-Cllvm-args=-static-func-strip-dirname-prefix={}",
-            builder.config.src.components().count()
-        ));
-    }
+    apply_pgo(builder, cargo, *build_compiler, &builder.config.rust_pgo);
 
     // The stage0 compiler changes infrequently and does not directly depend on code
     // in the current working directory. Therefore, caching it with sccache should be

src/bootstrap/src/core/build_steps/tool.rs

@@ -18,7 +18,8 @@ use crate::core::build_steps::toolstate::ToolState;
 use crate::core::build_steps::{compile, llvm};
 use crate::core::builder;
 use crate::core::builder::{
-    Builder, Cargo as CargoCommand, RunConfig, ShouldRun, Step, StepMetadata, cargo_profile_var,
+    Builder, Cargo as CargoCommand, RunConfig, ShouldRun, Step, StepMetadata, apply_pgo,
+    cargo_profile_var,
 };
 use crate::core::config::{DebuginfoLevel, RustcLto, TargetSelection};
 use crate::utils::exec::{BootstrapCommand, command};
@@ -137,6 +138,10 @@ impl Step for ToolBuild {
             }
         }
 
+        if self.path == "src/tools/rustdoc" {
+            apply_pgo(builder, &mut cargo, self.build_compiler, &builder.config.rustdoc_pgo);
+        }
+
         if !self.allow_features.is_empty() {
             cargo.allow_features(self.allow_features);
         }

src/bootstrap/src/core/builder/cargo.rs

@@ -6,6 +6,7 @@ use super::{Builder, Kind};
 use crate::core::build_steps::test;
 use crate::core::build_steps::tool::SourceType;
 use crate::core::config::flags::Color;
+use crate::core::config::toml::pgo::PgoConfig;
 use crate::core::config::{CompressDebuginfo, SplitDebuginfo};
 use crate::utils::build_stamp;
 use crate::utils::helpers::{self, LldThreads, check_cfg_arg, linker_flags};
@@ -1527,3 +1528,44 @@ pub fn cargo_profile_var(name: &str, config: &Config, mode: Mode) -> String {
     };
     format!("CARGO_PROFILE_{profile}_{name}")
 }
+
+/// Applies PGO compile flags to the given Cargo invocation based on the given PGO config.
+/// PGO flags are only applied when compiling a stage2 component.
+pub fn apply_pgo(
+    builder: &Builder<'_>,
+    cargo: &mut Cargo,
+    build_compiler: Compiler,
+    config: &PgoConfig,
+) {
+    let is_collecting = if let Some(path) = &config.generate_profile {
+        if build_compiler.stage == 1 {
+            cargo
+                .rustflag(&format!("-Cprofile-generate={}", path.to_str().expect("non-UTF8 path")));
+            // Apparently necessary to avoid overflowing the counters during
+            // a Cargo build profile
+            cargo.rustflag("-Cllvm-args=-vp-counters-per-site=4");
+            true
+        } else {
+            false
+        }
+    } else if let Some(path) = &config.use_profile {
+        if build_compiler.stage == 1 {
+            cargo.rustflag(&format!("-Cprofile-use={}", path.to_str().expect("non-UTF8 path")));
+            if builder.is_verbose() {
+                cargo.rustflag("-Cllvm-args=-pgo-warn-missing-function");
+            }
+            true
+        } else {
+            false
+        }
+    } else {
+        false
+    };
+    if is_collecting {
+        // Ensure paths to Rust sources are relative, not absolute.
+        cargo.rustflag(&format!(
+            "-Cllvm-args=-static-func-strip-dirname-prefix={}",
+            builder.config.src.components().count()
+        ));
+    }
+}

src/bootstrap/src/core/builder/mod.rs

@@ -13,7 +13,7 @@ use clap::ValueEnum;
 #[cfg(feature = "tracing")]
 use tracing::instrument;
 
-pub use self::cargo::{Cargo, cargo_profile_var};
+pub use self::cargo::{Cargo, apply_pgo, cargo_profile_var};
 pub use crate::Compiler;
 use crate::core::build_steps::compile::{Std, StdLink};
 use crate::core::build_steps::tool::RustcPrivateCompilers;

src/bootstrap/src/core/config/config.rs

@@ -235,6 +235,7 @@ pub struct Config {
     pub rust_parallel_frontend_threads: Option<u32>,
     pub rust_rustflags: Vec<String>,
     pub rust_pgo: PgoConfig,
+    pub rustdoc_pgo: PgoConfig,
 
     pub llvm_libunwind_default: Option<LlvmLibunwind>,
     pub enable_bolt_settings: bool,
@@ -656,7 +657,8 @@ impl Config {
             libgccjit_libs_dir: gcc_libgccjit_libs_dir,
         } = toml_gcc.unwrap_or_default();
 
-        let Pgo { rustc: pgo_rustc, llvm: pgo_llvm } = toml_pgo.unwrap_or_default();
+        let Pgo { rustc: pgo_rustc, llvm: pgo_llvm, rustdoc: pgo_rustdoc } =
+            toml_pgo.unwrap_or_default();
 
         // Backcompat: flags have priority over config
         if flags_rust_profile_use.is_some() || flags_rust_profile_generate.is_some() {
@@ -701,6 +703,11 @@ impl Config {
             panic!("Cannot use and generate LLVM PGO profiles at the same time");
         }
 
+        let pgo_rustdoc = pgo_rustdoc.unwrap_or_default();
+        if pgo_rustdoc.use_profile.is_some() && pgo_rustdoc.generate_profile.is_some() {
+            panic!("Cannot use and generate rustdoc PGO profiles at the same time");
+        }
+
         if rust_bootstrap_override_lld.is_some() && rust_bootstrap_override_lld_legacy.is_some() {
             panic!(
                 "Cannot use both `rust.use-lld` and `rust.bootstrap-override-lld`. Please use only `rust.bootstrap-override-lld`"
@@ -1567,6 +1574,7 @@ NOTE: Please add `--stage 2` to your command line, or if you're sure you want to
             rustc_debug_assertions: rust_rustc_debug_assertions.unwrap_or(rust_debug == Some(true)),
             rustc_default_linker: rust_default_linker,
             rustc_error_format: flags_rustc_error_format,
+            rustdoc_pgo: pgo_rustdoc,
             rustfmt_info,
             sanitizers: build_sanitizers.unwrap_or(false),
             save_toolstates: rust_save_toolstates.map(PathBuf::from),

src/bootstrap/src/core/config/toml/pgo.rs

@@ -26,6 +26,7 @@ define_config! {
     #[derive(Default)]
     struct Pgo {
         rustc: Option<PgoConfig> = "rustc",
+        rustdoc: Option<PgoConfig> = "rustdoc",
         llvm: Option<PgoConfig> = "llvm",
     }
 }

src/tools/opt-dist/src/exec.rs

@@ -7,7 +7,7 @@ use camino::{Utf8Path, Utf8PathBuf};
 use crate::environment::Environment;
 use crate::metrics::{load_metrics, record_metrics};
 use crate::timer::TimerSection;
-use crate::training::{BoltProfile, LlvmPGOProfile, RustcPGOProfile};
+use crate::training::{BoltProfile, LlvmPGOProfile, RustcPGOProfile, RustdocPGOProfile};
 use crate::utils::io::normalize_path;
 
 #[derive(Default)]
@@ -119,6 +119,11 @@ impl Bootstrap {
         Self { cmd, metrics_path }
     }
 
+    pub fn with_rustdoc(mut self) -> Self {
+        self.cmd = self.cmd.arg("rustdoc");
+        self
+    }
+
     pub fn dist(env: &Environment, dist_args: &[String]) -> Self {
         let metrics_path = env.build_root().join("metrics.json");
         let args = dist_args.iter().map(|arg| arg.as_str()).collect::<Vec<_>>();
@@ -158,6 +163,14 @@ impl Bootstrap {
         self
     }
 
+    pub fn rustdoc_pgo_instrument(mut self, profile_dir: &Utf8Path) -> Self {
+        self.cmd = self
+            .cmd
+            .arg("--set")
+            .arg(format!(r#"pgo.rustdoc.generate="{}""#, normalize_path(profile_dir).as_str()));
+        self
+    }
+
     pub fn without_llvm_lto(mut self) -> Self {
         self.cmd = self
             .cmd
@@ -176,6 +189,14 @@ impl Bootstrap {
         self
     }
 
+    pub fn rustdoc_pgo_optimize(mut self, profile: &RustdocPGOProfile) -> Self {
+        self.cmd = self
+            .cmd
+            .arg("--set")
+            .arg(format!(r#"pgo.rustdoc.use="{}""#, normalize_path(&profile.0).as_str()));
+        self
+    }
+
     pub fn with_llvm_bolt_ldflags(mut self) -> Self {
         self.cmd = self.cmd.arg("--set").arg("llvm.ldflags=-Wl,-q");
         self

src/tools/opt-dist/src/main.rs

@@ -10,8 +10,8 @@ use crate::exec::{Bootstrap, cmd};
 use crate::tests::run_tests;
 use crate::timer::Timer;
 use crate::training::{
-    gather_bolt_profiles, gather_llvm_profiles, gather_rustc_profiles, llvm_benchmarks,
-    rustc_benchmarks,
+    gather_bolt_profiles, gather_llvm_profiles, gather_rustc_profiles, gather_rustdoc_profiles,
+    llvm_benchmarks, rustc_benchmarks,
 };
 use crate::utils::artifact_size::print_binary_sizes;
 use crate::utils::io::{copy_directory, reset_directory};
@@ -239,39 +239,53 @@ fn execute_pipeline(
     // Stage 1: Build PGO instrumented rustc
     // We use a normal build of LLVM, because gathering PGO profiles for LLVM and `rustc` at the
     // same time can cause issues, because the host and in-tree LLVM versions can diverge.
-    let rustc_pgo_profile = timer.section("Stage 1 (Rustc PGO)", |stage| {
-        let rustc_profile_dir_root = env.artifact_dir().join("rustc-pgo");
-
-        stage.section("Build PGO instrumented rustc and LLVM", |section| {
-            let mut builder = Bootstrap::build(env).rustc_pgo_instrument(&rustc_profile_dir_root);
-
-            if env.supports_shared_llvm() {
-                // This first LLVM that we build will be thrown away after this stage, and it
-                // doesn't really need LTO. Without LTO, it builds in ~1 minute thanks to sccache,
-                // with LTO it takes almost 10 minutes. It makes the followup Rustc PGO
-                // instrumented/optimized build a bit slower, but it seems to be worth it.
-                builder = builder.without_llvm_lto();
-            }
+    let (rustc_pgo_profile, rustdoc_pgo_profile) =
+        timer.section("Stage 1 (Rustc + rustdoc PGO)", |stage| {
+            let rustc_profile_dir_root = env.artifact_dir().join("rustc-pgo");
+
+            stage.section("Build PGO instrumented rustc and LLVM", |section| {
+                // Rustc and rustdoc profiles are gathered together into the same directory, because
+                // they are executed together within a single rustc-perf invocation.
+                // Their profiles are then merged together into a single PGO profile.
+                let mut builder = Bootstrap::build(env)
+                    .with_rustdoc()
+                    .rustc_pgo_instrument(&rustc_profile_dir_root)
+                    .rustdoc_pgo_instrument(&rustc_profile_dir_root);
+
+                if env.supports_shared_llvm() {
+                    // This first LLVM that we build will be thrown away after this stage, and it
+                    // doesn't really need LTO. Without LTO, it builds in ~1 minute thanks to sccache,
+                    // with LTO it takes almost 10 minutes. It makes the followup Rustc PGO
+                    // instrumented/optimized build a bit slower, but it seems to be worth it.
+                    builder = builder.without_llvm_lto();
+                }
+
+                builder.run(section)
+            })?;
 
-            builder.run(section)
-        })?;
+            let rustc_profile = stage.section("Gather rustc profiles", |_| {
+                gather_rustc_profiles(env, &rustc_profile_dir_root)
+            })?;
+            let rustdoc_profile = stage.section("Gather rustdoc profiles", |_| {
+                gather_rustdoc_profiles(env, &rustc_profile_dir_root)
+            })?;
+            print_free_disk_space()?;
 
-        let profile = stage
-            .section("Gather profiles", |_| gather_rustc_profiles(env, &rustc_profile_dir_root))?;
-        print_free_disk_space()?;
+            stage.section("Build PGO optimized rustc", |section| {
+                let mut cmd = Bootstrap::build(env)
+                    .with_rustdoc()
+                    .rustc_pgo_optimize(&rustc_profile)
+                    .rustdoc_pgo_optimize(&rustdoc_profile);
+                if env.use_bolt() {
+                    cmd = cmd.with_rustc_bolt_ldflags();
+                }
 
-        stage.section("Build PGO optimized rustc", |section| {
-            let mut cmd = Bootstrap::build(env).rustc_pgo_optimize(&profile);
-            if env.use_bolt() {
-                cmd = cmd.with_rustc_bolt_ldflags();
-            }
+                cmd.run(section)
+            })?;
 
-            cmd.run(section)
+            Ok((rustc_profile, rustdoc_profile))
         })?;
 
-        Ok(profile)
-    })?;
-
     // Stage 2: Gather LLVM PGO profiles
     // Here we build a PGO instrumented LLVM, reusing the previously PGO optimized rustc.
     // Then we use the instrumented LLVM to gather LLVM PGO profiles.
@@ -386,7 +400,8 @@ fn execute_pipeline(
 
     let mut dist = Bootstrap::dist(env, &dist_args)
         .llvm_pgo_optimize(llvm_pgo_profile.as_ref())
-        .rustc_pgo_optimize(&rustc_pgo_profile);
+        .rustc_pgo_optimize(&rustc_pgo_profile)
+        .rustdoc_pgo_optimize(&rustdoc_pgo_profile);
 
     // if LLVM is not built we'll have PGO optimized rustc
     dist = if env.supports_shared_llvm() || !env.build_llvm() {

src/tools/opt-dist/src/training.rs

@@ -111,6 +111,10 @@ pub fn rustc_benchmarks(env: &Environment) -> CmdBuilder {
     init_compiler_benchmarks(env, &["Check", "Debug", "Opt"], &["All"], RUSTC_PGO_CRATES)
 }
 
+pub fn rustdoc_benchmarks(env: &Environment) -> CmdBuilder {
+    init_compiler_benchmarks(env, &["Doc"], &["Full"], RUSTC_PGO_CRATES)
+}
+
 pub struct LlvmPGOProfile(pub Utf8PathBuf);
 
 pub fn gather_llvm_profiles(
@@ -174,6 +178,44 @@ pub fn gather_rustc_profiles(
     Ok(RustcPGOProfile(merged_profile))
 }
 
+pub struct RustdocPGOProfile(pub Utf8PathBuf);
+
+pub fn gather_rustdoc_profiles(
+    env: &Environment,
+    profile_root: &Utf8Path,
+) -> anyhow::Result<RustdocPGOProfile> {
+    log::info!("Running benchmarks with PGO instrumented rustdoc");
+
+    // The profile data is written into a single filepath that is being repeatedly merged when each
+    // rustc invocation ends. Empirically, this can result in some profiling data being lost. That's
+    // why we override the profile path to include the PID. This will produce many more profiling
+    // files, but the resulting profile will produce a slightly faster rustc binary.
+    let profile_template = profile_root.join("default_%m_%p.profraw");
+
+    // Here we're profiling the `rustc` frontend, so we also include `Check`.
+    // The benchmark set includes various stress tests that put the frontend under pressure.
+    with_log_group("Running benchmarks", || {
+        rustdoc_benchmarks(env)
+            .env("LLVM_PROFILE_FILE", profile_template.as_str())
+            .run()
+            .context("Cannot gather rustdoc PGO profiles")
+    })?;
+
+    let merged_profile = env.artifact_dir().join("rustdoc-pgo.profdata");
+    log::info!("Merging Rustdoc PGO profiles to {merged_profile}");
+
+    let llvm_profdata = if env.build_llvm() { LlvmProfdata::Target } else { LlvmProfdata::Host };
+
+    merge_llvm_profiles(env, &merged_profile, profile_root, llvm_profdata)?;
+    log_profile_stats("Rustdoc", &merged_profile, profile_root)?;
+
+    // We don't need the individual .profraw files now that they have been merged
+    // into a final .profdata
+    delete_directory(profile_root)?;
+
+    Ok(RustdocPGOProfile(merged_profile))
+}
+
 pub struct BoltProfile(pub Utf8PathBuf);
 
 pub fn gather_bolt_profiles(