refactor(rocprofiler-sdk): migrate thread trace off the per-queue callback registry (step 1)

.github/workflows/rocprofiler-sdk-rocm_release_compatibility.yml

@@ -90,6 +90,14 @@ jobs:
         update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 10 --slave /usr/bin/g++ g++ /usr/bin/g++-11 --slave /usr/bin/gcov gcov /usr/bin/gcov-11
         update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 20 --slave /usr/bin/g++ g++ /usr/bin/g++-12 --slave /usr/bin/gcov gcov /usr/bin/gcov-12
         python3 -m pip install -U --user -r requirements.txt
+        # Held at 3.25 rather than the 4.x that requirements.txt resolves to: this
+        # branch's external submodules (json, pybind11, sqlite, yaml-cpp) still declare
+        # cmake_minimum_required below 3.5, which CMake 4 rejects outright. This must
+        # run after requirements.txt so the pin is not upgraded away.
+        python3 -m pip install -U --user cmake==3.25.0
+
+        # Add pip packages to path
+        echo "$HOME/.local/bin" >> "$GITHUB_PATH"
 
         # Add pip packages to path
         echo "$HOME/.local/bin" >> "$GITHUB_PATH"

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/context/context.cpp

@@ -295,6 +295,13 @@ start_context(rocprofiler_context_id_t context_id)
             // conflicting context
             return ROCPROFILER_STATUS_ERROR_CONTEXT_CONFLICT;
         }
+        else if(cfg->dispatch_thread_trace && itr->dispatch_thread_trace &&
+                cfg->dispatch_thread_trace->intersects(*itr->dispatch_thread_trace))
+        {
+            // Two dispatch ATT contexts can run concurrently as long as they target disjoint
+            // sets of GPU agents. Overlapping agent sets would cross-talk in post_kernel_call.
+            return ROCPROFILER_STATUS_ERROR_CONTEXT_CONFLICT;
+        }
     }
 
     uint64_t rocp_tot_contexts = get_registered_contexts_impl()->size();
@@ -376,25 +383,26 @@ stop_context(rocprofiler_context_id_t idx)
         const context* _expected = itr.load(std::memory_order_acquire);
         if(_expected && _expected->context_idx == idx.handle)
         {
+            // Stop queue-interposed services before clearing the active slot so
+            // disable_serialization() always runs before dispatches stop being instrumented.
+            if(_expected->dispatch_counter_collection)
+            {
+                rocprofiler::counters::stop_context(const_cast<context*>(_expected));
+            }
+
+            if(_expected->dispatch_spm)
+                rocprofiler::spm::stop_context(const_cast<context*>(_expected));
+
+            if(_expected->device_thread_trace) _expected->device_thread_trace->stop_context();
+            if(_expected->dispatch_thread_trace) _expected->dispatch_thread_trace->stop_context();
+
             bool success = itr.compare_exchange_strong(_expected, nullptr);
 
             if(success)
             {
                 auto nactive = get_num_active_contexts().load(std::memory_order_acquire);
                 if(nactive > 0) get_num_active_contexts().fetch_sub(1, std::memory_order_release);
 
-                if(_expected->dispatch_counter_collection)
-                {
-                    rocprofiler::counters::stop_context(const_cast<context*>(_expected));
-                }
-
-                if(_expected->dispatch_spm)
-                    rocprofiler::spm::stop_context(const_cast<context*>(_expected));
-
-                if(_expected->device_thread_trace) _expected->device_thread_trace->stop_context();
-                if(_expected->dispatch_thread_trace)
-                    _expected->dispatch_thread_trace->stop_context();
-
                 rocprofiler::hsa::queue_interposition::
                     notify_queue_interposition_consumer_context_stopped(_expected);
 

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/CMakeLists.txt

@@ -39,3 +39,4 @@ target_sources(rocprofiler-sdk-object-library PRIVATE ${ROCPROFILER_LIB_HSA_SOUR
                                                       ${ROCPROFILER_LIB_HSA_HEADERS})
 
 add_subdirectory(details)
+add_subdirectory(queue_hooks)

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp

@@ -45,6 +45,7 @@
 #include "lib/rocprofiler-sdk/pc_sampling/hsa_adapter.hpp"
 #include "lib/rocprofiler-sdk/pc_sampling/service.hpp"
 #include "lib/rocprofiler-sdk/registration.hpp"
+#include "lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp"
 #include "lib/rocprofiler-sdk/tracing/tracing.hpp"
 
 #include <rocprofiler-sdk/callback_tracing.h>
@@ -294,6 +295,15 @@ AsyncSignalHandler(hsa_signal_value_t /*signal_v*/, void* data)
             }
         });
 
