Avoid RMWs when polling empty remote queues (NVIDIA#2180)

* Avoid RMWs when polling empty remote queues

* Fix remote poll before worker sleep

* Harden remote queue sleep polling

* Consume remote poll notifications with acquire RMW

* Format static thread pool polling changes

* Clarify remote polling modes

* Strengthen remote polling coverage

* clang-format

---------

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

Author: fallintoplace

examples/benchmark/common.hpp

@@ -192,4 +192,4 @@ void my_main(int argc, char** argv, exec::numa_policy policy = exec::get_numa_po
   auto [dur_ms, ops_per_sec, avg, max, min, stddev] =
     compute_perf(starts, ends, warmup, nRuns - 1, total_scheds);
   std::cout << avg << " | " << max << " | " << min << " | " << stddev << "\n";
-}
+}

include/exec/static_thread_pool.hpp

@@ -142,6 +142,12 @@ namespace experimental::execution
       std::size_t index_{(std::numeric_limits<std::size_t>::max)()};
     };
 
+    enum class remote_poll_mode
+    {
+      speculative,
+      before_sleep
+    };
+
     struct remote_queue_list
     {
      private:
@@ -168,13 +174,18 @@ namespace experimental::execution
         }
       }
 
-      auto pop_all_reversed(std::size_t tid) noexcept -> __intrusive_queue<&task_base::next_>
+      auto pop_all_reversed(std::size_t tid, remote_poll_mode mode) noexcept
+        -> __intrusive_queue<&task_base::next_>
       {
         remote_queue*                        head = head_.load(__std::memory_order_acquire);
         __intrusive_queue<&task_base::next_> tasks{};
         while (head != nullptr)
         {
-          tasks.append(head->queues_[tid].pop_all_reversed());
+          auto& queue = head->queues_[tid];
+          if (mode == remote_poll_mode::before_sleep || !queue.empty())
+          {
+            tasks.append(queue.pop_all_reversed());
+          }
           head = head->next_;
         }
         return tasks;
@@ -648,7 +659,7 @@ namespace experimental::execution
         };
 
         auto try_pop() -> pop_result;
-        auto try_remote() -> pop_result;
+        auto try_remote(remote_poll_mode mode) -> pop_result;
         auto try_steal(std::span<workstealing_victim> victims) -> pop_result;
         auto try_steal_near() -> pop_result;
         auto try_steal_any() -> pop_result;
@@ -973,11 +984,11 @@ namespace experimental::execution
       tmp.clear();
     }
 
-    inline auto
-    _static_thread_pool::thread_state::try_remote() -> _static_thread_pool::thread_state::pop_result
+    inline auto _static_thread_pool::thread_state::try_remote(remote_poll_mode mode)
+      -> _static_thread_pool::thread_state::pop_result
     {
       pop_result                           result{.task = nullptr, .queue_index = index_};
-      __intrusive_queue<&task_base::next_> remotes = pool_->remotes_.pop_all_reversed(index_);
+      __intrusive_queue<&task_base::next_> remotes = pool_->remotes_.pop_all_reversed(index_, mode);
       pending_queue_.append(std::move(remotes));
       if (!pending_queue_.empty())
       {
@@ -997,7 +1008,7 @@ namespace experimental::execution
       {
         return result;
       }
-      return try_remote();
+      return try_remote(remote_poll_mode::speculative);
     }
 
     inline auto _static_thread_pool::thread_state::try_steal(std::span<workstealing_victim> victims)
@@ -1130,11 +1141,22 @@ namespace experimental::execution
           return result;
         }
         state expected = state::running;
