SWDEV-391594: Added macros for MPI env-vars

Change-Id: I373833a2ad997ba944a05f10ac8da6cfdf139fa3


[ROCm/rocprofiler commit: cd5456bf1c]

projects/rocprofiler/CHANGELOG.md

@@ -214,7 +214,6 @@ The resulting `a.out` will depend on
 ### Added
 - 'end_time' need to be disabled in roctx_trace.txt
 - support for hsa_amd_memory_async_copy_on_engine API function trace
-
 ### Fixed
 - rocprof in ROcm/5.4.0 gpu selector broken.
 - rocprof in ROCm/5.4.1 fails to generate kernel info.
@@ -229,6 +228,7 @@ The resulting `a.out` will depend on
 - Limiting file name sizes for ATT plugin.
 - Support for MI300 XCC modes for rocprof v2.
 - MI300 individual XCC counters dumped per-xcc as separate records but with same record-id and kernel dispatch info
+- Naming for MPI ranks. Filenames containing "%rank" are replaced by variables "MPI_RANK", "OMPI_COMM_WORLD_RANK" or "MV2_COMM_WORLD_RANK".
 ### Fixed
 - Samples are fixed to show the new usage of phases.
 - Plugin option validates the plugin names.

projects/rocprofiler/plugin/file/file.cpp

@@ -39,6 +39,7 @@
 #include <string>
 #include <hsa/hsa.h>
 #include <mutex>
+#include <unordered_map>
 
 #include "rocprofiler.h"
 #include "rocprofiler_plugin.h"
@@ -94,6 +95,8 @@ class file_plugin_t {
       }
 
       std::stringstream ss;
+      output_file_name = replace_MPI_macros(output_file_name);
+
       ss << output_file_name << GetPid() << "_" << name_;
       stream_.open(output_prefix / ss.str());
     }
@@ -101,6 +104,29 @@ class file_plugin_t {
     bool is_open() const { return stream_.is_open(); }
     bool fail() const { return stream_.fail(); }
 
+    // Returns a string with the MPI %macro replaced with the corresponding envvar
+    std::string replace_MPI_macros(std::string output_file_name) {
+      std::unordered_map<const char*, const char*> MPI_BUILTINS = {
+        {"MPI_RANK", "%rank"},
+        {"OMPI_COMM_WORLD_RANK", "%rank"},
+        {"MV2_COMM_WORLD_RANK", "%rank"}
+      };
+
+      for (const auto& [envvar, key] : MPI_BUILTINS) {
+        size_t key_find = output_file_name.rfind(key);
+        if (key_find == std::string::npos) continue; // Does not contain a %?rank var
+
+        const char* env_var_set = getenv(envvar);
+        if (env_var_set == nullptr) continue; // MPI_COMM_WORLD_x var is does not exist
+
+        int rank = atoi(env_var_set);
+        output_file_name =  output_file_name.substr(0, key_find) + std::to_string(rank)
+                          + output_file_name.substr(key_find + std::string(key).size());
+      }
+
+      return output_file_name;
+    }
+
    private:
     const std::string name_;
     std::ofstream stream_;

projects/rocprofiler/plugin/perfetto/perfetto.cpp

@@ -145,6 +145,7 @@ class perfetto_plugin_t {
     data_source_cfg->set_name("track_event");
     data_source_cfg->set_track_event_config_raw(track_event_cfg.SerializeAsString());
 
+    output_file_name = replace_MPI_macros(output_file_name);
     output_prefix_.append(output_file_name + std::to_string(GetPid()) + "_output.pftrace");
     file_descriptor_ = open(output_prefix_.string().c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600);
     if (file_descriptor_ == -1) rocprofiler::warning("Can't open output file\n");
@@ -167,9 +168,7 @@ class perfetto_plugin_t {
               track_counter_.fetch_add((1 + machine_id) * GetPid(), std::memory_order_acquire);
         }
       }
-      std::string thread_track_str =
-          rocprofiler::string_printf("Node: %s Process ID: %lu Thread ID:", hostname_, GetPid());
-      process_track_desc.mutable_process()->set_process_name(thread_track_str);
+      process_track_desc.mutable_process()->set_process_name(get_thread_track_str());
       perfetto::TrackEvent::SetTrackDescriptor(perfetto::ProcessTrack::Current(),
                                                process_track_desc);
       perfetto::ProcessTrack::Current().Serialize().set_uuid(track_id);
@@ -186,6 +185,40 @@ class perfetto_plugin_t {
     }
   }
 