+        // Thread trace completion is migrated off the callback registry (see WriteInterceptor);
+        // invoke it explicitly here.
+        thread_trace::signal_completion_hook(queue_info_session.queue,
+                                             packet.kernel_packet,
+                                             _session,
+                                             packet,
+                                             packet.instrumentation_packets,
+                                             dispatch_time);
+
         CHECK_NOTNULL(hsa::get_queue_controller())
             ->serializer(&queue_info_session.queue)
             .wlock([&](auto& serializer) {
@@ -428,8 +438,10 @@ WriteInterceptor(const void* packets,
 
     auto*      gls                 = ::rocprofiler::hip::graph::current_launch_state();
     const bool graph_launch_active = (gls != nullptr);
+    // Thread trace no longer registers a queue-controller callback, so it does not count toward
+    // get_notifiers(); detect it explicitly so an ATT-only run still enters the interceptor.
     const bool no_real_consumers =
-        (queue.get_notifiers() == 0 &&
+        (queue.get_notifiers() == 0 && !thread_trace::is_any_active() &&
          context::get_active_contexts(full_packet_instrumentation_context_filter).empty());
 
     const bool has_kernel_replay = kernel_replay::has_active_replay_contexts();
@@ -827,6 +839,18 @@ WriteInterceptor(const void* packets,
                 }
             });
 
+            // Thread trace is migrated off the per-queue callback registry: call its hook
+            // explicitly (the other services still flow through signal_callback above).
+            thread_trace::write_hook(queue,
+                                     kernel_packet,
+                                     kernel_id,
+                                     dispatch_id,
+                                     &_packet_data.user_data,
+                                     _packet_data.tracing_data.external_correlation_ids,
+                                     corr_id,
+                                     _packet_data.instrumentation_packets,
+                                     _packet_data.is_serialized);
+
             bool inserted_before = false;
             if(_packet_data.is_serialized)
             {
@@ -1195,6 +1219,9 @@ WriteInterceptor(const void* packets,
         }
     });
 
+    // Thread trace requires per-packet mode; it no longer participates in the registry above.
+    if(thread_trace::is_any_active()) should_batch_packets = false;
+
     if(should_batch_packets)
     {
         ROCP_TRACE_IF(pkt_count > 1) << fmt::format(

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp

@@ -677,21 +677,30 @@ common::Synchronized<hsa::profiler_serializer>&
 QueueController::serializer(const Queue* queue)
 {
     CHECK(queue);
+    const auto agent_id = queue->get_agent().get_rocp_agent()->id;
+
+    // Read the refcount before taking the serializer lock: update_serialization() acquires
+    // them in that order, and reversing it here would deadlock. A concurrent enable/disable
+    // for this agent can therefore race with the creation below, exactly as the previous
+    // _serialized_enabled load could, and is resolved the same way -- the next transition
+    // through update_serialization() reaches the entry once it is in the map.
+    const bool should_serialize = is_serialization_enabled(agent_id);
+
     common::Synchronized<hsa::profiler_serializer>* ret = nullptr;
     _profiler_serializer.ulock(
         [&](const auto& m) {
-            if(auto ptr = m.find(queue->get_agent().get_rocp_agent()->id); ptr != m.end())
+            if(auto ptr = m.find(agent_id); ptr != m.end())
             {
                 ret = ptr->second.get();
                 return true;
             }
             return false;
         },
         [&](auto& m) {
-            ret = m.emplace(queue->get_agent().get_rocp_agent()->id,
+            ret = m.emplace(agent_id,
                             std::make_shared<common::Synchronized<hsa::profiler_serializer>>())
                       .first->second.get();
-            if(_serialized_enabled.load() == true)
+            if(should_serialize)
             {
                 ret->wlock([&](auto& serializer) { serializer.enable({}); });
             }
@@ -715,47 +724,120 @@ per_dev_map(const QueueController::queue_map_t& _queues_v)
 };  // namespace
 
 void
-QueueController::disable_serialization()
+QueueController::update_serialization(const agent_handle_set_t& agents, bool enable)
 {
     _queues.rlock([&](const queue_map_t& _queues_v) {
-        _serialized_enabled.store(false);
         auto pd_map = per_dev_map(_queues_v);
-        _profiler_serializer.wlock([&](auto& m) {
-            for(auto& [k, v] : m)
+
+        // Agents whose effective state changed in this call; only those get their serializer
+        // toggled. Each entry stores the agent id together with its new effective state so
+        // that the serializer action is derived from the post-update refcount rather than
+        // from the caller's `enable` flag (which could be stale if another thread races
+        // between the refcount unlock and the serializer lock).
+        auto transitioned = std::vector<std::pair<rocprofiler_agent_id_t, bool>>{};
+        bool any_enabled  = false;
+
+        _serialization_refcount.wlock([&](auto& state) {
+            // An agent can be covered by `all` and by an explicit entry at the same time, so
+            // the transition has to be decided by comparing effective counts across the
+            // update rather than by watching a single counter hit zero.
+            auto was_enabled = std::unordered_map<rocprofiler_agent_id_t, bool>{};
+            _profiler_serializer.rlock([&](const auto& serializers) {
+                for(const auto& [agent_id, _] : serializers)
+                    was_enabled.emplace(agent_id, state.enabled(agent_id));
+            });
+
+            if(agents.empty())
+            {
+                if(enable)
+                    ++state.all;
+                else if(state.all > 0)
+                    --state.all;
+            }
+            else
             {
-                if(auto it = pd_map.find(k); it != pd_map.end())
+                for(const auto& agent_id : agents)
                 {
-                    v->wlock([&](auto& serializer) { serializer.disable(it->second); });
+                    auto& count = state.per_agent[agent_id];
+                    if(enable)
+                        ++count;
+                    else if(count > 0)
+                        --count;
                 }
-                else
+            }
+
+            _profiler_serializer.rlock([&](const auto& serializers) {
+                for(const auto& [agent_id, _] : serializers)
                 {
-                    v->wlock([&](auto& serializer) { serializer.disable({}); });
+                    auto itr = was_enabled.find(agent_id);
+                    if(itr == was_enabled.end()) continue;
+                    bool now_enabled = state.enabled(agent_id);
+                    if(itr->second != now_enabled) transitioned.emplace_back(agent_id, now_enabled);
                 }
+            });
+
+            any_enabled = state.any();
+        });
+
+        _serialized_enabled.store(any_enabled);
+
+        if(transitioned.empty()) return;
+
+        _profiler_serializer.wlock([&](auto& m) {
+            // Deliberately not a structured binding: the serializer lambda below captures
+            // now_enabled, and capturing a structured binding is a C++20 extension.
+            for(const auto& transition : transitioned)
+            {
+                const auto agent_id    = transition.first;
+                const bool now_enabled = transition.second;
+
+                auto itr = m.find(agent_id);
+                if(itr == m.end() || !itr->second) continue;
+
+                auto queues = hsa_barrier::queue_map_ptr_t{};
+                if(auto it = pd_map.find(agent_id); it != pd_map.end()) queues = it->second;
+
+                itr->second->wlock([&](auto& serializer) {
+                    if(now_enabled)
+                        serializer.enable(queues);
+                    else
+                        serializer.disable(queues);
+                });
             }
         });
     });
 }
 
+void
+QueueController::disable_serialization()
+{
+    update_serialization({}, false);
+}
+
 void
 QueueController::enable_serialization()
 {
-    _queues.rlock([&](const queue_map_t& _queues_v) {
-        _serialized_enabled.store(true);
-        auto pd_map = per_dev_map(_queues_v);
-        _profiler_serializer.wlock([&](auto& m) {
-            for(auto& [k, v] : m)
-            {
-                if(auto it = pd_map.find(k); it != pd_map.end())
-                {
-                    v->wlock([&](auto& serializer) { serializer.enable(it->second); });
-                }
-                else
-                {
-                    v->wlock([&](auto& serializer) { serializer.enable({}); });
-                }
-            }
-        });
-    });
+    update_serialization({}, true);
+}
+
+void
+QueueController::disable_serialization(const agent_handle_set_t& agents)
+{
+    update_serialization(agents, false);
+}
+
+void
+QueueController::enable_serialization(const agent_handle_set_t& agents)
+{
+    update_serialization(agents, true);
+}
+
+bool
+QueueController::is_serialization_enabled(rocprofiler_agent_id_t agent_id) const
+{
+    bool enabled = false;
+    _serialization_refcount.rlock([&](const auto& state) { enabled = state.enabled(agent_id); });
+    return enabled;
 }
 
 void