Barrier.h (2113B)
1 // This file is part of Eigen, a lightweight C++ template library 2 // for linear algebra. 3 // 4 // Copyright (C) 2018 Rasmus Munk Larsen <rmlarsen@google.com> 5 // 6 // This Source Code Form is subject to the terms of the Mozilla 7 // Public License v. 2.0. If a copy of the MPL was not distributed 8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 10 // Barrier is an object that allows one or more threads to wait until 11 // Notify has been called a specified number of times. 12 13 #ifndef EIGEN_CXX11_THREADPOOL_BARRIER_H 14 #define EIGEN_CXX11_THREADPOOL_BARRIER_H 15 16 namespace Eigen { 17 18 class Barrier { 19 public: 20 Barrier(unsigned int count) : state_(count << 1), notified_(false) { 21 eigen_plain_assert(((count << 1) >> 1) == count); 22 } 23 ~Barrier() { eigen_plain_assert((state_ >> 1) == 0); } 24 25 void Notify() { 26 unsigned int v = state_.fetch_sub(2, std::memory_order_acq_rel) - 2; 27 if (v != 1) { 28 // Clear the lowest bit (waiter flag) and check that the original state 29 // value was not zero. If it was zero, it means that notify was called 30 // more times than the original count. 31 eigen_plain_assert(((v + 2) & ~1) != 0); 32 return; // either count has not dropped to 0, or waiter is not waiting 33 } 34 std::unique_lock<std::mutex> l(mu_); 35 eigen_plain_assert(!notified_); 36 notified_ = true; 37 cv_.notify_all(); 38 } 39 40 void Wait() { 41 unsigned int v = state_.fetch_or(1, std::memory_order_acq_rel); 42 if ((v >> 1) == 0) return; 43 std::unique_lock<std::mutex> l(mu_); 44 while (!notified_) { 45 cv_.wait(l); 46 } 47 } 48 49 private: 50 std::mutex mu_; 51 std::condition_variable cv_; 52 std::atomic<unsigned int> state_; // low bit is waiter flag 53 bool notified_; 54 }; 55 56 // Notification is an object that allows a user to to wait for another 57 // thread to signal a notification that an event has occurred. 58 // 59 // Multiple threads can wait on the same Notification object, 60 // but only one caller must call Notify() on the object. 61 struct Notification : Barrier { 62 Notification() : Barrier(1){}; 63 }; 64 65 } // namespace Eigen 66 67 #endif // EIGEN_CXX11_THREADPOOL_BARRIER_H