rewrite pgo-branch-weights (1/2)

Author: Oneirical

src/tools/run-make-support/src/lib.rs

@@ -6,7 +6,7 @@
 pub mod cc;
 pub mod clang;
 pub mod diff;
-pub mod llvm_readobj;
+pub mod llvm;
 pub mod run;
 pub mod rustc;
 pub mod rustdoc;
@@ -23,7 +23,7 @@ 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 llvm::{llvm_profdata, llvm_readobj, LlvmProfdata, LlvmReadobj};
 pub use run::{run, run_fail};
 pub use rustc::{aux_build, rustc, Rustc};
 pub use rustdoc::{bare_rustdoc, rustdoc, Rustdoc};

src/tools/run-make-support/src/llvm.rs

@@ -0,0 +1,102 @@
+use std::env;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+
+use crate::handle_failed_output;
+
+/// 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()
+}
+
+/// A `llvm-readobj` invocation builder.
+#[derive(Debug)]
+pub struct LlvmReadobj {
+    cmd: Command,
+}
+
+/// A `llvm-profdata` invocation builder.
+#[derive(Debug)]
+pub struct LlvmProfdata {
+    cmd: Command,
+}
+
+crate::impl_common_helpers!(LlvmReadobj);
+
+/// Generate the path to the bin directory of LLVM.
+pub fn llvm_bin_dir() -> PathBuf {
+    let llvm_bin_dir = env::var("LLVM_BIN_DIR")
+    .expect("`LLVM_BIN_DIR` not specified, but this is required to find `llvm-readobj`");
+    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
+    }
+
+    /// Get the [`Output`][::std::process::Output] of the finished process.
+    #[track_caller]
+    pub fn command_output(&mut self) -> ::std::process::Output {
+        self.cmd.output().expect("failed to get output of finished process")
+    }
+}
+
+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("-o");
+        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(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
+    }
+
+    /// Get the [`Output`][::std::process::Output] of the finished process.
+    #[track_caller]
+    pub fn command_output(&mut self) -> ::std::process::Output {
+        self.cmd.output().expect("failed to get output of finished process")
+    }
+}

src/tools/run-make-support/src/llvm_readobj.rs

@@ -1,6 +1,9 @@
 use std::env;
 use std::path::{Path, PathBuf};
 use std::process::{Command, Output};
+use std::io::{BufReader, Read, Write};
+use std::fs::File;
+use std::ffi::OsStr;
 
 use crate::is_windows;
 
@@ -37,6 +40,38 @@ fn run_common(name: &str) -> (Command, Output) {
     (cmd, output)
 }
 
+fn read_bytes<R: Read>(mut reader: R) -> Result<Vec<u8>, std::io::Error> {
+    let mut buffer = Vec::new();
+    reader.read_to_end(&mut buffer)?;
+    Ok(buffer)
+}
+
+/// Run a command taking in standard input from a file, and passing an argument
+/// referring to a different file.
+pub fn stdin_command<P>(command: &str, stdin: P, argument: P)
+    where
+    P: AsRef<Path> + AsRef<OsStr> + std::convert::AsRef<std::ffi::OsStr>,
+{
+    let file = File::open(stdin).unwrap();
+    let reader = BufReader::new(file);
+
+    let mut child = Command::new(command)
+       .stdin(std::process::Stdio::piped())
+       .stdout(std::process::Stdio::piped())
+       .stderr(std::process::Stdio::piped())
+       .arg(argument)
+       .spawn()
+       .unwrap();
+
+    let mut child_stdin = child.stdin.take().unwrap();
+    let byte_vec = read_bytes(reader).expect("failed to read bytes of standard input");
+    let byte_slice = byte_vec.as_slice();
+    child_stdin.write_all(byte_slice).unwrap();
+    child_stdin.flush().unwrap();
+
+    child.wait_with_output().unwrap();
+}
+
 /// Run a built binary and make sure it succeeds.
 #[track_caller]
 pub fn run(name: &str) -> Output {

src/tools/run-make-support/src/run.rs

@@ -125,6 +125,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}"));

src/tools/run-make-support/src/rustc.rs

@@ -199,7 +199,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

src/tools/tidy/src/allowed_run_make_makefiles.txt

@@ -0,0 +1,67 @@
+// 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
+
+// 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.
+
+//@ needs-profiler-support
+//@ ignore-windows-gnu
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_profdata, rustc, rustdoc, stdin_command, target, tmp_dir};
+
+fn main() {
+    // For some very small programs GNU ld seems to not properly handle
+    // instrumentation sections correctly. Neither Gold nor LLD have that
+    // problem.
+    let link_args = (cfg!(target_os = "linux") && target().contains("x86"))
+        .then_some("-Clink-args=-fuse-ld=gold");
+
+    let path_prof_data_dir = tmp_dir().join("prof_data_dir");
+    let path_merged_profdata = path_prof_data_dir.join("merged.profdata");
+    let invoc = rustc().input("opaque.rs");
+    if let Some(link_args) = link_args {
+        invoc.arg(link_args);
+    }
+    invoc.run();
+    fs::create_dir_all(&path_prof_data_dir);
+    let invoc = rustc()
+        .input("interesting.rs")
+        .profile_generate(&path_prof_data_dir)
+        .opt()
+        .codegen_units(1);
+    if let Some(link_args) = link_args {
+        invoc.arg(link_args);
+    }
+    invoc.run();
+    let invoc = rustc().input("main.rs").profile_generate(&path_prof_data_dir).opt();
+    if let Some(link_args) = link_args {
+        invoc.arg(link_args);
+    }
+    invoc.run();
+    run("main aaaaaaaaaaaa2bbbbbbbbbbbb2bbbbbbbbbbbbbbbbcc");
+    llvm_profdata()
+        .merge()
+        .output(&path_merged_profdata)
+        .input(path_prof_data_dir)
+        .command_output();
+    let invoc = rustc()
+        .input("interesting.rs")
+        .profile_use(path_merged_profdata)
+        .opt()
+        .codegen_units(1)
+        .emit("llvm-ir");
+    if let Some(link_args) = link_args {
+        invoc.arg(link_args);
+    }
+    invoc.run();
+
+    let interesting_ll = tmp_dir().join("interesting.ll");
+    stdin_command(env::var("LLVM_FILECHECK"), interesting_ll, "filecheck-patterns.txt");
+}