diff --git a/src/common/platform/include/host/profiler_base.h b/src/common/platform/include/host/profiler_base.h index 197dfff0fe..cdb7ce2c6d 100644 --- a/src/common/platform/include/host/profiler_base.h +++ b/src/common/platform/include/host/profiler_base.h @@ -129,12 +129,17 @@ * 3. ... device execution ... * 4. stop() — atomically: * a) flips mgmt_running_, joins the mgmt thread(s); the drain thread's - * final-drain pass pushes the last L1→L2 entries before exiting. + * final-drain pass pushes the last device-ring entries into the host + * ready queue shard(s) before exiting. * b) execution_complete_ is set; each collector loop sees it on its * next idle tick, drains its host ready queue shard, and exits. * c) collector thread(s) joined. - * Caller is then guaranteed L1 and L2 are both empty and all collected - * data has been delivered to Derived::on_buffer_collected. + * Caller is then guaranteed the device-side ring and the host ready queue + * shard(s) are both empty and all collected data has been delivered to + * Derived::on_buffer_collected. + * + * quiesce() gives that same guarantee without (a)'s and (c)'s joins, so + * the threads survive it — see its own comment. * * SVM vs host-shadow paths (chosen at runtime by the collector's MemoryOps) * ------------------------------------------------------------------------- @@ -854,8 +859,8 @@ class ProfilerBase { * init() aborted before set_memory_context, or finalize() has cleared * the context) this is a no-op. * - * Order matters: mgmt is started before collectors because mgmt is the - * only writer to L2 (the ready queues) and collectors are the consumers. The + * Order matters: mgmt is started before collectors because mgmt is the only + * writer to the host ready queue shards and collectors are the consumers. The * register slot defaults to identity on the SVM path (copy_to_device_ * is null) or to a host-shadow malloc lambda on the non-SVM path * (copy_to_device_ installed) — so BufferPoolManager always has a @@ -925,6 +930,16 @@ class ProfilerBase { manager_.set_memory_context(std::move(ops), shm_dev_, shm_host_, shm_size_, device_id_); execution_complete_.store(false, std::memory_order_release); + // Reset the quiescence handshake so a restarted collector cannot see a + // previous run's acks. Safe to do unsynchronized: the std::thread + // constructors below are the synchronization point for the workers that + // read these. + drain_quiesce_epoch_.store(0, std::memory_order_relaxed); + collect_quiesce_epoch_.store(0, std::memory_order_relaxed); + for (int i = 0; i < Manager::kMaxCollectorShards; i++) { + drain_acked_[i].store(0, std::memory_order_relaxed); + collect_acked_[i].store(0, std::memory_order_relaxed); + } { DataHeader *header = Module::header_from_shm(manager_.shared_mem_host()); (void)ProfilerAlgorithms::proactive_replenish(manager_, header); @@ -966,17 +981,46 @@ class ProfilerBase { } } + /** + * Drain to a quiescent point without retiring the threads. On return the + * device-side ring and the host ready queue shard(s) are empty and + * Derived::on_buffer_collected has been called for every entry that was in + * either — the same guarantee stop() gives, minus the thread teardown. + * + * Precondition: the device-side producers have stopped. A drain worker + * reports its shard quiescent after one full sweep that found nothing, so a + * producer still writing could push a record in behind that report. + * Callers already satisfy this — the run is drained before teardown. + * + * Idempotent, and a no-op before start() or after stop(). Like stop(), it + * waits without a deadline: a wedged worker hangs the caller here exactly + * as it would hang the join in stop(). + */ + void quiesce() { + if (collector_threads_.empty()) return; + const int n = shard_count_; + + // Phase one: mgmt sweeps the device-side ring into the host shards. + const uint64_t epoch = drain_quiesce_epoch_.fetch_add(1, std::memory_order_acq_rel) + 1; + wait_for_epoch(drain_acked_, n, epoch); + + // Phase two: only now can a collector's "my shard is empty" mean the + // pipeline is empty rather than that mgmt has not pushed yet. + collect_quiesce_epoch_.store(epoch, std::memory_order_release); + wait_for_epoch(collect_acked_, n, epoch); + } + /** * Stop the drain/replenish mgmt threads, drain whatever the drain side * pushes during its final pass, and join the collector. Idempotent. Caller - * is guaranteed on return that mgmt's L1 ringbuffer and the host-side - * ready queue shard(s) are empty and Derived::on_buffer_collected has been - * called for every entry that was in either queue. Framework-owned buffers - * are NOT freed here — Derived's finalize() must do that. + * is guaranteed on return that mgmt's device-side ringbuffer and the + * host-side ready queue shard(s) are empty and Derived::on_buffer_collected + * has been called for every entry that was in either queue. Framework-owned + * buffers are NOT freed here — Derived's finalize() must do that. * * Order matters: stop+join mgmt first so its final-drain pass is fully - * landed in L2 BEFORE we tell poll to exit. Otherwise mgmt's last batch - * has no consumer. + * landed in the host shards BEFORE we tell poll to exit. Otherwise mgmt's + * last batch has no consumer. */ void stop() { mgmt_running_.store(false, std::memory_order_release); @@ -1121,6 +1165,17 @@ class ProfilerBase { } private: + // Teardown-path wait, so a sleep is permitted here: no task's latency + // passes through it (codestyle.md rule 5 exempts teardown). + template + static void wait_for_epoch(const Acks &acks, int n, uint64_t epoch) { + for (int i = 0; i < n; i++) { + while (acks[i].load(std::memory_order_acquire) != epoch) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + } + } + void mgmt_drain_loop(int queue_start, int queue_stride) { DataHeader *header = Module::header_from_shm(manager_.shared_mem_host()); using Alg = ProfilerAlgorithms; @@ -1140,6 +1195,17 @@ class ProfilerBase { idle_busy_polls = 0; } + // A full sweep that found nothing means this worker's slice of the + // device-side queues is empty. With producers stopped (quiesce()'s + // precondition) nothing can arrive behind this report, so it is the + // quiescent condition for phase one. + if (!found_any) { + const uint64_t requested = drain_quiesce_epoch_.load(std::memory_order_acquire); + if (drain_acked_[queue_start].load(std::memory_order_relaxed) != requested) { + drain_acked_[queue_start].store(requested, std::memory_order_release); + } + } + if (!found_any) { if (idle_busy_polls < kIdleBusyPollLoops) { idle_busy_polls++; @@ -1216,6 +1282,22 @@ class ProfilerBase { } break; } + // Phase two of the quiescence handshake. wait_pop_ready timed out, + // so this shard is empty; mgmt has already reported its own sweep + // done for this epoch, so nothing further can arrive. Placed above + // the has_seen_buffer guard below: a shard that never received a + // buffer is a valid run shape and still has to report, or quiesce() + // would wait on it forever. + { + const uint64_t requested = collect_quiesce_epoch_.load(std::memory_order_acquire); + if (collect_acked_[shard_index].load(std::memory_order_relaxed) != requested) { + while (manager_.try_pop_ready(info, shard_index)) { + consume(info, shard_index); + has_seen_buffer = true; + } + collect_acked_[shard_index].store(requested, std::memory_order_release); + } + } // A shard that has never seen a buffer is a valid run shape at any // shard count — a subsystem can legitimately emit nothing for a // whole run. execution_complete_ above is the exit path for that @@ -1256,6 +1338,19 @@ class ProfilerBase { std::vector mgmt_drain_threads_; std::thread mgmt_replenish_thread_; std::atomic mgmt_running_{false}; + + // Two-phase quiescence handshake. Each phase is a monotonic epoch the + // caller publishes and every worker of that phase echoes back once it has + // reached the quiescent condition for its own shard. + // + // The phases are ordered, not concurrent: a collector that reported its + // shard empty before mgmt finished its sweep would be reporting on a queue + // mgmt is still about to push into. So collect_quiesce_epoch_ is published + // only after every drain ack for that epoch has landed. + std::atomic drain_quiesce_epoch_{0}; + std::atomic collect_quiesce_epoch_{0}; + std::array, Manager::kMaxCollectorShards> drain_acked_{}; + std::array, Manager::kMaxCollectorShards> collect_acked_{}; }; } // namespace profiling_common diff --git a/tests/ut/cpp/common/test_profiler_base.cpp b/tests/ut/cpp/common/test_profiler_base.cpp index 24bf6e4e7c..03e40bd1e7 100644 --- a/tests/ut/cpp/common/test_profiler_base.cpp +++ b/tests/ut/cpp/common/test_profiler_base.cpp @@ -283,3 +283,75 @@ TEST(ProfilerBaseTest, CollectorStaysAliveAfterArmedIdleTimeout) { collector.stop(); EXPECT_EQ(collector.collected(), 2); } + +// quiesce() gives stop()'s drain guarantee without retiring the threads: on +// return every published buffer has reached on_buffer_collected, and the +// collector is still able to take more. This is what lets the collectors stay +// resident across runs instead of being started and joined per run. +TEST(ProfilerBaseTest, QuiesceDrainsWithoutRetiringThreads) { + constexpr int kThreads = PLATFORM_MAX_AICPU_THREADS; + + TestHeader header{}; + uint64_t first[kThreads]{}; + uint64_t second[kThreads]{}; + + TestCollector collector; + collector.init(kThreads, &header); + collector.start(nullptr); + + for (int q = 0; q < kThreads; q++) { + publish(collector, header, q, &first[q]); + } + collector.quiesce(); + // No wait_for_collected here on purpose: quiesce() must have delivered + // everything by the time it returns, so polling for it would hide a + // handshake that reports too early. + EXPECT_EQ(collector.collected(), kThreads); + + for (int q = 0; q < kThreads; q++) { + publish(collector, header, q, &second[q]); + } + collector.quiesce(); + EXPECT_EQ(collector.collected(), 2 * kThreads); + + collector.stop(); + EXPECT_EQ(collector.collected(), 2 * kThreads); +} + +// A subsystem that emitted nothing still has to complete the handshake. The +// collector loop skips its idle bookkeeping for a shard that has never seen a +// buffer, so an ack placed behind that guard would leave quiesce() waiting +// forever on a silent run — a hang, not a wrong count. +TEST(ProfilerBaseTest, QuiesceCompletesOnASilentCollector) { + TestHeader header{}; + TestCollector collector; + collector.init(2, &header); + collector.start(nullptr); + + collector.quiesce(); + EXPECT_EQ(collector.collected(), 0); + + // Still live afterwards. + uint64_t buffer = 0; + publish(collector, header, 1, &buffer); + collector.quiesce(); + EXPECT_EQ(collector.collected(), 1); + + collector.stop(); +} + +// quiesce() before start() and after stop() are both no-ops rather than hangs: +// there are no workers to answer the handshake in either state. +TEST(ProfilerBaseTest, QuiesceIsANoOpWithoutRunningThreads) { + TestHeader header{}; + TestCollector collector; + collector.init(2, &header); + + collector.quiesce(); // before start() + + collector.start(nullptr); + collector.stop(); + + collector.quiesce(); // after stop() + EXPECT_EQ(collector.collected(), 0); +}