Make static thread pool sleep when empty (#1476)

Co-authored-by: Eric Niebler <eniebler@nvidia.com>

Author: maikel

include/exec/static_thread_pool.hpp

@@ -608,6 +608,8 @@ namespace exec {
         void notify_one_sleeping();
         void set_stealing();
         void clear_stealing();
+        void set_sleeping();
+        void clear_sleeping();
 
         bwos::lifo_queue<task_base*, numa_allocator<task_base*>> local_queue_;
         __intrusive_queue<&task_base::next> pending_queue_{};
@@ -624,7 +626,7 @@ namespace exec {
       void run(std::uint32_t index) noexcept;
       void join() noexcept;
 
-      alignas(64) std::atomic<std::uint32_t> numThiefs_{};
+      alignas(64) std::atomic<std::uint32_t> numActive_{};
       alignas(64) remote_queue_list remotes_;
       std::uint32_t threadCount_;
       std::uint32_t maxSteals_{threadCount_ + 1};
@@ -695,6 +697,7 @@ namespace exec {
       threads_.reserve(threadCount);
 
       try {
+        numActive_.store(threadCount << 16, std::memory_order_relaxed);
         for (std::uint32_t i = 0; i < threadCount; ++i) {
           threads_.emplace_back([this, i] { run(i); });
         }
@@ -717,9 +720,9 @@ namespace exec {
     }
 
     inline void static_thread_pool_::run(std::uint32_t threadIndex) noexcept {
+      STDEXEC_ASSERT(threadIndex < threadCount_);
       // NOLINTNEXTLINE(bugprone-unused-return-value)
       numa_.bind_to_node(threadStates_[threadIndex]->numa_node());
-      STDEXEC_ASSERT(threadIndex < threadCount_);
       while (true) {
         // Make a blocking call to de-queue a task if we don't already have one.
         auto [task, queueIndex] = threadStates_[threadIndex]->pop();
@@ -938,12 +941,30 @@ namespace exec {
       pending_queue_.prepend(std::move(tasks));
     }
 
+    inline void static_thread_pool_::thread_state::set_sleeping() {
+      pool_->numActive_.fetch_sub(1 << 16, std::memory_order_relaxed);
+    }
+
+    // wakeup a worker thread and maintain the invariant that we always one active thief as long as a potential victim is awake
+    inline void static_thread_pool_::thread_state::clear_sleeping() {
+      const std::uint32_t numActive = pool_->numActive_.fetch_add(1 << 16, std::memory_order_relaxed);
+      if (numActive == 0) {
+        notify_one_sleeping();
+      }
+    }
+
     inline void static_thread_pool_::thread_state::set_stealing() {
-      pool_->numThiefs_.fetch_add(1, std::memory_order_relaxed);
+      const std::uint32_t diff = 1 - (1 << 16);
+      pool_->numActive_.fetch_add(diff, std::memory_order_relaxed);
     }
 
+    // put a thief to sleep but maintain the invariant that we always have one active thief as long as a potential victim is awake
     inline void static_thread_pool_::thread_state::clear_stealing() {
-      if (pool_->numThiefs_.fetch_sub(1, std::memory_order_relaxed) == 1) {
+      constexpr std::uint32_t diff = 1 - (1 << 16);
+      const std::uint32_t numActive = pool_->numActive_.fetch_sub(diff, std::memory_order_relaxed);
+      const std::uint32_t numVictims = numActive >> 16;
+      const std::uint32_t numThiefs = numActive & 0xffff;
+      if (numThiefs == 1 && numVictims != 0) {
         notify_one_sleeping();
       }
     }
@@ -995,9 +1016,14 @@ namespace exec {
           if (result.task) {
             return result;
           }
+          set_sleeping();
           cv_.wait(lock);
+          lock.unlock();
+          clear_sleeping();
+        }
+        if (lock.owns_lock()) {
+          lock.unlock();
         }
-        lock.unlock();
         state_.store(state::running, std::memory_order_relaxed);
         result = try_pop();
       }