amd-vkale · Oct 9, 2026 · Sep 22, 2026 · Sep 22, 2026 · Sep 23, 2026 · Sep 23, 2026
@@ -542,6 +542,14 @@ start_context(rocprofiler_context_id_t context_id)
                 // sets of GPU agents. A context with no agent restriction claims every agent.
                 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;
+            }
         }
 
         // try to find a nullptr slot first

@@ -247,11 +247,14 @@ class TraceControlAQLPacket : public AQLPacket
     {
         this->tracepool      = other.tracepool;
         this->packets        = other.packets;
+        this->owner          = other.owner;
         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; }
 
     void populate_before() override
     {
@@ -271,6 +274,7 @@ class TraceControlAQLPacket : public AQLPacket
 protected:
     std::shared_ptr<TraceMemoryPool>     tracepool;
     aqlprofile_att_control_aql_packets_t packets;
+    uint64_t                             owner = 0;
 
     std::unordered_map<code_object_id_t, std::shared_ptr<CodeobjMarkerAQLPacket>> loaded_codeobj;
 };

@@ -47,6 +47,7 @@
 #include "lib/rocprofiler-sdk/pc_sampling/service.hpp"
 #include "lib/rocprofiler-sdk/registration.hpp"
 #include "lib/rocprofiler-sdk/spm/queue_hooks.hpp"
+#include "lib/rocprofiler-sdk/thread_trace/queue_hooks.hpp"
 #include "lib/rocprofiler-sdk/tracing/tracing.hpp"
 
 #include <rocprofiler-sdk/callback_tracing.h>
@@ -314,6 +315,15 @@ AsyncSignalHandler(hsa_signal_value_t /*signal_v*/, void* data)
                                              packet.instrumentation_packets,
                                              dispatch_time);
 
+        // 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) {
@@ -455,8 +465,17 @@ WriteInterceptor(const void* packets,
     // dispatches that kernel_dispatch_phase_enter_hook() would filter out anyway.
     const bool spm_active =
         spm::is_active_on_agent(CHECK_NOTNULL(queue.get_agent().get_rocp_agent())->id);
+
+    // 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 && !spm_active &&
+        (queue.get_notifiers() == 0 && !spm_active && !thread_trace_active &&
          !pc_sampling::is_configured_on_agent(queue.get_agent().get_rocp_agent()->id) &&
          context::get_active_contexts(full_packet_instrumentation_context_filter).empty());
 
@@ -868,6 +887,19 @@ WriteInterceptor(const void* packets,
                 _packet_data.instrumentation_packets,
                 _packet_data.is_serialized);
 
+            // Thread trace is migrated off the per-queue callback registry: call its hook
+            // explicitly (only counter collection still flows 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)
             {
@@ -1238,6 +1270,8 @@ WriteInterceptor(const void* packets,
 
     // SPM requires per-packet mode; it no longer participates in the registry above.
     if(spm_active) should_batch_packets = false;
+    // 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)
     {

@@ -1,7 +1,8 @@
-set(ROCPROFILER_LIB_THREAD_TRACE_SOURCES core.cpp service.cpp code_object.cpp decode.cpp
-                                         dl.cpp threading.cpp hsa_util.cpp)
+set(ROCPROFILER_LIB_THREAD_TRACE_SOURCES
+    core.cpp service.cpp code_object.cpp decode.cpp dl.cpp threading.cpp hsa_util.cpp
+    queue_hooks.cpp)
 set(ROCPROFILER_LIB_THREAD_TRACE_HEADERS core.hpp code_object.hpp dl.hpp threading.hpp
-                                         hsa_util.hpp trace_decoder_api.h)
+                                         hsa_util.hpp trace_decoder_api.h queue_hooks.hpp)
 target_sources(
     rocprofiler-sdk-object-library PRIVATE ${ROCPROFILER_LIB_THREAD_TRACE_SOURCES}
                                            ${ROCPROFILER_LIB_THREAD_TRACE_HEADERS})

@@ -35,6 +35,7 @@
 #include "lib/rocprofiler-sdk/context/context.hpp"
 #include "lib/rocprofiler-sdk/hsa/queue.hpp"
 #include "lib/rocprofiler-sdk/hsa/queue_controller.hpp"
+#include "lib/rocprofiler-sdk/hsa/queue_hooks/client_ids.hpp"
 #include "lib/rocprofiler-sdk/internal_threading.hpp"
 #include "lib/rocprofiler-sdk/kernel_replay/local_context.hpp"
 #include "lib/rocprofiler-sdk/registration.hpp"
@@ -51,6 +52,7 @@
 
 #include <atomic>
 #include <cstdint>
+#include <memory>
 #include <mutex>
 #include <stdexcept>
 #include <thread>
@@ -82,10 +84,6 @@ struct cbdata_t
     aqlprofile_handle_t handle{};
 };
 
-// Keeps track of a single client registering for serialized thread trace
-// operations so we can gate new traces while one is active.
-common::Synchronized<std::optional<int64_t>> client;
-
 // True once the HSA runtime is registered. Gates start_context() so pre-init
 // start requests are deferred and replayed by start_active_contexts().
 std::atomic<bool>&
@@ -517,24 +515,32 @@ DispatchThreadTracer::resource_init()
     }
 }
 
