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

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

@@ -254,6 +254,88 @@ rocpd_post_processing::get_memory_allocate_callback() const
         };
     };
 }
+
+postprocessing_callback
+rocpd_post_processing::get_ompt_callback() const
+{
+    [[maybe_unused]] auto parse_args = []([[maybe_unused]] const std::string& arg_str) {
+#    if ROCPROFSYS_USE_ROCM > 0
+        rocprofiler_sdk::function_args_t args;
+        const std::string                delimiter = ";;";
+
+        auto split = [](const std::string& str, const std::string& _delimiter) {
+            std::vector<std::string> tokens;
+            size_t                   start = 0;
+            size_t                   end   = str.find(_delimiter);
+
+            while(end != std::string::npos)
+            {
+                tokens.push_back(str.substr(start, end - start));
+                start = end + _delimiter.length();
+                end   = str.find(_delimiter, start);
+            }
+
+            return tokens;
+        };
+
+        auto tokens = split(arg_str, delimiter);
+
+        // Ensure the number of tokens is a multiple of 4
+        if(tokens.size() % 4 != 0)
+        {
+            throw std::invalid_argument("Malformed argument string.");
+        }
+
+        for(auto it = tokens.begin(); it != tokens.end(); it += 4)
+        {
+            rocprofiler_sdk::argument_info arg = { static_cast<uint32_t>(std::stoi(*it)),
+                                                   *(it + 1), *(it + 2), *(it + 3) };
+            args.push_back(arg);
+        }
+
+        return args;
+#    endif
+    };
+
+    return [&]([[maybe_unused]] const storage_parsed_type_base& parsed) {
+#    if ROCPROFSYS_USE_ROCM > 0
+        auto  _ors           = static_cast<const struct ompt_region_sample&>(parsed);
+        auto& data_processor = get_data_processor();
+        auto& n_info         = node_info::get_instance();
+        auto  process        = m_metadata.get_process_info();
+        auto  thread_primary_key =
+            data_processor.map_thread_id_to_primary_key(_ors.thread_id);
+
+        auto callback_tracing_info = m_metadata.get_callback_tracing_info();
+        auto _name                 = _ors.name;
+        auto name_primary_key      = data_processor.insert_string(_name.c_str());
+
+        auto category_primary_key =
+            data_processor.insert_string(trait::name<category::rocm_ompt_api>::value);
+
+        size_t stack_id        = _ors.correlation_id_internal;
+        size_t parent_stack_id = _ors.correlation_id_ancestor;
+        size_t correlation_id  = 0;
+
+        auto event_primary_key =
+            data_processor.insert_event(category_primary_key, stack_id, parent_stack_id,
+                                        correlation_id, _ors.call_stack.c_str());
+
+        auto args = parse_args(_ors.args_str);
+        for(const auto& arg : args)
+        {
+            data_processor.insert_args(event_primary_key, arg.arg_number,
+                                       arg.arg_type.c_str(), arg.arg_name.c_str(),
+                                       arg.arg_value.c_str());
+        }
+
+        data_processor.insert_region(n_info.id, process.pid, thread_primary_key,
+                                     _ors.start_timestamp, _ors.end_timestamp,
+                                     name_primary_key, event_primary_key);
+#    endif
+    };
+}
+
 #endif
 
 postprocessing_callback
@@ -654,6 +736,7 @@ rocpd_post_processing::register_parser_callback([[maybe_unused]] storage_parser&
 #    if(ROCPROFILER_VERSION >= 600)
     parser.register_type_callback(entry_type::memory_alloc,
                                   get_memory_allocate_callback());
+    parser.register_type_callback(entry_type::ompt, get_ompt_callback());
 #    endif
     parser.register_type_callback(entry_type::in_time_sample,
                                   get_in_time_sample_callback());

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

@@ -48,6 +48,7 @@ class rocpd_post_processing
     postprocessing_callback get_memory_copy_callback() const;
 #if(ROCPROFILER_VERSION >= 600)
     postprocessing_callback get_memory_allocate_callback() const;
+    postprocessing_callback get_ompt_callback() const;
 #endif
     postprocessing_callback get_region_callback() const;
     postprocessing_callback get_in_time_sample_callback() const;

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

@@ -126,6 +126,45 @@ struct memory_allocate_sample : storage_parsed_type_base
     // Stream handle
     size_t stream_handle;
 };
+
+struct ompt_region_sample : storage_parsed_type_base
+{
+    ompt_region_sample() = default;
+    ompt_region_sample(std::string _name, uint64_t _thread_id, int32_t _operation,
+                       uint64_t _correlation_id_internal,
+                       uint64_t _correlation_id_ancestor, uint64_t _start_timestamp,
+                       uint64_t _end_timestamp, std::string _args_str,
+                       std::string _call_stack)
+    : name(std::move(_name))
+    , thread_id(_thread_id)
+    , operation(_operation)
+    , correlation_id_internal(_correlation_id_internal)
+    , correlation_id_ancestor(_correlation_id_ancestor)
+    , start_timestamp(_start_timestamp)
+    , end_timestamp(_end_timestamp)
+    , args_str(std::move(_args_str))
+    , call_stack(std::move(_call_stack))
+    {}
+
+    // Identification fields
+    std::string name;
+    uint64_t    thread_id;
+    int32_t     operation;
+
+    // Correlation fields
+    uint64_t correlation_id_internal;
+    uint64_t correlation_id_ancestor;
+
+    // Timing fields
+    uint64_t start_timestamp;
+    uint64_t end_timestamp;
+
+    // Callback arguments
+    std::string args_str;
+
+    // Additional fields
+    std::string call_stack;
+};
 #endif
 
 struct region_sample : storage_parsed_type_base
@@ -231,9 +270,10 @@ enum class entry_type : uint32_t
     memory_copy           = 0x0004,
 #if(ROCPROFSYS_USE_ROCM && ROCPROFILER_VERSION >= 600)
     memory_alloc = 0x0005,
+    ompt         = 0x0006,
 #endif
-    amd_smi_sample   = 0x0006,
-    cpu_freq_sample  = 0x0007,
+    amd_smi_sample   = 0x0007,
+    cpu_freq_sample  = 0x0008,
     fragmented_space = 0xFFFF
 };
 }  // namespace trace_cache

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

@@ -160,6 +160,20 @@ storage_parser::consume_storage()
                 invoke_callbacks(header.type, _memory_allocate_sample);
                 break;
             }
+            case entry_type::ompt:
+            {
+                ompt_region_sample _ompt_region_sample;
+                parse_data(sample.data(), _ompt_region_sample.name,
+                           _ompt_region_sample.thread_id, _ompt_region_sample.operation,
+                           _ompt_region_sample.correlation_id_internal,
+                           _ompt_region_sample.correlation_id_ancestor,
+                           _ompt_region_sample.start_timestamp,
+                           _ompt_region_sample.end_timestamp,
+                           _ompt_region_sample.args_str, _ompt_region_sample.call_stack);
+
+                invoke_callbacks(header.type, _ompt_region_sample);
+                break;
+            }
 #endif
             case entry_type::region:
             {