Auto merge of #125165 - Oneirical:pgo-branch-weights, r=jieyouxu

Migrate `run-make/pgo-branch-weights` to `rmake`

Part of #121876 and the associated [Google Summer of Code project](https://blog.rust-lang.org/2024/05/01/gsoc-2024-selected-projects.html).

This is a scary one and I expect things to break. Set as draft, because this isn't ready.

- [x] There is this comment here, which suggests the test is excluded from the testing process due to a platform specific issue? I can't see anything here that would cause this test to not run...
> // FIXME(mati865): MinGW GCC miscompiles compiler-rt profiling library but with Clang it works
// properly. Since we only have GCC on the CI ignore the test for now."

EDIT: This is specific to Windows-gnu.

- [x] The Makefile has this line:
```
ifneq (,$(findstring x86,$(TARGET)))
COMMON_FLAGS=-Clink-args=-fuse-ld=gold
```
I honestly can't tell whether this is checking if the target IS x86, or IS NOT. EDIT: It's checking if it IS x86.

- [x] I don't know why the Makefile was trying to pass an argument directly in the Makefile instead of setting that "aaaaaaaaaaaa2bbbbbbbbbbbb2bbbbbbbbbbbbbbbbcc" input as a variable in the Rust program directly. I changed that, let me know if that was wrong.

