[rocprofiler-systems] [ROCpd] Add OMPT callbacks to ROCpd

projects/rocprofiler-systems/source/lib/core/trace_cache/metadata_registry.cpp

@@ -210,7 +210,7 @@ metadata_registry::get_string_list() const
     return result;
 }
 
-#if ROCPROFSYS_USE_ROCM
+#if ROCPROFSYS_USE_ROCM > 0
 
 void
 metadata_registry::add_code_object(
@@ -278,6 +278,96 @@ metadata_registry::get_kernel_symbol_list() const
     return result;
 }
 
+void
+metadata_registry::overwrite_callback_names(
+    std::initializer_list<
+        std::pair<rocprofiler_callback_tracing_kind_t, callback_rename_map_t>>
+        rename_table)
+{
+    if(rename_table.size() == 0) return;
+
+    using callback_kind_t   = rocprofiler_callback_tracing_kind_t;
+    using operation_names_t = std::vector<std::string_view>;
+
+    auto category_names = std::vector<std::string_view>{};
+    auto modified_ops   = std::map<callback_kind_t, operation_names_t>{};
+
+    auto extract_operations = [&](callback_kind_t cat) -> operation_names_t {
+        auto        items           = m_callback_tracing_info.items();
+        const auto* target_category = items[static_cast<size_t>(cat)];
+
+        auto              operations_data = target_category->items();
+        operation_names_t operation_names;
+        operation_names.reserve(operations_data.size());
+
+        for(const auto& [op_idx, op_name] : operations_data)
+            operation_names.push_back(*op_name);
+
+        return operation_names;
+    };
+
+    // Store category names
+    category_names.resize(ROCPROFILER_CALLBACK_TRACING_LAST);
+    for(callback_kind_t i = ROCPROFILER_CALLBACK_TRACING_NONE;
+        i < ROCPROFILER_CALLBACK_TRACING_LAST;
+        i = static_cast<callback_kind_t>(static_cast<int>(i) + 1))
+    {
+        category_names[i] = m_callback_tracing_info.at(i);
+    }
+
+    // Process list
+    for(const auto& category_info : rename_table)
+    {
+        auto callback_kind = category_info.first;
+        // Store operations of all following categories
+        //  as they will be deleted
+        for(callback_kind_t i =
+                static_cast<callback_kind_t>(static_cast<int>(callback_kind) + 1);
+            i < ROCPROFILER_CALLBACK_TRACING_LAST;
+            i = static_cast<callback_kind_t>(static_cast<int>(i) + 1))
+        {
+            if(modified_ops.find(i) != modified_ops.end()) break;
+            modified_ops[i] = extract_operations(i);
+        }
+
+        ROCPROFSYS_CI_THROW(modified_ops.find(callback_kind) != modified_ops.end(),
+                            "Overwriting a previously overwritten entry is forbidden");
+
+        ROCPROFSYS_CI_THROW(!modified_ops.empty() &&
+                                callback_kind >= modified_ops.begin()->first,
+                            "Category must have a larger enum value than all previously "
+                            "modified_ops categories");
+
+        // Overwrite desired category
+        auto operation_names = extract_operations(callback_kind);
+        for(const auto& [index, new_value] : category_info.second)
+        {
+            ROCPROFSYS_CI_THROW(index < 0 ||
+                                    static_cast<size_t>(index) >= operation_names.size(),
+                                "Index is invalid");
+            operation_names[index] = new_value;
+        }
+        modified_ops[callback_kind] = std::move(operation_names);
+    }
+    if(modified_ops.empty()) return;
+
+    // Emplace the changed category operations
+    for(callback_kind_t i = modified_ops.begin()->first;
+        i < ROCPROFILER_CALLBACK_TRACING_LAST;
+        i = static_cast<callback_kind_t>(static_cast<int>(i) + 1))
+    {
+        auto        renaming_entry = modified_ops.find(i);
+        const auto& operations_vec = renaming_entry->second;
+        m_callback_tracing_info.emplace(i, category_names.at(i).data());
+        for(size_t op_idx = 0; op_idx < operations_vec.size(); ++op_idx)
+        {
+            m_callback_tracing_info.emplace(
+                i, static_cast<rocprofiler_tracing_operation_t>(op_idx),
+                operations_vec[op_idx].data());
+        }
+    }
+}
+
 rocprofiler::sdk::buffer_name_info_t<const char*>
 metadata_registry::get_buffer_name_info() const
 {
@@ -292,5 +382,20 @@ metadata_registry::get_callback_tracing_info() const
 
 #endif
 
+metadata_registry::metadata_registry()
+{
+#if ROCPROFSYS_USE_ROCM > 0
+    overwrite_callback_names({
+#    if(ROCPROFILER_VERSION >= 600)
+        { ROCPROFILER_CALLBACK_TRACING_OMPT,
+          { { ROCPROFILER_OMPT_ID_thread_begin, "omp_thread" },
+            { ROCPROFILER_OMPT_ID_thread_end, "omp_thread" },
+            { ROCPROFILER_OMPT_ID_parallel_begin, "omp_parallel" },
+            { ROCPROFILER_OMPT_ID_parallel_end, "omp_parallel" } } }
+#    endif
+    });
+#endif
+}
+
 }  // namespace trace_cache
 }  // namespace rocprofsys

projects/rocprofiler-systems/source/lib/core/trace_cache/metadata_registry.hpp

@@ -25,6 +25,7 @@
 #include "common/synchronized.hpp"
 #include "core/agent.hpp"
 #include "core/categories.hpp"
+#include "core/debug.hpp"
 
 #include <cassert>
 #include <cstdint>
@@ -35,6 +36,8 @@
 #    include <rocprofiler-sdk/callback_tracing.h>
 #    include <rocprofiler-sdk/cxx/name_info.hpp>
 #endif
+#include <initializer_list>
+#include <map>
 #include <set>
 #include <sstream>
 #include <stdint.h>
@@ -215,7 +218,7 @@ struct metadata_registry
 
 private:
     friend class cache_manager;
-    metadata_registry() = default;
+    metadata_registry();
     common::synchronized<info::process> m_process;
     common::synchronized<
         std::unordered_set<info::pmc, info::pmc_info_hash, info::pmc_info_equal>>
@@ -240,6 +243,14 @@ struct metadata_registry
     rocprofiler::sdk::callback_name_info_t<const char*> m_callback_tracing_info{
         rocprofiler::sdk::get_callback_tracing_names<const char*>()
     };
+
+    using callback_rename_map_t =
+        std::map<rocprofiler_tracing_operation_t, std::string_view>;
+
+    void overwrite_callback_names(
+        std::initializer_list<
+            std::pair<rocprofiler_callback_tracing_kind_t, callback_rename_map_t>>
+            rename_table);
 #endif
 };