Update PGO staging logic for static LLVM linkage

chriswailes · chriswailes · commit 257ec3b21478 · 2022-11-16T10:43:40.000-08:00
This commit updates the logic around PGO staging when linking LLVM
statically to ensure that the necessary runtime libraries are included.
In addition, several flags related to PGO can now be specified in the
config.toml file.
diff --git a/src/bootstrap/compile.rs b/src/bootstrap/compile.rs
@@ -666,8 +666,10 @@ impl Step for Rustc {
             cargo.rustflag("-Clink-args=-Wl,--icf=all");
         }
 
-        let is_collecting = if let Some(path) = &builder.config.rust_profile_generate {
-            if compiler.stage == 1 {
+        let use_relative_paths = if let Some(path) = &builder.config.rust_profile_generate {
+            if compiler.stage >= 1
+                || (builder.config.llvm_profile_generate.is_some() && !builder.llvm_link_shared())
+            {
                 cargo.rustflag(&format!("-Cprofile-generate={}", path));
                 // Apparently necessary to avoid overflowing the counters during
                 // a Cargo build profile
@@ -676,18 +678,28 @@ impl Step for Rustc {
             } else {
                 false
             }
-        } else if let Some(path) = &builder.config.rust_profile_use {
-            if compiler.stage == 1 {
-                cargo.rustflag(&format!("-Cprofile-use={}", path));
-                cargo.rustflag("-Cllvm-args=-pgo-warn-missing-function");
+        } else if let Some(path) = &builder.config.llvm_profile_generate {
+            // If libLLVM.a is instrumented it will need to be linked against
+            // the profiler's runtime environment.  The only way to ensure that
+            // occurs is to tell rustc to profile the compilation unit.
+            if !builder.llvm_link_shared() {
+                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.rust_profile_use {
+            cargo.rustflag(&format!("-Cprofile-use={}", path));
+            cargo.rustflag("-Cllvm-args=-pgo-warn-missing-function");
+            true
         } else {
             false
         };
-        if is_collecting {
+
+        if use_relative_paths {
             // Ensure paths to Rust sources are relative, not absolute.
             cargo.rustflag(&format!(
                 "-Cllvm-args=-static-func-strip-dirname-prefix={}",
@@ -821,7 +833,7 @@ pub fn rustc_cargo_env(builder: &Builder<'_>, cargo: &mut Cargo, target: TargetS
         // found. This is to avoid the linker errors about undefined references to
         // `__llvm_profile_instrument_memop` when linking `rustc_driver`.
         let mut llvm_linker_flags = String::new();
-        if builder.config.llvm_profile_generate && target.contains("msvc") {
+        if builder.config.llvm_profile_generate.is_some() && target.contains("msvc") {
             if let Some(ref clang_cl_path) = builder.config.llvm_clang_cl {
                 // Add clang's runtime library directory to the search path
                 let clang_rt_dir = get_clang_cl_resource_dir(clang_cl_path);
diff --git a/src/bootstrap/config.rs b/src/bootstrap/config.rs
@@ -140,6 +140,12 @@ pub struct Config {
     pub llvm_ldflags: Option<String>,
     pub llvm_use_libcxx: bool,
 
+    pub llvm_profile_use: Option<String>,
+    pub llvm_profile_generate: Option<String>,
+    pub llvm_libunwind_default: Option<LlvmLibunwind>,
+    pub llvm_bolt_profile_generate: bool,
+    pub llvm_bolt_profile_use: Option<String>,
+
     // rust codegen options
     pub rust_optimize: bool,
     pub rust_codegen_units: Option<u32>,
@@ -167,11 +173,6 @@ pub struct Config {
     pub rust_profile_use: Option<String>,
     pub rust_profile_generate: Option<String>,
     pub rust_lto: RustcLto,
-    pub llvm_profile_use: Option<String>,
-    pub llvm_profile_generate: bool,
-    pub llvm_libunwind_default: Option<LlvmLibunwind>,
-    pub llvm_bolt_profile_generate: bool,
-    pub llvm_bolt_profile_use: Option<String>,
 
     pub build: TargetSelection,
     pub hosts: Vec<TargetSelection>,
@@ -679,6 +680,10 @@ define_config! {
         clang: Option<bool> = "clang",
         download_ci_llvm: Option<StringOrBool> = "download-ci-llvm",
         build_config: Option<HashMap<String, String>> = "build-config",
+        profile_generate: Option<String> = "profile-generate",
+        profile_use: Option<String> = "profile-use",
+        bolt_profile_generate: Option<bool> = "bolt-profile-generate",
+        bolt_profile_use: Option<String> = "bolt-profile-use",
     }
 }
 
@@ -838,17 +843,6 @@ impl Config {
         if let Some(value) = flags.deny_warnings {
             config.deny_warnings = value;
         }
-        config.llvm_profile_use = flags.llvm_profile_use;
-        config.llvm_profile_generate = flags.llvm_profile_generate;
-        config.llvm_bolt_profile_generate = flags.llvm_bolt_profile_generate;
-        config.llvm_bolt_profile_use = flags.llvm_bolt_profile_use;
-
-        if config.llvm_bolt_profile_generate && config.llvm_bolt_profile_use.is_some() {
-            eprintln!(
-                "Cannot use both `llvm_bolt_profile_generate` and `llvm_bolt_profile_use` at the same time"
-            );
-            crate::detail_exit(1);
-        }
 
         // Infer the rest of the configuration.
 
@@ -1141,6 +1135,25 @@ impl Config {
                 // the link step) with each stage.
                 config.llvm_link_shared.set(Some(true));
             }
+
+            config.llvm_profile_generate = flags.llvm_profile_generate.or(llvm.profile_generate);
+            config.llvm_profile_use = flags.llvm_profile_use.or(llvm.profile_use);
+
+            config.llvm_bolt_profile_generate = flags.llvm_bolt_profile_generate;
+            config.llvm_bolt_profile_use = flags.llvm_bolt_profile_use.or(llvm.bolt_profile_use);
+        } else {
+            config.llvm_profile_generate = flags.llvm_profile_generate;
+            config.llvm_profile_use = flags.llvm_profile_use;
+
+            config.llvm_bolt_profile_generate = flags.llvm_bolt_profile_generate;
+            config.llvm_bolt_profile_use = flags.llvm_bolt_profile_use;
+        }
+
+        if config.llvm_bolt_profile_generate && config.llvm_bolt_profile_use.is_some() {
+            eprintln!(
+                "Cannot use both `llvm_bolt_profile_generate` and `llvm_bolt_profile_use` at the same time"
+            );
+            crate::detail_exit(1);
         }
 
         if let Some(rust) = toml.rust {
diff --git a/src/bootstrap/flags.rs b/src/bootstrap/flags.rs
@@ -69,15 +69,12 @@ pub struct Flags {
 
     pub llvm_skip_rebuild: Option<bool>,
 
-    pub rust_profile_use: Option<String>,
     pub rust_profile_generate: Option<String>,
+    pub rust_profile_use: Option<String>,
 
+    pub llvm_profile_generate: Option<String>,
     pub llvm_profile_use: Option<String>,
-    // LLVM doesn't support a custom location for generating profile
-    // information.
-    //
-    // llvm_out/build/profiles/ is the location this writes to.
-    pub llvm_profile_generate: bool,
+
     pub llvm_bolt_profile_generate: bool,
     pub llvm_bolt_profile_use: Option<String>,
 }
@@ -698,7 +695,7 @@ Arguments:
             rust_profile_use: matches.opt_str("rust-profile-use"),
             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"),
+            llvm_profile_generate: matches.opt_str("llvm-profile-generate"),
             llvm_bolt_profile_generate: matches.opt_present("llvm-bolt-profile-generate"),
             llvm_bolt_profile_use: matches.opt_str("llvm-bolt-profile-use"),
         }
diff --git a/src/bootstrap/native.rs b/src/bootstrap/native.rs
@@ -320,11 +320,9 @@ impl Step for Llvm {
         // This flag makes sure `FileCheck` is copied in the final binaries directory.
         cfg.define("LLVM_INSTALL_UTILS", "ON");
 
-        if builder.config.llvm_profile_generate {
+        if let Some(path) = &builder.config.llvm_profile_generate {
             cfg.define("LLVM_BUILD_INSTRUMENTED", "IR");
-            if let Ok(llvm_profile_dir) = std::env::var("LLVM_PROFILE_DIR") {
-                cfg.define("LLVM_PROFILE_DATA_DIR", llvm_profile_dir);
-            }
+            cfg.define("LLVM_PROFILE_DATA_DIR", &path);
             cfg.define("LLVM_BUILD_RUNTIME", "No");
         }
         if let Some(path) = builder.config.llvm_profile_use.as_ref() {
@@ -822,7 +820,7 @@ impl Step for Lld {
         // when doing PGO on CI, cmake or clang-cl don't automatically link clang's
         // profiler runtime in. In that case, we need to manually ask cmake to do it, to avoid
         // linking errors, much like LLVM's cmake setup does in that situation.
-        if builder.config.llvm_profile_generate && target.contains("msvc") {
+        if builder.config.llvm_profile_generate.is_some() && target.contains("msvc") {
             if let Some(clang_cl_path) = builder.config.llvm_clang_cl.as_ref() {
                 // Find clang's runtime library directory and push that as a search path to the
                 // cmake linker flags.
diff --git a/src/ci/pgo.sh b/src/ci/pgo.sh
@@ -84,14 +84,14 @@ LLVM_PROFILE_DIRECTORY_ROOT=$PGO_TMP/llvm-pgo
 # separate phases. This increases build time -- though not by a huge amount --
 # but prevents any problems from arising due to different profiling runtimes
 # being simultaneously linked in.
-# LLVM IR PGO does not respect LLVM_PROFILE_FILE, so we have to set the profiling file
-# path through our custom environment variable. We include the PID in the directory path
-# to avoid updates to profile files being lost because of race conditions.
-LLVM_PROFILE_DIR=${LLVM_PROFILE_DIRECTORY_ROOT}/prof-%p python3 $CHECKOUT/x.py build \
+# LLVM IR PGO does not respect LLVM_PROFILE_FILE, so we have to set the profiling path
+# through the LLVM build system in src/bootstrap/native.rs. We include the PID in the
+# directory path to avoid updates to profile files being lost because of race conditions.
+python3 $CHECKOUT/x.py build \
     --target=$PGO_HOST \
     --host=$PGO_HOST \
     --stage 2 library/std \
-    --llvm-profile-generate
+    --llvm-profile-generate=${LLVM_PROFILE_DIRECTORY_ROOT}/prof-%p
 
 # Compile rustc-perf:
 # - get the expected commit source code: on linux, the Dockerfile downloads a source archive before