- [x] Trying to rewrite `cat "$(TMPDIR)/interesting.ll" | "$(LLVM_FILECHECK)" filecheck-patterns.txt` resulted in some butchery. For starters, in `tools.mk`, LLVM_FILECHECK corrects its own backslashes on Windows distributions, but there is no further mention of it, so I assume this is a preset environment variable... but is it really? Then, the command itself uses a Standard Input and a passed input file as an argument simultaneously, according to the [documentation](https://llvm.org/docs/CommandGuide/FileCheck.html#synopsis).

try-job: aarch64-gnu

7 files changed, 215 insertions, 42 deletions

diff --git a/src/tools/run-make-support/src/lib.rs b/src/tools/run-make-support/src/lib.rs
index b920f9a07db..b7a936a1e11 100644
--- a/src/tools/run-make-support/src/lib.rs
+++ b/src/tools/run-make-support/src/lib.rs
@@ -9,7 +9,7 @@ mod command;
 pub mod diff;
 mod drop_bomb;
 pub mod fs_wrapper;
-pub mod llvm_readobj;
+pub mod llvm;
 pub mod run;
 pub mod rustc;
 pub mod rustdoc;
@@ -29,8 +29,10 @@ pub use wasmparser;
 pub use cc::{cc, extra_c_flags, extra_cxx_flags, Cc};
 pub use clang::{clang, Clang};
 pub use diff::{diff, Diff};
-pub use llvm_readobj::{llvm_readobj, LlvmReadobj};
-pub use run::{cmd, run, run_fail};
+pub use llvm::{
+    llvm_filecheck, llvm_profdata, llvm_readobj, LlvmFilecheck, LlvmProfdata, LlvmReadobj,
+};
+pub use run::{cmd, run, run_fail, run_with_args};
 pub use rustc::{aux_build, rustc, Rustc};
 pub use rustdoc::{bare_rustdoc, rustdoc, Rustdoc};
 
diff --git a/src/tools/run-make-support/src/llvm.rs b/src/tools/run-make-support/src/llvm.rs
new file mode 100644
index 00000000000..414251abda2
--- /dev/null
+++ b/src/tools/run-make-support/src/llvm.rs
@@ -0,0 +1,127 @@
+use std::path::{Path, PathBuf};
+
+use crate::{env_var, Command};
+
+/// Construct a new `llvm-readobj` invocation. This assumes that `llvm-readobj` is available
+/// at `$LLVM_BIN_DIR/llvm-readobj`.
+pub fn llvm_readobj() -> LlvmReadobj {
+    LlvmReadobj::new()
+}
+
+/// Construct a new `llvm-profdata` invocation. This assumes that `llvm-profdata` is available
+/// at `$LLVM_BIN_DIR/llvm-profdata`.
+pub fn llvm_profdata() -> LlvmProfdata {
+    LlvmProfdata::new()
+}
+
+/// Construct a new `llvm-filecheck` invocation. This assumes that `llvm-filecheck` is available
+/// at `$LLVM_FILECHECK`.
+pub fn llvm_filecheck() -> LlvmFilecheck {
+    LlvmFilecheck::new()
+}
+
+/// A `llvm-readobj` invocation builder.
+#[derive(Debug)]
+#[must_use]
+pub struct LlvmReadobj {
+    cmd: Command,
+}
+
+/// A `llvm-profdata` invocation builder.
+#[derive(Debug)]
+#[must_use]
+pub struct LlvmProfdata {
+    cmd: Command,
+}
+
+/// A `llvm-filecheck` invocation builder.
+#[derive(Debug)]
+#[must_use]
+pub struct LlvmFilecheck {
+    cmd: Command,
+}
+
+crate::impl_common_helpers!(LlvmReadobj);
+crate::impl_common_helpers!(LlvmProfdata);
+crate::impl_common_helpers!(LlvmFilecheck);
+
+/// Generate the path to the bin directory of LLVM.
+#[must_use]
+pub fn llvm_bin_dir() -> PathBuf {
+    let llvm_bin_dir = env_var("LLVM_BIN_DIR");
+    PathBuf::from(llvm_bin_dir)
+}
+
+impl LlvmReadobj {
+    /// Construct a new `llvm-readobj` invocation. This assumes that `llvm-readobj` is available
+    /// at `$LLVM_BIN_DIR/llvm-readobj`.
+    pub fn new() -> Self {
+        let llvm_readobj = llvm_bin_dir().join("llvm-readobj");
+        let cmd = Command::new(llvm_readobj);
+        Self { cmd }
+    }
+
+    /// Provide an input file.
+    pub fn input<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
+        self.cmd.arg(path.as_ref());
+        self
+    }
+
+    /// Pass `--file-header` to display file headers.
+    pub fn file_header(&mut self) -> &mut Self {
+        self.cmd.arg("--file-header");
+        self
+    }
+}
+
+impl LlvmProfdata {
+    /// Construct a new `llvm-profdata` invocation. This assumes that `llvm-profdata` is available
+    /// at `$LLVM_BIN_DIR/llvm-profdata`.
+    pub fn new() -> Self {
+        let llvm_profdata = llvm_bin_dir().join("llvm-profdata");
+        let cmd = Command::new(llvm_profdata);
+        Self { cmd }
+    }
+
+    /// Provide an input file.
+    pub fn input<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
+        self.cmd.arg(path.as_ref());
+        self
+    }
+
+    /// Specify the output file path.
+    pub fn output<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
+        self.cmd.arg("-o");
+        self.cmd.arg(path.as_ref());
+        self
+    }
+
+    /// Take several profile data files generated by PGO instrumentation and merge them
+    /// together into a single indexed profile data file.
+    pub fn merge(&mut self) -> &mut Self {
+        self.cmd.arg("merge");
+        self
+    }
+}
+
+impl LlvmFilecheck {
+    /// Construct a new `llvm-filecheck` invocation. This assumes that `llvm-filecheck` is available
+    /// at `$LLVM_FILECHECK`.
+    pub fn new() -> Self {
+        let llvm_filecheck = env_var("LLVM_FILECHECK");
+        let cmd = Command::new(llvm_filecheck);
+        Self { cmd }
+    }
+
+    /// Pipe a read file into standard input containing patterns that will be matched against the .patterns(path) call.
+    pub fn stdin<I: AsRef<[u8]>>(&mut self, input: I) -> &mut Self {
+        self.cmd.set_stdin(input.as_ref().to_vec().into_boxed_slice());
+        self
+    }
+
+    /// Provide the patterns that need to be matched.
+    pub fn patterns<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
+        self.cmd.arg(path.as_ref());
+        self
+    }
+}
diff --git a/src/tools/run-make-support/src/run.rs b/src/tools/run-make-support/src/run.rs
index 6fa1a75363c..54730bb7de7 100644
--- a/src/tools/run-make-support/src/run.rs
+++ b/src/tools/run-make-support/src/run.rs
@@ -9,12 +9,17 @@ use crate::{cwd, env_var, is_windows, set_host_rpath};
 use super::handle_failed_output;
 
 #[track_caller]
