1b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 3b04c9f48ce69377a1172fd58e8ab210b435f359eDerek MurrayLicensed under the Apache License, Version 2.0 (the "License"); 4b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murrayyou may not use this file except in compliance with the License. 5b04c9f48ce69377a1172fd58e8ab210b435f359eDerek MurrayYou may obtain a copy of the License at 6b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 7b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray http://www.apache.org/licenses/LICENSE-2.0 8b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 9b04c9f48ce69377a1172fd58e8ab210b435f359eDerek MurrayUnless required by applicable law or agreed to in writing, software 10b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murraydistributed under the License is distributed on an "AS IS" BASIS, 11b04c9f48ce69377a1172fd58e8ab210b435f359eDerek MurrayWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12b04c9f48ce69377a1172fd58e8ab210b435f359eDerek MurraySee the License for the specific language governing permissions and 13b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murraylimitations under the License. 14b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray==============================================================================*/ 15b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 16b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#ifndef TENSORFLOW_CORE_PLATFORM_DEFAULT_NOTIFICATION_H_ 17b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#define TENSORFLOW_CORE_PLATFORM_DEFAULT_NOTIFICATION_H_ 18b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 19b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#include <assert.h> 20ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower#include <atomic> // NOLINT 21b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#include <chrono> // NOLINT 22b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#include <condition_variable> // NOLINT 23b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 24b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#include "tensorflow/core/platform/mutex.h" 25b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#include "tensorflow/core/platform/types.h" 26b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 27b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murraynamespace tensorflow { 28b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 29b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murrayclass Notification { 30b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray public: 31ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower Notification() : notified_(0) {} 3228c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower ~Notification() { 3328c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower // In case the notification is being used to synchronize its own deletion, 3428c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower // force any prior notifier to leave its critical section before the object 3528c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower // is destroyed. 3628c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower mutex_lock l(mu_); 3728c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower } 38b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 39b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray void Notify() { 40b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray mutex_lock l(mu_); 41ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower assert(!HasBeenNotified()); 42ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower notified_.store(true, std::memory_order_release); 43b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray cv_.notify_all(); 44b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray } 45b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 46ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower bool HasBeenNotified() const { 47ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower return notified_.load(std::memory_order_acquire); 48b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray } 49b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 50b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray void WaitForNotification() { 5128c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower if (!HasBeenNotified()) { 5228c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower mutex_lock l(mu_); 5328c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower while (!HasBeenNotified()) { 5428c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower cv_.wait(l); 5528c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower } 56b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray } 57b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray } 58b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 59b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray private: 60b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray friend bool WaitForNotificationWithTimeout(Notification* n, 619d3eebb35ae339bfc8d58f56bb912336ca733e2dYuan Yu int64 timeout_in_us); 629d3eebb35ae339bfc8d58f56bb912336ca733e2dYuan Yu bool WaitForNotificationWithTimeout(int64 timeout_in_us) { 6328c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower bool notified = HasBeenNotified(); 6428c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower if (!notified) { 6528c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower mutex_lock l(mu_); 6628c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower do { 6728c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower notified = HasBeenNotified(); 6828c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower } while (!notified && 6928c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower cv_.wait_for(l, std::chrono::microseconds(timeout_in_us)) != 7028c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower std::cv_status::timeout); 71b48cfaea2aea3707a33e60c10385a87e37101b95A. Unique TensorFlower } 7228c32c380aa5c2dfe7070e22f8be6af5f67cd35cA. Unique TensorFlower return notified; 73b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray } 74b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 75ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower mutex mu_; // protects mutations of notified_ 76d0a5d885d61b837018cb931a4d577289acc826fcMartin Wicke condition_variable cv_; // signaled when notified_ becomes non-zero 77ead1d9ba76af1fc721ed16c11812a2783d41133fA. Unique TensorFlower std::atomic<bool> notified_; // mutations under mu_ 78b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray}; 79b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 80b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murrayinline bool WaitForNotificationWithTimeout(Notification* n, 819d3eebb35ae339bfc8d58f56bb912336ca733e2dYuan Yu int64 timeout_in_us) { 829d3eebb35ae339bfc8d58f56bb912336ca733e2dYuan Yu return n->WaitForNotificationWithTimeout(timeout_in_us); 83b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray} 84b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 85b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray} // namespace tensorflow 86b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray 87b04c9f48ce69377a1172fd58e8ab210b435f359eDerek Murray#endif // TENSORFLOW_CORE_PLATFORM_DEFAULT_NOTIFICATION_H_ 88