[rocprofiler-sdk] defer barrier retirement until transition packets drain

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

@@ -45,20 +45,19 @@ hsa_barrier::~hsa_barrier()
         return;
     }
 
-    // If this barrier is destroyed and has queued packets, a future signal creation may
-    // inadvertently block these packets if the same handle is re-used by HSA.
-    if(!complete())
-    {
-        _barrier_enqueued.rlock([&](auto& barrier_enqueued) {
-            if(!barrier_enqueued.empty())
-            {
-                ROCP_WARNING << "An incomplete hsa_barrier which was enqueued is being destroyed.";
-            }
-        });
-    }
+    bool outstanding = false;
+    _barrier_enqueued.rlock(
+        [&](const auto& barrier_enqueued) { outstanding = !barrier_enqueued.empty(); });
+
+    // retirement is gated on safe_to_destroy(), so a barrier should never be destroyed while a
+    // transition packet still references its signal; log an error (and still release) if it does.
+    ROCP_ERROR_IF(outstanding) << "hsa_barrier (handle: " << _barrier_signal.handle
+                               << ") destroyed with outstanding transition packets";
 
-    // Clear and destroy the barrier signal
+    // release the signal so any stale transition packet can pass
     clear_barrier();
+
+    // Destroy the barrier signal
     _core_api.hsa_signal_destroy_fn(_barrier_signal);
 }
 
@@ -85,15 +84,17 @@ hsa_barrier::set_barrier(const queue_map_ptr_t& q)
 }
 
 std::optional<rocprofiler_packet>
-hsa_barrier::enqueue_packet(const Queue* queue)
+hsa_barrier::enqueue_packet(int64_t queue_id, uint64_t dispatch_id)
 {
     if(complete()) return std::nullopt;
     bool return_block = false;
     _barrier_enqueued.wlock([&](auto& barrier_enqueued) {
-        if(barrier_enqueued.find(queue->get_id().handle) == barrier_enqueued.end())
+        if(barrier_enqueued.find(queue_id) == barrier_enqueued.end())
         {
             return_block = true;
-            barrier_enqueued.insert(queue->get_id().handle);
+            // record the dispatch id carrying this transition packet; it has executed once the
+            // queue's completed count reaches this id
+            barrier_enqueued.emplace(queue_id, dispatch_id);
         }
     });
 
@@ -103,13 +104,33 @@ hsa_barrier::enqueue_packet(const Queue* queue)
     barrier.barrier_and.header        = HSA_PACKET_TYPE_BARRIER_AND << HSA_PACKET_HEADER_TYPE;
     barrier.barrier_and.dep_signal[0] = _barrier_signal;
     ROCP_INFO << "Barrier (handle: " << _barrier_signal.handle
-              << ") added to queue (handle: " << queue->get_id().handle << ")";
+              << ") added to queue (handle: " << queue_id << ")";
     return barrier;
 }
 
