From 749dd05d53707a7e13f57547da297104123eb080 Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:55:11 -0700 Subject: [PATCH] Add: ProfilerBase::quiesce() to drain without retiring the threads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stop() couples two things: draining the pipeline and retiring the worker threads. Its drain guarantee is in fact *paid for* by the threads exiting — joining mgmt is what proves its final sweep landed in the host shards, and joining the collectors is what proves those shards were consumed. That coupling is why the DFX collectors are started and joined per run: there is no way to reach a known-empty state without tearing them down. quiesce() supplies the missing half — the same guarantee, threads intact — using a two-phase epoch handshake in place of the joins. The phases are ordered rather than concurrent. A collector reporting its shard empty before mgmt has finished sweeping would be reporting on a queue mgmt is about to push into, so collect_quiesce_epoch_ is published only once every drain ack for that epoch has landed. The caller having stopped the device-side producers is what makes a drain worker's "one full sweep found nothing" conclusive; that already holds where the DFX teardown runs. The collector-side ack sits above the has_seen_buffer guard on purpose: a shard that never received a buffer is a valid run shape, and an ack behind that guard would leave quiesce() waiting on a silent subsystem forever. The new silent-collector test hangs if it is moved. stop()'s implementation is unchanged, so no existing behavior moves. Also drops the L1/L2 shorthand from this file's comments. It denoted the device-side ring and the host ready queue shards, but L1 and L2 are already taken by the hierarchy model (die and chip runtime), so the same coordinates meant two unrelated things. --- .../platform/include/host/profiler_base.h | 117 ++++++++++++++++-- tests/ut/cpp/common/test_profiler_base.cpp | 72 +++++++++++ 2 files changed, 178 insertions(+), 11 deletions(-) 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); +}