SWDEV-391594: Added macros for MPI env-vars

Change-Id: I373833a2ad997ba944a05f10ac8da6cfdf139fa3

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_;

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