1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 // Class method definitions for HostStream, the Stream implementation for
17 // the HostExecutor implementation.
18 #include "tensorflow/compiler/xla/stream_executor/host/host_stream.h"
19
20 #include "absl/synchronization/notification.h"
21 #include "tensorflow/core/platform/denormal.h"
22 #include "tensorflow/core/platform/setround.h"
23
24 namespace stream_executor {
25 namespace host {
26
27 namespace {
28
GetThreadOptions(size_t stack_size_in_bytes)29 port::ThreadOptions GetThreadOptions(size_t stack_size_in_bytes) {
30 port::ThreadOptions options;
31 options.stack_size = stack_size_in_bytes;
32 return options;
33 }
34
35 } // namespace
36
HostStream(size_t stack_size_in_bytes)37 HostStream::HostStream(size_t stack_size_in_bytes)
38 : thread_(port::Env::Default()->StartThread(
39 GetThreadOptions(stack_size_in_bytes), "host_executor",
40 [this]() { WorkLoop(); })) {}
41
~HostStream()42 HostStream::~HostStream() {
43 {
44 absl::MutexLock lock(&mu_);
45 work_queue_.push(nullptr);
46 }
47 // thread_'s destructor blocks until the thread finishes running.
48 thread_.reset();
49 }
50
EnqueueTask(std::function<void ()> task)51 bool HostStream::EnqueueTask(std::function<void()> task) {
52 return EnqueueTaskWithStatus([task = std::move(task)]() {
53 task();
54 return ::tensorflow::OkStatus();
55 });
56 }
57
EnqueueTaskWithStatus(std::function<port::Status ()> task)58 bool HostStream::EnqueueTaskWithStatus(std::function<port::Status()> task) {
59 CHECK(task != nullptr);
60 absl::MutexLock lock(&mu_);
61 work_queue_.push(std::move(task));
62 return true;
63 }
64
WorkAvailable()65 bool HostStream::WorkAvailable() { return !work_queue_.empty(); }
66
WorkLoop()67 void HostStream::WorkLoop() {
68 // Set denormal and rounding behavior to match the default TF ThreadPool
69 // behavior.
70 // TODO(phawkins, jlebar): it's not clear this is the best place to set this.
71 tensorflow::port::ScopedFlushDenormal flush;
72 tensorflow::port::ScopedSetRound round(FE_TONEAREST);
73 while (true) {
74 std::queue<std::function<port::Status()>> queue;
75 {
76 absl::MutexLock lock(&mu_);
77 mu_.Await(absl::Condition(this, &HostStream::WorkAvailable));
78 std::swap(queue, work_queue_);
79 }
80 while (!queue.empty()) {
81 std::function<port::Status()>& fn = queue.front();
82 if (!fn) {
83 return;
84 }
85 status_.Update(fn());
86 queue.pop();
87 }
88 }
89 }
90
BlockUntilDone()91 port::Status HostStream::BlockUntilDone() {
92 absl::Notification done;
93 port::Status status;
94 EnqueueTask([&done, &status, this]() {
95 // This task is always executed synchronously before 'status_' is updated
96 // with the result of the task (always OK() in this case), so we don't need
97 // to worry about locking access to 'status_'.
98 status = status_;
99 status_ = ::tensorflow::OkStatus();
100 done.Notify();
101 });
102 done.WaitForNotification();
103 return status;
104 }
105
106 } // namespace host
107
108 } // namespace stream_executor
109