+uint64_t
+DispatchThreadTracer::allocate_tracer_id()
+{
+    static auto _counter = std::atomic<uint64_t>{1};
+    return _counter.fetch_add(1, std::memory_order_relaxed);
+}
+
 void
 DispatchThreadTracer::resource_deinit()
 {
-    auto was_enabled = enabled.exchange(false, std::memory_order_acq_rel);
-
-    if(auto* controller = hsa::get_queue_controller())
+    if(enabled.exchange(false, std::memory_order_acq_rel))
     {
-        client.wlock([&](auto& client_id) {
-            if(!client_id) return;
-            controller->remove_callback(*client_id);
-            client_id = std::nullopt;
-        });
-        if(was_enabled) controller->disable_serialization();
+        const auto serialization_agents = configured_agents();
+        if(auto* controller = hsa::get_queue_controller())
+            controller->disable_serialization(serialization_agents);
     }
 
     ROCP_TRACE << "Clearing agents";
-    auto lk = std::unique_lock{agents_map_mut};
-    agents.clear();
+    // Destroy the agents only after releasing agents_map_mut: a completion can be waiting for it
+    // in post_kernel_call on ROCr's async signal handler thread, and on the HSA backend an
+    // agent's destructor calls hsa_queue_destroy(), which waits on that same thread.
+    auto retired = decltype(agents){};
+    {
+        auto lk = std::unique_lock{agents_map_mut};
+        retired.swap(agents);
+    }
 }
 
 /**
@@ -557,8 +563,6 @@ DispatchThreadTracer::pre_kernel_call(const hsa::Queue&              queue,
     }
     // TODO: Get external
 
-    if(!enabled.load(std::memory_order_acquire)) return {nullptr, false};
-
     std::shared_lock<std::shared_mutex> lk(agents_map_mut);
 
     auto it = agents.find(queue.get_agent().get_rocp_agent()->id);
@@ -568,6 +572,12 @@ DispatchThreadTracer::pre_kernel_call(const hsa::Queue&              queue,
     auto&       agent      = *CHECK_NOTNULL(it->second);
     const auto& parameters = agent.params;
 
+    // stop_context() clears enabled before it drops the serialization reference, and the context
+    // stays in the active array until both are done. A dispatch in that window is treated like an
+    // untraced one: with SERIALIZE_ALL it must still take the barriers, or it overlaps a traced
+    // dispatch that is still running. Once serialization is disabled the serializer ignores it.
+    if(!enabled.load(std::memory_order_acquire)) return {nullptr, parameters.bSerialize};
+
     // Kernel-replay localized context control: a replay pass may disable this ATT context for the
     // pass -- skip the trace (but keep serialization) when it's forced off. No-op outside a replay
     // loop. See kernel_replay/local_context.hpp.
@@ -587,6 +597,7 @@ DispatchThreadTracer::pre_kernel_call(const hsa::Queue&              queue,
 
     auto packet = agent.get_start_packet();
     if(!packet) return {nullptr, parameters.bSerialize};
+    packet->SetOwner(tracer_id);
     post_move_data.fetch_add(1);
     packet->populate_before();
     packet->populate_after();
@@ -602,8 +613,11 @@ DispatchThreadTracer::post_kernel_call(DispatchThreadTracer::inst_pkt_t& aql,
 
     for(auto& aql_pkt : aql)
     {
+        if(aql_pkt.second != hsa::queue_hooks::THREAD_TRACE_CLIENT_ID) continue;
+
         auto* pkt = dynamic_cast<hsa::TraceControlAQLPacket*>(aql_pkt.first.get());
         if(!pkt) continue;
+        if(pkt->GetOwner() != tracer_id) continue;
 
         std::shared_lock<std::shared_mutex> lk(agents_map_mut);
         post_move_data.fetch_sub(1);
@@ -616,55 +630,62 @@ DispatchThreadTracer::post_kernel_call(DispatchThreadTracer::inst_pkt_t& aql,
     }
 }
 
+std::unordered_set<rocprofiler_agent_id_t>
+DispatchThreadTracer::configured_agents() const
+{
+    auto                                result = std::unordered_set<rocprofiler_agent_id_t>{};
+    std::shared_lock<std::shared_mutex> lk(agents_map_mut);
+    for(const auto& [agent_id, _] : params)
+        result.insert(agent_id);
+    return result;
+}
+
+bool
+DispatchThreadTracer::collects_on(rocprofiler_agent_id_t agent_id) const
+{
+    std::shared_lock<std::shared_mutex> lk(agents_map_mut);
+    return params.count(agent_id) > 0;
+}
+
+bool
+DispatchThreadTracer::intersects(const DispatchThreadTracer& rhs) const
+{
+    if(this == &rhs)
+    {
+        std::shared_lock<std::shared_mutex> lk(agents_map_mut);
+        return !params.empty();
+    }
+
+    auto lk     = std::shared_lock<std::shared_mutex>{agents_map_mut, std::defer_lock};
+    auto lk_rhs = std::shared_lock<std::shared_mutex>{rhs.agents_map_mut, std::defer_lock};
+    std::lock(lk, lk_rhs);
+    for(const auto& [agent_id, _] : params)
+    {
+        if(rhs.params.count(agent_id) > 0) return true;
+    }
+    return false;
+}
+
 void
 DispatchThreadTracer::start_context()
 {
-    using corr_id_map_t = hsa::queue_info_session_t::external_corr_id_map_t;
-
-    // Only installs queue-controller callbacks (cached and applied to queues as
-    // they are created), so this is safe to call before hsa_init.
-    CHECK_NOTNULL(hsa::get_queue_controller())->enable_serialization();
+    // An empty agent set means every agent.
+    const auto serialization_agents = configured_agents();
+    CHECK_NOTNULL(hsa::get_queue_controller())->enable_serialization(serialization_agents);
     enabled.store(true, std::memory_order_release);
-
-    // Only one thread should be attempting to enable/disable this context
-    client.wlock([&](auto& client_id) {
-        if(client_id) return;
-
-        auto&& _callbacks = hsa::queue_callbacks_t{
-            .batch_packets = []() { return false; },
-            .write_interceptor =
-                [=](const hsa::Queue& q,
-                    const hsa::rocprofiler_packet& /* kern_pkt */,
-                    rocprofiler_kernel_id_t   kernel_id,
-                    rocprofiler_dispatch_id_t dispatch_id,
-                    rocprofiler_user_data_t*  user_data,
-                    const corr_id_map_t& /* extern_corr_ids */,
-                    const context::correlation_id* corr_id) {
-                    return this->pre_kernel_call(q, kernel_id, dispatch_id, user_data, corr_id);
-                },
-            .signal_completion =
-                [=](const hsa::Queue& /* q */,
-                    hsa::rocprofiler_packet /* kern_pkt */,
-                    std::shared_ptr<hsa::queue_info_session_t>& session,
-                    hsa::packet_data_t&                         packet_data,
-                    inst_pkt_t&                                 aql,
-                    kernel_dispatch::profiling_time) {
-                    this->post_kernel_call(aql, *session, packet_data);
-                }};
-
-        client_id = CHECK_NOTNULL(hsa::get_queue_controller())
-                        ->add_callback(std::nullopt, std::move(_callbacks));
-    });
 }
 
 void