+  std::string replace_MPI_macros(std::string output_file_name) {
+    std::unordered_map<const char*, const char*> MPI_BUILTINS = {
+      {"MPI_RANK", "%rank"},
+      {"OMPI_COMM_WORLD_RANK", "%rank"},
+      {"MV2_COMM_WORLD_RANK", "%rank"}
+    };
+    bIsMPI = false;
+
+    for (const auto& [envvar, key] : MPI_BUILTINS) {
+      size_t key_find = output_file_name.rfind(key);
+      if (key_find == std::string::npos) continue; // Does not contain a %?rank var
+
+      const char* env_var_set = getenv(envvar);
+      if (env_var_set == nullptr) continue; // MPI_COMM_WORLD_x var is does not exist
+
+      int rank = atoi(env_var_set);
+      output_file_name =  output_file_name.substr(0, key_find) + std::to_string(rank)
+                        + output_file_name.substr(key_find + std::string(key).size());
+      if (!bIsMPI)
+        MPI_rank = rank;
+      bIsMPI = true;
+    }
+
+    return output_file_name;
+  }
+
+  std::string get_thread_track_str() {
+    return rocprofiler::string_printf("Node: %s Process ID: %lu Thread ID:", hostname_, GetPid());
+  }
+
+  std::string get_device_track_str() {
+      return rocprofiler::string_printf("Node: %s Device:", hostname_);
+  }
+
   const char* GetDomainName(rocprofiler_tracer_activity_domain_t domain) {
     switch (domain) {
       case ACTIVITY_DOMAIN_ROCTX:
@@ -232,8 +265,7 @@ class perfetto_plugin_t {
                 .first;
         auto gpu_desc = device_track_it->second.Serialize();
         gpu_desc.mutable_process()->set_pid(device_id);
-        std::string gpu_str = rocprofiler::string_printf("Node: %s Device:", hostname_);
-        gpu_desc.mutable_process()->set_process_name(gpu_str);
+        gpu_desc.mutable_process()->set_process_name(get_device_track_str());
         perfetto::TrackEvent::SetTrackDescriptor(device_track_it->second, gpu_desc);
         track_ids_used_.emplace_back(device_id + 1 + machine_id);
       }
@@ -370,8 +402,7 @@ class perfetto_plugin_t {
                   .first;
           auto gpu_desc = device_track_it->second.Serialize();
           gpu_desc.mutable_process()->set_pid(device_id);
-          std::string gpu_str = rocprofiler::string_printf("Node: %s Device:", hostname_);
-          gpu_desc.mutable_process()->set_process_name(gpu_str);
+          gpu_desc.mutable_process()->set_process_name(get_device_track_str());
           perfetto::TrackEvent::SetTrackDescriptor(device_track_it->second, gpu_desc);
           track_ids_used_.emplace_back(1 + machine_id + device_id);
         }
@@ -391,10 +422,8 @@ class perfetto_plugin_t {
         thread_track_it =
             thread_tracks_.emplace(thread_id, perfetto::ProcessTrack::Global(track_id)).first;
         auto thread_track_desc = thread_track_it->second.Serialize();
-        std::string thread_track_str =
-            rocprofiler::string_printf("Node: %s Process ID: %lu Thread ID:", hostname_, GetPid());
         thread_track_desc.mutable_process()->set_pid(thread_id);
-        thread_track_desc.mutable_process()->set_process_name(thread_track_str);
+        thread_track_desc.mutable_process()->set_process_name(get_thread_track_str());
         perfetto::TrackEvent::SetTrackDescriptor(thread_track_it->second, thread_track_desc);
       }
     }
@@ -714,6 +743,8 @@ class perfetto_plugin_t {
   std::unique_ptr<perfetto::TracingSession> tracing_session_;
   int file_descriptor_;
   bool is_valid_{false};
+  bool bIsMPI = false;
+  int MPI_rank = 0;
   size_t roctx_track_entries_{0};
 
   // Correlate stream id(s) with correlation id(s) to identify the stream id of every HIP activity

projects/rocprofiler/tests/featuretests/profiler/profiler_gtest.cpp

@@ -1346,7 +1346,6 @@ bool FilePluginTest::hasFileInDir(const std::string& filename, const char* direc
   return false;
 }
 
-
 void FilePluginTest::ProcessApplication(std::stringstream& ss) {
   FILE* handle = popen(ss.str().c_str(), "r");
   ASSERT_NE(handle, nullptr);
@@ -1367,7 +1366,7 @@ class VectorAddFileOnlyTest : public FilePluginTest {
   bool hasFile(){ return hasFileInDir("file_test_name", "."); }
 };
 
-TEST_F(VectorAddFileOnlyTest, WhenRunningProfilerWithOnlyOutputFilenameSetTest) {
+TEST_F(VectorAddFileOnlyTest, WhenRunningProfilerWithFilePluginTest) {
   EXPECT_EQ(hasFile(), true);
 }
 
@@ -1380,7 +1379,7 @@ class VectorAddFolderOnlyTest : public FilePluginTest {
   bool hasFile(){ return hasFileInDir("", "./plugin_test_folder_path"); }
 };
 
-TEST_F(VectorAddFolderOnlyTest, WhenRunningProfilerWithOnlyOutputFilenameSetTest) {
+TEST_F(VectorAddFolderOnlyTest, WhenRunningProfilerWithFilePluginTest) {
   EXPECT_EQ(hasFile(), true);
 }
 
@@ -1393,7 +1392,24 @@ class VectorAddFileAndFolderTest : public FilePluginTest {
   bool hasFile(){ return hasFileInDir("file_test_name", "./plugin_test_folder_path"); }
 };
 
-TEST_F(VectorAddFileAndFolderTest, WhenRunningProfilerWithOnlyOutputFilenameSetTest) {
+TEST_F(VectorAddFileAndFolderTest, WhenRunningProfilerWithFilePluginTest) {
+  EXPECT_EQ(hasFile(), true);
+}
+
+class VectorAddFilenameMPITest : public FilePluginTest {
+ protected:
+  virtual void SetUp() {
+    setenv("MPI_RANK", "7", true);
+    RunApplication("hip_vectoradd", "-d ./plugin_test_folder_path -o test_%rank_");
+  }
+  virtual void TearDown() {
+    std::experimental::filesystem::remove_all("./plugin_test_folder_path");
+    unsetenv("MPI_RANK");
+  }
+  bool hasFile(){ return hasFileInDir("test_7_", "./plugin_test_folder_path"); }
+};
+
+TEST_F(VectorAddFilenameMPITest, WhenRunningProfilerWithFilePluginTest) {
   EXPECT_EQ(hasFile(), true);
 }
 */