Use CS PGO for LLVM

Author: Kobzol

src/bootstrap/llvm.rs

@@ -331,15 +331,24 @@ impl Step for Llvm {
         // This flag makes sure `FileCheck` is copied in the final binaries directory.
         cfg.define("LLVM_INSTALL_UTILS", "ON");
 
+        let mut cxxflags = vec![];
         if 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);
+            if std::env::var("LLVM_USE_CS_PGO").is_ok() {
+                cfg.define("LLVM_BUILD_INSTRUMENTED", "CSIR");
+                if let Ok(llvm_profile_dir) = std::env::var("LLVM_PROFILE_DIR") {
+                    cfg.define("LLVM_CSPROFILE_DATA_DIR", llvm_profile_dir);
+                }
+            } else {
+                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_BUILD_RUNTIME", "No");
         }
         if let Some(path) = builder.config.llvm_profile_use.as_ref() {
-            cfg.define("LLVM_PROFDATA_FILE", &path);
+            // cfg.define("LLVM_PROFDATA_FILE", &path);
+            cxxflags.push(&format!("-fprofile-use={path}"));
         }
         if builder.config.llvm_bolt_profile_generate
             || builder.config.llvm_bolt_profile_use.is_some()

src/ci/stage-build.py

@@ -48,6 +48,7 @@
 
 LLVM_BOLT_CRATES = LLVM_PGO_CRATES
 
+
 class Pipeline:
     # Paths
     def checkout_path(self) -> Path:
@@ -86,6 +87,9 @@ def opt_artifacts(self) -> Path:
     def llvm_profile_dir_root(self) -> Path:
         return self.opt_artifacts() / "llvm-pgo"
 
+    def llvm_profile_merged_file_intermediate(self) -> Path:
+        return self.opt_artifacts() / "llvm-pgo-intermediate.profdata"
+
     def llvm_profile_merged_file(self) -> Path:
         return self.opt_artifacts() / "llvm-pgo.profdata"
 
@@ -450,6 +454,7 @@ def cmd(
             )
     return subprocess.run(args, env=environment, check=True)
 
+
 class BenchmarkRunner:
     def run_rustc(self, pipeline: Pipeline):
         raise NotImplementedError
@@ -460,6 +465,7 @@ def run_llvm(self, pipeline: Pipeline):
     def run_bolt(self, pipeline: Pipeline):
         raise NotImplementedError
 
+
 class DefaultBenchmarkRunner(BenchmarkRunner):
     def run_rustc(self, pipeline: Pipeline):
         # Here we're profiling the `rustc` frontend, so we also include `Check`.
@@ -473,6 +479,7 @@ def run_rustc(self, pipeline: Pipeline):
                 LLVM_PROFILE_FILE=str(pipeline.rustc_profile_template_path())
             )
         )
+
     def run_llvm(self, pipeline: Pipeline):
         run_compiler_benchmarks(
             pipeline,
@@ -489,6 +496,7 @@ def run_bolt(self, pipeline: Pipeline):
             crates=LLVM_BOLT_CRATES
         )
 
+
 def run_compiler_benchmarks(
         pipeline: Pipeline,
         profiles: List[str],
@@ -622,7 +630,7 @@ def gather_llvm_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
 
     runner.run_llvm(pipeline)
 
-    profile_path = pipeline.llvm_profile_merged_file()
+    profile_path = pipeline.llvm_profile_merged_file_intermediate()
     LOGGER.info(f"Merging LLVM PGO profiles to {profile_path}")
     cmd([
         pipeline.downloaded_llvm_dir() / "bin" / "llvm-profdata",
@@ -642,13 +650,37 @@ def gather_llvm_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
     delete_directory(pipeline.llvm_profile_dir_root())
 
 
+def gather_llvm_cs_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
+    LOGGER.info("Running benchmarks with CS PGO instrumented LLVM")
+
+    runner.run_llvm(pipeline)
+
+    profile_path = pipeline.llvm_profile_merged_file()
+    LOGGER.info(f"Merging LLVM CS PGO profiles to {profile_path}")
+    cmd([
+        pipeline.downloaded_llvm_dir() / "bin" / "llvm-profdata",
+        "merge",
+        "-o", profile_path,
+        pipeline.llvm_profile_dir_root(),
+        pipeline.llvm_profile_merged_file_intermediate()
+    ])
+
+    LOGGER.info("LLVM CS PGO statistics")
+    LOGGER.info(f"{profile_path}: {format_bytes(get_path_size(profile_path))}")
+    LOGGER.info(
+        f"{pipeline.llvm_profile_dir_root()}: {format_bytes(get_path_size(pipeline.llvm_profile_dir_root()))}")
+    LOGGER.info(f"Profile file count: {count_files(pipeline.llvm_profile_dir_root())}")
+
+    # We don't need the individual .profraw files now that they have been merged
+    # into a final .profdata
+    delete_directory(pipeline.llvm_profile_dir_root())
+
+
 def gather_rustc_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
     LOGGER.info("Running benchmarks with PGO instrumented rustc")
 
-
     runner.run_rustc(pipeline)
 
-
     profile_path = pipeline.rustc_profile_merged_file()
     LOGGER.info(f"Merging Rustc PGO profiles to {profile_path}")
     cmd([
@@ -787,6 +819,27 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, runner: BenchmarkRu
         print_free_disk_space(pipeline)
 
     clear_llvm_files(pipeline)
+
+    # Stage 1b: Build rustc + CS PGO instrumented LLVM
+    with timer.section("Stage 1b (LLVM CS PGO)") as stage1b:
+        with stage1.section("Build rustc and LLVM") as rustc_build:
+            build_rustc(pipeline, args=[
+                "--llvm-profile-generate",
+                "--llvm-profile-use",
+                pipeline.llvm_profile_merged_file_intermediate()
+            ], env=dict(
+                LLVM_USE_CS_PGO="1",
+                LLVM_PROFILE_DIR=str(pipeline.llvm_profile_dir_root() / "prof-%p"),
+
+            ))
+            record_metrics(pipeline, rustc_build)
+
+        with stage1b.section("Gather profiles"):
+            gather_llvm_cs_profiles(pipeline, runner)
+        print_free_disk_space(pipeline)
+
+    clear_llvm_files(pipeline)
+
     final_build_args += [
         "--llvm-profile-use",
         pipeline.llvm_profile_merged_file()
@@ -865,6 +918,7 @@ def run(runner: BenchmarkRunner):
 
     print_binary_sizes(pipeline)
 
+
 if __name__ == "__main__":
     runner = DefaultBenchmarkRunner()
     run(runner)