+void
+hsa_barrier::drain_queue(int64_t queue_id, uint64_t completed_id)
+{
+    // transition packet can't have executed until the barrier clears; skip the lock while armed
+    if(!complete()) return;
+    _barrier_enqueued.wlock([&](auto& barrier_enqueued) {
+        auto it = barrier_enqueued.find(queue_id);
+        if(it == barrier_enqueued.end()) return;
+        // packet executed once completions reach the carrying dispatch id (queue is in-order)
+        if(completed_id >= it->second)
+        {
+            barrier_enqueued.erase(it);
+        }
+    });
+}
+
 void
 hsa_barrier::remove_queue(const Queue* queue)
 {
+    // a torn-down queue can't execute its packet, so release its pin too
+    _barrier_enqueued.wlock(
+        [&](auto& barrier_enqueued) { barrier_enqueued.erase(queue->get_id().handle); });
+
     _queue_waiting.wlock([&](auto& queue_waiting) {
         if(queue_waiting.find(queue->get_id().handle) == queue_waiting.end()) return;
         queue_waiting.erase(queue->get_id().handle);
@@ -146,6 +167,16 @@ hsa_barrier::complete() const
     return _core_api.hsa_signal_load_scacquire_fn(_barrier_signal) == 0;
 }
 
+bool
+hsa_barrier::safe_to_destroy() const
+{
+    if(!complete()) return false;
+    bool no_outstanding = false;
+    _barrier_enqueued.rlock(
+        [&](const auto& barrier_enqueued) { no_outstanding = barrier_enqueued.empty(); });
+    return no_outstanding;
+}
+
 void
 hsa_barrier::clear_barrier()
 {

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/hsa_barrier.hpp

@@ -35,7 +35,6 @@
 #include <functional>
 #include <optional>
 #include <unordered_map>
-#include <unordered_set>
 
 namespace rocprofiler
 {
@@ -55,28 +54,37 @@ class hsa_barrier
     // packets, clears the barrier and marks it as complete.
     void set_barrier(const queue_map_ptr_t& q);
 
-    // Returns the packets needed to place this barrier in a queue.
-    // If the barrier has already been enqueued in the queue given or the barrier is complete, no
-    // packets are returned.
-    std::optional<rocprofiler_packet> enqueue_packet(const Queue* queue);
+    // Returns the transition packet to insert in the queue identified by queue_id, stamping it with
+    // dispatch_id (that queue's in-order serialized-dispatch id carrying the packet). Returns
+    // nothing if the barrier is already complete or the queue was already handed a packet.
+    std::optional<rocprofiler_packet> enqueue_packet(int64_t queue_id, uint64_t dispatch_id);
 
     // To be called when 1 async packet is finished in the given queue to decrement the remaining
     // packets counter. If all packets in all queues given in set_barrier are completed, clears the
     // barrier and marks it as complete. Returns true if the given queue had async packets waiting.
     bool register_completion(const Queue* queue);
 
+    // Drop the queue from the outstanding set once its completed count (completed_id) reaches the
+    // dispatch id that carried this barrier's transition packet (called on each serialized
+    // completion).
+    void drain_queue(int64_t queue_id, uint64_t completed_id);
+
     // Checks if this barrier is complete
     bool complete() const;
 
-    // Removes a queue from the barrier dependency list.
+    // complete() AND no queue still references the signal -> safe to destroy (no handle reuse).
+    bool safe_to_destroy() const;
+
+    // Removes a queue from the barrier dependency list (waiting + outstanding transition packets).
     // If this is the last queue waiting, clears the barrier and marks it as complete.
     void remove_queue(const Queue* queue);
 
 private:
     std::function<void()>                                      _barrier_finished = {};
     CoreApiTable                                               _core_api         = {};
     common::Synchronized<std::unordered_map<int64_t, int64_t>> _queue_waiting    = {};
-    common::Synchronized<std::unordered_set<int64_t>>          _barrier_enqueued = {};
+    // queue id -> serialized-dispatch id carrying this barrier's transition packet
+    common::Synchronized<std::unordered_map<int64_t, uint64_t>> _barrier_enqueued = {};
 
     void clear_barrier();
 

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

@@ -42,12 +42,24 @@ profiler_serializer_ready_signal_handler(hsa_signal_value_t /* signal_value */,
     return true;
 }
 
+void
+drain_barriers(std::deque<profiler_serializer::barrier_with_state>& barriers,
+               int64_t                                              queue_id,
+               uint64_t                                             completed_id)
+{
+    // let each barrier drop the queue if its transition packet on that queue has now executed
+    for(auto& barrier : barriers)
+        barrier.barrier->drain_queue(queue_id, completed_id);
+}
+
 void
 clear_complete_barriers(std::deque<profiler_serializer::barrier_with_state>& barriers)
 {
+    // Retire only when safe_to_destroy(): complete AND no queue still references the signal.
+    // Otherwise the barrier stays parked (signal at 0, stale packets pass) until its packets drain.
     while(!barriers.empty())
     {
-        if(barriers.front().barrier->complete())
+        if(barriers.front().barrier->safe_to_destroy())
         {
             barriers.pop_front();
         }
@@ -73,12 +85,23 @@ profiler_serializer::add_queue(hsa_queue_t** hsa_queues, const Queue& queue)
                                              profiler_serializer_ready_signal_handler,
                                              *hsa_queues);
     if(status != HSA_STATUS_SUCCESS) ROCP_FATAL << "hsa_amd_signal_async_handler failed";
+
+    // Track this queue's in-order serialized dispatch/completion ids. Created here (under wlock) so
+    // the dispatch path only ever increments an existing entry under the shared lock.
+    _serial.try_emplace(queue.get_id().handle);
 }
 
 void
 profiler_serializer::kernel_completion_signal(const Queue& completed)
 {
-    // We do not want to track kernel compleiton signals before we have reached the barrier
+    const auto qid          = completed.get_id().handle;
+    uint64_t   completed_id = 0;
+    if(auto it = _serial.find(qid); it != _serial.end())
+        completed_id = it->second.completed.fetch_add(1, std::memory_order_relaxed) + 1;
+
+    drain_barriers(_barrier, qid, completed_id);
+
+    // We do not want to track kernel completion signals before we have reached the barrier
     clear_complete_barriers(_barrier);
 
     // Find the state of this barrier
@@ -178,9 +201,15 @@ profiler_serializer::kernel_dispatch(const Queue& queue) const
         return barrier;
     };
 
+    // advance this queue's in-order dispatch id (counts every serialized dispatch)
+    const auto qid         = queue.get_id().handle;
+    uint64_t   dispatch_id = 0;
+    if(auto it = _serial.find(qid); it != _serial.end())
+        dispatch_id = it->second.dispatched.fetch_add(1, std::memory_order_relaxed) + 1;
+
     if(!_barrier.empty())
     {
-        if(auto maybe_barrier = _barrier.back().barrier->enqueue_packet(&queue))
+        if(auto maybe_barrier = _barrier.back().barrier->enqueue_packet(qid, dispatch_id))
         {
             ret.push_back(*maybe_barrier);
         }

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/profile_serializer.hpp

@@ -29,6 +29,8 @@
 #include <rocprofiler-sdk/fwd.h>
 #include <rocprofiler-sdk/hsa.h>
 
+#include <atomic>
+#include <cstdint>
 #include <deque>
 #include <unordered_map>
 
@@ -82,13 +84,26 @@ class profiler_serializer
 
     void destroy_queue(hsa_queue_t* id, const Queue& queue);
 
-    static void add_queue(hsa_queue_t** hsa_queues, const Queue& queue);
+    void add_queue(hsa_queue_t** hsa_queues, const Queue& queue);
 
 private:
-    const Queue*                   _dispatch_queue{nullptr};
-    std::deque<const Queue*>       _dispatch_ready;
-    std::atomic<Status>            _serializer_status{Status::DISABLED};
-    std::deque<barrier_with_state> _barrier;
+    // Per-queue in-order serialized dispatch/completion ids. hsa_barrier uses them to tell when a
+    // transition packet handed to a queue has executed (completed >= dispatched-at-enqueue).
+    // Relaxed atomics: ordering/visibility comes from the serializer lock; dispatch only takes a
+    // shared lock, so the values must be atomic and the map structure must stay stable. Entries are
+    // created in add_queue under the write lock and only ever incremented (never inserted) on the
+    // shared dispatch path.
+    struct serial_counts
+    {
+        std::atomic<uint64_t> dispatched{0};
+        std::atomic<uint64_t> completed{0};
+    };
+
+    const Queue*                                       _dispatch_queue{nullptr};
+    std::deque<const Queue*>                           _dispatch_ready;
+    std::atomic<Status>                                _serializer_status{Status::DISABLED};
+    std::deque<barrier_with_state>                     _barrier;
+    mutable std::unordered_map<int64_t, serial_counts> _serial;
 };
 
 }  // namespace hsa

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

@@ -170,6 +170,7 @@ class Queue
     {
         return _core_api.hsa_signal_load_scacquire_fn(_active_kernels);
     }
+
     void sync() const;
 
     void register_callback(ClientID id, queue_callbacks_t callbacks);

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/tests/hsa_barrier.cpp

@@ -143,7 +143,7 @@ inject_barriers(hsa_barrier& barrier, QueueController::queue_map_t& queues)
     for(auto& [hsa_queue, fq] : queues)
     {
         auto complete = barrier.complete();
-        auto _pkt     = barrier.enqueue_packet(fq.get());
+        auto _pkt     = barrier.enqueue_packet(fq->get_id().handle, 0);
         // If barrier is complete, no packets should be generated
         ASSERT_NE(complete, _pkt.has_value());
 
@@ -388,3 +388,90 @@ TEST(hsa_barrier, block_multi)
     registration::set_init_status(1);
     registration::finalize();
 }
+
+// Regression tests: a barrier must not be retired (its signal destroyed/recyclable) while a queue
+// still holds an un-executed transition packet referencing it.
+
+// complete-but-referenced barrier is not safe_to_destroy() until the carrying dispatch runs.
+TEST(hsa_barrier, safe_to_destroy_gated_on_transition_drain)
+{
+    ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS);
+    test_init();
+
+    registration::init_logging();
+    registration::set_init_status(-1);
+    context::push_client(1);
+
+    auto queues = create_queue_map(1);
+    ASSERT_EQ(queues.size(), 1u);
+    Queue* q = queues.begin()->second.get();
+
+    hsa_barrier::queue_map_ptr_t q_map;
+    for(const auto& [k, v] : queues)
+        q_map[k] = v.get();
+
+    const int64_t qid = q->get_id().handle;
+
+    q->async_started();  // pre-barrier dispatch (id 1) in flight
+
+    hsa_barrier barrier([]() {}, get_api_table());
+    barrier.set_barrier(q_map);
+    ASSERT_FALSE(barrier.complete());
+    ASSERT_FALSE(barrier.safe_to_destroy());
+
+    // carrying dispatch (id 2) gets the transition packet
+    auto pkt = barrier.enqueue_packet(qid, 2);
+    ASSERT_TRUE(pkt.has_value());
+
+    // dispatch 1 completes -> barrier clears, but carrier (id 2) has not run -> not retirable
+    q->async_complete();
+    barrier.register_completion(q);
+    barrier.drain_queue(qid, 1);  // completed(1) < token(2): no-op
+    ASSERT_TRUE(barrier.complete());
+    ASSERT_FALSE(barrier.safe_to_destroy());
+
+    // carrier (id 2) completes -> retirable
+    barrier.drain_queue(qid, 2);
+    ASSERT_TRUE(barrier.safe_to_destroy());
+
+    registration::set_init_status(1);
+    registration::finalize();
+}
+
+// teardown escape: removing an enqueued queue releases the barrier instead of pinning it forever.
+TEST(hsa_barrier, safe_to_destroy_after_enqueued_queue_removed)
+{
+    ASSERT_EQ(hsa_init(), HSA_STATUS_SUCCESS);
+    test_init();
+
+    registration::init_logging();
+    registration::set_init_status(-1);
+    context::push_client(1);
+
+    auto queues = create_queue_map(1);
+    ASSERT_EQ(queues.size(), 1u);
+    Queue* q = queues.begin()->second.get();
+
+    hsa_barrier::queue_map_ptr_t q_map;
+    for(const auto& [k, v] : queues)
+        q_map[k] = v.get();
+
+    const int64_t qid = q->get_id().handle;
+
+    q->async_started();
+    hsa_barrier barrier([]() {}, get_api_table());
+    barrier.set_barrier(q_map);
+
+    ASSERT_TRUE(barrier.enqueue_packet(qid, 1).has_value());
+
+    q->async_complete();
+    barrier.register_completion(q);
+    ASSERT_TRUE(barrier.complete());
+    ASSERT_FALSE(barrier.safe_to_destroy());  // transition packet still outstanding
+
+    barrier.remove_queue(q);
+    ASSERT_TRUE(barrier.safe_to_destroy());
+
+    registration::set_init_status(1);
+    registration::finalize();
+}