chore: import upstream snapshot with attribution

paddle/fluid/framework/new_executor/workqueue/event_count.h

@@ -0,0 +1,286 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2016 Dmitry Vyukov <dvyukov@google.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+// EventCount allows to wait for arbitrary predicates in non-blocking
+// algorithms. Think of condition variable, but wait predicate does not need to
+// be protected by a mutex. Usage:
+// Waiting thread does:
+//
+//   if (predicate)
+//     return act();
+//   EventCount::Waiter& w = waiters[my_index];
+//   ec.Prewait(&w);
+//   if (predicate) {
+//     ec.CancelWait(&w);
+//     return act();
+//   }
+//   ec.CommitWait(&w);
+//
+// Notifying thread does:
+//
+//   predicate = true;
+//   ec.Notify(true);
+//
+// Notify is cheap if there are no waiting threads. Prewait/CommitWait are not
+// cheap, but they are executed only if the preceding predicate check has
+// failed.
+//
+// Algorithm outline:
+// There are two main variables: predicate (managed by user) and state_.
+// Operation closely resembles Dekker mutual algorithm:
+// https://en.wikipedia.org/wiki/Dekker%27s_algorithm
+// Waiting thread sets state_ then checks predicate, Notifying thread sets
+// predicate then checks state_. Due to seq_cst fences in between these
+// operations it is guaranteed than either waiter will see predicate change
+// and won't block, or notifying thread will see state_ change and will unblock
+// the waiter, or both. But it can't happen that both threads don't see each
+// other changes, which would lead to deadlock.
+//
+// What changed by PaddlePaddle
+//   1. Allocate aligned storage for Waiters to get better performance.
+//   2. Replace Eigen utils with std utils.
+
+#pragma once
+
+#include <atomic>
+#include <cassert>
+#include <condition_variable>
+#include <cstdlib>
+#include <mutex>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/common/macros.h"
+
+namespace paddle {
+namespace framework {
+
+void* AlignedMalloc(size_t size, size_t alignment);
+void AlignedFree(void* memory_ptr);
+
+class EventCount {
+ public:
+  class Waiter;
+
+  explicit EventCount(size_t waiter_num) : state_(kStackMask) {
+    assert(waiter_num < (1 << kWaiterBits) - 1);
+    void* buffer = AlignedMalloc(sizeof(Waiter) * waiter_num, alignof(Waiter));
+    if (buffer == nullptr) {
+      return;
+    }
+    waiters_ = reinterpret_cast<Waiter*>(buffer);
+    waiter_num_ = waiter_num;
+    for (size_t i = 0; i < waiter_num_; ++i) {
+      new (&waiters_[i]) Waiter;
+    }
+  }
+
+  EventCount(const EventCount&) = delete;
+
+  void operator=(const EventCount&) = delete;
+
+  ~EventCount() {
+    // Ensure there are no waiters.
+    assert(state_.load() == kStackMask);
+    AlignedFree(waiters_);
+  }
+
+  Waiter* GetWaiter(size_t waiter_index) {
+    assert(waiter_index < waiter_num_);
+    return &waiters_[waiter_index];
+  }
+
+  // Prewait prepares for waiting.
+  // After calling Prewait, the thread must re-check the wait predicate
+  // and then call either CancelWait or CommitWait.
+  void Prewait() {
+    uint64_t state = state_.load(std::memory_order_relaxed);
+    for (;;) {
+      CheckState(state);
+      uint64_t newstate = state + kWaiterInc;
+      CheckState(newstate);
+      if (state_.compare_exchange_weak(
+              state, newstate, std::memory_order_seq_cst))
+        return;
+    }
+  }
+
+  // CommitWait commits waiting after Prewait.
+  void CommitWait(Waiter* w) {
+    assert((w->epoch & ~kEpochMask) == 0);
+    w->state = Waiter::kNotSignaled;
+    const uint64_t me = (w - &waiters_[0]) | w->epoch;
+    uint64_t state = state_.load(std::memory_order_seq_cst);
+    for (;;) {
+      CheckState(state, true);
+      uint64_t newstate;
+      if ((state & kSignalMask) != 0) {
+        // Consume the signal and return immediately.
+        newstate = state - kWaiterInc - kSignalInc;
+      } else {
+        // Remove this thread from pre-wait counter and add to the waiter stack.
+        newstate = ((state & kWaiterMask) - kWaiterInc) | me;
+        w->next.store(state & (kStackMask | kEpochMask),
+                      std::memory_order_relaxed);
+      }
+      CheckState(newstate);
+      if (state_.compare_exchange_weak(
+              state, newstate, std::memory_order_acq_rel)) {
+        if ((state & kSignalMask) == 0) {
+          w->epoch += kEpochInc;
+          Park(w);
+        }
+        return;
+      }
+    }
+  }
+
+  // CancelWait cancels effects of the previous Prewait call.
+  void CancelWait() {
+    uint64_t state = state_.load(std::memory_order_relaxed);
+    for (;;) {
+      CheckState(state, true);
+      uint64_t newstate = state - kWaiterInc;
+      // We don't know if the thread was also notified or not,
+      // so we should not consume a signal unconditionally.
+      // Only if number of waiters is equal to number of signals,
+      // we know that the thread was notified and we must take away the signal.
+      if (((state & kWaiterMask) >> kWaiterShift) ==
+          ((state & kSignalMask) >> kSignalShift))
+        newstate -= kSignalInc;
+      CheckState(newstate);
+      if (state_.compare_exchange_weak(
+              state, newstate, std::memory_order_acq_rel))
+        return;
+    }
+  }
+
+  // Notify wakes one or all waiting threads.
+  // Must be called after changing the associated wait predicate.
+  void Notify(bool notify_all) {
+    std::atomic_thread_fence(std::memory_order_seq_cst);
+    uint64_t state = state_.load(std::memory_order_acquire);
+    for (;;) {
+      CheckState(state);
+      const uint64_t waiters = (state & kWaiterMask) >> kWaiterShift;
+      const uint64_t signals = (state & kSignalMask) >> kSignalShift;
+      // Easy case: no waiters.
+      if ((state & kStackMask) == kStackMask && waiters == signals) return;
+      uint64_t newstate;
+      if (notify_all) {
+        // Empty wait stack and set signal to number of pre-wait threads.
+        newstate =
+            (state & kWaiterMask) | (waiters << kSignalShift) | kStackMask;
+      } else if (signals < waiters) {
+        // There is a thread in pre-wait state, unblock it.
+        newstate = state + kSignalInc;
+      } else {
+        // Pop a waiter from list and unpark it.
+        Waiter* w = &waiters_[state & kStackMask];
+        uint64_t next = w->next.load(std::memory_order_relaxed);
+        newstate = (state & (kWaiterMask | kSignalMask)) | next;
+      }
+      CheckState(newstate);
+      if (state_.compare_exchange_weak(
+              state, newstate, std::memory_order_acq_rel)) {
+        if (!notify_all && (signals < waiters))
+          return;  // unblocked pre-wait thread
+        if ((state & kStackMask) == kStackMask) return;
+        Waiter* w = &waiters_[state & kStackMask];
+        if (!notify_all) w->next.store(kStackMask, std::memory_order_relaxed);
+        Unpark(w);
+        return;
+      }
+    }
+  }
+
+  class Waiter {
+    friend class EventCount;
+    // Align to 128 byte boundary to prevent false sharing with other Waiter
+    // objects in the same vector.
+    alignas(128) std::atomic<uint64_t> next;
+    std::mutex mu;
+    std::condition_variable cv;
+    uint64_t epoch = 0;
+    unsigned state = kNotSignaled;
+    enum {
+      kNotSignaled,
+      kWaiting,
+      kSignaled,
+    };
+  };
+
+ private:
+  // State_ layout:
+  // - low kWaiterBits is a stack of waiters committed wait
+  //   (indexes in waiters_ array are used as stack elements,
+  //   kStackMask means empty stack).
+  // - next kWaiterBits is count of waiters in prewait state.
+  // - next kWaiterBits is count of pending signals.
+  // - remaining bits are ABA counter for the stack.
+  //   (stored in Waiter node and incremented on push).
+  static const uint64_t kWaiterBits = 14;
+  static const uint64_t kStackMask = (1ull << kWaiterBits) - 1;
+  static const uint64_t kWaiterShift = kWaiterBits;
+  static const uint64_t kWaiterMask = ((1ull << kWaiterBits) - 1)
+                                      << kWaiterShift;
+  static const uint64_t kWaiterInc = 1ull << kWaiterShift;
+  static const uint64_t kSignalShift = 2 * kWaiterBits;
+  static const uint64_t kSignalMask = ((1ull << kWaiterBits) - 1)
+                                      << kSignalShift;
+  static const uint64_t kSignalInc = 1ull << kSignalShift;
+  static const uint64_t kEpochShift = 3 * kWaiterBits;
+  static const uint64_t kEpochBits = 64 - kEpochShift;
+  static const uint64_t kEpochMask = ((1ull << kEpochBits) - 1) << kEpochShift;
+  static const uint64_t kEpochInc = 1ull << kEpochShift;
+  std::atomic<uint64_t> state_;
+  Waiter* waiters_{nullptr};
+  size_t waiter_num_{0};
+
+  static void CheckState(uint64_t state, bool waiter UNUSED = false) {
+    static_assert(kEpochBits >= 20, "not enough bits to prevent ABA problem");
+    const uint64_t waiters = (state & kWaiterMask) >> kWaiterShift;
+    const uint64_t signals = (state & kSignalMask) >> kSignalShift;
+    assert(waiters >= signals);
+    assert(waiters < (1 << kWaiterBits) - 1);
+    assert(!waiter || waiters > 0);
+    (void)waiters;
+    (void)signals;
+  }
+
+  void Park(Waiter* w) {
+    std::unique_lock<std::mutex> lock(w->mu);
+    while (w->state != Waiter::kSignaled) {
+      w->state = Waiter::kWaiting;
+      VLOG(10) << "Go to wait " << &(w->cv);
+      w->cv.wait(lock);
+    }
+  }
+
+  void Unpark(Waiter* w) {
+    for (Waiter* next; w; w = next) {
+      uint64_t wnext = w->next.load(std::memory_order_relaxed) & kStackMask;
+      next = wnext == kStackMask ? nullptr : &waiters_[wnext];
+      unsigned state;
+      {
+        std::unique_lock<std::mutex> lock(w->mu);
+        state = w->state;
+        w->state = Waiter::kSignaled;
+      }
+      // Avoid notifying if it wasn't waiting.
+      if (state == Waiter::kWaiting) {
+        VLOG(10) << "Go to notify " << &(w->cv);
+        w->cv.notify_one();
+      }
+    }
+  }
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/new_executor/workqueue/events_waiter.cc

@@ -0,0 +1,217 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/new_executor/workqueue/events_waiter.h"
+
+#include <glog/logging.h>
+
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle::framework {
+
+constexpr EventsWaiter::EventId kEmptyEventId = 0;
+
+EventsWaiter::EventsWaiter()
+    : trigger_event_(kEmptyEventId),
+      counter_(0),
+      eof_(true),
+      waiting_(false),
+      cv_(1) {}
+
+std::shared_ptr<EventsWaiter::EventNotifier> EventsWaiter::RegisterEvent(
+    const std::string& name, EventChecker checker) {
+  EventId id = kEmptyEventId;
+  EventInfo* evt = nullptr;
+  do {
+    auto counter = counter_.fetch_add(1);
+    id = std::hash<std::string>()(name + std::to_string(counter));
+    if (id == kEmptyEventId) {
+      continue;
+    }
+    std::lock_guard<paddle::memory::SpinLock> guard(events_lock_);
+    if (events_.count(id) > 0) {
+      continue;
+    }
+    evt = &(events_[id]);
+  } while (evt == nullptr);
+  evt->id = id;
+  evt->name = name;
+  evt->type = TriggerType::LevelTriggered;
+  evt->checker = std::move(checker);
+  eof_.store(false, std::memory_order_relaxed);
+  VLOG(10) << "Register event id:" << id << " name:" << name;
+  auto notifier = std::shared_ptr<EventNotifier>(new EventNotifier(id, this));
+  return notifier;
+}
+
+std::shared_ptr<EventsWaiter::EventNotifier> EventsWaiter::RegisterEvent(
+    const std::string& name) {
+  EventId id = kEmptyEventId;
+  EventInfo* evt = nullptr;
+  do {
+    auto counter = counter_.fetch_add(1);
+    id = std::hash<std::string>()(name + std::to_string(counter));
+    if (id == kEmptyEventId) {
+      continue;
+    }
+    std::lock_guard<paddle::memory::SpinLock> guard(events_lock_);
+    if (events_.count(id) > 0) {
+      continue;
+    }
+    evt = &(events_[id]);
+  } while (evt == nullptr);
+  evt->id = id;
+  evt->name = name;
+  evt->type = TriggerType::EdgeTriggered;
+  evt->checker = []() { return false; };
+  eof_.store(false, std::memory_order_relaxed);
+  VLOG(10) << "Register event id:" << id << " name:" << name;
+  auto notifier = std::shared_ptr<EventNotifier>(new EventNotifier(id, this));
+  return notifier;
+}
+
+void EventsWaiter::UnregisterEvent(const EventId& id) {
+  VLOG(10) << "Unregister event id:" << id;
+  {
+    std::lock_guard<paddle::memory::SpinLock> guard(events_lock_);
+    deleted_events_.insert(id);
+    if (deleted_events_.size() == events_.size()) {
+      eof_.store(true, std::memory_order_relaxed);
+    }
+  }
+  if (eof_.load(std::memory_order_relaxed)) {
+    cv_.Notify(true);
+  }
+}
+
+std::string EventsWaiter::WaitEvent() {
+  // only one user can wait at any time
+  bool waiting = false;
+  if (!waiting_.compare_exchange_strong(waiting,
+                                        true,
+                                        std::memory_order_seq_cst,
+                                        std::memory_order_relaxed)) {
+    PADDLE_THROW(
+        common::errors::ResourceExhausted("Another thread is waiting."));
+  }
+
+  auto w = cv_.GetWaiter(0);
+  EventId triggered = trigger_event_;
+  while (triggered == kEmptyEventId && !eof_) {
+    cv_.Prewait();
+
+    // double check
+    triggered = trigger_event_;
+    // checkers
+    if (triggered == kEmptyEventId) {
+      {
+        std::lock_guard<paddle::memory::SpinLock> guard(events_lock_);
+        for (auto& kv : events_) {
+          auto& evt = kv.second;
+          if (TriggerType::LevelTriggered == evt.type && evt.checker()) {
+            triggered = evt.id;
+            break;
+          }
+        }
+      }
+      if (triggered != kEmptyEventId) {
+        EventId prev = kEmptyEventId;
+        if (!trigger_event_.compare_exchange_strong(
+                prev,
+                triggered,
+                std::memory_order_seq_cst,
+                std::memory_order_relaxed)) {
+          triggered = prev;
+        }
+      }
+    }
+
+    if (triggered != kEmptyEventId || eof_) {
+      cv_.CancelWait();
+    } else {
+      cv_.CommitWait(w);
+      triggered = trigger_event_;
+    }
+  }
+
+  trigger_event_.store(kEmptyEventId, std::memory_order_relaxed);
+  std::string evt_name =
+      triggered == kEmptyEventId ? "NoEventNotifier" : GetEventName(triggered);
+  VLOG(10) << "Consume event id:" << triggered << ", name:" << evt_name;
+  // lazy deletion
+  {
+    triggered = trigger_event_;
+    std::lock_guard<paddle::memory::SpinLock> guard(events_lock_);
+    if (!deleted_events_.empty()) {
+      for (auto evt : deleted_events_) {
+        if (evt == triggered) {
+          continue;
+        }
+        events_.erase(evt);
+      }
+      deleted_events_.clear();
+    }
+  }
+  waiting_.store(false, std::memory_order_relaxed);
+  return evt_name;
+}
+
+int EventsWaiter::Clear() {
+  bool waiting = false;
+  if (!waiting_.compare_exchange_strong(waiting,
+                                        true,
+                                        std::memory_order_seq_cst,
+                                        std::memory_order_relaxed)) {
+    return -1;
+  }
+  trigger_event_.store(kEmptyEventId, std::memory_order_relaxed);
+  waiting_.store(false);
+  return 0;
+}
+
+void EventsWaiter::TriggerEvent(const EventId& id) {
+  VLOG(10) << "Try to trigger event id:" << id;
+  EventId prev = kEmptyEventId;
+  if (!trigger_event_.compare_exchange_strong(
+          prev, id, std::memory_order_seq_cst, std::memory_order_relaxed)) {
+    VLOG(10) << "Event id:" << prev << " is pending";
+    return;
+  }
+  VLOG(10) << "Triggered event id:" << id;
+  cv_.Notify(true);
+}
+
+void EventsWaiter::CancelEvent(const EventId& id) {
+  VLOG(10) << "Try to cancel event id:" << id;
+  EventId prev = id;
+  if (!trigger_event_.compare_exchange_strong(prev,
+                                              kEmptyEventId,
+                                              std::memory_order_seq_cst,
+                                              std::memory_order_relaxed)) {
+    VLOG(10) << "Event id:" << prev << " is pending";
+    return;
+  }
+  VLOG(10) << "Cancelled event id:" << id;
+}
+
+std::string EventsWaiter::GetEventName(const EventId& id) {
+  std::lock_guard<paddle::memory::SpinLock> guard(events_lock_);
+  auto iter = events_.find(id);
+  if (iter == events_.end()) {
+    return "Unregistered";
+  }
+  return iter->second.name;
+}
+
+}  // namespace paddle::framework

paddle/fluid/framework/new_executor/workqueue/events_waiter.h

@@ -0,0 +1,116 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+#include <functional>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "paddle/fluid/framework/new_executor/workqueue/event_count.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+
+namespace paddle {
+namespace framework {
+
+// A multiplexing waiter, be able to wait multiple kinds of events
+// simultaneously.
+// Multi-Producer single-consumer single-slot message-queue.
+class PADDLE_API EventsWaiter {
+ public:
+  using EventId = std::size_t;
+
+  using EventChecker = std::function<bool()>;
+
+  // Make sure EventsWaiter has a longer lifetime than EventNotifier.
+  class EventNotifier {
+   public:
+    ~EventNotifier() { waiter_.UnregisterEvent(id_); }
+
+    void NotifyEvent() { waiter_.TriggerEvent(id_); }
+
+    void CancelEvent() { waiter_.CancelEvent(id_); }
+
+    std::string GetEventName() { return waiter_.GetEventName(id_); }
+
+   private:
+    friend EventsWaiter;
+    EventNotifier(EventId id, EventsWaiter* waiter)
+        : id_(id), waiter_(*waiter) {}
+    EventNotifier(const EventNotifier&) = delete;
+    void operator=(const EventNotifier&) = delete;
+
+    EventId id_;
+    EventsWaiter& waiter_;
+  };
+
+  EventsWaiter();
+  EventsWaiter(const EventsWaiter&) = delete;
+  EventsWaiter& operator=(const EventsWaiter&) = delete;
+
+  // Register a level-triggered event. If the checker returns true or
+  // EventNotifier::NotifyEvent is called, the corresponding event will be
+  // distributed.
+  std::shared_ptr<EventNotifier> RegisterEvent(const std::string& name,
+                                               EventChecker checker);
+
+  // Register an edge-triggered event. The corresponding event will be
+  // distributed when EventNotifier::NotifyEvent is called.
+  std::shared_ptr<EventNotifier> RegisterEvent(const std::string& name);
+
+  void UnregisterEvent(const EventId& id);
+
+  // Blocking the calling thread to wait any of the registered events.
+  std::string WaitEvent();
+
+  // Nonblocking.
+  // Clear the slot, no matter whether there is an event.
+  // Return value:
+  //     -1 : another thread is waiting.
+  //      0 : succ.
+  int Clear();
+
+ private:
+  friend EventNotifier;
+
+  enum class TriggerType { LevelTriggered, EdgeTriggered };
+
+  struct EventInfo {
+    EventId id;
+    std::string name;
+    TriggerType type;
+    EventChecker checker;
+  };
+
+  void TriggerEvent(const EventId& id);
+
+  void CancelEvent(const EventId& id);
+
+  std::string GetEventName(const EventId& id);
+
+  std::unordered_map<EventId, EventInfo> events_;
+  std::unordered_set<EventId> deleted_events_;
+  paddle::memory::SpinLock events_lock_;
+  std::atomic<EventId> trigger_event_;
+  std::atomic<uint64_t> counter_;
+  std::atomic<bool> eof_;
+  std::atomic<bool> waiting_;
+  EventCount cv_;
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/new_executor/workqueue/nonblocking_threadpool.h

@@ -0,0 +1,475 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2016 Dmitry Vyukov <dvyukov@google.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include <atomic>
+#include <cstdlib>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/framework/new_executor/workqueue/event_count.h"
+#include "paddle/fluid/framework/new_executor/workqueue/run_queue.h"
+#include "paddle/fluid/framework/new_executor/workqueue/thread_environment.h"
+#include "paddle/phi/core/os_info.h"
+#include "paddle/phi/core/platform/profiler/event_tracing.h"
+
+namespace paddle {
+namespace framework {
+
+template <typename Environment>
+class ThreadPoolTempl {
+ public:
+  typedef typename Environment::Task Task;
+  typedef RunQueue<Task, 1024> Queue;
+
+  ThreadPoolTempl(const std::string& name,
+                  int num_threads,
+                  bool allow_spinning,
+                  bool always_spinning,
+                  Environment env = Environment())
+      : env_(env),
+        allow_spinning_(allow_spinning),
+        always_spinning_(always_spinning),
+        global_steal_partition_(EncodePartition(0, num_threads)),
+        blocked_(0),
+        done_(false),
+        cancelled_(false),
+        ec_(num_threads),
+        num_threads_(num_threads),
+        thread_data_(num_threads),
+        name_(name) {
+    // Calculate coprimes of all numbers [1, num_threads].
+    // Coprimes are used for random walks over all threads in Steal
+    // and NonEmptyQueueIndex. Iteration is based on the fact that if we take
+    // a random starting thread index t and calculate num_threads - 1 subsequent
+    // indices as (t + coprime) % num_threads, we will cover all threads without
+    // repetitions (effectively getting a presudo-random permutation of thread
+    // indices).
+    assert(num_threads_ >= 1 && num_threads_ < kMaxThreads);
+    all_coprimes_.reserve(num_threads_);
+    for (int i = 1; i <= num_threads_; ++i) {
+      all_coprimes_.emplace_back(i);
+      ComputeCoprimes(i, &(all_coprimes_.back()));
+    }
+    for (int i = 0; i < num_threads_; i++) {
+      SetStealPartition(i, EncodePartition(0, num_threads_));
+      thread_data_[i].thread.reset(
+          env_.CreateThread([this, i]() { WorkerLoop(i); }));
+    }
+  }
+
+  ~ThreadPoolTempl() {
+    done_ = true;
+
+    // Now if all threads block without work, they will start exiting.
+    // But note that threads can continue to work arbitrary long,
+    // block, submit new work, unblock and otherwise live full life.
+    if (!cancelled_) {
+      ec_.Notify(true);
+    } else {
+      // Since we were cancelled, there might be entries in the queues.
+      // Empty them to prevent their destructor from asserting.
+      for (size_t i = 0; i < thread_data_.size(); i++) {
+        thread_data_[i].queue.Flush();
+      }
+    }
+    // Join threads explicitly (by destroying) to avoid destruction order within
+    // this class.
+    for (size_t i = 0; i < thread_data_.size(); ++i) {
+      thread_data_[i].thread.reset();
+    }
+  }
+
+  void SetStealPartitions(
+      const std::vector<std::pair<unsigned, unsigned>>& partitions) {
+    assert(partitions.size() == static_cast<std::size_t>(num_threads_));
+
+    // Pass this information to each thread queue.
+    for (int i = 0; i < num_threads_; i++) {
+      const auto& pair = partitions[i];
+      unsigned start = pair.first, end = pair.second;
+      AssertBounds(start, end);
+      unsigned val = EncodePartition(start, end);
+      SetStealPartition(i, val);
+    }
+  }
+
+  void AddTask(std::function<void()> fn) {
+    AddTaskWithHint(std::move(fn), 0, num_threads_);
+  }
+
+  void AddTaskWithHint(std::function<void()> fn, int start, int limit) {
+    Task t = env_.CreateTask(std::move(fn));
+    PerThread* pt = GetPerThread();
+    if (pt->pool == this) {
+      // Worker thread of this pool, push onto the thread's queue.
+      Queue& q = thread_data_[pt->thread_id].queue;
+      t = q.PushFront(std::move(t));
+    } else {
+      // A free-standing thread (or worker of another pool), push onto a random
+      // queue.
+      assert(start < limit);
+      assert(limit <= num_threads_);
+      int num_queues = limit - start;
+      int rnd = Rand(&pt->rand) % num_queues;
+      assert(start + rnd < limit);
+      Queue& q = thread_data_[start + rnd].queue;
+      t = q.PushBack(std::move(t));
+    }
+
+    // Note: below we touch this after making w available to worker threads.
+    // Strictly speaking, this can lead to a racy-use-after-free. Consider that
+    // Schedule is called from a thread that is neither main thread nor a worker
+    // thread of this pool. Then, execution of w directly or indirectly
+    // completes overall computations, which in turn leads to destruction of
+    // this. We expect that such scenario is prevented by program, that is,
+    // this is kept alive while any threads can potentially be in Schedule.
+    if (!t.f) {
+      // Allow 'false positive' which makes a redundant notification.
+      VLOG(6) << "Add task, Notify";
+      ec_.Notify(false);
+    } else {
+      env_.ExecuteTask(t);  // Push failed, execute directly.
+    }
+  }
+
+  void Cancel() {
+    cancelled_ = true;
+    done_ = true;
+
+    // Wake up the threads without work to let them exit on their own.
+    ec_.Notify(true);
+  }
+
+  void WaitThreadsExit() {
+    for (size_t i = 0; i < thread_data_.size(); ++i) {
+      thread_data_[i].thread->WaitExit();
+    }
+  }
+
+  size_t NumThreads() const { return num_threads_; }
+
+  int CurrentThreadId() const {
+    const PerThread* pt = const_cast<ThreadPoolTempl*>(this)->GetPerThread();
+    if (pt->pool == this) {
+      return pt->thread_id;
+    } else {
+      return -1;
+    }
+  }
+
+ private:
+  // Create a single atomic<int> that encodes start and limit information for
+  // each thread.
+  // We expect num_threads_ < 65536, so we can store them in a single
+  // std::atomic<unsigned>.
+  // Exposed publicly as static functions so that external callers can reuse
+  // this encode/decode logic for maintaining their own thread-safe copies of
+  // scheduling and steal domain(s).
+  static const int kMaxPartitionBits = 16;
+  static const int kMaxThreads = 1 << kMaxPartitionBits;
+
+  inline unsigned EncodePartition(unsigned start, unsigned limit) {
+    return (start << kMaxPartitionBits) | limit;
+  }
+
+  inline void DecodePartition(unsigned val, unsigned* start, unsigned* limit) {
+    *limit = val & (kMaxThreads - 1);
+    val >>= kMaxPartitionBits;
+    *start = val;
+  }
+
+  void AssertBounds(int start, int end) {
+    assert(start >= 0);
+    assert(start < end);  // non-zero sized partition
+    assert(end <= num_threads_);
+  }
+
+  inline void SetStealPartition(size_t i, unsigned val) {
+    thread_data_[i].steal_partition.store(val, std::memory_order_relaxed);
+  }
+
+  inline unsigned GetStealPartition(int i) {
+    return thread_data_[i].steal_partition.load(std::memory_order_relaxed);
+  }
+
+  inline void ComputeCoprimes(int n, std::vector<unsigned>* coprimes) {
+    for (int i = 1; i <= n; i++) {
+      unsigned a = i;
+      unsigned b = n;
+      // If GCD(a, b) == 1, then a and b are coprimes.
+      while (b != 0) {
+        unsigned tmp = a;
+        a = b;
+        b = tmp % b;
+      }
+      if (a == 1) {
+        coprimes->push_back(i);
+      }
+    }
+  }
+
+  typedef typename Environment::EnvThread Thread;
+
+  struct PerThread {
+    constexpr PerThread() : pool(NULL), rand(0), thread_id(-1) {}
+    ThreadPoolTempl* pool;  // Parent pool, or null for normal threads.
+    uint64_t rand;          // Random generator state.
+    int thread_id;          // Worker thread index in pool.
+  };
+
+  struct ThreadData {
+    constexpr ThreadData() : thread(), steal_partition(0), queue() {}
+    std::unique_ptr<Thread> thread;
+    std::atomic<unsigned> steal_partition;
+    Queue queue;
+  };
+
+  Environment env_;
+  const bool allow_spinning_;
+  const bool always_spinning_;
+  std::vector<std::vector<unsigned>> all_coprimes_;
+  unsigned global_steal_partition_;
+  std::atomic<unsigned> blocked_;
+  std::atomic<bool> done_;
+  std::atomic<bool> cancelled_;
+  EventCount ec_;
+  const int num_threads_;
+  std::vector<ThreadData> thread_data_;
+  std::string name_;
+
+  // Main worker thread loop.
+  void WorkerLoop(int thread_id) {
+    std::string thr_name = name_ + "_thread_" + std::to_string(thread_id);
+    VLOG(1) << thr_name << " started ";
+    phi::SetCurrentThreadName(thr_name);
+    PerThread* pt = GetPerThread();
+    pt->pool = this;
+    pt->rand = GlobalThreadIdHash();
+    pt->thread_id = thread_id;
+    Queue& q = thread_data_[thread_id].queue;
+    EventCount::Waiter* waiter = ec_.GetWaiter(thread_id);
+    // TODO(dvyukov,rmlarsen): The time spent in NonEmptyQueueIndex() is
+    // proportional to num_threads_ and we assume that new work is scheduled at
+    // a constant rate, so we set spin_count to 5000 / num_threads_. The
+    // constant was picked based on a fair dice roll, tune it.
+    const int spin_count =
+        allow_spinning_ && num_threads_ > 0 ? 5000 / num_threads_ : 0;
+    if (num_threads_ == 1) {
+      // For num_threads_ == 1 there is no point in going through the expensive
+      // steal loop. Moreover, since NonEmptyQueueIndex() calls PopBack() on the
+      // victim queues it might reverse the order in which ops are executed
+      // compared to the order in which they are added, which tends to be
+      // counter-productive for the types of I/O workloads the single thread
+      // pools tend to be used for.
+      while (!cancelled_) {
+        Task t = q.PopFront();
+        for (int i = 0; i < spin_count && !t.f; i++) {
+          if (!cancelled_.load(std::memory_order_relaxed)) {
+            t = q.PopFront();
+          }
+        }
+        if (!t.f) {
+          if (!WaitForWork(waiter, &t)) {
+            return;
+          }
+        }
+        if (t.f) {
+          env_.ExecuteTask(t);
+        }
+      }
+    } else {
+      while (!cancelled_) {
+        Task t = q.PopFront();
+        if (!t.f) {
+          t = LocalSteal();
+          if (!t.f) {
+            t = GlobalSteal();
+            if (!t.f) {
+              if (allow_spinning_) {
+                for (int i = 0; i < spin_count && !t.f; i++) {
+                  if (!cancelled_.load(std::memory_order_relaxed)) {
+                    t = GlobalSteal();
+                  } else {
+                    return;
+                  }
+                }
+              }
+              if (!t.f) {
+                if (!WaitForWork(waiter, &t)) {
+                  return;
+                }
+              }
+            }
+          }
+        }
+        if (t.f) {
+          env_.ExecuteTask(t);
+        }
+      }
+    }
+  }
+
+  // Steal tries to steal work from other worker threads in the range [start,
+  // limit) in best-effort manner.
+  Task Steal(unsigned start, unsigned limit) {
+    PerThread* pt = GetPerThread();
+    const size_t size = limit - start;
+    unsigned r = Rand(&pt->rand);
+    // Reduce r into [0, size) range, this utilizes trick from
+    // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
+    assert(all_coprimes_[size - 1].size() < (1 << 30));
+    unsigned victim = ((uint64_t)r * (uint64_t)size) >> 32;
+    unsigned index =
+        ((uint64_t)all_coprimes_[size - 1].size() * (uint64_t)r) >> 32;
+    unsigned inc = all_coprimes_[size - 1][index];
+
+    for (unsigned i = 0; i < size; i++) {
+      assert(start + victim < limit);
+      Task t = thread_data_[start + victim].queue.PopBack();
+      if (t.f) {
+        return t;
+      }
+      victim += inc;
+      if (victim >= size) {
+        victim -= size;
+      }
+    }
+    return Task();
+  }
+
+  // Steals work within threads belonging to the partition.
+  Task LocalSteal() {
+    PerThread* pt = GetPerThread();
+    unsigned partition = GetStealPartition(pt->thread_id);
+    // If thread steal partition is the same as global partition, there is no
+    // need to go through the steal loop twice.
+    if (global_steal_partition_ == partition) return Task();
+    unsigned start, limit;
+    DecodePartition(partition, &start, &limit);
+    AssertBounds(start, limit);
+
+    return Steal(start, limit);
+  }
+
+  // Steals work from any other thread in the pool.
+  Task GlobalSteal() { return Steal(0, num_threads_); }
+
+  // WaitForWork blocks until new work is available (returns true), or if it is
+  // time to exit (returns false). Can optionally return a task to execute in t
+  // (in such case t.f != nullptr on return).
+  bool WaitForWork(EventCount::Waiter* waiter, Task* t) {
+    assert(t != nullptr && !t->f);
+    // We already did best-effort emptiness check in Steal, so prepare for
+    // blocking.
+    ec_.Prewait();
+    if (cancelled_) {
+      ec_.CancelWait();
+      return false;
+    }
+
+    // Number of blocked threads is used as notification condition.
+    // We must increase the counter before the emptiness check.
+    blocked_++;
+
+    // Now do a reliable emptiness check.
+    int victim = NonEmptyQueueIndex();
+    if (victim != -1) {
+      ec_.CancelWait();
+      *t = thread_data_[victim].queue.PopBack();
+      blocked_--;
+      return true;
+    }
+
+    // Number of blocked threads is used as termination condition.
+    // If we are shutting down and all worker threads blocked without work,
+    // that's we are done.
+    if (done_ && blocked_ == static_cast<unsigned>(num_threads_)) {
+      ec_.CancelWait();
+      // Almost done, but need to re-check queues.
+      // Consider that all queues are empty and all worker threads are preempted
+      // right after incrementing blocked_ above. Now a free-standing thread
+      // submits work and calls destructor (which sets done_). If we don't
+      // re-check queues, we will exit leaving the work unexecuted.
+      if (NonEmptyQueueIndex() != -1) {
+        // Note: we must not pop from queues before we decrement blocked_,
+        // otherwise the following scenario is possible. Consider that instead
+        // of checking for emptiness we popped the only element from queues.
+        // Now other worker threads can start exiting, which is bad if the
+        // work item submits other work. So we just check emptiness here,
+        // which ensures that all worker threads exit at the same time.
+        blocked_--;
+        return true;
+      }
+      // Reached stable termination state.
+      ec_.Notify(true);
+      return false;
+    }
+
+    // Cancel wait if always_spinning_
+    if (always_spinning_) {
+      ec_.CancelWait();
+      blocked_--;
+      return true;
+    }
+
+    // Wait for work
+    phi::RecordEvent record(
+        "WaitForWork", phi::TracerEventType::UserDefined, 10);
+    ec_.CommitWait(waiter);
+    blocked_--;
+    return true;
+  }
+
+  int NonEmptyQueueIndex() {
+    PerThread* pt = GetPerThread();
+    // We intentionally design NonEmptyQueueIndex to steal work from
+    // anywhere in the queue so threads don't block in WaitForWork() forever
+    // when all threads in their partition go to sleep. Steal is still local.
+    const size_t size = thread_data_.size();
+    unsigned r = Rand(&pt->rand);
+    unsigned inc = all_coprimes_[size - 1][r % all_coprimes_[size - 1].size()];
+    unsigned victim = r % size;
+    for (unsigned i = 0; i < size; i++) {
+      if (!thread_data_[victim].queue.Empty()) {
+        return victim;
+      }
+      victim += inc;
+      if (victim >= size) {
+        victim -= size;
+      }
+    }
+    return -1;
+  }
+
+  static inline uint64_t GlobalThreadIdHash() {
+    return std::hash<std::thread::id>()(std::this_thread::get_id());
+  }
+
+  inline PerThread* GetPerThread() {
+    static thread_local PerThread per_thread_;
+    PerThread* pt = &per_thread_;
+    return pt;
+  }
+
+  static inline unsigned Rand(uint64_t* state) {
+    uint64_t current = *state;
+    // Update the internal state
+    *state = current * 6364136223846793005ULL + 0xda3e39cb94b95bdbULL;
+    // Generate the random output (using the PCG-XSH-RS scheme)
+    return static_cast<unsigned>((current ^ (current >> 22)) >>
+                                 (22 + (current >> 61)));
+  }
+};
+
+using NonblockingThreadPool = ThreadPoolTempl<StlThreadEnvironment>;
+
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/new_executor/workqueue/run_queue.h

@@ -0,0 +1,270 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2016 Dmitry Vyukov <dvyukov@google.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+// RunQueue is a fixed-size, partially non-blocking deque or Work items.
+// Operations on front of the queue must be done by a single thread (owner),
+// operations on back of the queue can be done by multiple threads concurrently.
+//
+// Algorithm outline:
+// All remote threads operating on the queue back are serialized by a mutex.
+// This ensures that at most two threads access state: owner and one remote
+// thread (Size aside). The algorithm ensures that the occupied region of the
+// underlying array is logically continuous (can wraparound, but no stray
+// occupied elements). Owner operates on one end of this region, remote thread
+// operates on the other end. Synchronization between these threads
+// (potential consumption of the last element and take up of the last empty
+// element) happens by means of state variable in each element. States are:
+// empty, busy (in process of insertion of removal) and ready. Threads claim
+// elements (empty->busy and ready->busy transitions) by means of a CAS
+// operation. The finishing transition (busy->empty and busy->ready) are done
+// with plain store as the element is exclusively owned by the current thread.
+//
+// Note: we could permit only pointers as elements, then we would not need
+// separate state variable as null/non-null pointer value would serve as state,
+// but that would require malloc/free per operation for large, complex values
+// (and this is designed to store std::function<()>).
+//
+// What changed by PaddlePaddle
+//   1. Use paddle::memory::SpinLock instead of std::mutex to protect back_.
+//   2. Make front_/back_ aligned to get better performance.
+//   3. Replace Eigen utils with std utils.
+
+#pragma once
+
+#include <atomic>
+#include <cassert>
+#include <cstdint>
+#include <mutex>
+#include <vector>
+
+#include "paddle/fluid/framework/new_executor/workqueue/workqueue_utils.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+
+namespace paddle {
+namespace framework {
+
+template <typename Work, unsigned kSize>
+class RunQueue {
+ public:
+  RunQueue() : front_(0), back_(0) {
+    // require power-of-two for fast masking
+    static_assert((kSize & (kSize - 1)) == 0,
+                  "need to be a power of two for fast masking");
+    static_assert(
+        kSize > 2,
+        "need to be in [4, 65536] range to leave enough space for counter");
+    static_assert(
+        kSize <= (64 << 10),
+        "need to be in [4, 65536] range to leave enough space for counter");
+    for (unsigned i = 0; i < kSize; i++)
+      array_[i].state.store(kEmpty, std::memory_order_relaxed);
+  }
+
+  RunQueue(const RunQueue&) = delete;
+  void operator=(const RunQueue&) = delete;
+
+  ~RunQueue() { assert(Size() == 0); }
+
+  // PushFront inserts w at the beginning of the queue.
+  // If queue is full returns w, otherwise returns default-constructed Work.
+  Work PushFront(Work w) {
+    unsigned front = front_.load(std::memory_order_relaxed);
+    Elem* e = &array_[front & kMask];
+    uint8_t s = e->state.load(std::memory_order_relaxed);
+    if (s != kEmpty || !e->state.compare_exchange_strong(
+                           s, kBusy, std::memory_order_acquire)) {
+      return w;
+    }
+    front_.store(front + 1 + (kSize << 1), std::memory_order_relaxed);
+    e->w = std::move(w);
+    e->state.store(kReady, std::memory_order_release);
+    return Work();
+  }
+
+  // PopFront removes and returns the first element in the queue.
+  // If the queue was empty returns default-constructed Work.
+  Work PopFront() {
+    unsigned front = front_.load(std::memory_order_relaxed);
+    Elem* e = &array_[(front - 1) & kMask];
+    uint8_t s = e->state.load(std::memory_order_relaxed);
+    if (s != kReady || !e->state.compare_exchange_strong(
+                           s, kBusy, std::memory_order_acquire)) {
+      return Work();
+    }
+    Work w = std::move(e->w);
+    e->state.store(kEmpty, std::memory_order_release);
+    front = ((front - 1) & kMask2) | (front & ~kMask2);
+    front_.store(front, std::memory_order_relaxed);
+    return w;
+  }
+
+  // PushBack adds w at the end of the queue.
+  // If queue is full returns w, otherwise returns default-constructed Work.
+  Work PushBack(Work w) {
+    std::unique_lock<paddle::memory::SpinLock> lock(mutex_);
+    unsigned back = back_.load(std::memory_order_relaxed);
+    Elem* e = &array_[(back - 1) & kMask];
+    uint8_t s = e->state.load(std::memory_order_relaxed);
+    if (s != kEmpty || !e->state.compare_exchange_strong(
+                           s, kBusy, std::memory_order_acquire)) {
+      return w;
+    }
+    back = ((back - 1) & kMask2) | (back & ~kMask2);
+    back_.store(back, std::memory_order_relaxed);
+    e->w = std::move(w);
+    e->state.store(kReady, std::memory_order_release);
+    return Work();
+  }
+
+  // PopBack removes and returns the last elements in the queue.
+  Work PopBack() {
+    if (Empty()) {
+      return Work();
+    }
+
+    std::unique_lock<paddle::memory::SpinLock> lock(mutex_);
+    unsigned back = back_.load(std::memory_order_relaxed);
+    Elem* e = &array_[back & kMask];
+    uint8_t s = e->state.load(std::memory_order_relaxed);
+    if (s != kReady || !e->state.compare_exchange_strong(
+                           s, kBusy, std::memory_order_acquire)) {
+      return Work();
+    }
+    Work w = std::move(e->w);
+    e->state.store(kEmpty, std::memory_order_release);
+    back_.store(back + 1 + (kSize << 1), std::memory_order_relaxed);
+    return w;
+  }
+
+  // PopBackHalf removes and returns half last elements in the queue.
+  // Returns number of elements removed.
+  unsigned PopBackHalf(std::vector<Work>* result) {
+    if (Empty()) {
+      return 0;
+    }
+
+    std::unique_lock<paddle::memory::SpinLock> lock(mutex_);
+    unsigned back = back_.load(std::memory_order_relaxed);
+    unsigned size = Size();
+    unsigned mid = back;
+    if (size > 1) mid = back + (size - 1) / 2;
+    unsigned n = 0;
+    unsigned start = 0;
+    for (; static_cast<int>(mid - back) >= 0; mid--) {
+      Elem* e = &array_[mid & kMask];
+      uint8_t s = e->state.load(std::memory_order_relaxed);
+      if (n == 0) {
+        if (s != kReady || !e->state.compare_exchange_strong(
+                               s, kBusy, std::memory_order_acquire))
+          continue;
+        start = mid;
+      } else {
+        // Note: no need to store temporal kBusy, we exclusively own these
+        // elements.
+        assert(s == kReady);
+      }
+      result->push_back(std::move(e->w));
+      e->state.store(kEmpty, std::memory_order_release);
+      n++;
+    }
+    if (n != 0) {
+      back_.store(start + 1 + (kSize << 1), std::memory_order_relaxed);
+    }
+    return n;
+  }
+
+  // Size returns current queue size.
+  // Can be called by any thread at any time.
+  unsigned Size() const { return SizeOrNotEmpty<true>(); }
+
+  // Empty tests whether container is empty.
+  // Can be called by any thread at any time.
+  bool Empty() const { return SizeOrNotEmpty<false>() == 0; }
+
+  // Delete all the elements from the queue.
+  void Flush() {
+    while (!Empty()) {
+      PopFront();
+    }
+  }
+
+ private:
+  static const unsigned kMask = kSize - 1;
+  static const unsigned kMask2 = (kSize << 1) - 1;
+  struct alignas(64) Elem {
+    std::atomic<uint8_t> state;
+    Work w;
+  };
+  enum {
+    kEmpty,
+    kBusy,
+    kReady,
+  };
+
+  // Low log(kSize) + 1 bits in front_ and back_ contain rolling index of
+  // front/back, respectively. The remaining bits contain modification counters
+  // that are incremented on Push operations. This allows us to (1) distinguish
+  // between empty and full conditions (if we would use log(kSize) bits for
+  // position, these conditions would be indistinguishable); (2) obtain
+  // consistent snapshot of front_/back_ for Size operation using the
+  // modification counters.
+  alignas(64) std::atomic<unsigned> front_;
+  alignas(64) std::atomic<unsigned> back_;
+  paddle::memory::SpinLock mutex_;
+  Elem array_[kSize];
+
+  // SizeOrNotEmpty returns current queue size; if NeedSizeEstimate is false,
+  // only whether the size is 0 is guaranteed to be correct.
+  // Can be called by any thread at any time.
+  template <bool NeedSizeEstimate>
+  unsigned SizeOrNotEmpty() const {
+    // Emptiness plays critical role in thread pool blocking. So we go to great
+    // effort to not produce false positives (claim non-empty queue as empty).
+    unsigned front = front_.load(std::memory_order_acquire);
+    for (;;) {
+      // Capture a consistent snapshot of front/tail.
+      unsigned back = back_.load(std::memory_order_acquire);
+      unsigned front1 = front_.load(std::memory_order_relaxed);
+      if (front != front1) {
+        front = front1;
+        std::atomic_thread_fence(std::memory_order_acquire);
+        continue;
+      }
+      if (NeedSizeEstimate) {
+        return CalculateSize(front, back);
+      } else {
+        // This value will be 0 if the queue is empty, and undefined otherwise.
+        unsigned maybe_zero = ((front ^ back) & kMask2);
+        // Queue size estimate must agree with maybe zero check on the queue
+        // empty/non-empty state.
+        assert((CalculateSize(front, back) == 0) == (maybe_zero == 0));
+        return maybe_zero;
+      }
+    }
+  }
+
+  inline unsigned CalculateSize(unsigned front, unsigned back) const {
+    int size = (front & kMask2) - (back & kMask2);
+    // Fix overflow.
+    if (size < 0) {
+      size += 2 * kSize;
+    }
+    // Order of modification in push/pop is crafted to make the queue look
+    // larger than it is during concurrent modifications. E.g. push can
+    // increment size before the corresponding pop has decremented it.
+    // So the computed size can be up to kSize + 1, fix it.
+    if (size > static_cast<int>(kSize)) {
+      size = kSize;
+    }
+    return static_cast<unsigned>(size);
+  }
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/new_executor/workqueue/thread_environment.h

@@ -0,0 +1,51 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include <functional>
+#include <thread>
+
+namespace paddle {
+namespace framework {
+
+struct StlThreadEnvironment {
+  struct Task {
+    std::function<void()> f;
+  };
+
+  // EnvThread constructor must start the thread,
+  // destructor must join the thread.
+  class EnvThread {
+   public:
+    explicit EnvThread(std::function<void()> f) : thr_(std::move(f)) {}
+    void WaitExit() {
+      if (thr_.joinable()) {
+        thr_.join();
+      }
+    }
+    ~EnvThread() {
+      if (thr_.joinable()) {
+        thr_.join();
+      }
+    }
+
+   private:
+    std::thread thr_;
+  };
+
+  EnvThread* CreateThread(std::function<void()> f) {
+    return new EnvThread(std::move(f));
+  }
+  Task CreateTask(std::function<void()> f) { return Task{std::move(f)}; }
+  void ExecuteTask(const Task& t) { t.f(); }
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/new_executor/workqueue/workqueue.cc

@@ -0,0 +1,251 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#include "paddle/fluid/framework/new_executor/workqueue/workqueue.h"
+
+#include "paddle/fluid/framework/new_executor/workqueue/nonblocking_threadpool.h"
+#include "paddle/fluid/framework/new_executor/workqueue/workqueue_utils.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/phi/core/platform/profiler/event_tracing.h"
+
+namespace paddle::framework {
+
+void WorkQueueOptions::Validate() const {
+  PADDLE_ENFORCE_GT(name.size(),
+                    0,
+                    common::errors::InvalidArgument(
+                        "WorkQueueOptions.name must be nonempty"));
+  PADDLE_ENFORCE_EQ(
+      name.find('_'),
+      std::string::npos,
+      common::errors::InvalidArgument(
+          "WorkQueueOptions.name shouldn't contain an underline"));
+  PADDLE_ENFORCE_EQ(
+      allow_spinning == false && always_spinning == true,
+      false,
+      common::errors::InvalidArgument("WorkQueueOptions.allow_spinning must "
+                                      "be true when always_spinning is set"));
+}
+
+namespace {
+
+using TaskTracker = TaskTracker<EventsWaiter::EventNotifier>;
+
+class WorkQueueImpl : public WorkQueue {
+ public:
+  explicit WorkQueueImpl(const WorkQueueOptions& options) : WorkQueue(options) {
+    if (options_.track_task && options.events_waiter != nullptr) {
+      empty_notifier_ = options.events_waiter->RegisterEvent(kQueueEmptyEvent);
+      void* storage = AlignedMalloc(sizeof(TaskTracker), alignof(TaskTracker));
+      tracker_ = new (storage) TaskTracker(*empty_notifier_);
+    }
+    if (options_.detached == false && options.events_waiter != nullptr) {
+      destruct_notifier_ =
+          options.events_waiter->RegisterEvent(kQueueDestructEvent);
+    }
+    queue_ = new NonblockingThreadPool(options_.name,
+                                       static_cast<int>(options_.num_threads),
+                                       options_.allow_spinning,
+                                       options_.always_spinning);
+  }
+
+  ~WorkQueueImpl() override {
+    delete queue_;
+    if (tracker_ != nullptr) {
+      tracker_->~TaskTracker();
+      AlignedFree(tracker_);
+    }
+    if (destruct_notifier_) {
+      destruct_notifier_->NotifyEvent();
+    }
+  }
+
+  void AddTask(std::function<void()> fn) override {
+    phi::RecordEvent record(
+        "WorkQueue::AddTask", phi::TracerEventType::UserDefined, 10 /*level*/);
+    if (tracker_ != nullptr) {
+      fn = [task = std::move(fn),
+            raii = CounterGuard<TaskTracker>(tracker_)]() mutable { task(); };
+    }
+    queue_->AddTask(std::move(fn));
+  }
+
+  void Cancel() override {
+    queue_->Cancel();
+    queue_->WaitThreadsExit();
+  }
+
+  size_t NumThreads() const override { return queue_->NumThreads(); }
+
+ private:
+  NonblockingThreadPool* queue_{nullptr};
+  TaskTracker* tracker_{nullptr};
+  std::shared_ptr<EventsWaiter::EventNotifier> empty_notifier_;
+  std::shared_ptr<EventsWaiter::EventNotifier> destruct_notifier_;
+};
+
+class WorkQueueGroupImpl : public WorkQueueGroup {
+ public:
+  explicit WorkQueueGroupImpl(
+      const std::vector<WorkQueueOptions>& queue_options);
+
+  ~WorkQueueGroupImpl() override;
+
+  void AddTask(size_t queue_idx, std::function<void()> fn) override;
+
+  size_t QueueNumThreads(size_t queue_idx) const override;
+
+  size_t QueueGroupNumThreads() const override;
+
+  void Cancel() override;
+
+ private:
+  std::vector<NonblockingThreadPool*> queues_;
+  NonblockingThreadPool* queues_storage_;
+  TaskTracker* tracker_;
+  std::shared_ptr<EventsWaiter::EventNotifier> empty_notifier_;
+  std::shared_ptr<EventsWaiter::EventNotifier> destruct_notifier_;
+};
+
+WorkQueueGroupImpl::WorkQueueGroupImpl(
+    const std::vector<WorkQueueOptions>& queues_options)
+    : WorkQueueGroup(queues_options),
+      queues_storage_(nullptr),
+      tracker_(nullptr) {
+  size_t num_queues = queues_options_.size();
+  queues_.resize(num_queues);
+  void* buffer = malloc(sizeof(NonblockingThreadPool) * num_queues);  // NOLINT
+  queues_storage_ = reinterpret_cast<NonblockingThreadPool*>(buffer);
+
+  for (size_t idx = 0; idx < num_queues; ++idx) {
+    const auto& options = queues_options_[idx];
+    if (options.num_threads == 0) {
+      queues_[idx] = nullptr;
+      continue;
+    }
+    if (options.track_task && tracker_ == nullptr &&
+        options.events_waiter != nullptr) {
+      empty_notifier_ = options.events_waiter->RegisterEvent(kQueueEmptyEvent);
+      void* storage = AlignedMalloc(sizeof(TaskTracker), alignof(TaskTracker));
+      tracker_ = new (storage) TaskTracker(*empty_notifier_);
+    }
+    if (options.detached == false && options.events_waiter != nullptr &&
+        !destruct_notifier_) {
+      destruct_notifier_ =
+          options.events_waiter->RegisterEvent(kQueueDestructEvent);
+    }
+    queues_[idx] = new (&queues_storage_[idx])
+        NonblockingThreadPool(options.name,
+                              static_cast<int>(options.num_threads),
+                              options.allow_spinning,
+                              options.always_spinning);
+  }
+}
+
+WorkQueueGroupImpl::~WorkQueueGroupImpl() {
+  for (auto queue : queues_) {
+    if (queue) {
+      queue->~NonblockingThreadPool();
+    }
+  }
+  if (tracker_ != nullptr) {
+    tracker_->~TaskTracker();
+    AlignedFree(tracker_);
+  }
+  free(queues_storage_);  // NOLINT
+  if (destruct_notifier_) {
+    destruct_notifier_->NotifyEvent();
+  }
+}
+
+void WorkQueueGroupImpl::AddTask(size_t queue_idx, std::function<void()> fn) {
+  phi::RecordEvent record(
+      "WorkQueue::AddTask", phi::TracerEventType::UserDefined, 10 /*level*/);
+  assert(queue_idx < queues_.size());
+  PADDLE_ENFORCE_NOT_NULL(
+      queues_.at(queue_idx),
+      common::errors::NotFound("Workqueue of index %d is not initialized.",
+                               queue_idx));
+  if (queues_options_.at(queue_idx).track_task) {
+    fn = [task = std::move(fn),
+          raii = CounterGuard<TaskTracker>(tracker_)]() mutable { task(); };
+  }
+  queues_[queue_idx]->AddTask(std::move(fn));
+}
+
+size_t WorkQueueGroupImpl::QueueNumThreads(size_t queue_idx) const {
+  assert(queue_idx < queues_.size());
+  if (!queues_.at(queue_idx)) {
+    return 0;
+  }
+  return queues_.at(queue_idx)->NumThreads();
+}
+
+size_t WorkQueueGroupImpl::QueueGroupNumThreads() const {
+  size_t total_num = 0;
+  for (auto queue : queues_) {
+    total_num += queue->NumThreads();
+  }
+  return total_num;
+}
+
+void WorkQueueGroupImpl::Cancel() {
+  for (auto queue : queues_) {
+    if (queue) {
+      queue->Cancel();
+    }
+  }
+  for (auto queue : queues_) {
+    if (queue) {
+      queue->WaitThreadsExit();
+    }
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<WorkQueue> CreateSingleThreadedWorkQueue(
+    const WorkQueueOptions& options) {
+  options.Validate();
+  // extra check
+  PADDLE_ENFORCE_EQ(options.num_threads,
+                    1u,
+                    common::errors::InvalidArgument(
+                        "For a SingleThreadedWorkQueue, "
+                        "WorkQueueOptions.num_threads must equals to 1."));
+  std::unique_ptr<WorkQueue> ptr(new WorkQueueImpl(options));
+  return ptr;
+}
+
+std::unique_ptr<WorkQueue> CreateMultiThreadedWorkQueue(
+    const WorkQueueOptions& options) {
+  options.Validate();
+  // extra check
+  PADDLE_ENFORCE_GT(
+      options.num_threads,
+      1u,
+      common::errors::InvalidArgument("For a MultiThreadedWorkQueue, "
+                                      "WorkQueueOptions.num_threads must be "
+                                      "greater than 1."));
+  std::unique_ptr<WorkQueue> ptr(new WorkQueueImpl(options));
+  return ptr;
+}
+
+std::unique_ptr<WorkQueueGroup> CreateWorkQueueGroup(
+    const std::vector<WorkQueueOptions>& queues_options) {
+  PADDLE_ENFORCE_GT(queues_options.size(),
+                    1u,
+                    common::errors::InvalidArgument(
+                        "For a WorkQueueGroup, the number of WorkQueueOptions "
+                        "must be greater than 1."));
+  for (const auto& opts : queues_options) {
+    opts.Validate();
+  }
+  std::unique_ptr<WorkQueueGroup> ptr(new WorkQueueGroupImpl(queues_options));
+  return ptr;
+}
+
+}  // namespace paddle::framework

paddle/fluid/framework/new_executor/workqueue/workqueue.h

@@ -0,0 +1,197 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <functional>
+#include <future>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace framework {
+
+constexpr const char* kQueueEmptyEvent = "QueueEmpty";
+constexpr const char* kQueueDestructEvent = "QueueDestruct";
+
+// For std::function
+// https://stackoverflow.com/questions/25421346/how-to-create-an-stdfunction-from-a-move-capturing-lambda-expression
+template <typename OnlyMovable>
+class FakeCopyable {
+ public:
+  explicit FakeCopyable(OnlyMovable&& obj) : obj_(std::move(obj)) {
+    static_assert(std::is_copy_constructible<OnlyMovable>::value == false,
+                  "Need not to use FakeCopyable");
+  }
+
+  FakeCopyable(FakeCopyable&& other) : obj_(std::move(other.obj_)) {}
+
+  FakeCopyable(const FakeCopyable& other) {
+    PADDLE_THROW(common::errors::Unavailable(
+        "Never use the copy constructor of FakeCopyable."));
+  }
+
+  OnlyMovable& Get() { return obj_; }
+
+ private:
+  OnlyMovable obj_;
+};
+
+class EventsWaiter;
+
+struct WorkQueueOptions {
+  WorkQueueOptions(const std::string& name,
+                   size_t num_threads,
+                   bool allow_spinning,
+                   bool track_task)
+      : name(name),
+        num_threads(num_threads),
+        allow_spinning(allow_spinning),
+        track_task(track_task) {
+    Validate();
+  }
+
+  WorkQueueOptions(const std::string& name,
+                   size_t num_threads,
+                   bool allow_spinning,
+                   bool always_spinning,
+                   bool track_task,
+                   bool detached,
+                   EventsWaiter* waiter)
+      : name(name),
+        num_threads(num_threads),
+        allow_spinning(allow_spinning),
+        always_spinning(always_spinning),
+        track_task(track_task),
+        detached(detached),
+        events_waiter(waiter) {
+    Validate();
+  }
+
+  // throw an exception if there is an invalid option
+  PADDLE_API void Validate() const;
+
+  std::string name;
+  size_t num_threads;
+  // Worker threads will spin for a while if this flag is set.
+  bool allow_spinning;
+  // Worker threads will never sleep if this flag is set.
+  // Better performance vs. higher CPU utilization.
+  bool always_spinning{false};
+  // If you need to blocking the calling  thread to wait "queue empty", set
+  // track_task = true and set events_waiter. EventsWaiter::WaitEvent will
+  // block the calling thread until any of events (including "queue empty")
+  // occurred.
+  bool track_task;
+  // If you need to be noticed when a WorkQueue Destruct() , set detached =
+  // false and set events_waiter.
+  bool detached{true};
+  EventsWaiter* events_waiter{nullptr};  // not owned
+};
+
+class WorkQueue {
+ public:
+  explicit WorkQueue(const WorkQueueOptions& options) : options_(options) {}
+
+  WorkQueue(const WorkQueue&) = delete;
+
+  WorkQueue& operator=(const WorkQueue&) = delete;
+
+  virtual ~WorkQueue() = default;
+
+  virtual void AddTask(std::function<void()> fn) = 0;
+
+  // Higher cost than AddTask
+  template <typename F, typename... Args>
+  std::future<std::invoke_result_t<F, Args...>> AddAwaitableTask(
+      F&& f, Args&&... args) {
+    using ReturnType = std::invoke_result_t<F, Args...>;
+    std::function<ReturnType()> task =
+        std::bind(std::forward<F>(f), std::forward<Args>(args)...);
+    std::promise<ReturnType> prom;
+    std::future<ReturnType> res = prom.get_future();
+    AddTask([t = std::move(task),
+             p = FakeCopyable<std::promise<ReturnType>>(
+                 std::move(prom))]() mutable { p.Get().set_value(t()); });
+    return res;
+  }
+
+  // See WorkQueueOptions.track_task for details
+  // virtual void WaitQueueEmpty() = 0;
+
+  virtual size_t NumThreads() const = 0;
+
+  virtual void Cancel() = 0;
+
+ protected:
+  WorkQueueOptions options_;
+};
+
+class WorkQueueGroup {
+ public:
+  explicit WorkQueueGroup(const std::vector<WorkQueueOptions>& queues_options)
+      : queues_options_(queues_options) {}
+
+  WorkQueueGroup(const WorkQueueGroup&) = delete;
+
+  WorkQueueGroup& operator=(const WorkQueueGroup&) = delete;
+
+  virtual ~WorkQueueGroup() = default;
+
+  virtual void AddTask(size_t queue_idx, std::function<void()> fn) = 0;
+
+  // Higher cost than AddTask
+  template <typename F, typename... Args>
+  std::future<std::invoke_result_t<F, Args...>> AddAwaitableTask(
+      size_t queue_idx, F&& f, Args&&... args) {
+    using ReturnType = std::invoke_result_t<F, Args...>;
+    std::function<ReturnType()> task =
+        std::bind(std::forward<F>(f), std::forward<Args>(args)...);
+    std::promise<ReturnType> prom;
+    std::future<ReturnType> res = prom.get_future();
+    AddTask(queue_idx,
+            [t = std::move(task),
+             p = FakeCopyable<std::promise<ReturnType>>(
+                 std::move(prom))]() mutable { p.Get().set_value(t()); });
+    return res;
+  }
+
+  // See WorkQueueOptions.track_task for details
+  // virtual void WaitQueueGroupEmpty() = 0;
+
+  virtual size_t QueueNumThreads(size_t queue_idx) const = 0;
+
+  virtual size_t QueueGroupNumThreads() const = 0;
+
+  virtual void Cancel() = 0;
+
+ protected:
+  std::vector<WorkQueueOptions> queues_options_;
+};
+
+PADDLE_API std::unique_ptr<WorkQueue> CreateSingleThreadedWorkQueue(
+    const WorkQueueOptions& options);
+
+PADDLE_API std::unique_ptr<WorkQueue> CreateMultiThreadedWorkQueue(
+    const WorkQueueOptions& options);
+
+PADDLE_API std::unique_ptr<WorkQueueGroup> CreateWorkQueueGroup(
+    const std::vector<WorkQueueOptions>& queues_options);
+
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/new_executor/workqueue/workqueue_utils.cc

@@ -0,0 +1,58 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/new_executor/workqueue/workqueue_utils.h"
+
+#include <cstdint>
+#include <cstdlib>
+
+namespace paddle::framework {
+
+void* AlignedMalloc(size_t size, size_t alignment) {
+  assert(alignment >= sizeof(void*) && (alignment & (alignment - 1)) == 0);
+  size = (size + alignment - 1) / alignment * alignment;
+#if defined(_POSIX_C_SOURCE) && _POSIX_C_SOURCE >= 200112L
+  void* aligned_mem = nullptr;
+  if (posix_memalign(&aligned_mem, alignment, size) != 0) {
+    aligned_mem = nullptr;
+  }
+  return aligned_mem;
+#elif defined(_WIN32)
+  return _aligned_malloc(size, alignment);
+#else
+  void* mem = malloc(size + alignment);  // NOLINT
+  if (mem == nullptr) {
+    return nullptr;
+  }
+  size_t adjust = alignment - reinterpret_cast<uint64_t>(mem) % alignment;
+  void* aligned_mem = reinterpret_cast<char*>(mem) + adjust;
+  *(reinterpret_cast<void**>(aligned_mem) - 1) = mem;
+  assert(reinterpret_cast<uint64_t>(aligned_mem) % alignment == 0);
+  return aligned_mem;
+#endif
+}
+
+void AlignedFree(void* mem_ptr) {
+#if defined(_POSIX_C_SOURCE) && _POSIX_C_SOURCE >= 200112L
+  free(mem_ptr);
+#elif defined(_WIN32)
+  _aligned_free(mem_ptr);
+#else
+  if (mem_ptr) {
+    free(*(reinterpret_cast<void**>(mem_ptr) - 1));  // NOLINT
+  }
+#endif
+}
+
+}  // namespace paddle::framework

paddle/fluid/framework/new_executor/workqueue/workqueue_utils.h

@@ -0,0 +1,109 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+#include <cassert>
+#include <cstddef>
+#include <cstdlib>
+#include <memory>
+#include <set>
+#include <string>
+
+#include "paddle/fluid/framework/new_executor/workqueue/events_waiter.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace framework {
+
+template <typename Holder>
+class CounterGuard {
+ public:
+  explicit CounterGuard(Holder* holder) : counter_holder_(holder) {
+    assert(holder != nullptr);
+    counter_holder_->AddCounter();
+  }
+
+  ~CounterGuard() {
+    if (counter_holder_ != nullptr) {
+      counter_holder_->SubCounter();
+    }
+  }
+
+  CounterGuard(CounterGuard&& other) : counter_holder_(other.counter_holder_) {
+    other.counter_holder_ = nullptr;
+  }
+
+  CounterGuard& operator=(CounterGuard&& other) {
+    counter_holder_ = other.counter_holder_;
+    other.counter_holder_ = nullptr;
+    return *this;
+  }
+
+  // copy constructor deleted, we define this for std::function
+  // never use it directly
+  CounterGuard(const CounterGuard& other) {
+    PADDLE_THROW(common::errors::Unavailable(
+        "Never use the copy constructor of CounterGuard."));
+  }
+
+  CounterGuard& operator=(const CounterGuard&) = delete;
+
+ private:
+  Holder* counter_holder_{nullptr};
+};
+
+void* AlignedMalloc(size_t size, size_t alignment);
+
+void AlignedFree(void* memory_ptr);
+
+template <typename Notifier>
+class TaskTracker {
+ public:
+  TaskTracker() = default;
+
+  explicit TaskTracker(Notifier& notifier) : notifier_(&notifier) {}
+
+  TaskTracker(const TaskTracker&) = delete;
+
+  TaskTracker& operator=(const TaskTracker&) = delete;
+
+  ~TaskTracker() = default;
+
+  void AddCounter() {
+    if (0 == num_tasks_.fetch_add(1, std::memory_order_relaxed)) {
+      if (notifier_ != nullptr) {
+        notifier_->CancelEvent();
+      }
+    }
+  }
+
+  void SubCounter() {
+    if (1 == num_tasks_.fetch_sub(1, std::memory_order_relaxed)) {
+      if (notifier_ != nullptr) {
+        notifier_->NotifyEvent();
+      }
+    }
+  }
+
+  uint64_t PendingTaskNum() { return num_tasks_.load(); }
+
+ private:
+  alignas(64) std::atomic<uint64_t> num_tasks_{0};
+  Notifier* notifier_{nullptr};
+};
+
+}  // namespace framework
+}  // namespace paddle