-fn run_common(name: &str) -> Command {
+fn run_common(name: &str, args: Option<&[&str]>) -> Command {
     let mut bin_path = PathBuf::new();
     bin_path.push(cwd());
     bin_path.push(name);
     let ld_lib_path_envvar = env_var("LD_LIB_PATH_ENVVAR");
     let mut cmd = Command::new(bin_path);
+    if let Some(args) = args {
+        for arg in args {
+            cmd.arg(arg);
+        }
+    }
     cmd.env(&ld_lib_path_envvar, {
         let mut paths = vec![];
         paths.push(cwd());
@@ -43,7 +48,19 @@ fn run_common(name: &str) -> Command {
 #[track_caller]
 pub fn run(name: &str) -> CompletedProcess {
     let caller = panic::Location::caller();
-    let mut cmd = run_common(name);
+    let mut cmd = run_common(name, None);
+    let output = cmd.run();
+    if !output.status().success() {
+        handle_failed_output(&cmd, output, caller.line());
+    }
+    output
+}
+
+/// Run a built binary with one or more argument(s) and make sure it succeeds.
+#[track_caller]
+pub fn run_with_args(name: &str, args: &[&str]) -> CompletedProcess {
+    let caller = panic::Location::caller();
+    let mut cmd = run_common(name, Some(args));
     let output = cmd.run();
     if !output.status().success() {
         handle_failed_output(&cmd, output, caller.line());
@@ -55,7 +72,7 @@ pub fn run(name: &str) -> CompletedProcess {
 #[track_caller]
 pub fn run_fail(name: &str) -> CompletedProcess {
     let caller = panic::Location::caller();
-    let mut cmd = run_common(name);
+    let mut cmd = run_common(name, None);
     let output = cmd.run_fail();
     if output.status().success() {
         handle_failed_output(&cmd, output, caller.line());
diff --git a/src/tools/run-make-support/src/rustc.rs b/src/tools/run-make-support/src/rustc.rs
index 13cec0cb6be..c27af966f39 100644
--- a/src/tools/run-make-support/src/rustc.rs
+++ b/src/tools/run-make-support/src/rustc.rs
@@ -152,6 +152,24 @@ impl Rustc {
         self
     }
 
+    /// Specify directory path used for profile generation
+    pub fn profile_generate<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
+        let mut arg = OsString::new();
+        arg.push("-Cprofile-generate=");
+        arg.push(path.as_ref());
+        self.cmd.arg(&arg);
+        self
+    }
+
+    /// Specify directory path used for profile usage
+    pub fn profile_use<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
+        let mut arg = OsString::new();
+        arg.push("-Cprofile-use=");
+        arg.push(path.as_ref());
+        self.cmd.arg(&arg);
+        self
+    }
+
     /// Specify error format to use
     pub fn error_format(&mut self, format: &str) -> &mut Self {
         self.cmd.arg(format!("--error-format={format}"));
diff --git a/src/tools/tidy/src/allowed_run_make_makefiles.txt b/src/tools/tidy/src/allowed_run_make_makefiles.txt
index 87185eef952..d4c3a970c70 100644
--- a/src/tools/tidy/src/allowed_run_make_makefiles.txt
+++ b/src/tools/tidy/src/allowed_run_make_makefiles.txt
@@ -168,7 +168,6 @@ run-make/pass-linker-flags/Makefile
 run-make/pass-non-c-like-enum-to-c/Makefile
 run-make/pdb-alt-path/Makefile
 run-make/pdb-buildinfo-cl-cmd/Makefile
-run-make/pgo-branch-weights/Makefile
 run-make/pgo-gen-lto/Makefile
 run-make/pgo-gen-no-imp-symbols/Makefile
 run-make/pgo-gen/Makefile
diff --git a/tests/run-make/pgo-branch-weights/Makefile b/tests/run-make/pgo-branch-weights/Makefile
deleted file mode 100644
index 4c9f8b2493a..00000000000
--- a/tests/run-make/pgo-branch-weights/Makefile
+++ /dev/null
@@ -1,35 +0,0 @@
-# needs-profiler-support
-# ignore-windows-gnu
-# ignore-cross-compile
-
-# FIXME(mati865): MinGW GCC miscompiles compiler-rt profiling library but with Clang it works
-# properly. Since we only have GCC on the CI ignore the test for now.
-
-include ../tools.mk
-
-# For some very small programs GNU ld seems to not properly handle
-# instrumentation sections correctly. Neither Gold nor LLD have that problem.
-ifeq ($(UNAME),Linux)
-ifneq (,$(findstring x86,$(TARGET)))
-COMMON_FLAGS=-Clink-args=-fuse-ld=gold
-endif
-endif
-
-
-all:
-	# We don't compile `opaque` with either optimizations or instrumentation.
-	$(RUSTC) $(COMMON_FLAGS) opaque.rs || exit 1
-	# Compile the test program with instrumentation
-	mkdir -p "$(TMPDIR)/prof_data_dir" || exit 1
-	$(RUSTC) $(COMMON_FLAGS) interesting.rs \
-		-Cprofile-generate="$(TMPDIR)/prof_data_dir" -O -Ccodegen-units=1 || exit 1
-	$(RUSTC) $(COMMON_FLAGS) main.rs -Cprofile-generate="$(TMPDIR)/prof_data_dir" -O || exit 1
-	# The argument below generates to the expected branch weights
-	$(call RUN,main aaaaaaaaaaaa2bbbbbbbbbbbb2bbbbbbbbbbbbbbbbcc) || exit 1
-	"$(LLVM_BIN_DIR)/llvm-profdata" merge \
-		-o "$(TMPDIR)/prof_data_dir/merged.profdata" \
-		"$(TMPDIR)/prof_data_dir" || exit 1
-	$(RUSTC) $(COMMON_FLAGS) interesting.rs \
-		-Cprofile-use="$(TMPDIR)/prof_data_dir/merged.profdata" -O \
-		-Ccodegen-units=1 --emit=llvm-ir || exit 1
-	cat "$(TMPDIR)/interesting.ll" | "$(LLVM_FILECHECK)" filecheck-patterns.txt
diff --git a/tests/run-make/pgo-branch-weights/rmake.rs b/tests/run-make/pgo-branch-weights/rmake.rs
new file mode 100644
index 00000000000..55f6e7e56c5
--- /dev/null
+++ b/tests/run-make/pgo-branch-weights/rmake.rs
@@ -0,0 +1,45 @@
+// This test generates an instrumented binary - a program which
+// will keep track of how many times it calls each function, a useful
+// feature for optimization. Then, an argument (aaaaaaaaaaaa2bbbbbbbbbbbb2bbbbbbbbbbbbbbbbcc)
+// is passed into the instrumented binary, which should react with a number of function calls
+// fully known in advance. (For example, the letter 'a' results in calling f1())
+
+// If the test passes, the expected function call count was added to the use-phase LLVM-IR.
+// See https://github.com/rust-lang/rust/pull/66631
+
+//@ needs-profiler-support
+//@ ignore-cross-compile
+
+// FIXME(Oneirical): This test has problems generating profdata on mingw.
+// For more information, see https://github.com/rust-lang/rust/pull/122613
+//@ ignore-windows-gnu
+
+use run_make_support::{fs_wrapper, llvm_filecheck, llvm_profdata, run_with_args, rustc};
+use std::path::Path;
+
+fn main() {
+    let path_prof_data_dir = Path::new("prof_data_dir");
+    let path_merged_profdata = path_prof_data_dir.join("merged.profdata");
+    rustc().input("opaque.rs").run();
+    fs_wrapper::create_dir_all(&path_prof_data_dir);
+    rustc()
+        .input("interesting.rs")
+        .profile_generate(&path_prof_data_dir)
+        .opt()
+        .codegen_units(1)
+        .run();
+    rustc().input("main.rs").profile_generate(&path_prof_data_dir).opt().run();
+    run_with_args("main", &["aaaaaaaaaaaa2bbbbbbbbbbbb2bbbbbbbbbbbbbbbbcc"]);
+    llvm_profdata().merge().output(&path_merged_profdata).input(path_prof_data_dir).run();
+    rustc()
+        .input("interesting.rs")
+        .profile_use(path_merged_profdata)
+        .opt()
+        .codegen_units(1)
+        .emit("llvm-ir")
+        .run();
+    llvm_filecheck()
+        .patterns("filecheck-patterns.txt")
+        .stdin(fs_wrapper::read("interesting.ll"))
+        .run();
+}