Use PGO for libstd

Author: Kobzol

compiler/rustc_metadata/src/creader.rs

@@ -772,6 +772,11 @@ impl<'a> CrateLoader<'a> {
 
         info!("loading profiler");
 
+        let local_crate_name = self.local_crate_name.as_str();
+        if std::env::var("RUST_LIBSTD_PGO").is_ok() && local_crate_name != "std" {
+            return;
+        }
+
         let name = Symbol::intern(&self.sess.opts.debugging_opts.profiler_runtime);
         if name == sym::profiler_builtins && self.sess.contains_name(&krate.attrs, sym::no_core) {
             self.sess.err(

src/bootstrap/compile.rs

@@ -108,6 +108,42 @@ impl Step for Std {
         let mut cargo = builder.cargo(compiler, Mode::Std, SourceType::InTree, target, "build");
         std_cargo(builder, target, compiler.stage, &mut cargo);
 
+        if builder.config.libstd_profile_use.is_some()
+            && builder.config.libstd_profile_generate.is_some()
+        {
+            panic!("Cannot use and generate PGO profiles for libstd at the same time");
+        }
+
+        let is_collecting = if let Some(path) = &builder.config.libstd_profile_generate {
+            if compiler.stage == 1 {
+                cargo.env("RUST_LIBSTD_PGO", "1");
+                cargo.rustflag(&format!("-Cprofile-generate={}", 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.libstd_profile_use {
+            if compiler.stage == 1 {
+                cargo.rustflag(&format!("-Cprofile-use={}", path));
+                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()
+            ));
+        }
+
         builder.info(&format!(
             "Building stage{} std artifacts ({} -> {})",
             compiler.stage, &compiler.host, target

src/bootstrap/config.rs

@@ -156,6 +156,8 @@ pub struct Config {
     pub rust_new_symbol_mangling: Option<bool>,
     pub rust_profile_use: Option<String>,
     pub rust_profile_generate: Option<String>,
+    pub libstd_profile_use: Option<String>,
+    pub libstd_profile_generate: Option<String>,
     pub llvm_profile_use: Option<String>,
     pub llvm_profile_generate: bool,
     pub llvm_libunwind_default: Option<LlvmLibunwind>,
@@ -797,6 +799,8 @@ impl Config {
         }
         config.llvm_profile_use = flags.llvm_profile_use;
         config.llvm_profile_generate = flags.llvm_profile_generate;
+        config.libstd_profile_use = flags.libstd_profile_use;
+        config.libstd_profile_generate = flags.libstd_profile_generate;
 
         let stage0_json = t!(std::fs::read(&config.src.join("src").join("stage0.json")));
         config.stage0_metadata = t!(serde_json::from_slice::<Stage0Metadata>(&stage0_json));

src/bootstrap/flags.rs

@@ -79,6 +79,9 @@ pub struct Flags {
     //
     // llvm_out/build/profiles/ is the location this writes to.
     pub llvm_profile_generate: bool,
+
+    pub libstd_profile_use: Option<String>,
+    pub libstd_profile_generate: Option<String>,
 }
 
 #[cfg_attr(test, derive(Clone))]
@@ -247,6 +250,13 @@ To learn more about a subcommand, run `./x.py <subcommand> -h`",
         opts.optopt("", "rust-profile-use", "use PGO profile for rustc build", "PROFILE");
         opts.optflag("", "llvm-profile-generate", "generate PGO profile with llvm built for rustc");
         opts.optopt("", "llvm-profile-use", "use PGO profile for llvm build", "PROFILE");
+        opts.optopt(
+            "",
+            "libstd-profile-generate",
+            "generate PGO profile with libstd build",
+            "PROFILE",
+        );
+        opts.optopt("", "libstd-profile-use", "use PGO profile for libstd build", "PROFILE");
 
         // We can't use getopt to parse the options until we have completed specifying which
         // options are valid, but under the current implementation, some options are conditional on
@@ -676,6 +686,8 @@ Arguments:
             rust_profile_generate: matches.opt_str("rust-profile-generate"),
             llvm_profile_use: matches.opt_str("llvm-profile-use"),
             llvm_profile_generate: matches.opt_present("llvm-profile-generate"),
+            libstd_profile_use: matches.opt_str("libstd-profile-use"),
+            libstd_profile_generate: matches.opt_str("libstd-profile-generate"),
         }
     }
 }

src/ci/pgo.sh

@@ -120,6 +120,26 @@ du -sh ${RUSTC_PROFILE_DIRECTORY_ROOT}
 echo "Profile file count"
 find ${RUSTC_PROFILE_DIRECTORY_ROOT} -type f | wc -l
 
+rm -r ./build/$PGO_HOST/llvm ./build/$PGO_HOST/lld
+
+python3 ../x.py build --target=$PGO_HOST --host=$PGO_HOST \
+    --stage 2 library/std \
+    --rust-profile-use=/tmp/rustc-pgo.profdata \
+    --libstd-profile-generate=/tmp/libstd-pgo
+
+# 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.
+LLVM_PROFILE_FILE=/tmp/libstd-pgo/default_%m_%p.profraw gather_profiles "Check,Debug,Opt" "All" \
+  "externs,ctfe-stress-5,cargo-0.60.0,token-stream-stress,match-stress,tuple-stress"
+
+# Merge the profile data we gathered
+./build/$PGO_HOST/llvm/bin/llvm-profdata \
+    merge -o /tmp/libstd-pgo.profdata /tmp/libstd-pgo
+
+echo "libstd PGO statistics"
+./build/$PGO_HOST/llvm/bin/llvm-profdata \
+    show /tmp/libstd-pgo.profdata --all-functions
+
 # Rustbuild currently doesn't support rebuilding LLVM when PGO options
 # change (or any other llvm-related options); so just clear out the relevant
 # directories ourselves.
@@ -129,4 +149,5 @@ rm -r ./build/$PGO_HOST/llvm ./build/$PGO_HOST/lld
 # collected profiling data.
 $@ \
     --rust-profile-use=${RUSTC_PROFILE_MERGED_FILE} \
-    --llvm-profile-use=${LLVM_PROFILE_MERGED_FILE}
+    --llvm-profile-use=${LLVM_PROFILE_MERGED_FILE} \
+    --libstd-profile-use=/tmp/libstd-pgo.profdata