-DispatchThreadTracer::stop_context()  // NOLINT(readability-convert-member-functions-to-static)
+DispatchThreadTracer::stop_context()
 {
-    // Stop injecting ATT packets before transitioning serialization. Completion callbacks remain
-    // registered so packets already in the queues can drain through the serializer transition.
+    // Stop injecting ATT packets before transitioning serialization. Completion hooks continue
+    // to route already-tagged packets via kernel_dispatch_phase_exit_hook even after the context
+    // stops.
     if(!enabled.exchange(false, std::memory_order_acq_rel)) return;
 
-    if(auto* controller = hsa::get_queue_controller()) controller->disable_serialization();
+    const auto serialization_agents = configured_agents();
+    if(auto* controller = hsa::get_queue_controller())
+        controller->disable_serialization(serialization_agents);
 }
 
 DeviceThreadTracer::DeviceThreadTracer()

@@ -45,6 +45,7 @@
 #include <string>
 #include <tuple>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace rocprofiler
@@ -164,6 +165,9 @@ class DispatchThreadTracer
     void resource_init();
     void resource_deinit();
 
+    bool collects_on(rocprofiler_agent_id_t agent_id) const;
+    bool intersects(const DispatchThreadTracer& rhs) const;
+
     void add_agent(rocprofiler_agent_id_t agent, thread_trace_parameter_pack pack)
     {
         auto lk       = std::unique_lock{agents_map_mut};
@@ -181,13 +185,23 @@ class DispatchThreadTracer
                                  const hsa::packet_data_t&        packet_data);
     const auto& get_agents() const { return agents; }
 
+    /// Number of injected trace packets still awaiting post_kernel_call.
+    int pending_post_moves() const { return post_move_data.load(); }
+
+    std::unordered_set<rocprofiler_agent_id_t> configured_agents() const;
+
 private:
+    /// Hands out a value that is never reused, so a packet outliving its tracer cannot be
+    /// mistaken for one produced by a later tracer that happens to occupy the same address.
+    static uint64_t allocate_tracer_id();
+
     std::unordered_map<rocprofiler_agent_id_t, std::unique_ptr<ThreadTracerAgent>> agents{};
     std::unordered_map<rocprofiler_agent_id_t, thread_trace_parameter_pack>        params{};
 
-    std::shared_mutex agents_map_mut{};
-    std::atomic<int>  post_move_data{0};
-    std::atomic<bool> enabled{false};
+    const uint64_t            tracer_id = allocate_tracer_id();
+    mutable std::shared_mutex agents_map_mut{};
+    std::atomic<int>          post_move_data{0};
+    std::atomic<bool>         enabled{false};
 };
 
 class DeviceThreadTracer