PGO: Add test case for branch weights and function call counts recording.

5 files changed, 122 insertions, 0 deletions

diff --git a/src/test/run-make-fulldeps/pgo-branch-weights/Makefile b/src/test/run-make-fulldeps/pgo-branch-weights/Makefile
new file mode 100644
index 00000000000..edb4006313a
--- /dev/null
+++ b/src/test/run-make-fulldeps/pgo-branch-weights/Makefile
@@ -0,0 +1,35 @@
+# needs-profiler-support
+
+-include ../tools.mk
+
+# This test makes sure that instrumented binaries record the right counts for
+# functions being called and branches being taken. We run an instrumented binary
+# with an argument that causes a know path through the program and then check
+# that the expected counts get added to the use-phase LLVM IR.
+
+# LLVM doesn't support instrumenting binaries that use SEH:
+# https://github.com/rust-lang/rust/issues/61002
+#
+# Things work fine with -Cpanic=abort though.
+ifdef IS_MSVC
+COMMON_FLAGS=-Cpanic=abort
+endif
+
+all:
+	# We don't compile `opaque` with either optimizations or instrumentation.
+	# We don't compile `opaque` with either optimizations or instrumentation.
+	$(RUSTC) $(COMMON_FLAGS) opaque.rs
+	# Compile the test program with instrumentation
+	mkdir -p "$(TMPDIR)"/prof_data_dir
+	$(RUSTC) $(COMMON_FLAGS) interesting.rs \
+		-Cprofile-generate="$(TMPDIR)"/prof_data_dir -O -Ccodegen-units=1
+	$(RUSTC) $(COMMON_FLAGS) main.rs -Cprofile-generate="$(TMPDIR)"/prof_data_dir -O
+	# 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
+	$(RUSTC) $(COMMON_FLAGS) interesting.rs \
+		-Cprofile-use="$(TMPDIR)"/prof_data_dir/merged.profdata -O \
+		-Ccodegen-units=1 --emit=llvm-ir
+	cat "$(TMPDIR)"/interesting.ll | "$(LLVM_FILECHECK)" filecheck-patterns.txt
diff --git a/src/test/run-make-fulldeps/pgo-branch-weights/filecheck-patterns.txt b/src/test/run-make-fulldeps/pgo-branch-weights/filecheck-patterns.txt
new file mode 100644
index 00000000000..e68b35c0c23
--- /dev/null
+++ b/src/test/run-make-fulldeps/pgo-branch-weights/filecheck-patterns.txt
@@ -0,0 +1,24 @@
+
+# First, establish that certain !prof labels are attached to the expected
+# functions and branching instructions.
+
+CHECK: define void @function_called_twice(i32 %c) {{.*}} !prof !29 {
+CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !prof !30
+
+CHECK: define void @function_called_42_times(i32 %c) {{.*}} !prof !31 {
+CHECK:      switch i32 %c, label {{.*}} [
+CHECK-NEXT:     i32 97, label {{.*}}
+CHECK-NEXT:     i32 98, label {{.*}}
+CHECK-NEXT: ], !prof !32
+
+CHECK: define void @function_called_never(i32 {{.*}} !prof !33 {
+
+
+
+# Now check that those !prof tags hold the expected counts
+
+CHECK: !29 = !{!"function_entry_count", i64 2}
+CHECK: !30 = !{!"branch_weights", i32 2, i32 0}
+CHECK: !31 = !{!"function_entry_count", i64 42}
+CHECK: !32 = !{!"branch_weights", i32 2, i32 12, i32 28}
+CHECK: !33 = !{!"function_entry_count", i64 0}
diff --git a/src/test/run-make-fulldeps/pgo-branch-weights/interesting.rs b/src/test/run-make-fulldeps/pgo-branch-weights/interesting.rs
new file mode 100644
index 00000000000..a26d6fd69d1
--- /dev/null
+++ b/src/test/run-make-fulldeps/pgo-branch-weights/interesting.rs
@@ -0,0 +1,40 @@
+#![crate_name="interesting"]
+#![crate_type="rlib"]
+
+extern crate opaque;
+
+#[no_mangle]
+#[inline(never)]
+pub fn function_called_twice(c: char) {
+    if c == '2' {
+        // This branch is taken twice
+        opaque::f1();
+    } else {
+        // This branch is never taken
+        opaque::f2();
+    }
+}
+
+#[no_mangle]
+#[inline(never)]
+pub fn function_called_42_times(c: char) {
+    if c == 'a' {
+        // This branch is taken 12 times
+        opaque::f1();
+    } else {
+
+        if c == 'b' {
+            // This branch is taken 28 times
+            opaque::f2();
+        } else {
+            // This branch is taken 2 times
+            opaque::f3();
+        }
+    }
+}
+
+#[no_mangle]
+#[inline(never)]
+pub fn function_called_never(_: char) {
+    opaque::f1();
+}
diff --git a/src/test/run-make-fulldeps/pgo-branch-weights/main.rs b/src/test/run-make-fulldeps/pgo-branch-weights/main.rs
new file mode 100644
index 00000000000..619cf9c698c
--- /dev/null
+++ b/src/test/run-make-fulldeps/pgo-branch-weights/main.rs
@@ -0,0 +1,17 @@
+extern crate interesting;
+
+fn main() {
+    let arg = std::env::args().skip(1).next().unwrap();
+
+    for c in arg.chars() {
+        if c == '2' {
+            interesting::function_called_twice(c);
+        } else {
+            interesting::function_called_42_times(c);
+        }
+
+        if c == '0' {
+            interesting::function_called_never(c);
+        }
+    }
+}
diff --git a/src/test/run-make-fulldeps/pgo-branch-weights/opaque.rs b/src/test/run-make-fulldeps/pgo-branch-weights/opaque.rs
new file mode 100644
index 00000000000..72f93c9feab
--- /dev/null
+++ b/src/test/run-make-fulldeps/pgo-branch-weights/opaque.rs
@@ -0,0 +1,6 @@
+#![crate_name="opaque"]
+#![crate_type="rlib"]
+
+pub fn f1() {}
+pub fn f2() {}
+pub fn f3() {}