xref: /aosp_15_r20/external/eigen/unsupported/Eigen/CXX11/src/ThreadPool/Barrier.h (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2018 Rasmus Munk Larsen <[email protected]>
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:
Barrier(unsigned int count)20   Barrier(unsigned int count) : state_(count << 1), notified_(false) {
21     eigen_plain_assert(((count << 1) >> 1) == count);
22   }
~Barrier()23   ~Barrier() { eigen_plain_assert((state_ >> 1) == 0); }
24 
Notify()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 
Wait()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 {
NotificationNotification62   Notification() : Barrier(1){};
63 };
64 
65 }  // namespace Eigen
66 
67 #endif  // EIGEN_CXX11_THREADPOOL_BARRIER_H
68