[memprof] Add flag to control profile dump at exit

[ellishg]

Add the dump_at_exit flag to control whether or not profiles should be dumped when the program exits. Since we can call __memprof_profile_dump() directly, we don't necessarily need to dump profiles at exit.

[llvmbot]

@llvm/pr-subscribers-pgo

Author: Ellis Hoag (ellishg)

Changes

Add the dump_at_exit flag to control whether or not profiles should be dumped when the program exits. Since we can call __memprof_profile_dump() directly, we don't necessarily need to dump profiles at exit.

---

Full diff: https://github.com/llvm/llvm-project/pull/119452.diff

3 Files Affected:

• (modified) compiler-rt/lib/memprof/memprof_allocator.cpp (+2-1)
• (modified) compiler-rt/lib/memprof/memprof_flags.inc (+3-1)
• (added) compiler-rt/test/memprof/TestCases/dump_at_exit.cpp (+16)

diff --git a/compiler-rt/lib/memprof/memprof_allocator.cpp b/compiler-rt/lib/memprof/memprof_allocator.cpp
index 19b2b901068246..c3448c22532bb8 100644
--- a/compiler-rt/lib/memprof/memprof_allocator.cpp
+++ b/compiler-rt/lib/memprof/memprof_allocator.cpp
@@ -301,7 +301,8 @@ struct Allocator {
 
   ~Allocator() {
     atomic_store_relaxed(&destructing, 1);
-    FinishAndWrite();
+    if (flags()->dump_at_exit)
+      FinishAndWrite();
   }
 
   static void PrintCallback(const uptr Key, LockedMemInfoBlock *const &Value,
diff --git a/compiler-rt/lib/memprof/memprof_flags.inc b/compiler-rt/lib/memprof/memprof_flags.inc
index 7c5dc091f79351..3a40ba0ab5498b 100644
--- a/compiler-rt/lib/memprof/memprof_flags.inc
+++ b/compiler-rt/lib/memprof/memprof_flags.inc
@@ -38,4 +38,6 @@ MEMPROF_FLAG(bool, allocator_frees_and_returns_null_on_realloc_zero, true,
 MEMPROF_FLAG(bool, print_text, false,
   "If set, prints the heap profile in text format. Else use the raw binary serialization format.")
 MEMPROF_FLAG(bool, print_terse, false,
-             "If set, prints memory profile in a terse format. Only applicable if print_text = true.")
\ No newline at end of file
+             "If set, prints memory profile in a terse format. Only applicable if print_text = true.")
+MEMPROF_FLAG(bool, dump_at_exit, true,
+             "If set, dump profiles when the program terminates.")
diff --git a/compiler-rt/test/memprof/TestCases/dump_at_exit.cpp b/compiler-rt/test/memprof/TestCases/dump_at_exit.cpp
new file mode 100644
index 00000000000000..426849d1cea01e
--- /dev/null
+++ b/compiler-rt/test/memprof/TestCases/dump_at_exit.cpp
@@ -0,0 +1,16 @@
+// RUN: %clangxx_memprof %s -o %t
+
+// RUN: %env_memprof_opts=print_text=true:log_path=stdout:dump_at_exit=false %run %t | count 0
+// RUN: %env_memprof_opts=print_text=true:log_path=stdout:dump_at_exit=true %run %t | FileCheck %s
+
+#include <stdlib.h>
+#include <string.h>
+
+int main() {
+  char *x = (char *)malloc(10);
+  memset(x, 0, 10);
+  free(x);
+  return 0;
+}
+
+// CHECK: Recorded MIBs

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[snehasish]

Can you describe the scenario where this is useful?

[snehasish]

What is count 0 here?

[ellishg]

count N will pass if it reads N lines from stdin. So this just checks that there is no output.

[snehasish]

Ah count is a utility we bundle with LLVM for testing. TIL.

[ellishg]

Can you describe the scenario where this is useful?

If we manually call __memprof_profile_dump(), we don't want a second profile to be automatically dumped when the program exits.

[snehasish]

lgtm