From f4e27d8aeec3595130362fcb0be99bf86aa74ce8 Mon Sep 17 00:00:00 2001
From: "Jonathan R. Madsen" <jrmadsen@users.noreply.github.com>
Date: Wed, 23 Mar 2022 00:28:13 -0500
Subject: [PATCH] Fixes for ROCM-SMI + MPI (#34)

---
 source/lib/omnitrace/src/library.cpp          |  1 +
 .../src/library/components/rocm_smi.cpp       |  8 +++++--
 source/lib/omnitrace/src/library/config.cpp   | 21 ++++++++++++-------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/source/lib/omnitrace/src/library.cpp b/source/lib/omnitrace/src/library.cpp
index e8950d7d44..d0783137e7 100644
--- a/source/lib/omnitrace/src/library.cpp
+++ b/source/lib/omnitrace/src/library.cpp
@@ -1060,6 +1060,7 @@ omnitrace_finalize_hidden(void)
         }
     }
 
+    OMNITRACE_DEBUG_F("Post-processing the system-level samples...\n");
     thread_sampler::post_process();
 
     if(get_use_critical_trace())
diff --git a/source/lib/omnitrace/src/library/components/rocm_smi.cpp b/source/lib/omnitrace/src/library/components/rocm_smi.cpp
index 9645ec215b..592f3baf43 100644
--- a/source/lib/omnitrace/src/library/components/rocm_smi.cpp
+++ b/source/lib/omnitrace/src/library/components/rocm_smi.cpp
@@ -215,8 +215,6 @@ data::post_process(uint32_t _dev_id)
     using component::sampling_gpu_memory;
     using component::sampling_gpu_power;
     using component::sampling_gpu_temp;
-    using bundle_t = tim::lightweight_tuple<sampling_gpu_busy, sampling_gpu_temp,
-                                            sampling_gpu_power, sampling_gpu_memory>;
 
     if(device_count < _dev_id) return;
 
@@ -255,6 +253,11 @@ data::post_process(uint32_t _dev_id)
 
     if(!get_use_timemory()) return;
 
+#if !defined(TIMEMORY_USE_MPI)
+    // timemory + MPI here causes hangs for some reason. it is unclear why
+    using bundle_t = tim::lightweight_tuple<sampling_gpu_busy, sampling_gpu_temp,
+                                            sampling_gpu_power, sampling_gpu_memory>;
+
     for(auto& itr : _rocm_smi)
     {
         using entry_t = critical_trace::entry;
@@ -282,6 +285,7 @@ data::post_process(uint32_t _dev_id)
             _v.pop();
         }
     }
+#endif
 }
 
 //--------------------------------------------------------------------------------------//
diff --git a/source/lib/omnitrace/src/library/config.cpp b/source/lib/omnitrace/src/library/config.cpp
index f50cee88ec..5617c28dc0 100644
--- a/source/lib/omnitrace/src/library/config.cpp
+++ b/source/lib/omnitrace/src/library/config.cpp
@@ -371,7 +371,7 @@ configure_settings()
     settings::suppress_config()   = true;
     settings::use_output_suffix() = _config->get<bool>("OMNITRACE_USE_PID");
 #if !defined(TIMEMORY_USE_MPI) && defined(TIMEMORY_USE_MPI_HEADERS)
-    if(tim::mpi::is_initialized()) settings::default_process_suffix() = tim::mpi::rank();
+    if(tim::dmp::is_initialized()) settings::default_process_suffix() = tim::dmp::rank();
 #endif
     OMNITRACE_CONDITIONAL_BASIC_PRINT(get_verbose_env() > 0, "configuration complete\n");
 
@@ -865,10 +865,11 @@ get_backend()
 std::string&
 get_perfetto_output_filename()
 {
-    static auto  _v = get_config()->find("OMNITRACE_OUTPUT_FILE");
-    static auto& _t = static_cast<tim::tsettings<std::string>&>(*_v->second);
-    if(_t.get().empty())
-    {
+    static auto  _v        = get_config()->find("OMNITRACE_OUTPUT_FILE");
+    static auto& _t        = static_cast<tim::tsettings<std::string>&>(*_v->second);
+    static auto  _generate = []() {
+        if(tim::dmp::is_initialized())
+            settings::default_process_suffix() = tim::dmp::rank();
         // default name: perfetto-trace.<pid>.proto or perfetto-trace.<rank>.proto
         auto _default_fname =
             settings::compose_output_filename("perfetto-trace", "proto", get_use_pid());
@@ -879,10 +880,16 @@ get_perfetto_output_filename()
             _default_fname =
                 _default_fname.replace(_dpos, _pid_patch.length(), "/perfetto-trace");
         // have the default display the full path to the output file
-        _t.set(tim::get_env<std::string>(
+        return tim::get_env<std::string>(
             "OMNITRACE_OUTPUT_FILE",
             JOIN('/', tim::get_env<std::string>("PWD", ".", false), _default_fname),
-            false));
+            false);
+    };
+    static auto _generated = _generate();
+    if(_t.get().empty() || _t.get() == _generated)
+    {
+        _t.set(_generate());
+        _generated = _t.get();
     }
     return _t.get();
 }