support PGO on custom project

src/ci/stage-build.py

@@ -48,7 +48,6 @@
 
 LLVM_BOLT_CRATES = LLVM_PGO_CRATES
 
-
 class Pipeline:
     # Paths
     def checkout_path(self) -> Path:
@@ -451,6 +450,44 @@ def cmd(
             )
     return subprocess.run(args, env=environment, check=True)
 
+class BenchmarkRunner:
+    def run_rustc(self, pipeline: Pipeline):
+        raise NotImplementedError
+
+    def run_llvm(self, pipeline: Pipeline):
+        raise NotImplementedError
+
+    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`.
+        # The benchmark set includes various stress tests that put the frontend under pressure.
+        run_compiler_benchmarks(
+            pipeline,
+            profiles=["Check", "Debug", "Opt"],
+            scenarios=["All"],
+            crates=RUSTC_PGO_CRATES,
+            env=dict(
+                LLVM_PROFILE_FILE=str(pipeline.rustc_profile_template_path())
+            )
+        )
+    def run_llvm(self, pipeline: Pipeline):
+        run_compiler_benchmarks(
+            pipeline,
+            profiles=["Debug", "Opt"],
+            scenarios=["Full"],
+            crates=LLVM_PGO_CRATES
+        )
+
+    def run_bolt(self, pipeline: Pipeline):
+        run_compiler_benchmarks(
+            pipeline,
+            profiles=["Check", "Debug", "Opt"],
+            scenarios=["Full"],
+            crates=LLVM_BOLT_CRATES
+        )
 
 def run_compiler_benchmarks(
         pipeline: Pipeline,
@@ -580,14 +617,10 @@ def create_pipeline() -> Pipeline:
         raise Exception(f"Optimized build is not supported for platform {sys.platform}")
 
 
-def gather_llvm_profiles(pipeline: Pipeline):
+def gather_llvm_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
     LOGGER.info("Running benchmarks with PGO instrumented LLVM")
-    run_compiler_benchmarks(
-        pipeline,
-        profiles=["Debug", "Opt"],
-        scenarios=["Full"],
-        crates=LLVM_PGO_CRATES
-    )
+
+    runner.run_llvm(pipeline)
 
     profile_path = pipeline.llvm_profile_merged_file()
     LOGGER.info(f"Merging LLVM PGO profiles to {profile_path}")
@@ -609,20 +642,12 @@ def gather_llvm_profiles(pipeline: Pipeline):
     delete_directory(pipeline.llvm_profile_dir_root())
 
 
-def gather_rustc_profiles(pipeline: Pipeline):
+def gather_rustc_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
     LOGGER.info("Running benchmarks with PGO instrumented rustc")
 
-    # 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.
-    run_compiler_benchmarks(
-        pipeline,
-        profiles=["Check", "Debug", "Opt"],
-        scenarios=["All"],
-        crates=RUSTC_PGO_CRATES,
-        env=dict(
-            LLVM_PROFILE_FILE=str(pipeline.rustc_profile_template_path())
-        )
-    )
+
+    runner.run_rustc(pipeline)
+
 
     profile_path = pipeline.rustc_profile_merged_file()
     LOGGER.info(f"Merging Rustc PGO profiles to {profile_path}")
@@ -644,14 +669,10 @@ def gather_rustc_profiles(pipeline: Pipeline):
     delete_directory(pipeline.rustc_profile_dir_root())
 
 
-def gather_llvm_bolt_profiles(pipeline: Pipeline):
+def gather_llvm_bolt_profiles(pipeline: Pipeline, runner: BenchmarkRunner):
     LOGGER.info("Running benchmarks with BOLT instrumented LLVM")
-    run_compiler_benchmarks(
-        pipeline,
-        profiles=["Check", "Debug", "Opt"],
-        scenarios=["Full"],
-        crates=LLVM_BOLT_CRATES
-    )
+
+    runner.run_bolt(pipeline)
 
     merged_profile_path = pipeline.llvm_bolt_profile_merged_file()
     profile_files_path = Path("/tmp/prof.fdata")
@@ -744,7 +765,7 @@ def record_metrics(pipeline: Pipeline, timer: Timer):
     log_metrics(metrics)
 
 
-def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: List[str]):
+def execute_build_pipeline(timer: Timer, pipeline: Pipeline, runner: BenchmarkRunner, final_build_args: List[str]):
     # Clear and prepare tmp directory
     shutil.rmtree(pipeline.opt_artifacts(), ignore_errors=True)
     os.makedirs(pipeline.opt_artifacts(), exist_ok=True)
@@ -762,7 +783,7 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: L
             record_metrics(pipeline, rustc_build)
 
         with stage1.section("Gather profiles"):
-            gather_llvm_profiles(pipeline)
+            gather_llvm_profiles(pipeline, runner)
         print_free_disk_space(pipeline)
 
     clear_llvm_files(pipeline)
@@ -781,7 +802,7 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: L
             record_metrics(pipeline, rustc_build)
 
         with stage2.section("Gather profiles"):
-            gather_rustc_profiles(pipeline)
+            gather_rustc_profiles(pipeline, runner)
         print_free_disk_space(pipeline)
 
     clear_llvm_files(pipeline)
@@ -804,7 +825,7 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: L
                 record_metrics(pipeline, rustc_build)
 
             with stage3.section("Gather profiles"):
-                gather_llvm_bolt_profiles(pipeline)
+                gather_llvm_bolt_profiles(pipeline, runner)
 
         # LLVM is not being cleared here, we want to reuse the previous build
         print_free_disk_space(pipeline)
@@ -819,7 +840,7 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: L
         record_metrics(pipeline, stage4)
 
 
-if __name__ == "__main__":
+def run(runner: BenchmarkRunner):
     logging.basicConfig(
         level=logging.DEBUG,
         format="%(name)s %(levelname)-4s: %(message)s",
@@ -832,8 +853,9 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: L
 
     timer = Timer()
     pipeline = create_pipeline()
+
     try:
-        execute_build_pipeline(timer, pipeline, build_args)
+        execute_build_pipeline(timer, pipeline, runner, build_args)
     except BaseException as e:
         LOGGER.error("The multi-stage build has failed")
         raise e
@@ -842,3 +864,7 @@ def execute_build_pipeline(timer: Timer, pipeline: Pipeline, final_build_args: L
         print_free_disk_space(pipeline)
 
     print_binary_sizes(pipeline)
+
+if __name__ == "__main__":
+    runner = DefaultBenchmarkRunner()
+    run(runner)