amd-vkale · Aug 18, 2026 · Jul 30, 2026 · Jul 30, 2026 · Jul 31, 2026 · Jul 31, 2026
@@ -368,25 +368,30 @@ 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.
+            // Clearing the slot first opens the opposite window, in which write_hook sees no
+            // active context and a dispatch is submitted without serializer packets while the
+            // serializer is still enabled.
+            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);
 

@@ -290,6 +290,27 @@ stop_context(const context::context* ctx)
 
     if(controller)
     {
+        // Drain in-flight dispatches, then disable serialization. The review of #8887 asked for
+        // this sync to be kept and for SPM to stay visible to the enter hook and to serialization
+        // through a "draining" window until both this and disable_serialization() complete.
+        //
+        // The visibility half comes from ordering rather than a separate flag:
+        // context::stop_context calls this function while the context is still in the active list,
+        // so for the duration of the drain the enter hook still reaches pre_kernel_call, whose
+        // disabled path deliberately returns serialize=true to coordinate the
+        // serialized->unserialized transition. Only once this function returns does the active slot
+        // get cleared. Introducing a draining flag as well would duplicate that guarantee, but the
+        // ordering is load-bearing: if the slot were cleared first, the enter hook would go blind
+        // before disable_serialization() ran and a dispatch could be submitted unserialized while
+        // the serializer was still enabled.
+        //
+        // A removal of this sync was considered on the grounds that its two original purposes --
+        // avoiding dangling callback pointers once the per-queue callbacks were unregistered, and
+        // letting in-flight dispatches complete -- are both addressed by the migration and by
+        // routing completions over registered contexts. The first is genuinely gone. The second is
+        // weaker than it looks: provenance routing ensures a completion that arrives is delivered,
+        // whereas the drain bounds when completions arrive at all, which is what the rest of the
+        // teardown depends on. Keeping it.
         hsa::queue_controller_sync();
         controller->disable_serialization();
         // No per-queue callback to remove; spm::write_hook no-ops once dispatch_spm is

@@ -31,7 +31,7 @@ namespace spm
 namespace
 {
 auto
-active_spm_contexts_filter()
+spm_contexts_filter()
 {
     return [](const context::context* ctx) -> bool { return ctx && ctx->dispatch_spm != nullptr; };
 }
@@ -48,7 +48,7 @@ write_hook(const hsa::Queue&                                        queue,
            hsa::inst_pkt_t&                                         inst_pkt,
            bool&                                                    is_serialized)
 {
-    auto active = context::get_active_contexts(active_spm_contexts_filter());
+    auto active = context::get_active_contexts(spm_contexts_filter());
     for(const auto* ctx : active)
     {
         for(auto& cb : ctx->dispatch_spm->callbacks)
@@ -76,8 +76,18 @@ signal_completion_hook(const hsa::Queue& /*queue*/,
                        hsa::inst_pkt_t&                inst_pkt,
                        kernel_dispatch::profiling_time dispatch_time)
 {
-    auto active = context::get_active_contexts(active_spm_contexts_filter());
-    for(const auto* ctx : active)
+    // Route by packet provenance, not current activeness: post_kernel_call self-filters via
+    // packet_return_map, so in-flight dispatches still complete after stop_context removes the
+    // context from the active list. Without this, a dispatch in flight at stop skips the
+    // DISPATCH_END record, kfd_stop() and the barrier-signal cleanup, and leaks its map entry.
+    //
+    // Iterating registered contexts widens the candidate set, but routing stays per-origin because
+    // each context's post_kernel_call only claims packets present in its own packet_return_map.
+    // That matters because there is no conflict guard rejecting a second concurrent dispatch_spm
+    // context, unlike dispatch_counter_collection, so more than one SPM context can be registered
+    // at once.
+    auto contexts = context::get_registered_contexts(spm_contexts_filter());
+    for(const auto* ctx : contexts)
     {
         for(auto& cb : ctx->dispatch_spm->callbacks)
         {
@@ -89,7 +99,7 @@ signal_completion_hook(const hsa::Queue& /*queue*/,
 bool
 is_any_active()
 {
-    return !context::get_active_contexts(active_spm_contexts_filter()).empty();
+    return !context::get_active_contexts(spm_contexts_filter()).empty();
 }
 }  // namespace spm
 }  // namespace rocprofiler
@@ -49,8 +49,10 @@ write_hook(const hsa::Queue&                                        queue,
            hsa::inst_pkt_t&                                         inst_pkt,
            bool&                                                    is_serialized);
 
-// Explicit replacement for the SPM completion callback. Iterates active
-// dispatch_spm contexts and calls each callback's post_kernel_call.
+// Explicit replacement for the SPM completion callback. Iterates registered
+// dispatch_spm contexts (not only active ones) and calls each callback's
+// post_kernel_call; post_kernel_call self-filters via packet_return_map so
+// in-flight dispatches still complete after stop_context.
 void
 signal_completion_hook(const hsa::Queue&                           queue,
                        const hsa::rocprofiler_packet&              kernel_packet,

@@ -29,10 +29,12 @@
 #include "lib/rocprofiler-sdk/hsa/hsa.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/kernel_dispatch/profiling_time.hpp"
 #include "lib/rocprofiler-sdk/registration.hpp"
 #include "lib/rocprofiler-sdk/spm/decode.hpp"
 #include "lib/rocprofiler-sdk/spm/dispatch_handlers.hpp"
+#include "lib/rocprofiler-sdk/spm/queue_hooks.hpp"
 
 #include <rocprofiler-sdk/dispatch_counting_service.h>
 #include <rocprofiler-sdk/experimental/spm.h>
@@ -928,7 +930,7 @@ TEST(spm_core, stop_context_removes_callbacks)
     ASSERT_TRUE(ctx.dispatch_spm);
     ASSERT_EQ(ctx.dispatch_spm->callbacks.size(), 1);
 
-    // Stop exercises queue_controller_sync + context teardown.
+    // Stop disables collection, drains in-flight dispatches, then tears the context down.
     ROCPROFILER_CALL(rocprofiler_stop_context(get_client_ctx()), "stop context");
 
     bool enabled = true;
@@ -969,7 +971,7 @@ TEST(spm_core, stop_context_sync_and_restart)
 
     ROCPROFILER_CALL(rocprofiler_stop_context(get_client_ctx()), "stop context");
 
-    // Restart to verify queue controller is in a clean state after sync + teardown
+    // Restart to verify the queue controller is in a clean state after drain + teardown
     ROCPROFILER_CALL(rocprofiler_start_context(get_client_ctx()), "restart context");
 
     bool enabled = false;
@@ -983,3 +985,129 @@ TEST(spm_core, stop_context_sync_and_restart)
     context::pop_client(1);
     set_client_ctx(get_client_ctx());
 }
+
+// Regression for callback-registry removal (#8887): a dispatch enqueued while the SPM context
+// is active must still complete (and drain packet_return_map) when stop_context runs before the
+// GPU completion callback. Exercises the queue_hooks routing layer, not pre/post_kernel_call
+// directly.
+TEST(spm_queue_hooks, stop_context_in_flight_completion_routes_via_hook_path)
+{
+    rocprofiler::common::set_env("ROCPROFILER_SPM_BETA_ENABLED", true);
+    ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS);
+    test_init();
+
+    registration::init_logging();
+    registration::set_init_status(-1);
+    context::push_client(1);
+
+    ROCPROFILER_CALL(rocprofiler_create_context(&get_client_ctx()), "context creation failed");
+    ROCPROFILER_CALL(
+        rocprofiler_spm_configure_callback_dispatch_service(
+            get_client_ctx(), user_dispatch_cb, nullptr, null_record_callback, nullptr),
+        "Could not setup SPM service");
+    ROCPROFILER_CALL(rocprofiler_start_context(get_client_ctx()), "start context");
+
+    auto* ctx_p = context::get_mutable_registered_context(get_client_ctx());
+    ASSERT_TRUE(ctx_p);
+    ASSERT_TRUE(ctx_p->dispatch_spm);
+    auto cb_info = ctx_p->dispatch_spm->callbacks.front();
+
+    ASSERT_TRUE(hsa::get_queue_controller() != nullptr);
+    auto agents = hsa::get_queue_controller()->get_supported_agents();
+    ASSERT_GT(agents.size(), 0);
+
+    bool found_spm_agent = false;
+    for(const auto& [_, agent] : agents)
+    {
+        auto rocp_agent = CHECK_NOTNULL(agent.get_rocp_agent());
+        if(!rocp_agent->runtime_visibility.hsa || !rocp_agent->runtime_visibility.hip) continue;
+        if(!is_spm_supported_arch(agent)) continue;
+
+        found_spm_agent = true;
+        auto metrics    = getSPMMetrics(agent);
+        for(auto& metric : metrics)
+        {
+            if(!metric.expression().empty()) continue;
+
+            expected_dispatch                          expected = {};
+            rocprofiler_counter_id_t                   id       = {.handle = metric.id()};
+            std::vector<rocprofiler_spm_parameters_t*> input_params{};
+            rocprofiler_spm_parameters_t               param{
+                .size  = sizeof(rocprofiler_spm_parameters_t),
+                .type  = ROCPROFILER_SPM_PARAMETER_TYPE_SAMPLE_INTERVAL_SCLK_CYCLES,
+                .value = 1200};
+            input_params.push_back(&param);
+
+            ROCPROFILER_CALL(rocprofiler_spm_create_counter_config(agent.get_rocp_agent()->id,
+                                                                   &id,
+                                                                   1,
+                                                                   input_params.data(),
+                                                                   input_params.size(),
+                                                                   &expected.id),
+                             "Unable to create profile");
+            cb_info->callback_args = &expected;
+
+            rocprofiler_queue_id_t  qid = {.handle = 1};
+            hsa::FakeQueue          fq(agent, qid);
+            hsa::rocprofiler_packet pkt{};
+            context::correlation_id corr_id{.internal = 99};
+
+            expected.queue_id    = qid;
+            expected.agent_id    = agent.get_rocp_agent()->id;
+            expected.kernel_id   = 42;
+            expected.dispatch_id = 7;
+
+            auto user_data = rocprofiler_user_data_t{.value = corr_id.internal};
+            auto ret_pkt   = spm::pre_kernel_call(ctx_p,
+                                                cb_info,
+                                                fq,
+                                                pkt,
+                                                expected.kernel_id,
+                                                expected.dispatch_id,
+                                                &user_data,
+                                                {},
+                                                &corr_id);
+            if(!ret_pkt.packet) continue;
+
+            size_t map_size_before_stop = 0;
+            cb_info->packet_return_map.rlock(
+                [&](const auto& data) { map_size_before_stop = data.size(); });
+            if(map_size_before_stop == 0) continue;
+
+            ROCPROFILER_CALL(rocprofiler_stop_context(get_client_ctx()), "stop context");
+            EXPECT_FALSE(rocprofiler::spm::is_any_active());
+
+            auto sess =
+                std::make_shared<hsa::queue_info_session_t>(hsa::queue_info_session_t{.queue = fq});
+            sess->correlation_id = &corr_id;
+
+            spm::inst_pkt_t inst_pkt;
+            inst_pkt.emplace_back(
+                std::make_pair(std::move(ret_pkt.packet), hsa::queue_hooks::SPM_CLIENT_ID));
+
+            spm::signal_completion_hook(fq,
+                                        pkt,
+                                        sess,
+                                        sess->packet_data.emplace_back(),
+                                        inst_pkt,
+                                        kernel_dispatch::profiling_time{});
+
+            size_t map_size_after_completion = 1;
+            cb_info->packet_return_map.rlock(
+                [&](const auto& data) { map_size_after_completion = data.size(); });
+            EXPECT_EQ(map_size_after_completion, 0U)
+                << "packet_return_map must drain via signal_completion_hook after stop_context";
+
+            ROCPROFILER_CALL(rocprofiler_spm_destroy_counter_config(expected.id),
+                             "Could not delete profile id");
+            registration::set_init_status(1);
+            registration::finalize();
+            context::pop_client(1);
+            set_client_ctx(get_client_ctx());
+            return;
+        }
+    }
+
+    if(!found_spm_agent) GTEST_SKIP() << "SPM unavailable";
+    FAIL() << "Could not exercise SPM in-flight completion hook path";
+}