src/ThreadPool/Barrier.h

0001 // This file is part of Eigen, a lightweight C++ template library
0002 // for linear algebra.
0003 //
0004 // Copyright (C) 2018 Rasmus Munk Larsen <rmlarsen@google.com>
0005 //
0006 // This Source Code Form is subject to the terms of the Mozilla
0007 // Public License v. 2.0. If a copy of the MPL was not distributed
0008 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
0009
0010 // Barrier is an object that allows one or more threads to wait until
0011 // Notify has been called a specified number of times.
0012
0013 #ifndef EIGEN_CXX11_THREADPOOL_BARRIER_H
0014 #define EIGEN_CXX11_THREADPOOL_BARRIER_H
0015
0016 namespace Eigen {
0017
0018 class Barrier {
0019  public:
0020   Barrier(unsigned int count) : state_(count << 1), notified_(false) {
0021     eigen_plain_assert(((count << 1) >> 1) == count);
0022   }
0023   ~Barrier() { eigen_plain_assert((state_ >> 1) == 0); }
0024
0025   void Notify() {
0026     unsigned int v = state_.fetch_sub(2, std::memory_order_acq_rel) - 2;
0027     if (v != 1) {
0028       // Clear the lowest bit (waiter flag) and check that the original state
0029       // value was not zero. If it was zero, it means that notify was called
0030       // more times than the original count.
0031       eigen_plain_assert(((v + 2) & ~1) != 0);
0032       return;  // either count has not dropped to 0, or waiter is not waiting
0033     }
0034     std::unique_lock<std::mutex> l(mu_);
0035     eigen_plain_assert(!notified_);
0036     notified_ = true;
0037     cv_.notify_all();
0038   }
0039
0040   void Wait() {
0041     unsigned int v = state_.fetch_or(1, std::memory_order_acq_rel);
0042     if ((v >> 1) == 0) return;
0043     std::unique_lock<std::mutex> l(mu_);
0044     while (!notified_) {
0045       cv_.wait(l);
0046     }
0047   }
0048
0049  private:
0050   std::mutex mu_;
0051   std::condition_variable cv_;
0052   std::atomic<unsigned int> state_;  // low bit is waiter flag
0053   bool notified_;
0054 };
0055
0056 // Notification is an object that allows a user to to wait for another
0057 // thread to signal a notification that an event has occurred.
0058 //
0059 // Multiple threads can wait on the same Notification object,
0060 // but only one caller must call Notify() on the object.
0061 struct Notification : Barrier {
0062   Notification() : Barrier(1){};
0063 };
0064
0065 }  // namespace Eigen
0066
0067 #endif  // EIGEN_CXX11_THREADPOOL_BARRIER_H