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

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

@@ -35,13 +35,17 @@
 
 #include <rocprofiler-sdk/fwd.h>
 #include <rocprofiler-sdk/registration.h>
+#include <rocprofiler-sdk/cxx/hash.hpp>
+#include <rocprofiler-sdk/cxx/operators.hpp>
 
 #include <array>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <unordered_map>
+#include <unordered_set>
 
 namespace rocprofiler
 {
@@ -196,6 +200,40 @@ stop_context(rocprofiler_context_id_t id);
 using context_array_t    = common::container::small_vector<const context*>;
 using context_id_array_t = common::container::small_vector<rocprofiler_context_id_t>;
 
+// defined in context.cpp
+struct registered_context_storage;
+
+/// \brief An immutable view of the registered contexts, taken at the point of construction.
+///
+/// The registry itself may only be walked with get_contexts_mutex() held: allocate_context()
+/// appends to it while dispatch and completion paths read it. The completion path cannot take
+/// that mutex, because stop_context() waits on the GPU drain and the drain waits on completions,
+/// so it reads a snapshot instead. A snapshot also owns the contexts it names, which is what makes
+/// it safe against deregister_client_contexts() destroying a context mid-iteration.
+class registered_contexts_snapshot
+{
+public:
+    registered_contexts_snapshot() = default;
+
+    const context* const* begin() const { return m_begin; }
+    const context* const* end() const { return m_end; }
+    size_t                size() const { return m_end - m_begin; }
+    bool                  empty() const { return m_begin == m_end; }
+
+private:
+    friend registered_contexts_snapshot get_registered_contexts_snapshot();
+
+    explicit registered_contexts_snapshot(std::shared_ptr<const registered_context_storage>&&);
+
+    std::shared_ptr<const registered_context_storage> m_data  = {};
+    const context* const*                             m_begin = nullptr;
+    const context* const*                             m_end   = nullptr;
+};
+
+/// \brief Returns the currently published snapshot of the registered contexts.
+registered_contexts_snapshot
+get_registered_contexts_snapshot();
+
 context*
 get_mutable_registered_context(rocprofiler_context_id_t id);
 
@@ -207,6 +245,11 @@ using context_filter_t = bool (*)(const context*);
 inline bool
 default_context_filter(const context* val);
 
+// These read the published snapshot and then drop it, so the pointers they return outlive it. The
+// contexts themselves stay alive regardless: deregister_client_contexts() retires a context out of
+// the registry without destroying it, so a pointer obtained here stays dereferenceable until
+// static teardown. Which contexts are reachable is still a snapshot-in-time answer, so a caller
+// that needs a stable set across a long operation should hold a registered_contexts_snapshot.
 context_array_t&
 get_registered_contexts(context_array_t& data, context_filter_t filter = default_context_filter);
 
@@ -242,6 +285,20 @@ deactivate_client_contexts(rocprofiler_client_id_t id);
 void
 deregister_client_contexts(rocprofiler_client_id_t id);
 
+// The mutex that serialises start_context / stop_context and any mutation that must be
+// atomic with activation state (e.g. setting the agent restriction on a context).
+std::mutex&
+get_contexts_mutex();
+
+// Blocks until no context stop is in progress. stop_context() releases get_contexts_mutex()
+// across the GPU drain, so holding the mutex alone is not enough to exclude a stop: for that
+// window the context is still in the active array with nothing held. Anything that reads
+// activation state to decide whether a context is running has to wait here first.
+//
+// _lk must own get_contexts_mutex(); it is released while waiting and re-acquired on return.
+void
+wait_for_stopping_contexts(std::unique_lock<std::mutex>& _lk);
+
 inline bool
 default_context_filter(const context* val)
 {

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/aql_packet.hpp

@@ -247,11 +247,20 @@ class TraceControlAQLPacket : public AQLPacket
     {
         this->tracepool      = other.tracepool;
         this->packets        = other.packets;
+        this->owner          = other.owner;
+        this->owner_state    = other.owner_state;
         this->loaded_codeobj = other.loaded_codeobj;
     }
 
     aqlprofile_handle_t    GetHandle() const { return tracepool->handle; }
     rocprofiler_agent_id_t GetAgent() const { return tracepool->agent_id; }
+    uint64_t               GetOwner() const { return owner; }
+    void                   SetOwner(uint64_t value) { owner = value; }
+
+    // Holds a reference to the per-agent state this packet will deliver into, so a completion
+    // arriving after the owning tracer has torn down its agent map can still be serviced.
+    const std::shared_ptr<void>& GetOwnerState() const { return owner_state; }
+    void SetOwnerState(std::shared_ptr<void> value) { owner_state = std::move(value); }
 
     void populate_before() override
     {
@@ -271,6 +280,8 @@ class TraceControlAQLPacket : public AQLPacket
 protected:
     std::shared_ptr<TraceMemoryPool>     tracepool;
     aqlprofile_att_control_aql_packets_t packets;
+    uint64_t                             owner = 0;
+    std::shared_ptr<void>                owner_state{};
 
     std::unordered_map<code_object_id_t, std::shared_ptr<CodeobjMarkerAQLPacket>> loaded_codeobj;
 };

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::kernel_dispatch_phase_exit_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,16 @@ 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.
+    //
+    // Scoped to this queue's agent: a tracer is configured per agent, so queues on agents it was
+    // never configured for must stay on the fast path instead of paying interception and losing
+    // batching for dispatches that kernel_dispatch_phase_enter_hook() would filter out anyway.
+    const bool thread_trace_active =
+        thread_trace::is_active_on_agent(CHECK_NOTNULL(queue.get_agent().get_rocp_agent())->id);
     const bool no_real_consumers =
-        (queue.get_notifiers() == 0 &&
+        (queue.get_notifiers() == 0 && !thread_trace_active &&
          context::get_active_contexts(full_packet_instrumentation_context_filter).empty());
 
     const bool has_kernel_replay = kernel_replay::has_active_replay_contexts();
@@ -827,6 +845,19 @@ 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::kernel_dispatch_phase_enter_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 +1226,9 @@ WriteInterceptor(const void* packets,
         }
     });
 
+    // Thread trace requires per-packet mode; it no longer participates in the registry above.
+    if(thread_trace_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,26 +677,36 @@ common::Synchronized<hsa::profiler_serializer>&
 QueueController::serializer(const Queue* queue)
 {
     CHECK(queue);
+    const auto agent_id = queue->get_agent().get_rocp_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())
-            {
-                ret = ptr->second.get();
+    // Hold the refcount read lock across the lookup/insertion so that a concurrent
+    // update_serialization() cannot run between reading the refcount and inserting the
+    // entry: if it did, it would find no serializer to transition and the new entry would
+    // be created in the wrong state.  update_serialization() holds the refcount write lock
+    // then the serializer write lock, so taking them in the same order here is deadlock-free.
+    _serialization_refcount.rlock([&](const auto& state) {
+        const bool should_serialize = state.enabled(agent_id);
+        _profiler_serializer.ulock(
+            [&](const auto& m) {
+                if(auto ptr = m.find(agent_id); ptr != m.end())
+                {
+                    ret = ptr->second.get();
+                    return true;
+                }
+                return false;
+            },
+            [&](auto& m) {
+                ret = m.emplace(agent_id,
+                                std::make_shared<common::Synchronized<hsa::profiler_serializer>>())
+                          .first->second.get();
+                if(should_serialize)
+                {
+                    ret->wlock([&](auto& serializer) { serializer.enable({}); });
+                }
                 return true;
-            }
-            return false;
-        },
-        [&](auto& m) {
-            ret = m.emplace(queue->get_agent().get_rocp_agent()->id,
-                            std::make_shared<common::Synchronized<hsa::profiler_serializer>>())
-                      .first->second.get();
-            if(_serialized_enabled.load() == true)
-            {
-                ret->wlock([&](auto& serializer) { serializer.enable({}); });
-            }
-            return true;
-        });
+            });
+    });
     return *ret;
 }
 
@@ -715,47 +725,107 @@ 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 refcount crossed zero in this call; only those get their
+        // serializer toggled. Both refcount mutation and serializer transitions happen
+        // together inside the refcount lock to prevent concurrent calls from interleaving.
+        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.wlock([&](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;
+                    const bool now_enabled = state.enabled(agent_id);
+                    if(itr->second == now_enabled) continue;
+
+                    auto queues = hsa_barrier::queue_map_ptr_t{};
+                    if(auto it = pd_map.find(agent_id); it != pd_map.end()) queues = it->second;
+
+                    auto ser_itr = serializers.find(agent_id);
+                    if(ser_itr == serializers.end() || !ser_itr->second) continue;
+                    // Derived from the post-update refcount rather than the caller's `enable`
+                    // flag, so the applied state cannot disagree with the count that produced
+                    // the transition.
+                    ser_itr->second->wlock([&](auto& serializer) {
+                        if(now_enabled)
+                            serializer.enable(queues);
+                        else
+                            serializer.disable(queues);
+                    });
                 }
-            }
+            });
+
+            any_enabled = state.any();
         });
+
+        _serialized_enabled.store(any_enabled);
     });
 }
 
+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