-        if (state_.compare_exchange_weak(expected, state::sleeping, __std::memory_order_relaxed))
-        {
-          result = try_remote();
+        if (state_.compare_exchange_weak(expected,
+                                         state::sleeping,
+                                         __std::memory_order_relaxed,
+                                         __std::memory_order_relaxed))
+        {
+          // The relaxed empty probe is safe during normal polling, but the
+          // running-to-sleeping boundary must perform the CAS dequeue so work
+          // published before the transition cannot be missed.
+          result = try_remote(remote_poll_mode::before_sleep);
           if (result.task)
           {
+            state expected_sleeping = state::sleeping;
+            state_.compare_exchange_strong(expected_sleeping,
+                                           state::running,
+                                           __std::memory_order_relaxed,
+                                           __std::memory_order_relaxed);
             return result;
           }
           set_sleeping();
@@ -1146,15 +1168,15 @@ namespace experimental::execution
         {
           lock.unlock();
         }
-        state_.store(state::running, __std::memory_order_relaxed);
+        state_.exchange(state::running, __std::memory_order_acquire);
         result = try_pop();
       }
       return result;
     }
 
     inline auto _static_thread_pool::thread_state::notify() -> bool
     {
-      if (state_.exchange(state::notified, __std::memory_order_relaxed) == state::sleeping)
+      if (state_.exchange(state::notified, __std::memory_order_release) == state::sleeping)
       {
         {
           std::lock_guard lock{mut_};

test/exec/test_static_thread_pool.cpp

@@ -16,17 +16,22 @@
 
 #include <exec/sequence/ignore_all_values.hpp>
 #include <exec/sequence/transform_each.hpp>
+#include <exec/start_detached.hpp>
 #include <exec/static_thread_pool.hpp>
 #include <stdexec/execution.hpp>
 #include <test_common/catch2.hpp>  // IWYU pragma: keep
 
 #include <atomic>
+#include <chrono>
 #include <exception>
+#include <latch>
 #include <mutex>
+#include <optional>
 #include <ranges>
 #include <stdexcept>
 #include <thread>
 #include <unordered_set>
+#include <vector>
 namespace ex = STDEXEC;
 
 namespace
@@ -229,3 +234,107 @@ TEST_CASE("bulk on static_thread_pool executes on multiple threads, take 2",
   ex::sync_wait(std::move(sender));
   REQUIRE(thread_ids.size() == num_of_threads);
 }
+
+namespace
+{
+  void run_remote_poll_stress(bool separate_schedulers)
+  {
+    constexpr std::size_t num_producers = 4;
+    constexpr std::size_t rounds        = 10'000;
+
+    std::latch                            ready{num_producers};
+    std::atomic<bool>                     start{false};
+    std::atomic<bool>                     stop{false};
+    std::vector<std::atomic<std::size_t>> completed(num_producers);
+    std::vector<std::thread>              producers;
+    producers.reserve(num_producers);
+    for (auto& count: completed)
+    {
+      count.store(0, std::memory_order_relaxed);
+    }
+
+    exec::static_thread_pool pool{1};
+    using scheduler_t = decltype(pool.get_scheduler());
+    std::optional<scheduler_t> shared_scheduler;
+    if (!separate_schedulers)
+    {
+      shared_scheduler.emplace(pool.get_scheduler());
+    }
+
+    for (std::size_t producer = 0; producer < num_producers; ++producer)
+    {
+      producers.emplace_back(
+        [&, producer]
+        {
+          auto scheduler = separate_schedulers ? pool.get_scheduler() : *shared_scheduler;
+          ready.count_down();
+          while (!start.load(std::memory_order_acquire))
+          {
+            std::this_thread::yield();
+          }
+
+          auto* const producer_completed = &completed[producer];
+          std::size_t expected           = 0;
+          for (std::size_t round = 0; round < rounds && !stop.load(std::memory_order_relaxed);
+               ++round)
+          {
+            std::size_t const batch_size = (round % 4 == 0) ? 2 : 1;
+            expected += batch_size;
+            for (std::size_t i = 0; i < batch_size; ++i)
+            {
+              exec::start_detached(
+                ex::schedule(scheduler)
+                | ex::then([producer_completed]
+                           { producer_completed->fetch_add(1, std::memory_order_relaxed); }));
+            }
+
+            while (!stop.load(std::memory_order_relaxed)
+                   && producer_completed->load(std::memory_order_relaxed) < expected)
+            {
+              std::this_thread::yield();
+            }
+            std::this_thread::yield();
+          }
+        });
+    }
+
+    ready.wait();
+    start.store(true, std::memory_order_release);
+
+    auto const expected        = num_producers * rounds + num_producers * ((rounds + 3) / 4);
+    auto const deadline        = std::chrono::steady_clock::now() + std::chrono::seconds(10);
+    auto       completed_total = [&]
+    {
+      std::size_t result = 0;
+      for (auto const & count: completed)
+      {
+        result += count.load(std::memory_order_relaxed);
+      }
+      return result;
+    };
+
+    while (completed_total() < expected && std::chrono::steady_clock::now() < deadline)
+    {
+      std::this_thread::yield();
+    }
+    stop.store(true, std::memory_order_release);
+    for (auto& producer: producers)
+    {
+      producer.join();
+    }
+
+    CHECK(completed_total() == expected);
+  }
+}  // namespace
+
+TEST_CASE("static_thread_pool drains remote work from a shared scheduler",
+          "[types][static_thread_pool][stress]")
+{
+  run_remote_poll_stress(false);
+}
+
+TEST_CASE("static_thread_pool drains remote work from producer schedulers",
+          "[types][static_thread_pool][stress]")
+{
+  run_remote_poll_stress(true);
+}

test/rrd/CMakeLists.txt

@@ -54,7 +54,7 @@ function(add_relacy_test target_name)
 endfunction()
 
 set(relacy_tests async_scope bwos_lifo_queue intrusive_mpsc_queue split
-                 sync_wait)
+                 static_thread_pool_remote_poll sync_wait)
 
 foreach(test ${relacy_tests})
   add_relacy_test(${test})