amd-vkale · Aug 25, 2026 · Sep 18, 2026 · Sep 18, 2026 · Sep 19, 2026 · Sep 19, 2026
@@ -217,6 +217,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);
 
@@ -228,6 +262,10 @@ 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 the
+// guarantee that the contexts are alive. That is fine for callers that cannot overlap with
+// deregister_client_contexts(); anything on a dispatch or completion path should hold a
+// registered_contexts_snapshot for as long as it uses the pointers.
 context_array_t&
 get_registered_contexts(context_array_t& data, context_filter_t filter = default_context_filter);
 
@@ -268,6 +306,15 @@ deregister_client_contexts(rocprofiler_client_id_t id);
 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)
 {

@@ -226,6 +226,12 @@ set_dispatch_agents(rocprofiler_context_id_t      context_id,
     // set for serialization) and then race with the unordered-set assignment below while
     // dispatch hooks read it.
     auto _lk = std::unique_lock<std::mutex>{context::get_contexts_mutex()};
+
+    // A context that is mid-stop is still in the active array while its GPU drain runs, and that
+    // drain happens with the mutex released, so the lock alone does not exclude it. Without this
+    // the scan below would read a context that is on its way out as a running one.
+    context::wait_for_stopping_contexts(_lk);
+
     for(const auto* itr : context::get_active_contexts())
     {
         if(itr && itr->context_idx == ctx_p->context_idx)

@@ -698,7 +698,17 @@ device_counting_service_hsa_registration()
 
 agent_callback_data::~agent_callback_data()
 {
-    if(completion.handle != 0) hsa::get_core_table()->hsa_signal_destroy_fn(completion);
+    if(completion.handle == 0) return;
+
+    // Guarded the same way as hsa::active_signals::destroy(). The owning context is not
+    // guaranteed to be destroyed while HSA is still up -- teardown can reach here after the
+    // runtime has been shut down, and the table pointer is also null under unit test -- so
+    // reaching through the table unconditionally faults instead of leaking a signal.
+    const auto* _core_table = hsa::get_core_table();
+    if(hsa::get_hsa_ref_count() > 0 && _core_table && _core_table->hsa_signal_destroy_fn)
+    {
+        _core_table->hsa_signal_destroy_fn(completion);
+    }
 }
 }  // namespace counters
 }  // namespace rocprofiler
@@ -84,7 +84,7 @@ kernel_dispatch_phase_enter_hook(
 }
 
 void
-kernel_dispatch_phase_exit_hook(const hsa::Queue& /*queue*/,
+kernel_dispatch_phase_exit_hook(const hsa::Queue* /*queue*/,
                                 const hsa::rocprofiler_packet& /*kernel_packet*/,
                                 std::shared_ptr<hsa::queue_info_session_t>& session,
                                 hsa::packet_data_t&                         packet,
@@ -107,9 +107,17 @@ kernel_dispatch_phase_exit_hook(const hsa::Queue& /*queue*/,
     // context from the active list. This is what guarantees a completion that arrives is delivered;
     // the drain in stop_context is what bounds when completions arrive. Kernel replay needs the
     // same property, since each pass completes separately.
-    auto contexts = context::get_registered_contexts(counter_contexts_filter());
+    //
+    // Read through a snapshot rather than the registry itself. This runs on the HSA async signal
+    // handler, so it cannot take get_contexts_mutex() -- stop_context() waits on the drain that
+    // this completion is part of -- and the registry is not safe to walk unlocked while
+    // allocate_context() appends to it. Holding the snapshot also keeps every context it names
+    // alive for the duration of the loop.
+    auto contexts = context::get_registered_contexts_snapshot();
     for(const auto* ctx : contexts)
     {
+        if(!counter_contexts_filter()(ctx)) continue;
+
         for(auto& cb : ctx->dispatch_counter_collection->callbacks)
         {
             completed_cb(ctx, cb, session, packet, inst_pkt, dispatch_time);

@@ -55,8 +55,10 @@ kernel_dispatch_phase_enter_hook(
 // registered dispatch_counter_collection contexts (not only active ones) and calls
 // each callback's completed_cb; completed_cb self-filters via packet_return_map so
 // in-flight dispatches still complete after stop_context.
+// queue is nullable: the hook never dereferences it, and tests exercise the
+// early-return path without an HSA runtime to build a queue from.
 void
-kernel_dispatch_phase_exit_hook(const hsa::Queue&                           queue,
+kernel_dispatch_phase_exit_hook(const hsa::Queue*                           queue,
                                 const hsa::rocprofiler_packet&              kernel_packet,
                                 std::shared_ptr<hsa::queue_info_session_t>& session,
                                 hsa::packet_data_t&                         packet,

@@ -176,25 +176,19 @@ TEST(counters_queue_hooks, is_any_active_false_when_no_context_active)
 
 TEST(counters_queue_hooks, exit_hook_skips_when_inst_pkt_has_no_counter_client_id)
 {
-    ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS);
-    test_init();
-
-    auto agents = hsa::get_queue_controller()->get_supported_agents();
-    ASSERT_FALSE(agents.empty());
-    const auto& [_, agent] = *agents.begin();
-    hsa::QueueHooksFakeQueue fq(agent, {.handle = 1});
-
     hsa::inst_pkt_t inst_pkt;
     inst_pkt.emplace_back(std::make_pair(std::make_unique<rocprofiler::hsa::EmptyAQLPacket>(),
                                          rocprofiler::hsa::queue_hooks::THREAD_TRACE_CLIENT_ID));
 
-    auto sess = std::make_shared<hsa::queue_info_session_t>(hsa::queue_info_session_t{.queue = fq});
+    // The hook returns before it dereferences the session, so a null one keeps this test free of
+    // any HSA runtime: queue_info_session_t holds a Queue& and cannot be default constructed.
+    auto sess   = std::shared_ptr<rocprofiler::hsa::queue_info_session_t>{};
     auto packet = rocprofiler::hsa::packet_data_t{};
     auto fq_pkt = rocprofiler::hsa::rocprofiler_packet{};
 
     // Must return without touching registered counter contexts.
     rocprofiler::counters::kernel_dispatch_phase_exit_hook(
-        fq, fq_pkt, sess, packet, inst_pkt, rocprofiler::kernel_dispatch::profiling_time{});
+        nullptr, fq_pkt, sess, packet, inst_pkt, rocprofiler::kernel_dispatch::profiling_time{});
     SUCCEED();
 }
 
@@ -353,7 +347,7 @@ TEST(counters_queue_hooks, stop_context_in_flight_completion_routes_via_hook_pat
         inst_pkt.emplace_back(std::move(in_flight[i]), hsa::queue_hooks::COUNTERS_CLIENT_ID);
 
         rocprofiler::counters::kernel_dispatch_phase_exit_hook(
-            fq, pkt, sess, packet_data, inst_pkt, rocprofiler::kernel_dispatch::profiling_time{});
+            &fq, pkt, sess, packet_data, inst_pkt, rocprofiler::kernel_dispatch::profiling_time{});
 
         size_t remaining = num_dispatches;
         cb_info->packet_return_map.rlock([&](const auto& data) { remaining = data.size(); });

@@ -297,7 +297,7 @@ AsyncSignalHandler(hsa_signal_value_t /*signal_v*/, void* data)
 
         // Counter collection completion is migrated off the callback registry (see
         // WriteInterceptor); invoke it explicitly here.
-        counters::kernel_dispatch_phase_exit_hook(queue_info_session.queue,
+        counters::kernel_dispatch_phase_exit_hook(&queue_info_session.queue,
                                                   packet.kernel_packet,
                                                   _session,
                                                   packet,

@@ -769,15 +769,19 @@ QueueController::update_serialization(const agent_handle_set_t& agents, bool ena
                 {
                     auto itr = was_enabled.find(agent_id);
                     if(itr == was_enabled.end()) continue;
-                    if(itr->second == state.enabled(agent_id)) 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(enable)
+                        if(now_enabled)
                             serializer.enable(queues);
                         else
                             serializer.disable(queues);

@@ -30,30 +30,35 @@ namespace hsa
 {
 namespace queue_hooks
 {
-// Tags identifying the producer subsystem of each inst_pkt_t entry.
-// Only distinctness matters; values are fixed for test-order stability.
-// Services are migrated off the per-queue callback registry one at a time; the
-// unused ids are reserved so the tag values stay stable across those PRs.
+// Tags identifying the producer subsystem of each inst_pkt_t entry, stored in the
+// inst_pkt_t pair as an hsa::ClientID.
 //
-// The underlying type is fixed to int64_t to match hsa::ClientID, which is how the
-// tag is stored in inst_pkt_t. The enum is deliberately unscoped rather than an
-// enum class: unscoped keeps the implicit conversion to ClientID, so the emplace
-// sites need no cast, while still grouping the values under one type.
+// The underlying type is fixed to int64_t to match hsa::ClientID. The enum is
+// deliberately unscoped rather than an enum class: unscoped keeps the implicit
+// conversion to ClientID, so the emplace sites need no cast, while still grouping
+// the values under one type.
 //
-// These values numerically overlap the per-queue registry's auto-incrementing
-// ClientID, which also starts at 1 and is still in use by the services that have
-// not been migrated yet. That is inert only because no consumer routes on the tag:
-// counters' completed_cb and thread trace's post_kernel_call both identify their
-// own packets by pointer lookup or dynamic_cast. Before adding a consumer that
-// dispatches on these values, either finish migrating the remaining services off
-// the registry or move this enum to a range the registry cannot produce.
+// The values are negative because they share the inst_pkt_t tag space with the
+// per-queue callback registry, whose QueueController::add_callback hands out
+// strictly positive ids counting up from 1. Migrated services (counters, thread
+// trace, SPM) select their own packets by comparing against these tags, so a tag
+// drawn from the registry's range could match a packet the service does not own.
+// Services that have not migrated yet still receive registry ids, so the two
+// ranges must stay disjoint until the registry is gone. Any tag added here must
+// likewise be negative; -1 is left unused because QueueController::add_callback
+// uses it as a local "unassigned" initializer.
 enum client_id : int64_t
 {
-    COUNTERS_CLIENT_ID     = 1,
-    THREAD_TRACE_CLIENT_ID = 2,
-    PC_SAMPLING_CLIENT_ID  = 3,
-    SPM_CLIENT_ID          = 4,
+    COUNTERS_CLIENT_ID     = -1001,
+    THREAD_TRACE_CLIENT_ID = -1002,
+    PC_SAMPLING_CLIENT_ID  = -1003,
+    SPM_CLIENT_ID          = -1004,
 };
+
+static_assert(COUNTERS_CLIENT_ID < 0 && THREAD_TRACE_CLIENT_ID < 0 && PC_SAMPLING_CLIENT_ID < 0 &&
+                  SPM_CLIENT_ID < 0,
+              "queue hook tags must stay disjoint from the positive ids handed out by "
+              "QueueController::add_callback");
 }  // namespace queue_hooks
 }  // namespace hsa
 }  // namespace rocprofiler
@@ -91,6 +91,10 @@ rocprofiler_context_id_t ctx_all_second  = {.handle = 0};
 // reorder HIP ordinals.
 std::vector<rocprofiler_agent_id_t> gpu_agents = {};
 
+// HIP ordinal that could not be matched to a rocprofiler agent, or -1 when every device
+// mapped. Reported by main() so a skip is distinguishable from a machine with one GPU.
+int unmapped_device = -1;
+
 // Intermediate storage used to build the HIP-to-rocprofiler mapping: keeps the
 // location_id (BDF) and PCI domain alongside the agent id until the mapping is done.
 struct agent_info_t
@@ -254,22 +258,36 @@ tool_init(rocprofiler_client_finalize_t, void*)
     int device_count = 0;
     if(hipGetDeviceCount(&device_count) != hipSuccess || device_count < 2) return 0;
 
+    // A device with no matching agent must not be skipped silently: dropping it would shift
+    // every later ordinal, so gpu_agents[i] would stop being HIP device i and the contexts
+    // would be scoped to a different physical GPU than the workload runs on. Bail instead and
+    // let main() report the skip.
     for(int dev = 0; dev < device_count; ++dev)
     {
         hipDeviceProp_t prop{};
         if(hipGetDeviceProperties(&prop, dev) != hipSuccess) return 0;
         const auto bus = static_cast<uint32_t>(prop.pciBusID);
         const auto did = static_cast<uint32_t>(prop.pciDeviceID);
         const auto dom = static_cast<uint32_t>(prop.pciDomainID);
+
+        bool matched = false;
         for(const auto& info : all_gpu_agent_info)
         {
             if(info.domain == dom && (info.location_id >> 8) == bus &&
                ((info.location_id >> 3) & 0x1Fu) == did)
             {
                 gpu_agents.push_back(info.id);
+                matched = true;
                 break;
             }
         }
+
+        // Only the first two ordinals are used, so an unmapped device beyond them is harmless.
+        if(!matched)
+        {
+            unmapped_device = dev;
+            break;
+        }
     }
 
     if(gpu_agents.size() < 2) return 0;
@@ -371,7 +389,12 @@ main()
 
     if(device_count < 2 || gpu_agents.size() < 2)
     {
-        printf("SKIP: per-agent scoping requires at least 2 GPUs (found %d)\n", device_count);
+        if(unmapped_device >= 0 && unmapped_device < 2)
+            printf("SKIP: HIP device %d has no rocprofiler agent with matching PCI "
+                   "coordinates, so per-agent scoping cannot be checked on this machine\n",
+                   unmapped_device);
+        else
+            printf("SKIP: per-agent scoping requires at least 2 GPUs (found %d)\n", device_count);
         return 0;
     }
 
@@ -406,6 +429,17 @@ main()
            scoped_ms,
            kNumStreams,
            kKernelsPer);
+    // This test only runs on a multi-GPU machine, so print the mapping and the raw counts the
+    // assertions below read: a control failure is otherwise indistinguishable from a scoping bug.
+    printf("hip device 0 -> agent %lu, hip device 1 -> agent %lu\n",
+           static_cast<unsigned long>(agent_0.handle),
+           static_cast<unsigned long>(agent_1.handle));
+    printf("unrestricted: dispatches(gpu0)=%zu | scoped: records(gpu1)=%zu dispatches(gpu0)=%zu "
+           "records(gpu0)=%zu\n",
+           unrestricted_dispatches_gpu0,
+           scoped_records_gpu1,
+           scoped_dispatches_gpu0,
+           scoped_records_gpu0);
 
     // Exact assertions. These hold regardless of how fast the machine is.
     check(unrestricted_dispatches_gpu0 > 0,