SWDEV-408509: Added MPI awareness, dispatch ID and buffer size to ATT API. Fixed wrong dispatch ID on kernel.txt. Readme fixes.

Change-Id: Ib1e2d515d14419097d73dc28648e47ca752d95a6

src/core/hsa/packets/packets_generator.cpp

@@ -479,7 +479,6 @@ hsa_ven_amd_aqlprofile_profile_t* InitializeDeviceProfilingAqlPackets(
 }
 
 // ATT
-uint32_t g_output_buffer_size = 0x40000000;  // 1GB
 bool g_output_buffer_local = true;
 
 // Allocate system memory accessible by both CPU and GPU
@@ -511,11 +510,11 @@ uint8_t* AllocateLocalMemory(size_t size, hsa_amd_memory_pool_t* gpu_pool) {
   return ptr;
 }
 
-hsa_status_t Allocate(hsa_agent_t gpu_agent, hsa_ven_amd_aqlprofile_profile_t* profile) {
+hsa_status_t Allocate(hsa_agent_t gpu_agent, hsa_ven_amd_aqlprofile_profile_t* profile, size_t att_buffer_size) {
   Agent::AgentInfo& agentInfo = rocprofiler::hsa_support::GetAgentInfo(gpu_agent.handle);
   profile->command_buffer.ptr =
       AllocateSysMemory(gpu_agent, profile->command_buffer.size, &agentInfo.cpu_pool);
-  profile->output_buffer.size = g_output_buffer_size;
+  profile->output_buffer.size = att_buffer_size;
   profile->output_buffer.ptr = (g_output_buffer_local)
       ? AllocateLocalMemory(profile->output_buffer.size, &agentInfo.gpu_pool)
       : AllocateSysMemory(gpu_agent, profile->output_buffer.size, &agentInfo.cpu_pool);
@@ -555,7 +554,7 @@ att_mem_pools_map_t* GetAttMemPoolsMap() {
 hsa_ven_amd_aqlprofile_profile_t* GenerateATTPackets(
     hsa_agent_t cpu_agent, hsa_agent_t gpu_agent,
     std::vector<hsa_ven_amd_aqlprofile_parameter_t>& att_params, packet_t* start_packet,
-    packet_t* stop_packet) {
+    packet_t* stop_packet, size_t att_buffer_size) {
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wconversion-null"
   // Preparing the profile structure to get the packets
@@ -576,7 +575,7 @@ hsa_ven_amd_aqlprofile_profile_t* GenerateATTPackets(
   // Allocate command and output buffers
   // command buffer -> from CPU memory pool
   // output buffer -> from GPU memory pool
-  status = Allocate(gpu_agent, profile);
+  status = Allocate(gpu_agent, profile, att_buffer_size);
   CHECK_HSA_STATUS("Error: Att Buffers Allocation", status);
 
   // Generate start/stop/read profiling packets

src/core/hsa/packets/packets_generator.h

@@ -58,7 +58,7 @@ hsa_amd_memory_pool_t& GetOutputPool();
 hsa_ven_amd_aqlprofile_profile_t* GenerateATTPackets(
     hsa_agent_t cpu_agent, hsa_agent_t gpu_agent,
     std::vector<hsa_ven_amd_aqlprofile_parameter_t>& att_params, packet_t* start_packet,
-    packet_t* stop_packet);
+    packet_t* stop_packet, size_t att_buffer_size);
 
 
 uint8_t* AllocateSysMemory(hsa_agent_t gpu_agent, size_t size, hsa_amd_memory_pool_t* cpu_pool);

src/core/hsa/queues/queue.cpp

@@ -51,6 +51,8 @@
 
 #define __NR_gettid 186
 
+#define DEFAULT_ATT_BUFFER_SIZE 0x40000000
+
 std::mutex sessions_pending_signal_lock;
 
 namespace rocprofiler {
@@ -514,10 +516,10 @@ bool AsyncSignalHandlerATT(hsa_signal_value_t /* signal */, void* data) {
       if (/*pending.counters_count > 0 && */ pending.profile) {
         AddAttRecord(&record, queue_info_session->agent, pending);
       }
-      // Kernel Descriptor is the right record id generated in the WriteInterceptor function and
-      // will be used to handle the kernel name of that dispatch
+      // July/01/2023 -> Changed this to writer ID so we can correlate to dispatches
+      // kernel_id already has the descriptor.
       record.header = {ROCPROFILER_ATT_TRACER_RECORD,
-                       rocprofiler_record_id_t{pending.kernel_descriptor}};
+                       rocprofiler_record_id_t{queue_info_session->writer_id}};
 
       if (pending.session_id.handle == 0) {
         pending.session_id = GetROCProfilerSingleton()->GetCurrentSessionId();
@@ -602,6 +604,7 @@ bool is_pc_sampling_collection_mode = false;
 std::vector<rocprofiler_att_parameter_t> att_parameters_data;
 uint32_t replay_mode_count = 0;
 std::vector<std::string> kernel_profile_names;
+std::vector<uint64_t> kernel_profile_dispatch_ids;
 std::vector<std::string> att_counters_names;
 
 rocprofiler::Session* session = nullptr;
@@ -649,6 +652,8 @@ void CheckNeededProfileConfigs() {
         att_counters_names = filter->GetCounterData();
         kernel_profile_names = std::get<std::vector<std::string>>(
             filter->GetProperty(ROCPROFILER_FILTER_KERNEL_NAMES));
+        kernel_profile_dispatch_ids = std::get<std::vector<uint64_t>>(
+            filter->GetProperty(ROCPROFILER_FILTER_DISPATCH_IDS));
       } else if (session && session->FindFilterWithKind(ROCPROFILER_PC_SAMPLING_COLLECTION)) {
         is_pc_sampling_collection_mode = true;
       }
@@ -658,6 +663,121 @@ void CheckNeededProfileConfigs() {
 
 static int KernelInterceptCount = 0;
 std::atomic<uint32_t> WRITER_ID{0};
+
+std::pair<std::vector<bool>, bool> GetAllowedProfilesList(const void* packets, int pkt_count) {
+  std::vector<bool> can_profile_packet;
+  bool b_can_profile_anypacket = false;
+  can_profile_packet.reserve(pkt_count);
+
+  std::lock_guard<std::mutex> lock(ksymbol_map_lock);
+  assert(ksymbols);
+
+  uint32_t current_writer_id = WRITER_ID.load(std::memory_order_relaxed);
+
+  for (int i = 0; i < pkt_count; ++i) {
+    auto& original_packet = static_cast<const hsa_barrier_and_packet_t*>(packets)[i];
+    bool b_profile_this_object = false;
+
+    // Skip packets other than kernel dispatch packets.
+    if (bit_extract(original_packet.header, HSA_PACKET_HEADER_TYPE,
+                    HSA_PACKET_HEADER_TYPE + HSA_PACKET_HEADER_WIDTH_TYPE - 1) ==
+        HSA_PACKET_TYPE_KERNEL_DISPATCH) {
+      auto& kdispatch = static_cast<const hsa_kernel_dispatch_packet_s*>(packets)[i];
+
+      // If Dispatch IDs specified, profile based on dispatch ID
+      for (auto id : kernel_profile_dispatch_ids)
+        b_profile_this_object |= id == current_writer_id;
+      try {
+        // Can throw
+        const std::string& kernel_name = ksymbols->at(kdispatch.kernel_object);
+
+        // If no filters specified, auto profile this kernel
+        if (kernel_profile_names.size() == 0 &&
+            kernel_profile_dispatch_ids.size() == 0 &&
+            kernel_name.find("__amd_rocclr_") == std::string::npos)
+            b_profile_this_object = true;
+
+        // Try to match the mangled kernel name with given matches in input.txt
+        // We want to initiate att profiling if a match exists
+        for (const std::string& kernel_matches : kernel_profile_names)
+          if (kernel_name.find(kernel_matches) != std::string::npos)
+            b_profile_this_object = true;
+      } catch (...) {
+        printf("Warning: Unknown name for object %lu\n", kdispatch.kernel_object);
+      }
+      current_writer_id += 1;
+    }
+    b_can_profile_anypacket |= b_profile_this_object;
+    can_profile_packet.push_back(b_profile_this_object);
+  }
+  // If we're going to skip all packets, need to update writer ID
+  if (!b_can_profile_anypacket)
+    WRITER_ID.store(current_writer_id, std::memory_order_release);
+  return {can_profile_packet, b_can_profile_anypacket};
+}
+
+hsa_ven_amd_aqlprofile_profile_t* ProcessATTParams(
+  Packet::packet_t& start_packet,
+  Packet::packet_t& stop_packet,
+  Queue& queue_info,
+  Agent::AgentInfo& agentInfo
+) {
+  std::vector<hsa_ven_amd_aqlprofile_parameter_t> att_params;
+  int num_att_counters = 0;
+  uint32_t att_buffer_size = DEFAULT_ATT_BUFFER_SIZE;
+
+  for (rocprofiler_att_parameter_t& param : att_parameters_data) {
+    switch (param.parameter_name) {
+      case ROCPROFILER_ATT_PERFCOUNTER_NAME:
+        break;
+      case ROCPROFILER_ATT_BUFFER_SIZE:
+        att_buffer_size = std::max(96l<<10l, std::min(int64_t(param.value)<<20l, (1l<<32l)-(3l<<20)));
+        break; // Clip to [96KB, 4GB)
+      case ROCPROFILER_ATT_PERFCOUNTER:
+        num_att_counters += 1;
+        break;
+      default:
+        att_params.push_back(
+        {static_cast<hsa_ven_amd_aqlprofile_parameter_name_t>(int(param.parameter_name)),
+        param.value});
+    }
+  }
+
+  if (att_counters_names.size() > 0) {
+    MetricsDict* metrics_dict_ = MetricsDict::Create(&agentInfo);
+
+    for (const std::string& counter_name : att_counters_names) {
+      const Metric* metric = metrics_dict_->Get(counter_name);
+      const BaseMetric* base = dynamic_cast<const BaseMetric*>(metric);
+      if (!base) {
+        printf("Invalid base metric value: %s\n", counter_name.c_str());
+        exit(1);
+      }
+      std::vector<const counter_t*> counters;
+      base->GetCounters(counters);
+      hsa_ven_amd_aqlprofile_event_t event = counters[0]->event;
+      if (event.block_name != HSA_VEN_AMD_AQLPROFILE_BLOCK_NAME_SQ) {
+        printf("Only events from the SQ block can be selected for ATT.");
+        exit(1);
+      }
+      att_params.push_back({static_cast<hsa_ven_amd_aqlprofile_parameter_name_t>(
+                                int(ROCPROFILER_ATT_PERFCOUNTER)),
+                            event.counter_id | (event.counter_id ? (0xF << 24) : 0)});
+      num_att_counters += 1;
+    }
+
+    hsa_ven_amd_aqlprofile_parameter_t zero_perf = {
+        static_cast<hsa_ven_amd_aqlprofile_parameter_name_t>(int(ROCPROFILER_ATT_PERFCOUNTER)),
+        0};
+
+    // Fill other perfcounters with 0's
+    for (; num_att_counters < 16; num_att_counters++) att_params.push_back(zero_perf);
+  }
+  // Get the PM4 Packets using packets_generator
+  return Packet::GenerateATTPackets(queue_info.GetCPUAgent(), queue_info.GetGPUAgent(),
+                                      att_params, &start_packet, &stop_packet, att_buffer_size);
+}
+
 /**
  * @brief This function is a queue write interceptor. It intercepts the
  * packet write function. Creates an instance of packet class with the raw
@@ -817,40 +937,7 @@ void WriteInterceptor(const void* packets, uint64_t pkt_count, uint64_t user_pkt
 
     bool can_profile_anypacket = false;
     std::vector<bool> can_profile_packet;
-
-    for (size_t i = 0; i < pkt_count; ++i) {
-      auto& original_packet = static_cast<const hsa_barrier_and_packet_t*>(packets)[i];
-      bool b_profile_this_object = false;
-
-      // Skip packets other than kernel dispatch packets.
-      if (bit_extract(original_packet.header, HSA_PACKET_HEADER_TYPE,
-                      HSA_PACKET_HEADER_TYPE + HSA_PACKET_HEADER_WIDTH_TYPE - 1) ==
-          HSA_PACKET_TYPE_KERNEL_DISPATCH) {
-        auto& kdispatch = static_cast<const hsa_kernel_dispatch_packet_s*>(packets)[i];
-        uint64_t kernel_object = kdispatch.kernel_object;
-
-        // Try to match the mangled kernel name with given matches in input.txt
-        try {
-          std::lock_guard<std::mutex> lock(ksymbol_map_lock);
-          assert(ksymbols);
-          const std::string& kernel_name = ksymbols->at(kernel_object);
-
-          // We want to initiate att profiling only if a match exists
-          for (const std::string& kernel_matches : kernel_profile_names) {
-            if (kernel_name.find(kernel_matches) != std::string::npos) {
-              b_profile_this_object = true;
-              break;
-            }
-          }
-          if (!b_profile_this_object) printf("Skipping: %s\n", kernel_name.c_str());
-        } catch (...) {
-          printf("Warning: Unknown name for object %lu\n", kernel_object);
-        }
-      }
-
-      if (b_profile_this_object) can_profile_anypacket = true;
-      can_profile_packet.push_back(b_profile_this_object);
-    }
+    std::tie(can_profile_packet, can_profile_anypacket) = GetAllowedProfilesList(packets, pkt_count);
 
     if (!can_profile_anypacket) {
       /* Write the original packets to the hardware if no patch will be profiled */
@@ -863,67 +950,27 @@ void WriteInterceptor(const void* packets, uint64_t pkt_count, uint64_t user_pkt
     Packet::packet_t stop_packet{};
     hsa_ven_amd_aqlprofile_profile_t* profile = nullptr;
 
-    if (att_parameters_data.size() > 0 && is_att_collection_mode) {
-      // TODO sauverma: convert att_parameters_data to pass to generateattPackets
-      std::vector<hsa_ven_amd_aqlprofile_parameter_t> att_params;
-      int num_att_counters = 0;
-
-      for (rocprofiler_att_parameter_t& param : att_parameters_data) {
-        att_params.push_back(
-            {static_cast<hsa_ven_amd_aqlprofile_parameter_name_t>(int(param.parameter_name)),
-             param.value});
-        num_att_counters += param.parameter_name == ROCPROFILER_ATT_PERFCOUNTER;
-      }
-
-      if (att_counters_names.size() > 0) {
-        MetricsDict* metrics_dict_ = MetricsDict::Create(&agentInfo);
-
-        for (const std::string& counter_name : att_counters_names) {
-          const Metric* metric = metrics_dict_->Get(counter_name);
-          const BaseMetric* base = dynamic_cast<const BaseMetric*>(metric);
-          if (!base) {
-            printf("Invalid base metric value: %s\n", counter_name.c_str());
-            exit(1);
-          }
-          std::vector<const counter_t*> counters;
-          base->GetCounters(counters);
-          hsa_ven_amd_aqlprofile_event_t event = counters[0]->event;
-          if (event.block_name != HSA_VEN_AMD_AQLPROFILE_BLOCK_NAME_SQ) {
-            printf("Only events from the SQ block can be selected for ATT.");
-            exit(1);
-          }
-          att_params.push_back({static_cast<hsa_ven_amd_aqlprofile_parameter_name_t>(
-                                    int(ROCPROFILER_ATT_PERFCOUNTER)),
-                                event.counter_id | (event.counter_id ? (0xF << 24) : 0)});
-          num_att_counters += 1;
-        }
-
-        hsa_ven_amd_aqlprofile_parameter_t zero_perf = {
-            static_cast<hsa_ven_amd_aqlprofile_parameter_name_t>(int(ROCPROFILER_ATT_PERFCOUNTER)),
-            0};
-
-        // Fill other perfcounters with 0's
-        for (; num_att_counters < 16; num_att_counters++) att_params.push_back(zero_perf);
-      }
-
-      // Get the PM4 Packets using packets_generator
-      profile = Packet::GenerateATTPackets(queue_info.GetCPUAgent(), queue_info.GetGPUAgent(),
-                                           att_params, &start_packet, &stop_packet);
-    }
+    if (att_parameters_data.size() > 0 && is_att_collection_mode)
+      profile = ProcessATTParams(start_packet, stop_packet, queue_info, agentInfo);
 
     // Searching across all the packets given during this write
     for (size_t i = 0; i < pkt_count; ++i) {
       auto& original_packet = static_cast<const hsa_barrier_and_packet_t*>(packets)[i];
+      uint32_t writer_id = 0;
 
       // Skip all packets marked with !can_profile
       if (i >= can_profile_packet.size() || can_profile_packet[i] == false) {
         transformed_packets.emplace_back(packets_arr[i]);
+
+        // increment writer ID for every packet
+        if (bit_extract(original_packet.header, HSA_PACKET_HEADER_TYPE,
+          HSA_PACKET_HEADER_TYPE+HSA_PACKET_HEADER_WIDTH_TYPE-1) == HSA_PACKET_TYPE_KERNEL_DISPATCH)
+            writer_id = WRITER_ID.fetch_add(1, std::memory_order_release);
+
         continue;
       }
       KernelInterceptCount += 1;
-
-      uint32_t writer_id = WRITER_ID.fetch_add(1, std::memory_order_release);
-
+      writer_id = WRITER_ID.fetch_add(1, std::memory_order_release);
 
       if (att_parameters_data.size() > 0 && is_att_collection_mode && profile) {
         // Adding start packet and its barrier with a dummy signal

src/core/session/att/att.cpp

@@ -20,6 +20,7 @@
 
 #include "att.h"
 #include <cassert>
+#include <atomic>
 
 namespace rocprofiler {
 
@@ -42,9 +43,11 @@ void AttTracer::AddPendingSignals(uint32_t writer_id, uint64_t kernel_object,
   sessions_pending_signals_.at(writer_id).emplace_back(
       att_pending_signal_t{kernel_object, original_completion_signal, new_completion_signal, session_id_, buffer_id, profile,
                             kernel_properties, thread_id, queue_index});
+  std::atomic_thread_fence(std::memory_order_release);
 }
 
 const std::vector<att_pending_signal_t>& AttTracer::GetPendingSignals(uint32_t writer_id) {
+  std::atomic_thread_fence(std::memory_order_acquire);
   std::lock_guard<std::mutex> lock(sessions_pending_signals_lock_);
   assert(sessions_pending_signals_.find(writer_id) != sessions_pending_signals_.end() &&
          "writer_id is not found in the pending_signals");

src/core/session/filter.cpp

@@ -192,6 +192,11 @@ void Filter::SetProperty(rocprofiler_filter_property_t property) {
       }
       break;
     }
+    case ROCPROFILER_FILTER_DISPATCH_IDS:
+      dispatch_id_filter_.clear();
+      for (uint32_t j = 0; j < property.data_count; j++)
+          dispatch_id_filter_.emplace_back(property.dispatch_ids[j]);
+      break;
     default:
       break;
       // TODO(aelwazir): Check for empty property
@@ -200,9 +205,9 @@ void Filter::SetProperty(rocprofiler_filter_property_t property) {
       //     "profiler mode!\n");
   }
 }
-std::variant<std::vector<std::string>, uint32_t*> Filter::GetProperty(
+std::variant<std::vector<std::string>, uint32_t*, std::vector<uint64_t>> Filter::GetProperty(
     rocprofiler_filter_property_kind_t kind) {
-  std::variant<std::vector<std::string>, uint32_t*> property;
+  std::variant<std::vector<std::string>, uint32_t*, std::vector<uint64_t>> property;
   switch (kind) {
     case ROCPROFILER_FILTER_GPU_NAME: {
       property = agent_names_;
@@ -224,6 +229,10 @@ std::variant<std::vector<std::string>, uint32_t*> Filter::GetProperty(
       property = hip_tracer_api_calls_;
       break;
     }
+    case ROCPROFILER_FILTER_DISPATCH_IDS: {
+      property = dispatch_id_filter_;
+      break;
+    }
     default:
       fatal(
           "Error: ROCProfiler filter specified is not supported for the given "
@@ -261,6 +270,9 @@ size_t Filter::GetPropertiesCount(rocprofiler_filter_property_kind_t kind) {
     case ROCPROFILER_FILTER_HIP_TRACER_API_FUNCTIONS: {
       return hip_tracer_api_calls_.size();
     }
+    case ROCPROFILER_FILTER_DISPATCH_IDS: {
+      return dispatch_id_filter_.size();
+    }
   }
   fatal(
       "Error: ROCProfiler filter specified is not supported for the given "

src/core/session/filter.h

@@ -53,8 +53,12 @@ class Filter {
   bool HasCallback();
 
   void SetProperty(rocprofiler_filter_property_t property);
-  std::variant<std::vector<std::string>, uint32_t*> GetProperty(
-      rocprofiler_filter_property_kind_t kind);
+  std::variant<
+    std::vector<std::string>,
+    uint32_t*,
+    std::vector<uint64_t>
+  > GetProperty(rocprofiler_filter_property_kind_t kind);
+
   size_t GetPropertiesCount(rocprofiler_filter_property_kind_t kind);
   rocprofiler_spm_parameter_t* GetSpmParameterData();
   rocprofiler_counters_sampler_parameters_t GetCountersSamplerParameterData();
@@ -75,6 +79,7 @@ class Filter {
   rocprofiler_spm_parameter_t* spm_parameter_;                     // spm parameter
   std::vector<rocprofiler_att_parameter_t> att_parameters_;  // ATT Parameters
   rocprofiler_counters_sampler_parameters_t counters_sampler_parameters_;        // sampled counters parameters
+  std::vector<uint64_t> dispatch_id_filter_;
 
   bool has_sync_callback_{false};
   rocprofiler_sync_callback_t callback_;

src/core/session/profiler/profiler.cpp

@@ -133,7 +133,7 @@ void Profiler::AddPendingSignals(
   std::lock_guard<std::mutex> lock(sessions_pending_signals_lock_);
   if (sessions_pending_signals_->find(writer_id) == sessions_pending_signals_->end())
     sessions_pending_signals_->emplace(writer_id, std::vector<pending_signal_t*>());
-  sessions_pending_signals_->at(writer_id).emplace_back(new pending_signal_t{
+    sessions_pending_signals_->at(writer_id).emplace_back(new pending_signal_t{
       kernel_object, original_completion_signal, new_completion_signal, session_id_, buffer_id, context, session_data_count,
       profile, kernel_properties, thread_id, queue_index, correlation_id});
 }

src/tools/tool.cpp

@@ -252,15 +252,28 @@ std::vector<std::string> GetCounterNames() {
   return counters;
 }
 
-typedef std::tuple<std::vector<std::pair<rocprofiler_att_parameter_name_t, uint32_t>>,
-                   std::vector<std::string>, std::vector<std::string>>
-    att_parsed_input_t;
+typedef std::tuple<
+  std::vector<std::pair<rocprofiler_att_parameter_name_t, uint32_t>>,
+  std::vector<std::string>,
+  std::vector<std::string>,
+  std::vector<uint64_t>
+> att_parsed_input_t;
+
+static int GetMpRank() {
+  std::vector<const char*> mpivars = {"MPI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK"};
+  for (const char* envvar : mpivars)
+    if (const char* env = getenv(envvar))
+      return atoi(env);
+  return -1;
+}
 
 att_parsed_input_t GetATTParams() {
   std::vector<std::pair<rocprofiler_att_parameter_name_t, uint32_t>> parameters;
   std::vector<std::string> kernel_names;
   std::vector<std::string> counters_names;
+  std::vector<uint64_t> dispatch_ids;
   const char* path = getenv("COUNTERS_PATH");
+  if (!path) return {{}, {}, {}, {}};
 
   // List of parameters the user can set. Maxvalue is unused.
   std::unordered_map<std::string, rocprofiler_att_parameter_name_t> ATT_PARAM_NAMES{};
@@ -268,6 +281,7 @@ att_parsed_input_t GetATTParams() {
   ATT_PARAM_NAMES["att: TARGET_CU"] = ROCPROFILER_ATT_COMPUTE_UNIT_TARGET;
   ATT_PARAM_NAMES["SE_MASK"] = ROCPROFILER_ATT_SE_MASK;
   ATT_PARAM_NAMES["SIMD_MASK"] = ROCPROFILER_ATT_MAXVALUE;
+  ATT_PARAM_NAMES["BUFFER_SIZE"] = ROCPROFILER_ATT_BUFFER_SIZE;
   ATT_PARAM_NAMES["PERFCOUNTER_ID"] = ROCPROFILER_ATT_PERFCOUNTER;
   ATT_PARAM_NAMES["PERFCOUNTER"] = ROCPROFILER_ATT_PERFCOUNTER_NAME;
   ATT_PARAM_NAMES["PERFCOUNTERS_COL_PERIOD"] = ROCPROFILER_ATT_MAXVALUE;
@@ -276,50 +290,62 @@ att_parsed_input_t GetATTParams() {
 
   // Default values used for token generation.
   std::unordered_map<std::string, uint32_t> default_params = {
-      {"ATT_MASK", 0x3F01}, {"TOKEN_MASK", 0x344B}, {"TOKEN_MASK2", 0xFFFFFFF}};
+    {"ATT_MASK", 0x3F01}, {"TOKEN_MASK", 0x344B},
+    {"TOKEN_MASK2", 0xFFFFFFF}, {"SE_MASK", 0x111111}
+  };
 
-  bool started_att_counters = false;
-
-  if (!path) return {parameters, kernel_names, counters_names};
-
-  std::string line;
   std::ifstream trace_file(path);
   if (!trace_file.is_open()) {
     std::cout << "Unable to open att trace file." << std::endl;
-    return {parameters, kernel_names, counters_names};
+    return {{}, {}, {}, {}};
   }
 
+  int MPI_RANK = GetMpRank();
+
+  bool started_att_counters = false;
+  std::string line;
   while (getline(trace_file, line)) {
     if (line.find("//") != std::string::npos)
       line = line.substr(0, line.find("//"));  // Remove comments
 
-    auto pos = line.find('=');
-    if (pos == std::string::npos) continue;
+    std::string param_name;
+    {
+      auto pos = line.find('=');
+      if (pos == std::string::npos) continue;
 
-    std::string param_name = line.substr(0, pos);
-    uint32_t param_value;
+      param_name = line.substr(0, pos);
+      line = line.substr(pos+1);
+    }
 
     if (param_name == "att: TARGET_CU") started_att_counters = true;
     if (!started_att_counters) continue;
 
     if (param_name == "KERNEL") {
-      kernel_names.push_back(line.substr(pos + 1));
+      kernel_names.push_back(line);
       continue;
     } else if (param_name == "PERFCOUNTER") {
-      counters_names.push_back(line.substr(pos + 1));
+      counters_names.push_back(line);
+      continue;
+    } else if (param_name == "DISPATCH") {
+      size_t comma = line.find(',');
+      int id = stoi(line.substr(0, comma));
+      int rank = (comma < line.size()-1) ? stoi(line.substr(comma+1)) : 0;
+
+      if (MPI_RANK < 0 || rank == MPI_RANK) // Only add ID if rank matches the one in input.txt
+        dispatch_ids.push_back(id);
+      continue;
+    }
+    // param_value is a number
+    uint32_t param_value;
+    try {
+      auto hexa_pos = line.find("0x");  // Is it hex?
+      if (hexa_pos != std::string::npos)
+        param_value = stoi(line.substr(hexa_pos + 2), 0, 16);  // hexadecimal
+      else
+        param_value = stoi(line, 0, 10);  // decimal
+    } catch (...) {
+      printf("Error: Invalid parameter value %s\n", line.c_str());
       continue;
-    } else {  // param_value is a number
-      try {
-        auto hexa_pos = line.find("0x", pos);  // Is it hex?
-        if (hexa_pos != std::string::npos)
-          param_value = stoi(line.substr(hexa_pos + 2), 0, 16);  // hexadecimal
-        else
-          param_value = stoi(line.substr(pos + 1), 0, 10);  // decimal
-      } catch (...) {
-        printf("Error: Invalid parameter value %s - (%s)\n",
-               line.substr(pos + 1, line.size()).c_str(), line.c_str());
-        exit(1);
-      }
     }
 
     if (param_name == "PERFCOUNTERS_COL_PERIOD") {
@@ -355,7 +381,7 @@ att_parsed_input_t GetATTParams() {
   }
   trace_file.close();
 
-  if (!started_att_counters) return {parameters, kernel_names, counters_names};
+  if (!started_att_counters) return {{}, {}, {}, {}};
 
   ATT_PARAM_NAMES["ATT_MASK"] = ROCPROFILER_ATT_MASK;
   ATT_PARAM_NAMES["TOKEN_MASK"] = ROCPROFILER_ATT_TOKEN_MASK;
@@ -364,11 +390,7 @@ att_parsed_input_t GetATTParams() {
   for (auto& param : default_params)
     parameters.push_back(std::make_pair(ATT_PARAM_NAMES[param.first], param.second));
 
-  // If no kernel names were provided, collect them all.
-  // Empty string always returns true for "str.find()".
-  if (kernel_names.size() == 0) kernel_names.push_back("");
-
-  return {parameters, kernel_names, counters_names};
+  return {parameters, kernel_names, counters_names, dispatch_ids};
 }
 
 void finish() {
@@ -651,7 +673,8 @@ ROCPROFILER_EXPORT bool OnLoad(void* table, uint64_t runtime_version, uint64_t f
   std::vector<std::pair<rocprofiler_att_parameter_name_t, uint32_t>> params;
   std::vector<std::string> kernel_names;
   std::vector<std::string> att_counters_names;
-  std::tie(params, kernel_names, att_counters_names) = GetATTParams();
+  std::vector<uint64_t> dispatch_ids;
+  std::tie(params, kernel_names, att_counters_names, dispatch_ids) = GetATTParams();
 
   for (auto& kv_pair : params)
     parameters.emplace_back(rocprofiler_att_parameter_t{kv_pair.first, kv_pair.second});
@@ -756,20 +779,26 @@ ROCPROFILER_EXPORT bool OnLoad(void* table, uint64_t runtime_version, uint64_t f
             1 << 20, &buffer_id));
         buffer_ids.emplace_back(buffer_id);
         printf("Enabling ATT Tracing\n");
+
         rocprofiler_filter_id_t filter_id;
-
-        std::vector<const char*> kernel_names_c;
-        for (auto& name : kernel_names) kernel_names_c.push_back(name.data());
-
         rocprofiler_filter_property_t property = {};
-        property.kind = ROCPROFILER_FILTER_KERNEL_NAMES;
-        property.data_count = kernel_names_c.size();
-        property.name_regex = kernel_names_c.data();
+        std::vector<const char*> kernel_names_c;
 
+        if (dispatch_ids.size()) { // Correlation ID filter
+          property.kind = ROCPROFILER_FILTER_DISPATCH_IDS;
+          property.data_count = dispatch_ids.size();
+          property.dispatch_ids = dispatch_ids.data();
+        } else { // Kernel names filter
+          for (auto& name : kernel_names) kernel_names_c.push_back(name.data());
+
+          property.kind = ROCPROFILER_FILTER_KERNEL_NAMES;
+          property.data_count = kernel_names_c.size();
+          property.name_regex = kernel_names_c.data();
+        }
         CHECK_ROCPROFILER(
-            rocprofiler_create_filter(session_id, ROCPROFILER_ATT_TRACE_COLLECTION,
-                                      rocprofiler_filter_data_t{.att_parameters = &parameters[0]},
-                                      parameters.size(), &filter_id, property));
+              rocprofiler_create_filter(session_id, ROCPROFILER_ATT_TRACE_COLLECTION,
+                                        rocprofiler_filter_data_t{.att_parameters = &parameters[0]},
+                                        parameters.size(), &filter_id, property));
         CHECK_ROCPROFILER(rocprofiler_set_filter_buffer(session_id, filter_id, buffer_id));
         filter_ids.emplace_back(filter_id);
         break;