1 /* Copyright 2015 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 #include "tensorflow/core/common_runtime/executor.h"
17
18 #include <atomic>
19 #include <deque>
20 #include <memory>
21 #include <string>
22 #include <unordered_map>
23 #include <vector>
24
25 #include "tensorflow/core/common_runtime/costmodel_manager.h"
26 #include "tensorflow/core/common_runtime/executor_factory.h"
27 #include "tensorflow/core/common_runtime/pending_counts.h"
28 #include "tensorflow/core/common_runtime/step_stats_collector.h"
29 #include "tensorflow/core/framework/allocation_description.pb.h"
30 #include "tensorflow/core/framework/allocator.h"
31 #include "tensorflow/core/framework/cancellation.h"
32 #include "tensorflow/core/framework/collective.h"
33 #include "tensorflow/core/framework/control_flow.h"
34 #include "tensorflow/core/framework/device_attributes.pb.h"
35 #include "tensorflow/core/framework/graph.pb.h"
36 #include "tensorflow/core/framework/log_memory.h"
37 #include "tensorflow/core/framework/node_def_util.h"
38 #include "tensorflow/core/framework/op_kernel.h"
39 #include "tensorflow/core/framework/op_segment.h"
40 #include "tensorflow/core/framework/step_stats.pb.h"
41 #include "tensorflow/core/framework/tensor.h"
42 #include "tensorflow/core/framework/tensor_reference.h"
43 #include "tensorflow/core/framework/types.h"
44 #include "tensorflow/core/framework/types.pb.h"
45 #include "tensorflow/core/graph/edgeset.h"
46 #include "tensorflow/core/lib/core/errors.h"
47 #include "tensorflow/core/lib/core/notification.h"
48 #include "tensorflow/core/lib/core/stringpiece.h"
49 #include "tensorflow/core/lib/core/threadpool.h"
50 #include "tensorflow/core/lib/gtl/flatmap.h"
51 #include "tensorflow/core/lib/gtl/flatset.h"
52 #include "tensorflow/core/lib/gtl/inlined_vector.h"
53 #include "tensorflow/core/lib/gtl/manual_constructor.h"
54 #include "tensorflow/core/lib/gtl/stl_util.h"
55 #include "tensorflow/core/lib/hash/hash.h"
56 #include "tensorflow/core/lib/strings/str_util.h"
57 #include "tensorflow/core/lib/strings/stringprintf.h"
58 #include "tensorflow/core/platform/context.h"
59 #include "tensorflow/core/platform/env.h"
60 #include "tensorflow/core/platform/logging.h"
61 #include "tensorflow/core/platform/macros.h"
62 #include "tensorflow/core/platform/mutex.h"
63 #include "tensorflow/core/platform/profile_utils/cpu_utils.h"
64 #include "tensorflow/core/platform/thread_annotations.h"
65 #include "tensorflow/core/platform/tracing.h"
66 #include "tensorflow/core/platform/types.h"
67 #include "tensorflow/core/util/tensor_slice_reader_cache.h"
68
69 namespace tensorflow {
70 namespace {
71
72 // 1-D, 0 element tensor.
73 static const Tensor* const kEmptyTensor = new Tensor;
74
IsInitializationOp(const Node * node)75 bool IsInitializationOp(const Node* node) {
76 return node->op_def().allows_uninitialized_input();
77 }
78
79 // Helper routines for collecting step stats.
80 namespace nodestats {
NowInNsec()81 inline int64 NowInNsec() { return Env::Default()->NowNanos(); }
82
SetScheduled(NodeExecStatsInterface * stats,int64 micros)83 void SetScheduled(NodeExecStatsInterface* stats, int64 micros) {
84 if (!stats) return;
85 stats->SetScheduled(micros * EnvTime::kMicrosToNanos);
86 }
87
SetAllStart(NodeExecStatsInterface * stats)88 void SetAllStart(NodeExecStatsInterface* stats) {
89 if (!stats) return;
90 stats->RecordExecutorStarted();
91 }
92
SetOpStart(NodeExecStatsInterface * stats)93 void SetOpStart(NodeExecStatsInterface* stats) {
94 if (!stats) return;
95 stats->RecordComputeStarted();
96 }
97
SetOpEnd(NodeExecStatsInterface * stats)98 void SetOpEnd(NodeExecStatsInterface* stats) {
99 if (!stats) return;
100 stats->RecordComputeEnded();
101 }
102
SetAllEnd(NodeExecStatsInterface * stats)103 void SetAllEnd(NodeExecStatsInterface* stats) {
104 if (!stats) return;
105 stats->RecordExecutorEnded();
106 }
107
SetOutput(NodeExecStatsInterface * stats,int slot,const Tensor * v)108 void SetOutput(NodeExecStatsInterface* stats, int slot, const Tensor* v) {
109 if (!stats) return;
110 stats->SetOutput(slot, v);
111 }
112
SetMemory(NodeExecStatsInterface * stats,OpKernelContext * ctx)113 void SetMemory(NodeExecStatsInterface* stats, OpKernelContext* ctx) {
114 if (!stats) return;
115 stats->SetMemory(ctx);
116 }
117
SetReferencedTensors(NodeExecStatsInterface * stats,const TensorReferenceVector & tensors)118 void SetReferencedTensors(NodeExecStatsInterface* stats,
119 const TensorReferenceVector& tensors) {
120 if (!stats) return;
121 stats->SetReferencedTensors(tensors);
122 }
123
124 } // namespace nodestats
125
126 class ExecutorImpl;
127 class GraphView;
128
129 struct EdgeInfo {
130 int dst_id;
131 int output_slot : 31;
132 // true if this is the last info for output_slot in the EdgeInfo list.
133 bool is_last : 1;
134 int input_slot;
135 };
136
137 // Time the execution of kernels (in CPU cycles). Used to dynamically identify
138 // inexpensive kernels which can be dispatched inline.
139 struct KernelTimer {
140 uint64 start_cycles = profile_utils::CpuUtils::GetCurrentClockCycle();
141
ElapsedCyclestensorflow::__anon675242740111::KernelTimer142 uint64 ElapsedCycles() {
143 return profile_utils::CpuUtils::GetCurrentClockCycle() - start_cycles;
144 }
145 };
146
147 struct NodeItem {
NodeItemtensorflow::__anon675242740111::NodeItem148 NodeItem() {}
149
150 // A graph node.
151 const Node* node = nullptr;
152
153 // The kernel for this node.
154 OpKernel* kernel = nullptr;
155
156 bool kernel_is_async : 1; // True iff kernel->AsAsync() != nullptr
157 bool is_merge : 1; // True iff IsMerge(node)
158 bool is_enter : 1; // True iff IsEnter(node)
159 bool is_constant_enter : 1; // True iff IsEnter(node) and
160 // node->GetAttr("is_constant") == true.
161 bool is_exit : 1; // True iff IsExit(node)
162 bool is_control_trigger : 1; // True iff IsControlTrigger(node)
163 bool is_sink : 1; // True iff IsSink(node)
164 // True iff IsEnter(node) || IsExit(node) || IsNextIteration(node)
165 bool is_enter_exit_or_next_iter : 1;
166
167 // Cached values of node->num_inputs() and node->num_outputs(), to
168 // avoid levels of indirection.
169 int num_inputs;
170 int num_outputs;
171
172 // ExecutorImpl::tensors_[input_start] is the 1st positional input
173 // for this node.
174 int input_start = 0;
175
176 // Number of output edges.
177 size_t num_output_edges;
178
179 PendingCounts::Handle pending_id;
180
output_edge_listtensorflow::__anon675242740111::NodeItem181 const EdgeInfo* output_edge_list() const { return output_edge_base(); }
182
183 // ith output edge.
output_edgetensorflow::__anon675242740111::NodeItem184 const EdgeInfo& output_edge(int i) const {
185 DCHECK_GE(i, 0);
186 DCHECK_LT(i, num_output_edges);
187 return output_edge_base()[i];
188 }
189
input_typetensorflow::__anon675242740111::NodeItem190 DataType input_type(int i) const {
191 DCHECK_LT(i, num_inputs);
192 return static_cast<DataType>(input_type_base()[i]);
193 }
output_typetensorflow::__anon675242740111::NodeItem194 DataType output_type(int i) const {
195 DCHECK_LT(i, num_outputs);
196 return static_cast<DataType>(output_type_base()[i]);
197 }
198
199 // Return array of per-output allocator attributes.
output_attrstensorflow::__anon675242740111::NodeItem200 const AllocatorAttributes* output_attrs() const { return output_attr_base(); }
201
202 // Return array of expected input index from which each output should
203 // be forwarded:
204 // kNeverForward (-2) for DO NOT FORWARD (must allocate).
205 // kNoReservation (-1) for no expected forwarding.
206 // 0... for forward from that input.
forward_fromtensorflow::__anon675242740111::NodeItem207 const int* forward_from() const { return forward_from_base(); }
208
209 private:
210 friend class GraphView;
211
212 // Variable length section starts immediately after *this
213 // (uint8 is enough for DataType).
214 // EdgeInfo out_edges[num_out_edges];
215 // AllocatorAttributes output_attr[num_outputs];
216 // int forward_from[num_outputs];
217 // uint8 input_type[num_inputs];
218 // uint8 output_type[num_outputs];
219
220 // Return pointer to variable length section.
vartensorflow::__anon675242740111::NodeItem221 char* var() const {
222 return const_cast<char*>(reinterpret_cast<const char*>(this) +
223 sizeof(NodeItem));
224 }
225
output_edge_basetensorflow::__anon675242740111::NodeItem226 EdgeInfo* output_edge_base() const {
227 return reinterpret_cast<EdgeInfo*>(var());
228 }
output_attr_basetensorflow::__anon675242740111::NodeItem229 AllocatorAttributes* output_attr_base() const {
230 return reinterpret_cast<AllocatorAttributes*>(var() + sizeof(EdgeInfo) *
231 num_output_edges);
232 }
forward_from_basetensorflow::__anon675242740111::NodeItem233 int* forward_from_base() const {
234 return reinterpret_cast<int*>(var() + sizeof(EdgeInfo) * num_output_edges +
235 sizeof(AllocatorAttributes) * num_outputs);
236 }
input_type_basetensorflow::__anon675242740111::NodeItem237 uint8* input_type_base() const {
238 return reinterpret_cast<uint8*>(
239 var() + sizeof(EdgeInfo) * num_output_edges +
240 sizeof(AllocatorAttributes) * num_outputs + sizeof(int) * num_outputs);
241 }
output_type_basetensorflow::__anon675242740111::NodeItem242 uint8* output_type_base() const {
243 return reinterpret_cast<uint8*>(
244 var() + sizeof(EdgeInfo) * num_output_edges +
245 sizeof(AllocatorAttributes) * num_outputs + sizeof(int) * num_outputs +
246 sizeof(uint8) * num_inputs);
247 }
248
249 TF_DISALLOW_COPY_AND_ASSIGN(NodeItem);
250 };
251
252 typedef gtl::InlinedVector<TensorValue, 4> TensorValueVec;
253 typedef gtl::InlinedVector<DeviceContext*, 4> DeviceContextVec;
254 typedef gtl::InlinedVector<AllocatorAttributes, 4> AllocatorAttributeVec;
255
256 // Immutable view of a Graph organized for efficient execution.
257 class GraphView {
258 public:
GraphView()259 GraphView() : space_(nullptr) {}
260 ~GraphView();
261
262 void Initialize(const Graph* g);
263 Status SetAllocAttrs(const Graph* g, const Device* device);
264 void SetScopedAllocatorAttrs(const std::vector<const Node*>& sa_nodes);
265
node(size_t id) const266 NodeItem* node(size_t id) const {
267 DCHECK_GE(id, 0);
268 DCHECK_LT(id, num_nodes_);
269 uint32 offset = node_offsets_[id];
270 return ((offset == kuint32max)
271 ? nullptr
272 : reinterpret_cast<NodeItem*>(space_ + node_offsets_[id]));
273 }
274
275 private:
276 char* InitializeNode(char* ptr, const Node* n);
277 size_t NodeItemBytes(const Node* n);
278
279 int32 num_nodes_ = 0;
280 uint32* node_offsets_ = nullptr; // array of size "graph_.num_node_ids()"
281 // node_offsets_[id] holds the byte offset for node w/ "id" in space_
282
283 char* space_; // NodeItem objects are allocated here
284
285 TF_DISALLOW_COPY_AND_ASSIGN(GraphView);
286 };
287
288 class ExecutorImpl : public Executor {
289 public:
ExecutorImpl(const LocalExecutorParams & p,std::unique_ptr<const Graph> g)290 ExecutorImpl(const LocalExecutorParams& p, std::unique_ptr<const Graph> g)
291 : params_(p), graph_(std::move(g)), gview_() {
292 CHECK(p.create_kernel != nullptr);
293 CHECK(p.delete_kernel != nullptr);
294 }
295
~ExecutorImpl()296 ~ExecutorImpl() override {
297 for (int i = 0; i < graph_->num_node_ids(); i++) {
298 NodeItem* item = gview_.node(i);
299 if (item != nullptr) {
300 params_.delete_kernel(item->kernel);
301 }
302 }
303 for (auto fiter : frame_info_) {
304 delete fiter.second;
305 }
306 }
307
308 Status Initialize();
309
310 // Process all Nodes in the current graph, attempting to infer the
311 // memory allocation attributes to be used wherever they may allocate
312 // a tensor buffer.
313 Status SetAllocAttrs();
314
315 void RunAsync(const Args& args, DoneCallback done) override;
316
317 private:
318 friend class ExecutorState;
319
320 struct ControlFlowInfo {
321 gtl::FlatSet<string> unique_frame_names;
322 std::vector<string> frame_names;
323 };
324
325 struct FrameInfo {
FrameInfotensorflow::__anon675242740111::ExecutorImpl::FrameInfo326 FrameInfo()
327 : input_count(0),
328 total_inputs(0),
329 pending_counts(nullptr),
330 nodes(nullptr) {}
331
332 // The total number of inputs to a frame.
333 int input_count;
334
335 // The total number of input tensors of a frame.
336 // == sum(nodes[*].num_inputs()) where nodes are the nodes in the frame.
337 int total_inputs;
338
339 // Used to determine the next place to allocate space in the
340 // pending_counts data structure we'll eventually construct
341 PendingCounts::Layout pending_counts_layout;
342
343 // Each frame has its own PendingCounts only for the nodes in the frame.
344 PendingCounts* pending_counts; // Owned
345
346 // The nodes in a frame. Used only for debugging.
347 std::vector<const Node*>* nodes; // Owned
348
~FrameInfotensorflow::__anon675242740111::ExecutorImpl::FrameInfo349 ~FrameInfo() {
350 delete pending_counts;
351 delete nodes;
352 }
353 };
354
355 static Status BuildControlFlowInfo(const Graph* graph,
356 ControlFlowInfo* cf_info);
357 void InitializePending(const Graph* graph, const ControlFlowInfo& cf_info);
358
EnsureFrameInfo(const string & fname)359 FrameInfo* EnsureFrameInfo(const string& fname) {
360 auto slot = &frame_info_[fname];
361 if (*slot == nullptr) {
362 *slot = new FrameInfo;
363 }
364 return *slot;
365 }
366
367 // Owned.
368 LocalExecutorParams params_;
369 std::unique_ptr<const Graph> graph_;
370 GraphView gview_;
371
372 // A cached value of params_
373 bool device_record_tensor_accesses_ = false;
374
375 // Root nodes (with no in edges) that should form the initial ready queue
376 std::vector<const Node*> root_nodes_;
377
378 // Mapping from frame name to static information about the frame.
379 // TODO(yuanbyu): We could cache it along with the graph so to avoid
380 // the overhead of constructing it for each executor instance.
381 gtl::FlatMap<string, FrameInfo*> frame_info_;
382
383 TF_DISALLOW_COPY_AND_ASSIGN(ExecutorImpl);
384 };
385
386 // Infer memory allocation attributes of a node n's output,
387 // based on its use node dst. Note that dst might not be directly
388 // connected to n by a single edge, but might be a downstream
389 // consumer of n's output by reference. *attr is updated with any
390 // necessary attributes.
391 Status InferAllocAttr(const Node* n, const Node* dst,
392 const DeviceNameUtils::ParsedName& local_dev_name,
393 AllocatorAttributes* attr);
394
~GraphView()395 GraphView::~GraphView() {
396 static_assert(std::is_trivially_destructible<AllocatorAttributes>::value,
397 "Update code if AllocatorAttributes gains a destructor");
398 static_assert(std::is_trivially_destructible<EdgeInfo>::value,
399 "Update code if EdgeInfo gains a destructor");
400 for (int i = 0; i < num_nodes_; i++) {
401 NodeItem* n = node(i);
402 if (n != nullptr) {
403 n->NodeItem::~NodeItem();
404 // Memory for "n" itself is held in space_ & gets cleaned up below
405 }
406 }
407 delete[] node_offsets_;
408 delete[] space_;
409 }
410
NodeItemBytes(const Node * n)411 size_t GraphView::NodeItemBytes(const Node* n) {
412 const size_t num_output_edges = n->out_edges().size();
413 const int num_inputs = n->num_inputs();
414 const int num_outputs = n->num_outputs();
415
416 // Compute number of bytes needed for NodeItem and variable length data.
417 // We do not subtract sizeof(var) since num_inputs/num_outputs might
418 // both be zero.
419 const size_t raw_bytes =
420 sizeof(NodeItem) // Fixed
421 + num_output_edges * sizeof(EdgeInfo) // output_edges[...]
422 + num_outputs * sizeof(AllocatorAttributes) // output_attr[...]
423 + num_outputs * sizeof(int) // forward_from[num_outputs]
424 + num_inputs * sizeof(uint8) // input_type[num_inputs]
425 + num_outputs * sizeof(uint8); // output_type[num_outputs]
426 static constexpr size_t kItemAlignment = sizeof(NodeItem*);
427 static_assert(kItemAlignment % alignof(NodeItem) == 0,
428 "NodeItem must be aligned with kItemAlignment");
429 static_assert(kItemAlignment % alignof(EdgeInfo) == 0,
430 "EdgeInfo must be aligned with kItemAlignment");
431 static_assert(kItemAlignment % alignof(AllocatorAttributes) == 0,
432 "AllocatorAttributes must be aligned with kItemAlignment");
433 static_assert(sizeof(NodeItem) % alignof(EdgeInfo) == 0,
434 "NodeItem must be aligned with EdgeInfo");
435 static_assert(sizeof(NodeItem) % alignof(AllocatorAttributes) == 0,
436 "NodeItem must be aligned with AllocatorAttributes");
437 static_assert(sizeof(EdgeInfo) % alignof(AllocatorAttributes) == 0,
438 "EdgeInfo must be aligned with AllocatorAttributes");
439 const size_t bytes =
440 ((raw_bytes + kItemAlignment - 1) / kItemAlignment) * kItemAlignment;
441 return bytes;
442 }
443
InitializeNode(char * ptr,const Node * n)444 char* GraphView::InitializeNode(char* ptr, const Node* n) {
445 const int id = n->id();
446 CHECK(node_offsets_[id] == kuint32max); // Initial value in constructor
447
448 const size_t bytes = NodeItemBytes(n);
449 constexpr size_t kItemAlignment = sizeof(NodeItem*);
450 CHECK_EQ(reinterpret_cast<uintptr_t>(ptr) % kItemAlignment, 0);
451 NodeItem* item = reinterpret_cast<NodeItem*>(ptr);
452
453 // We store a 32-bit offset relative to the beginning of space_, so that we
454 // only need an array of 32-bit values to map from node id to the NodeItem*,
455 // (versus 64 bits on most machines if we just stored an array of NodeItem*
456 // pointers). Casting to int64 is needed on 32bit CPU to avoid comparing
457 // values as "int" vs "size_t" in CHECK_LE.
458 CHECK_LE(static_cast<int64>(ptr - space_), kuint32max);
459 const uint32 offset = static_cast<uint32>(ptr - space_);
460 node_offsets_[id] = offset;
461 ptr += bytes;
462
463 const size_t num_output_edges = n->out_edges().size();
464 const int num_inputs = n->num_inputs();
465 const int num_outputs = n->num_outputs();
466
467 new (item) NodeItem();
468 item->num_inputs = num_inputs;
469 item->num_outputs = num_outputs;
470 item->num_output_edges = num_output_edges;
471
472 // Fill output edges.
473 // Keep track of the last EdgeInfo in the EdgeInfo array that references
474 // a given output slot. For all but the last, we need to do a copy of the
475 // Tensor when propagating results downstream in the graph, but for the
476 // last one, we can just do a move of the Tensor object to propagate it.
477 gtl::InlinedVector<EdgeInfo*, 4> last_indices(num_outputs, nullptr);
478 EdgeInfo* dst_edge = item->output_edge_base();
479 for (auto e : n->out_edges()) {
480 dst_edge->dst_id = e->dst()->id();
481 CHECK_LE(e->src_output(), 0x3FFFFFFF); // Must fit in 31 bits
482 dst_edge->output_slot = e->src_output();
483 dst_edge->is_last = false;
484 const int output_slot = dst_edge->output_slot;
485 if (output_slot >= 0) {
486 last_indices[output_slot] = dst_edge;
487 }
488 dst_edge->input_slot = e->dst_input();
489 dst_edge++;
490 }
491 for (EdgeInfo* edge_info : last_indices) {
492 if (edge_info != nullptr) {
493 edge_info->is_last = true;
494 }
495 }
496
497 AllocatorAttributes* output_attrs = item->output_attr_base();
498 for (int i = 0; i < num_outputs; i++) {
499 new (&output_attrs[i]) AllocatorAttributes();
500 }
501
502 DCHECK_LT(DataType_MAX, 255); // Must fit in uint8
503 uint8* input_types = item->input_type_base();
504 for (int i = 0; i < num_inputs; i++) {
505 input_types[i] = static_cast<uint8>(n->input_type(i));
506 DCHECK_EQ(item->input_type(i), n->input_type(i));
507 }
508
509 // Check ScopedAllocatorAttrs and forward_from. Also assign output_types.
510 {
511 std::vector<int> forward_input;
512 Status fwd_status =
513 GetNodeAttr(n->attrs(), "_forward_input", &forward_input);
514 std::vector<int> scoped_allocator_attrs;
515 Status sa_status =
516 GetNodeAttr(n->attrs(), "_scoped_allocator", &scoped_allocator_attrs);
517
518 int* forward_from = item->forward_from_base();
519 uint8* output_types = item->output_type_base();
520 for (int i = 0; i < num_outputs; ++i) {
521 output_types[i] = static_cast<uint8>(n->output_type(i));
522 DCHECK_EQ(item->output_type(i), n->output_type(i));
523
524 forward_from[i] = OpKernelContext::Params::kNoReservation;
525 if (sa_status.ok()) {
526 for (int j = 0; j < scoped_allocator_attrs.size(); j += 2) {
527 if (scoped_allocator_attrs[j] == i) {
528 // This output slot must be explicitly allocated from a
529 // ScopedAllocator.
530 forward_from[i] = OpKernelContext::Params::kNeverForward;
531 DCHECK_EQ(output_attrs[i].scope_id, 0);
532 output_attrs[i].scope_id = scoped_allocator_attrs[j + 1];
533 }
534 }
535 }
536 if (fwd_status.ok() &&
537 forward_from[i] == OpKernelContext::Params::kNoReservation) {
538 DCHECK_EQ(forward_input.size() % 2, 0);
539 for (int j = 0; j < forward_input.size(); j += 2) {
540 if (forward_input[j + 1] == i) {
541 DCHECK_EQ(forward_from[i], OpKernelContext::Params::kNoReservation);
542 forward_from[i] = forward_input[j];
543 break;
544 }
545 }
546 }
547 }
548 }
549
550 return ptr;
551 }
552
Initialize(const Graph * g)553 void GraphView::Initialize(const Graph* g) {
554 CHECK(node_offsets_ == nullptr);
555 const int num_nodes = g->num_node_ids();
556 num_nodes_ = num_nodes;
557 size_t total_bytes = 0;
558 for (const Node* n : g->nodes()) {
559 total_bytes += NodeItemBytes(n);
560 }
561
562 node_offsets_ = new uint32[num_nodes];
563 for (int i = 0; i < num_nodes; i++) {
564 node_offsets_[i] = kuint32max;
565 }
566
567 space_ = new char[total_bytes]; // NodeItem objects are allocated here
568 char* ptr = space_;
569 for (const Node* n : g->nodes()) {
570 ptr = InitializeNode(ptr, n);
571 }
572 CHECK_EQ(ptr, space_ + total_bytes);
573 }
574
GetMaxPendingCounts(const Node * n,size_t * max_pending,size_t * max_dead_count)575 void GetMaxPendingCounts(const Node* n, size_t* max_pending,
576 size_t* max_dead_count) {
577 const size_t num_in_edges = n->in_edges().size();
578 size_t initial_count;
579 if (IsMerge(n)) {
580 // merge waits all control inputs so we initialize the pending
581 // count to be the number of control edges.
582 int32 num_control_edges = 0;
583 for (const Edge* edge : n->in_edges()) {
584 if (edge->IsControlEdge()) {
585 num_control_edges++;
586 }
587 }
588 // Use bit 0 to indicate if we are waiting for a ready live data input.
589 initial_count = 1 + (num_control_edges << 1);
590 } else {
591 initial_count = num_in_edges;
592 }
593
594 *max_pending = initial_count;
595 *max_dead_count = num_in_edges;
596 }
597
Initialize()598 Status ExecutorImpl::Initialize() {
599 gview_.Initialize(graph_.get());
600
601 // Build the information about frames in this subgraph.
602 ControlFlowInfo cf_info;
603 TF_RETURN_IF_ERROR(BuildControlFlowInfo(graph_.get(), &cf_info));
604
605 // Cache this value so we make this virtual function call once, rather
606 // that O(# steps * # nodes per step) times.
607 device_record_tensor_accesses_ =
608 params_.device->RequiresRecordingAccessedTensors();
609
610 for (auto& it : cf_info.unique_frame_names) {
611 EnsureFrameInfo(it)->nodes = new std::vector<const Node*>;
612 }
613
614 // Preprocess every node in the graph to create an instance of op
615 // kernel for each node.
616 for (const Node* n : graph_->nodes()) {
617 const int id = n->id();
618 const string& frame_name = cf_info.frame_names[id];
619 FrameInfo* frame_info = EnsureFrameInfo(frame_name);
620
621 // See if this node is a root node, and if so, add to root_nodes_.
622 if (n->in_edges().empty()) {
623 root_nodes_.push_back(n);
624 }
625
626 NodeItem* item = gview_.node(id);
627 item->node = n;
628
629 item->input_start = frame_info->total_inputs;
630 frame_info->total_inputs += n->num_inputs();
631
632 Status s = params_.create_kernel(n->def(), &item->kernel);
633 if (!s.ok()) {
634 item->kernel = nullptr;
635 s = AttachDef(s, *n);
636 LOG(ERROR) << "Executor failed to create kernel. " << s;
637 return s;
638 }
639 CHECK(item->kernel);
640 item->kernel_is_async = (item->kernel->AsAsync() != nullptr);
641 item->is_merge = IsMerge(n);
642 item->is_enter = IsEnter(n);
643 if (item->is_enter) {
644 bool is_constant_enter;
645 TF_RETURN_IF_ERROR(
646 GetNodeAttr(n->attrs(), "is_constant", &is_constant_enter));
647 item->is_constant_enter = is_constant_enter;
648 } else {
649 item->is_constant_enter = false;
650 }
651 item->is_exit = IsExit(n);
652 item->is_control_trigger = IsControlTrigger(n);
653 item->is_sink = IsSink(n);
654 item->is_enter_exit_or_next_iter =
655 (IsEnter(n) || IsExit(n) || IsNextIteration(n));
656
657 // Compute the maximum values we'll store for this node in the
658 // pending counts data structure, and allocate a handle in
659 // that frame's pending counts data structure that has enough
660 // space to store these maximal count values.
661 size_t max_pending, max_dead;
662 GetMaxPendingCounts(n, &max_pending, &max_dead);
663 item->pending_id =
664 frame_info->pending_counts_layout.CreateHandle(max_pending, max_dead);
665
666 // Initialize static information about the frames in the graph.
667 frame_info->nodes->push_back(n);
668 if (IsEnter(n)) {
669 string enter_name;
670 TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "frame_name", &enter_name));
671 EnsureFrameInfo(enter_name)->input_count++;
672 }
673 }
674
675 // Initialize PendingCounts only after item->pending_id is initialized for
676 // all nodes.
677 InitializePending(graph_.get(), cf_info);
678
679 return gview_.SetAllocAttrs(graph_.get(), params_.device);
680 }
681
682 // If a Node has been marked to use a ScopedAllocator x for output i, then
683 // sc_attr will contain the subsequence (i, x) at an even offset. This function
684 // extracts and transfers that ScopedAllocator id to alloc_attr. For now, we
685 // only allow one ScopedAllocator use per Node.
ExtractScopedAllocatorAttr(const std::vector<int> & sc_attr,int output_index,AllocatorAttributes * alloc_attr)686 bool ExtractScopedAllocatorAttr(const std::vector<int>& sc_attr,
687 int output_index,
688 AllocatorAttributes* alloc_attr) {
689 DCHECK_LE(2, sc_attr.size());
690 for (int i = 0; i < sc_attr.size(); i += 2) {
691 if (sc_attr[i] == output_index) {
692 CHECK_EQ(alloc_attr->scope_id, 0);
693 alloc_attr->scope_id = sc_attr[i + 1];
694 return true;
695 }
696 }
697 return false;
698 }
699
SetScopedAllocatorAttrs(const std::vector<const Node * > & sa_nodes)700 void GraphView::SetScopedAllocatorAttrs(
701 const std::vector<const Node*>& sa_nodes) {
702 for (const Node* sa : sa_nodes) {
703 NodeItem* sa_item = node(sa->id());
704 AllocatorAttributes* sa_attrs = sa_item->output_attr_base();
705 // Control edges out of the ScopedAllocator should be use instances, but may
706 // include a few other nodes.
707 for (const auto& e : sa->out_edges()) {
708 if (!e->IsControlEdge()) {
709 continue;
710 }
711 Node* use_node = e->dst();
712 NodeItem* item = node(use_node->id());
713 AllocatorAttributes* use_attrs = item->output_attr_base();
714 std::vector<int> scoped_allocator_attrs;
715 Status s = GetNodeAttr(use_node->attrs(), "_scoped_allocator",
716 &scoped_allocator_attrs);
717 if (!s.ok()) {
718 VLOG(2) << "Failed to find expected ScopedAllocator attr on "
719 << use_node->name();
720 continue;
721 }
722 // There can be more than one output using ScopedAllocation, but this
723 // analysis assumes they use the same ScopedAllocator.
724 for (const auto& e : use_node->out_edges()) {
725 if (!e->IsControlEdge()) {
726 AllocatorAttributes attr;
727 if (ExtractScopedAllocatorAttr(scoped_allocator_attrs,
728 e->src_output(), &attr)) {
729 // Set the scope_id on this use instance node.
730 (use_attrs + e->src_output())->Merge(attr);
731 // Propagate the other attributes of this node back to the SA node.
732 attr = *(use_attrs + e->src_output());
733 attr.scope_id = 0;
734 sa_attrs->Merge(attr);
735 }
736 }
737 }
738 }
739 }
740 }
741
SetAllocAttrs(const Graph * g,const Device * device)742 Status GraphView::SetAllocAttrs(const Graph* g, const Device* device) {
743 Status s;
744 DeviceNameUtils::ParsedName local_dev_name = device->parsed_name();
745
746 std::vector<const Node*> scoped_allocator_instances;
747 for (const Node* n : g->nodes()) {
748 NodeItem* item = node(n->id());
749 AllocatorAttributes* attrs = item->output_attr_base();
750 if (IsScopedAllocator(n)) {
751 scoped_allocator_instances.push_back(n);
752 }
753
754 // Examine the out edges of each node looking for special use
755 // cases that may affect memory allocation attributes.
756 for (const auto& e : n->out_edges()) {
757 if (!e->IsControlEdge()) {
758 AllocatorAttributes attr;
759 s = InferAllocAttr(n, e->dst(), local_dev_name, &attr);
760 if (!s.ok()) return s;
761 if (attr.value != 0 || attr.scope_id != 0) {
762 attrs[e->src_output()].Merge(attr);
763 }
764 }
765 }
766
767 for (int out = 0; out < n->num_outputs(); out++) {
768 const OpKernel* op_kernel = item->kernel;
769 DCHECK_LT(out, op_kernel->output_memory_types().size());
770 bool on_host = op_kernel->output_memory_types()[out] == HOST_MEMORY;
771 if (on_host) {
772 AllocatorAttributes h;
773 h.set_on_host(on_host);
774 attrs[out].Merge(h);
775 }
776 }
777 }
778 SetScopedAllocatorAttrs(scoped_allocator_instances);
779 return s;
780 }
781
InferAllocAttr(const Node * n,const Node * dst,const DeviceNameUtils::ParsedName & local_dev_name,AllocatorAttributes * attr)782 Status InferAllocAttr(const Node* n, const Node* dst,
783 const DeviceNameUtils::ParsedName& local_dev_name,
784 AllocatorAttributes* attr) {
785 Status s;
786 // Note that it's possible for *n to be a Recv and *dst to be a Send,
787 // so these two cases are not mutually exclusive.
788 if (IsRecv(n)) {
789 string src_name;
790 s = GetNodeAttr(n->attrs(), "send_device", &src_name);
791 if (!s.ok()) return s;
792 DeviceNameUtils::ParsedName parsed_src_name;
793 if (!DeviceNameUtils::ParseFullName(src_name, &parsed_src_name)) {
794 s = errors::Internal("Bad send_device attr '", src_name, "' in node ",
795 n->name());
796 return s;
797 }
798 if (!DeviceNameUtils::IsSameAddressSpace(parsed_src_name, local_dev_name)) {
799 // Value is going to be the sink of an RPC.
800 attr->set_nic_compatible(true);
801 VLOG(2) << "node " << n->name() << " is the sink of an RPC in";
802 } else if ((local_dev_name.type == "CPU" || n->IsHostRecv()) &&
803 parsed_src_name.type != "CPU") {
804 // Value is going to be the sink of a local DMA from GPU to CPU (or
805 // other types of accelerators).
806 attr->set_gpu_compatible(true);
807 VLOG(2) << "node " << n->name() << " is the sink of a gpu->cpu copy";
808 } else {
809 VLOG(2) << "default alloc case local type " << local_dev_name.type
810 << " remote type " << parsed_src_name.type;
811 }
812 }
813 if (IsSend(dst)) {
814 string dst_name;
815 s = GetNodeAttr(dst->attrs(), "recv_device", &dst_name);
816 if (!s.ok()) return s;
817 DeviceNameUtils::ParsedName parsed_dst_name;
818 if (!DeviceNameUtils::ParseFullName(dst_name, &parsed_dst_name)) {
819 s = errors::Internal("Bad recv_device attr '", dst_name, "' in node ",
820 n->name());
821 return s;
822 }
823 if (!DeviceNameUtils::IsSameAddressSpace(parsed_dst_name, local_dev_name)) {
824 // Value is going to be the source of an RPC.
825 attr->set_nic_compatible(true);
826 VLOG(2) << "node " << n->name() << " is the source of an RPC out";
827 } else if ((local_dev_name.type == "CPU" || dst->IsHostSend()) &&
828 parsed_dst_name.type != "CPU") {
829 // Value is going to be the source of a local DMA from CPU to GPU (or
830 // other types of accelerators).
831 // Note that this does not cover the case where the allocation of the
832 // output tensor is not generated by the src: n.
833 attr->set_gpu_compatible(true);
834 VLOG(2) << "node " << n->name() << " is the source of a cpu->gpu copy";
835 } else {
836 VLOG(2) << "default alloc case local type " << local_dev_name.type
837 << " remote type " << parsed_dst_name.type;
838 }
839 }
840 if (n->IsCollective()) {
841 // We'll make the sweeping assumption that any collective op is going
842 // to be involved in network i/o.
843 attr->set_nic_compatible(true);
844 }
845 return s;
846 }
847
848 // The state associated with one invocation of ExecutorImpl::Run.
849 // ExecutorState dispatches nodes when they become ready and keeps
850 // track of how many predecessors of a node have not done (pending_).
851 class ExecutorState {
852 public:
853 ExecutorState(const Executor::Args& args, ExecutorImpl* impl);
854 ~ExecutorState();
855
856 void RunAsync(Executor::DoneCallback done);
857
858 private:
859 // Either a tensor pointer (pass-by-reference) or a tensor (pass-by-value).
860 // TODO(yuanbyu): A better way to do "has_value"?
861 struct Entry {
Entrytensorflow::__anon675242740111::ExecutorState::Entry862 Entry() {}
Entrytensorflow::__anon675242740111::ExecutorState::Entry863 Entry(const Entry& other)
864 : ref(other.ref),
865 ref_mu(other.ref_mu),
866 has_value(other.has_value),
867 val_field_is_set(other.val_field_is_set),
868 alloc_attr(other.alloc_attr),
869 device_context(other.device_context) {
870 if (val_field_is_set) {
871 val.Init(*other.val);
872 }
873 }
~Entrytensorflow::__anon675242740111::ExecutorState::Entry874 ~Entry() {
875 if (val_field_is_set) val.Destroy();
876 }
877
operator =tensorflow::__anon675242740111::ExecutorState::Entry878 Entry& operator=(const Entry& other) {
879 if (val_field_is_set) {
880 val.Destroy();
881 }
882 ref = other.ref;
883 ref_mu = other.ref_mu;
884 has_value = other.has_value;
885 val_field_is_set = other.val_field_is_set;
886 alloc_attr = other.alloc_attr;
887 device_context = other.device_context;
888 if (val_field_is_set) {
889 val.Init(*other.val);
890 }
891 return *this;
892 }
893
operator =tensorflow::__anon675242740111::ExecutorState::Entry894 Entry& operator=(Entry&& other) {
895 if (val_field_is_set) {
896 val.Destroy();
897 }
898 ref = other.ref;
899 ref_mu = other.ref_mu;
900 has_value = other.has_value;
901 val_field_is_set = other.val_field_is_set;
902 alloc_attr = other.alloc_attr;
903 device_context = other.device_context;
904 if (val_field_is_set) {
905 val.Init(std::move(*other.val));
906 }
907 return *this;
908 }
909
910 // Clears the <val> field.
ClearValtensorflow::__anon675242740111::ExecutorState::Entry911 void ClearVal() {
912 if (val_field_is_set) {
913 val.Destroy();
914 val_field_is_set = false;
915 has_value = false;
916 }
917 }
918
919 // A tensor value, if val_field_is_set.
920 ManualConstructor<Tensor> val;
921
922 Tensor* ref = nullptr; // A tensor reference.
923 mutex* ref_mu = nullptr; // mutex for *ref if ref is not nullptr.
924
925 // Whether the value exists, either in <val> or <ref>.
926 bool has_value = false;
927
928 bool val_field_is_set = false;
929
930 // The attributes of the allocator that creates the tensor.
931 AllocatorAttributes alloc_attr;
932
933 // Every entry carries an optional DeviceContext containing
934 // Device-specific information about how the Tensor was produced.
935 DeviceContext* device_context = nullptr;
936 };
937
938 // Contains a value for [node->id()] for the device context assigned by the
939 // device at the beginning of a step.
940 DeviceContextMap device_context_map_;
941
942 struct TaggedNode;
943 typedef gtl::InlinedVector<TaggedNode, 8> TaggedNodeSeq;
944 typedef gtl::InlinedVector<Entry, 4> EntryVector;
945
946 struct IterationState {
IterationStatetensorflow::__anon675242740111::ExecutorState::IterationState947 explicit IterationState(const PendingCounts* pending_counts,
948 int total_input_tensors)
949 : input_tensors(new Entry[total_input_tensors]),
950 outstanding_ops(0),
951 outstanding_frame_count(0),
952 counts_(*pending_counts) { // Initialize with copy of *pending_counts
953 }
954
955 // The state of an iteration.
956
957 // One copy per iteration. For iteration k, i-th node's j-th input is in
958 // input_tensors[k][impl_->nodes[i].input_start + j]. An entry is either
959 // a tensor pointer (pass-by-reference) or a tensor (pass-by-value).
960 //
961 // NOTE: No need to protect input_tensors[i] by any locks because it
962 // is resized once. Each element of tensors_ is written once by the
963 // source node of an edge and is cleared by the destination of the same
964 // edge. The latter node is never run concurrently with the former node.
965 Entry* input_tensors;
966
967 // The number of outstanding ops for each iteration.
968 size_t outstanding_ops;
969
970 // The number of outstanding frames for each iteration.
971 int outstanding_frame_count;
pendingtensorflow::__anon675242740111::ExecutorState::IterationState972 int pending(PendingCounts::Handle h) { return counts_.pending(h); }
decrement_pendingtensorflow::__anon675242740111::ExecutorState::IterationState973 int decrement_pending(PendingCounts::Handle h, int v) {
974 return counts_.decrement_pending(h, v);
975 }
976 // Mark a merge node as live
977 // REQUIRES: Node corresponding to "h" is a merge node
mark_livetensorflow::__anon675242740111::ExecutorState::IterationState978 void mark_live(PendingCounts::Handle h) { counts_.mark_live(h); }
979 // Mark a node to show that processing has started.
mark_startedtensorflow::__anon675242740111::ExecutorState::IterationState980 void mark_started(PendingCounts::Handle h) { counts_.mark_started(h); }
981 // Mark a node to show that processing has completed.
mark_completedtensorflow::__anon675242740111::ExecutorState::IterationState982 void mark_completed(PendingCounts::Handle h) { counts_.mark_completed(h); }
node_statetensorflow::__anon675242740111::ExecutorState::IterationState983 PendingCounts::NodeState node_state(PendingCounts::Handle h) {
984 return counts_.node_state(h);
985 }
986
dead_counttensorflow::__anon675242740111::ExecutorState::IterationState987 int dead_count(PendingCounts::Handle h) { return counts_.dead_count(h); }
increment_dead_counttensorflow::__anon675242740111::ExecutorState::IterationState988 void increment_dead_count(PendingCounts::Handle h) {
989 counts_.increment_dead_count(h);
990 }
adjust_for_activationtensorflow::__anon675242740111::ExecutorState::IterationState991 void adjust_for_activation(PendingCounts::Handle h, bool increment_dead,
992 int* pending_result, int* dead_result) {
993 counts_.adjust_for_activation(h, increment_dead, pending_result,
994 dead_result);
995 }
996
~IterationStatetensorflow::__anon675242740111::ExecutorState::IterationState997 ~IterationState() { delete[] input_tensors; }
998
999 private:
1000 PendingCounts counts_;
1001 };
1002
1003 struct FrameState {
FrameStatetensorflow::__anon675242740111::ExecutorState::FrameState1004 explicit FrameState(const ExecutorImpl* impl, int parallel_iters)
1005 : executor(impl),
1006 max_parallel_iterations(parallel_iters),
1007 num_outstanding_iterations(1) {}
1008
1009 // A new frame is created for each loop. Execution starts at iteration 0.
1010 // When a value at iteration 0 passes through a NextIteration node,
1011 // iteration 1 is created and starts running. Note that iteration 0 may
1012 // still be running so multiple iterations may run in parallel. The
1013 // frame maintains the state of iterations in several data structures
1014 // such as pending_count and input_tensors. When iteration 0 completes,
1015 // we garbage collect the state of iteration 0.
1016 //
1017 // A frame instance is considered "done" and can be garbage collected
1018 // if all its inputs have entered and all its iterations are "done".
1019 //
1020 // A frame manages the live iterations of an iterative computation.
1021 // Iteration i is considered "done" when there are no outstanding ops,
1022 // frames at iteration i are done, all recvs for this iteration are
1023 // completed, and iteration i-1 is done. For iteration 0, we instead
1024 // wait for there to be no more pending inputs of the frame.
1025 //
1026 // Frames and iterations are garbage collected once they are done.
1027 // The state we need to keep around is highly dependent on the
1028 // parallelism enabled by the scheduler. We may want to have the
1029 // scheduler dynamically control the outstanding number of live
1030 // parallel frames and iterations. To reduce the state space, the
1031 // scheduler might want to schedule ops in inner frames first and
1032 // lower iterations first.
1033 //
1034 // This frame state is mostly initialized lazily on demand so we
1035 // don't introduce unnecessary overhead.
1036
1037 // The executor the frame is in.
1038 const ExecutorImpl* executor = nullptr;
1039
1040 // The name of this frame, which is the concatenation of its parent
1041 // frame name, the iteration of the parent frame when this frame was
1042 // created, and the value of the attr 'frame_name'.
1043 string frame_name;
1044
1045 // The unique id for this frame. Generated by fingerprinting
1046 // frame_name.
1047 uint64 frame_id;
1048
1049 // The iteration id of its parent frame when this frame is created.
1050 // -1 if there is no parent frame. The frame_name/parent_iter pair
1051 // uniquely identifies this FrameState.
1052 int64 parent_iter = -1;
1053
1054 // The FrameState of its parent frame.
1055 FrameState* parent_frame = nullptr;
1056
1057 // The maximum allowed number of parallel iterations.
1058 const int max_parallel_iterations;
1059
1060 // The number of inputs this frame is still waiting.
1061 int num_pending_inputs = 0;
1062
1063 // The highest iteration number we have reached so far in this frame.
1064 int64 iteration_count GUARDED_BY(mu) = 0;
1065
1066 // The number of outstanding iterations.
1067 int num_outstanding_iterations GUARDED_BY(mu) = 1;
1068
1069 // The active iteration states of this frame.
1070 gtl::InlinedVector<IterationState*, 12> iterations;
1071
1072 // The NextIteration nodes to enter a new iteration. If the number of
1073 // outstanding iterations reaches the limit, we will defer the start of
1074 // the next iteration until the number of outstanding iterations falls
1075 // below the limit.
1076 std::vector<std::pair<const Node*, Entry>> next_iter_roots GUARDED_BY(mu);
1077
1078 // The values of the loop invariants for this loop. They are added into
1079 // this list as they "enter" the frame. When a loop invariant enters,
1080 // we make it available to all active iterations. When the frame starts
1081 // a new iteration, we make all the current loop invariants available
1082 // to the new iteration.
1083 std::vector<std::pair<const Node*, Entry>> inv_values GUARDED_BY(mu);
1084
1085 // The list of dead exit nodes for the current highest iteration. We
1086 // will only "execute" the dead exits of the final iteration.
1087 std::vector<const Node*> dead_exits GUARDED_BY(mu);
1088
1089 // Static information specific to this frame.
1090 PendingCounts* pending_counts = nullptr;
1091 int total_input_tensors = 0;
1092 std::vector<const Node*>* nodes = nullptr;
1093
1094 // Lock ordering: ExecutorState.mu_ < mu;
1095 // during structured traversal: parent_frame->mu < mu.
1096 mutex mu;
1097
InitializeFrameInfotensorflow::__anon675242740111::ExecutorState::FrameState1098 void InitializeFrameInfo(const string& enter_name) {
1099 auto it_frame_info = executor->frame_info_.find(enter_name);
1100 DCHECK(it_frame_info != executor->frame_info_.end());
1101 ExecutorImpl::FrameInfo* finfo = it_frame_info->second;
1102 pending_counts = finfo->pending_counts;
1103 total_input_tensors = finfo->total_inputs;
1104 num_pending_inputs = finfo->input_count;
1105 nodes = finfo->nodes;
1106 }
1107
GetIterationtensorflow::__anon675242740111::ExecutorState::FrameState1108 inline IterationState* GetIteration(int64 iter)
1109 EXCLUSIVE_LOCKS_REQUIRED(mu) {
1110 size_t index = iter % iterations.size();
1111 return iterations[index];
1112 }
1113
SetIterationtensorflow::__anon675242740111::ExecutorState::FrameState1114 inline void SetIteration(int64 iter, IterationState* state)
1115 EXCLUSIVE_LOCKS_REQUIRED(mu) {
1116 size_t index = iter % iterations.size();
1117 DCHECK(state == nullptr || iterations[index] == nullptr);
1118 iterations[index] = state;
1119 }
1120
1121 // Decrement the outstanding op count and clean up the iterations in the
1122 // frame. Return true iff the execution of the frame is done.
DecrementOutstandingOpstensorflow::__anon675242740111::ExecutorState::FrameState1123 inline bool DecrementOutstandingOps(const GraphView* gview, int64 iter,
1124 TaggedNodeSeq* ready) {
1125 mutex_lock l(mu);
1126 return DecrementOutstandingOpsLocked(gview, iter, ready);
1127 }
1128
1129 // Decrement the outstanding op count and clean up the iterations in the
1130 // frame. Return true iff the execution of the frame is done.
DecrementOutstandingOpsLockedtensorflow::__anon675242740111::ExecutorState::FrameState1131 inline bool DecrementOutstandingOpsLocked(const GraphView* gview,
1132 int64 iter, TaggedNodeSeq* ready)
1133 EXCLUSIVE_LOCKS_REQUIRED(mu) {
1134 IterationState* istate = GetIteration(iter);
1135 istate->outstanding_ops--;
1136 if (istate->outstanding_ops != 0) {
1137 return false;
1138 } else {
1139 return CleanupIterations(gview, iter, ready);
1140 }
1141 }
1142
1143 // Returns true if the computation in the frame is completed.
IsFrameDonetensorflow::__anon675242740111::ExecutorState::FrameState1144 inline bool IsFrameDone() EXCLUSIVE_LOCKS_REQUIRED(mu) {
1145 return (num_pending_inputs == 0 && num_outstanding_iterations == 0);
1146 }
1147
1148 // Returns true if the iteration of the frame is completed.
1149 bool IsIterationDone(int64 iter) EXCLUSIVE_LOCKS_REQUIRED(mu);
1150
1151 // Increments the iteration id. If this is a new iteration, initialize it.
1152 void IncrementIteration(const GraphView* gview, TaggedNodeSeq* ready)
1153 EXCLUSIVE_LOCKS_REQUIRED(mu);
1154
1155 // Activate all the deferred NextIteration nodes in a new iteration.
1156 void ActivateNexts(const GraphView* gview, int64 iter, TaggedNodeSeq* ready)
1157 EXCLUSIVE_LOCKS_REQUIRED(mu);
1158
1159 // Activate all the current loop invariants in a new iteration.
1160 void ActivateLoopInvs(const GraphView* gview, int64 iter,
1161 TaggedNodeSeq* ready) EXCLUSIVE_LOCKS_REQUIRED(mu);
1162
1163 // Add a new loop invariant and make it available to all active
1164 // iterations.
1165 void AddLoopInv(const NodeItem* item, const Entry& value,
1166 TaggedNodeSeq* ready) EXCLUSIVE_LOCKS_REQUIRED(mu);
1167
1168 // Activate the successors of a node. Contents of *outputs are left in an
1169 // indeterminate state after returning from this method.
1170 void ActivateNodes(const NodeItem* item, const bool is_dead, int64 iter,
1171 EntryVector* outputs, TaggedNodeSeq* ready)
1172 EXCLUSIVE_LOCKS_REQUIRED(mu);
1173
1174 // Cleanup iterations of this frame starting from iteration iter.
1175 bool CleanupIterations(const GraphView* gview, int64 iter,
1176 TaggedNodeSeq* ready) EXCLUSIVE_LOCKS_REQUIRED(mu);
1177
~FrameStatetensorflow::__anon675242740111::ExecutorState::FrameState1178 ~FrameState() {
1179 for (size_t i = 0; i < iterations.size(); ++i) {
1180 delete iterations[i];
1181 iterations[i] = nullptr;
1182 }
1183 }
1184 };
1185
1186 // A tagged node: <frame*, iter, node*>.
1187 struct TaggedNode {
1188 const Node* node = nullptr;
1189 FrameState* input_frame = nullptr;
1190 int64 input_iter = -1;
1191 bool is_dead = false;
1192
TaggedNodetensorflow::__anon675242740111::ExecutorState::TaggedNode1193 TaggedNode(const Node* t_node, FrameState* in_frame, int64 in_iter,
1194 bool dead) {
1195 node = t_node;
1196 input_frame = in_frame;
1197 input_iter = in_iter;
1198 is_dead = dead;
1199 }
1200 };
1201
1202 // A drop-in replacement for std::deque<TaggedNode>. We typically don't
1203 // have that many nodes in the ready queue, so we just use a vector and
1204 // don't free up memory from the queue as we consume nodes.
1205 class TaggedNodeReadyQueue {
1206 public:
TaggedNodeReadyQueue()1207 TaggedNodeReadyQueue() : front_index_(0) {}
1208
push_back(TaggedNode node)1209 void push_back(TaggedNode node) { ready_.push_back(node); }
front() const1210 TaggedNode front() const {
1211 DCHECK_LT(front_index_, ready_.size());
1212 return ready_[front_index_];
1213 }
pop_front()1214 void pop_front() {
1215 DCHECK_LT(front_index_, ready_.size());
1216 front_index_++;
1217 if ((front_index_ == ready_.size()) || (front_index_ > 16384)) {
1218 if (front_index_ == ready_.size()) {
1219 ready_.clear();
1220 } else {
1221 // Lots of unused entries at beginning of vector: move everything
1222 // down to start of vector.
1223 ready_.erase(ready_.begin(), ready_.begin() + front_index_);
1224 }
1225 front_index_ = 0;
1226 }
1227 }
empty() const1228 bool empty() const { return ready_.empty(); }
begin() const1229 const TaggedNode* begin() const { return ready_.begin() + front_index_; }
end() const1230 const TaggedNode* end() const { return ready_.end(); }
1231
1232 private:
1233 gtl::InlinedVector<TaggedNode, 16> ready_;
1234 int front_index_;
1235 };
1236
1237 struct AsyncState;
1238
1239 const bool vlog_; // true if VLOG_IS_ON(1). Used to check vlog cheaply.
1240
1241 // true if LogMemory::IsEnabled(). Used to check memory enabled cheaply.
1242 const bool log_memory_;
1243
1244 int64 step_id_;
1245 // Not owned.
1246 Rendezvous* rendezvous_;
1247 CollectiveExecutor* collective_executor_ = nullptr;
1248 SessionState* session_state_;
1249 string session_handle_;
1250 TensorStore* tensor_store_;
1251 // Step-local container.
1252 ScopedStepContainer* step_container_;
1253 StepStatsCollectorInterface* const stats_collector_;
1254 const tracing::EventCollector* const event_collector_;
1255 Context context_;
1256
1257 // QUESTION: Make it a checkpoint::TensorSliceReaderCacheWrapper
1258 // instead of a pointer? (avoids having to delete).
1259 checkpoint::TensorSliceReaderCacheWrapper* slice_reader_cache_;
1260 CallFrameInterface* call_frame_;
1261 const ExecutorImpl* impl_;
1262 CancellationManager* cancellation_manager_;
1263 Executor::Args::Runner runner_;
1264 bool sync_on_finish_;
1265 const bool trace_using_annotations_;
1266
1267 // Owned.
1268
1269 // A flag that is set on error after the frame state has been
1270 // dumped for diagnostic purposes.
1271 bool dumped_on_error_ = false;
1272
1273 // The root frame in which the execution of this step is started.
1274 FrameState* root_frame_;
1275
1276 // Invoked when the execution finishes.
1277 Executor::DoneCallback done_cb_;
1278
1279 std::atomic_int_fast32_t num_outstanding_ops_;
1280
1281 // Available via OpKernelContext to every OpKernel invocation.
1282 mutex num_deferred_ops_mu_;
1283 condition_variable no_deferred_ops_cv_;
1284 int64 num_deferred_ops_ GUARDED_BY(num_deferred_ops_mu_) = 0;
1285
1286 mutex mu_;
1287 Status status_ GUARDED_BY(mu_);
1288
1289 // Mapping from frame name to outstanding frames. A new frame is created
1290 // at some iteration of an active frame. So the unique key for the new
1291 // child frame is composed of the name of the parent frame, the iteration
1292 // number at which the parent frame is creating the new frame, and the
1293 // name of the new frame from nodedef.
1294 gtl::FlatMap<string, FrameState*> outstanding_frames_ GUARDED_BY(mu_);
1295
1296 // The unique name of a frame.
MakeFrameName(FrameState * frame,int64 iter_id,const string & name)1297 inline string MakeFrameName(FrameState* frame, int64 iter_id,
1298 const string& name) {
1299 return strings::StrCat(frame->frame_name, ";", iter_id, ";", name);
1300 }
1301
1302 // Find an existing or create a new child frame in the frame 'frame' at
1303 // iteration 'iter'.
1304 void FindOrCreateChildFrame(FrameState* frame, int64 iter, const Node* node,
1305 FrameState** child);
1306
1307 // Delete a frame. Called when the frame is done.
1308 void DeleteFrame(FrameState* frame, TaggedNodeSeq* ready);
1309
1310 // Cleanup frames and iterations starting from frame/iter. Called when
1311 // a child frame is done.
1312 void CleanupFramesIterations(FrameState* frame, int64 iter,
1313 TaggedNodeSeq* ready);
1314
1315 // Process a ready node in current thread.
1316 void Process(TaggedNode node, int64 scheduled_nsec);
1317
1318 // Before invoking item->kernel, fills in its "inputs".
1319 Status PrepareInputs(const NodeItem& item, Entry* first_input,
1320 TensorValueVec* inputs,
1321 DeviceContextVec* input_device_contexts,
1322 AllocatorAttributeVec* input_alloc_attrs,
1323 bool* is_input_dead);
1324
1325 // After item->kernel computation is done, processes its outputs.
1326 Status ProcessOutputs(const NodeItem& item, OpKernelContext* ctx,
1327 EntryVector* outputs, NodeExecStatsInterface* stats);
1328
1329 // After processing the outputs, propagates the outputs to their dsts.
1330 // Contents of *outputs are left in an indeterminate state after
1331 // returning from this method.
1332 void PropagateOutputs(const TaggedNode& tagged_node, const NodeItem* item,
1333 EntryVector* outputs, TaggedNodeSeq* ready);
1334
1335 // "node" just finishes. Takes ownership of "stats". Returns true if
1336 // execution has completed.
1337 bool NodeDone(const Status& s, const Node* node, const TaggedNodeSeq& ready,
1338 NodeExecStatsInterface* stats,
1339 TaggedNodeReadyQueue* inline_ready);
1340
1341 // Schedule all the expensive nodes in 'ready', and put all the inexpensive
1342 // nodes in 'ready' into 'inline_ready'.
1343 void ScheduleReady(const TaggedNodeSeq& ready,
1344 TaggedNodeReadyQueue* inline_ready);
1345
1346 // For debugging/logging only.
1347 inline void MaybeMarkCompleted(FrameState* frame, int64 iter, int64 id);
1348
1349 // Provide debugging output about an outstanding node in the executor.
1350 void DumpPendingNodeState(const int node_id, const Entry* input_vector,
1351 bool show_nodes_with_no_ready_inputs);
1352 void DumpActiveNodeState(const int node_id, const Entry* input_vector);
1353
1354 // Provide debugging output about an outstanding iteration in the executor.
1355 void DumpIterationState(const FrameState* frame, IterationState* iteration);
1356
1357 // Provide debugging output of the state of the executor.
1358 void DumpState();
1359 const Tensor* GetTensorValueForDump(const Entry& input);
1360
1361 // Clean up when this executor is done.
1362 void Finish();
1363 // Schedule Finish() on a separate thread if it needs to wait for deferred
1364 // async ops to complete; otherwise run it on the current thread.
1365 void ScheduleFinish();
1366
1367 // A standalone routine for this expression so that we can express
1368 // that we don't want thread safety analysis on this reference (it's
1369 // safe to do without the lock because the iterations array never
1370 // resizes and this particular iteration's array element will not
1371 // be changed out from under us because the iteration is still alive).
GetInputTensors(FrameState * input_frame,int64 input_iter) const1372 Entry* GetInputTensors(FrameState* input_frame,
1373 int64 input_iter) const NO_THREAD_SAFETY_ANALYSIS {
1374 return input_frame->GetIteration(input_iter)->input_tensors;
1375 }
1376 };
1377
ExecutorState(const Executor::Args & args,ExecutorImpl * impl)1378 ExecutorState::ExecutorState(const Executor::Args& args, ExecutorImpl* impl)
1379 : vlog_(VLOG_IS_ON(1)),
1380 log_memory_(LogMemory::IsEnabled()),
1381 step_id_(args.step_id),
1382 rendezvous_(args.rendezvous),
1383 collective_executor_(args.collective_executor),
1384 session_state_(args.session_state),
1385 session_handle_(args.session_handle),
1386 tensor_store_(args.tensor_store),
1387 step_container_(args.step_container),
1388 stats_collector_(args.stats_collector),
1389 event_collector_(
1390 tracing::GetEventCollector(tracing::EventCategory::kCompute)),
1391 context_(ContextKind::kThread),
1392 slice_reader_cache_(new checkpoint::TensorSliceReaderCacheWrapper),
1393 call_frame_(args.call_frame),
1394 impl_(impl),
1395 cancellation_manager_(args.cancellation_manager),
1396 runner_(args.runner),
1397 sync_on_finish_(args.sync_on_finish),
1398 trace_using_annotations_(impl->params_.device->TraceUsingAnnotations()),
1399 num_outstanding_ops_(0) {
1400 // We start the entire execution in iteration 0 of the root frame
1401 // so let us create the root frame and the state for iteration 0.
1402 // We assume root_frame_->frame_name.empty().
1403 root_frame_ = new FrameState(impl_, 1);
1404 root_frame_->frame_id = 0; // must be 0
1405 root_frame_->InitializeFrameInfo(root_frame_->frame_name);
1406
1407 // Initialize iteration 0.
1408 root_frame_->iterations.resize(root_frame_->max_parallel_iterations);
1409 root_frame_->iterations[0] = new IterationState(
1410 root_frame_->pending_counts, root_frame_->total_input_tensors);
1411
1412 outstanding_frames_.insert({root_frame_->frame_name, root_frame_});
1413 }
1414
~ExecutorState()1415 ExecutorState::~ExecutorState() {
1416 for (auto name_frame : outstanding_frames_) {
1417 delete name_frame.second;
1418 }
1419 for (auto it : device_context_map_) {
1420 it->Unref();
1421 }
1422 delete slice_reader_cache_;
1423 }
1424
BuildControlFlowInfo(const Graph * g,ControlFlowInfo * cf_info)1425 Status ExecutorImpl::BuildControlFlowInfo(const Graph* g,
1426 ControlFlowInfo* cf_info) {
1427 const int num_nodes = g->num_node_ids();
1428 cf_info->frame_names.resize(num_nodes);
1429 std::vector<Node*> parent_nodes;
1430 parent_nodes.resize(num_nodes);
1431 std::vector<bool> visited;
1432 visited.resize(num_nodes);
1433
1434 string frame_name;
1435 std::deque<Node*> ready;
1436
1437 // Initialize with the root nodes.
1438 for (Node* n : g->nodes()) {
1439 if (n->in_edges().empty()) {
1440 visited[n->id()] = true;
1441 cf_info->unique_frame_names.insert(frame_name);
1442 ready.push_back(n);
1443 }
1444 }
1445
1446 while (!ready.empty()) {
1447 Node* curr_node = ready.front();
1448 int curr_id = curr_node->id();
1449 ready.pop_front();
1450
1451 Node* parent = nullptr;
1452 if (IsEnter(curr_node)) {
1453 // Enter a child frame.
1454 TF_RETURN_IF_ERROR(
1455 GetNodeAttr(curr_node->attrs(), "frame_name", &frame_name));
1456 parent = curr_node;
1457 } else if (IsExit(curr_node)) {
1458 // Exit to the parent frame.
1459 parent = parent_nodes[curr_id];
1460 frame_name = cf_info->frame_names[parent->id()];
1461 parent = parent_nodes[parent->id()];
1462 } else {
1463 parent = parent_nodes[curr_id];
1464 frame_name = cf_info->frame_names[curr_id];
1465 }
1466
1467 for (const Edge* out_edge : curr_node->out_edges()) {
1468 Node* out = out_edge->dst();
1469 const int out_id = out->id();
1470
1471 // Add to ready queue if not visited.
1472 bool is_visited = visited[out_id];
1473 if (!is_visited) {
1474 ready.push_back(out);
1475 visited[out_id] = true;
1476
1477 // Process the node 'out'.
1478 cf_info->frame_names[out_id] = frame_name;
1479 parent_nodes[out_id] = parent;
1480 cf_info->unique_frame_names.insert(frame_name);
1481 }
1482 }
1483 }
1484
1485 return Status::OK();
1486 }
1487
InitializePending(const Graph * graph,const ControlFlowInfo & cf_info)1488 void ExecutorImpl::InitializePending(const Graph* graph,
1489 const ControlFlowInfo& cf_info) {
1490 for (auto& it : cf_info.unique_frame_names) {
1491 FrameInfo* finfo = EnsureFrameInfo(it);
1492 PendingCounts* counts = new PendingCounts(finfo->pending_counts_layout);
1493 DCHECK_EQ(finfo->pending_counts, nullptr);
1494 finfo->pending_counts = counts;
1495 }
1496 for (const Node* n : graph->nodes()) {
1497 const int id = n->id();
1498 const string& name = cf_info.frame_names[id];
1499 size_t max_pending, max_dead;
1500 GetMaxPendingCounts(n, &max_pending, &max_dead);
1501 const NodeItem* item = gview_.node(id);
1502 PendingCounts* counts = EnsureFrameInfo(name)->pending_counts;
1503 counts->set_initial_count(item->pending_id, max_pending);
1504 }
1505 }
1506
RunAsync(Executor::DoneCallback done)1507 void ExecutorState::RunAsync(Executor::DoneCallback done) {
1508 const Graph* graph = impl_->graph_.get();
1509 TaggedNodeSeq ready;
1510
1511 // Ask the device to fill in the device context map.
1512 Device* device = impl_->params_.device;
1513 const Status fill_status =
1514 device->FillContextMap(graph, &device_context_map_);
1515 if (!fill_status.ok()) {
1516 delete this;
1517 done(fill_status);
1518 return;
1519 }
1520
1521 // Initialize the ready queue.
1522 for (const Node* n : impl_->root_nodes_) {
1523 DCHECK_EQ(n->in_edges().size(), 0);
1524 ready.push_back(TaggedNode{n, root_frame_, 0, false});
1525 }
1526 if (ready.empty()) {
1527 delete this;
1528 done(Status::OK());
1529 } else {
1530 num_outstanding_ops_ = ready.size();
1531 root_frame_->iterations[0]->outstanding_ops = ready.size();
1532 done_cb_ = std::move(done);
1533 // Schedule to run all the ready ops in thread pool.
1534 ScheduleReady(ready, nullptr);
1535 }
1536 }
1537
1538 // State kept alive for executing an asynchronous node in another
1539 // thread. NOTE: We need to make a copy of p.input,
1540 // p.input_device_contexts, and p.input_alloc_attrs for asynchronous
1541 // kernels because OpKernelContext methods like input_type(i) needs
1542 // the param points to valid input type vector. It's not an issue for
1543 // sync kernels because these vectors are kept on the stack.
1544 struct ExecutorState::AsyncState {
AsyncStatetensorflow::__anon675242740111::ExecutorState::AsyncState1545 AsyncState(const OpKernelContext::Params& p, const TaggedNode& _tagged_node,
1546 const NodeItem* _item, Entry* _first_input,
1547 NodeExecStatsInterface* _stats)
1548 : saved_inputs(*p.inputs),
1549 saved_input_device_contexts(*p.input_device_contexts),
1550 saved_input_alloc_attrs(*p.input_alloc_attrs),
1551 params(p),
1552 tagged_node(_tagged_node),
1553 item(_item),
1554 first_input(_first_input),
1555 // ParamsButClearingEigenGPUDevice does equivalent of
1556 // params.eigen_gpu_device = nullptr;
1557 ctx(ParamsButClearingEigenGPUDevice(¶ms), item->num_outputs),
1558 stats(_stats) {
1559 params.inputs = &saved_inputs;
1560 params.input_device_contexts = &saved_input_device_contexts;
1561 params.input_alloc_attrs = &saved_input_alloc_attrs;
1562 }
1563
1564 TensorValueVec saved_inputs;
1565 DeviceContextVec saved_input_device_contexts;
1566 AllocatorAttributeVec saved_input_alloc_attrs;
1567 OpKernelContext::Params params;
1568 TaggedNode tagged_node;
1569 const NodeItem* item;
1570 Entry* first_input;
1571 OpKernelContext ctx;
1572 NodeExecStatsInterface* stats;
1573
1574 private:
ParamsButClearingEigenGPUDevicetensorflow::__anon675242740111::ExecutorState::AsyncState1575 OpKernelContext::Params* ParamsButClearingEigenGPUDevice(
1576 OpKernelContext::Params* p) {
1577 // Ensure OpKernelContext constructor will make a new eigen GPU device if
1578 // necessary.
1579 p->eigen_gpu_device = nullptr; // Force allocation
1580 return p;
1581 }
1582 };
1583
1584 // Returns true if `item` might be traced by the given trace and event
1585 // collectors. Returns false only if `item` definitely will not be traced.
MightTrace(const NodeItem & item,const tracing::EventCollector * event_collector,bool using_annotations)1586 bool MightTrace(const NodeItem& item,
1587 const tracing::EventCollector* event_collector,
1588 bool using_annotations) {
1589 // Tracing will only be enabled if either `event_collector` is non null,
1590 // or `trace_collector` is non-null and enabled for this particular kernel.
1591 // Although `tracing::ScopedActivity`,
1592 // `tracing::ScopedAnnotation`, and `tracing::ScopedRegion` check subsets of
1593 // these properties internally in their constructors, the cost of passing the
1594 // necessary arguments to them can be significant, so we avoid constructing
1595 // them in the common case (when we know they will not be used).
1596 if (event_collector != nullptr) {
1597 return true;
1598 }
1599 auto* trace_collector = tracing::GetTraceCollector();
1600 if (trace_collector) {
1601 if (using_annotations) {
1602 return trace_collector->IsEnabledForAnnotations();
1603 } else {
1604 return trace_collector->IsEnabledForActivities(
1605 item.kernel->IsExpensive());
1606 }
1607 }
1608 return false;
1609 }
1610
Process(TaggedNode tagged_node,int64 scheduled_nsec)1611 void ExecutorState::Process(TaggedNode tagged_node, int64 scheduled_nsec) {
1612 WithContext wc(context_);
1613 const GraphView& gview = impl_->gview_;
1614 TaggedNodeSeq ready;
1615 TaggedNodeReadyQueue inline_ready;
1616
1617 // Parameters passed to OpKernel::Compute.
1618 TensorValueVec inputs;
1619 DeviceContextVec input_device_contexts;
1620 AllocatorAttributeVec input_alloc_attrs;
1621
1622 OpKernelContext::Params params;
1623 params.step_id = step_id_;
1624 Device* device = impl_->params_.device;
1625 params.device = device;
1626 params.log_memory = log_memory_;
1627 params.record_tensor_accesses = impl_->device_record_tensor_accesses_;
1628 params.rendezvous = rendezvous_;
1629 params.collective_executor = collective_executor_;
1630 params.session_state = session_state_;
1631 params.session_handle = session_handle_;
1632 params.tensor_store = tensor_store_;
1633 params.cancellation_manager = cancellation_manager_;
1634 params.call_frame = call_frame_;
1635 params.function_library = impl_->params_.function_library;
1636 params.resource_manager = device->resource_manager();
1637 params.step_container = step_container_;
1638 params.slice_reader_cache = slice_reader_cache_;
1639 params.inputs = &inputs;
1640 params.input_device_contexts = &input_device_contexts;
1641 params.input_alloc_attrs = &input_alloc_attrs;
1642 params.runner = &runner_;
1643 params.stats_collector = stats_collector_;
1644 params.inc_num_deferred_ops_function = [this]() {
1645 mutex_lock lock(num_deferred_ops_mu_);
1646 num_deferred_ops_++;
1647 };
1648 params.dec_num_deferred_ops_function = [this]() {
1649 mutex_lock lock(num_deferred_ops_mu_);
1650 num_deferred_ops_--;
1651 if (num_deferred_ops_ == 0) {
1652 no_deferred_ops_cv_.notify_all();
1653 }
1654 };
1655
1656 Status s;
1657 NodeExecStatsInterface* stats = nullptr;
1658
1659 EntryVector outputs;
1660 bool completed = false;
1661 inline_ready.push_back(tagged_node);
1662 while (!inline_ready.empty()) {
1663 tagged_node = inline_ready.front();
1664 inline_ready.pop_front();
1665 const Node* node = tagged_node.node;
1666 FrameState* input_frame = tagged_node.input_frame;
1667 const int64 input_iter = tagged_node.input_iter;
1668 const int id = node->id();
1669 const NodeItem& item = *gview.node(id);
1670
1671 // TODO(misard) Replace with a finer-grain enabling flag once we
1672 // add better optional debugging support.
1673 if (vlog_ && VLOG_IS_ON(1)) {
1674 mutex_lock l(input_frame->mu);
1675 input_frame->GetIteration(input_iter)->mark_started(item.pending_id);
1676 }
1677
1678 // Set the device_context for this node id, if it exists.
1679 if (id < device_context_map_.size()) {
1680 params.op_device_context = device_context_map_[id];
1681 }
1682
1683 params.track_allocations = false;
1684 stats = nullptr;
1685 if (stats_collector_ && !tagged_node.is_dead) {
1686 stats = stats_collector_->CreateNodeExecStats(node);
1687 // Track allocations if and only if we are collecting statistics, and
1688 // `stats` object is expecting allocations to be tracked.
1689 params.track_allocations = stats ? stats->TrackAllocations() : false;
1690 nodestats::SetScheduled(stats, scheduled_nsec);
1691 nodestats::SetAllStart(stats);
1692 }
1693
1694 if (vlog_) {
1695 VLOG(1) << "Process node: " << id << " step " << params.step_id << " "
1696 << SummarizeNode(*node) << (tagged_node.is_dead ? " is dead" : "")
1697 << " device: " << device->name();
1698 }
1699
1700 Entry* input_tensors = GetInputTensors(input_frame, input_iter);
1701 Entry* first_input = input_tensors + item.input_start;
1702 outputs.clear();
1703
1704 TensorReferenceVector accessed_tensors;
1705 DeviceContext* device_context = nullptr;
1706 // Only execute this node if it is not dead or it is a send/recv
1707 // transfer node. For transfer nodes, we need to propagate the "dead"
1708 // bit even when the node is dead.
1709 bool launched_asynchronously = false;
1710 if (tagged_node.is_dead && !IsTransferNode(node)) {
1711 outputs.resize(item.num_outputs);
1712 } else {
1713 // Prepares inputs.
1714 bool is_input_dead = false;
1715 s = PrepareInputs(item, first_input, &inputs, &input_device_contexts,
1716 &input_alloc_attrs, &is_input_dead);
1717 if (!s.ok()) {
1718 // Clear inputs.
1719 int num_inputs = item.num_inputs;
1720 for (int i = 0; i < num_inputs; ++i) {
1721 (first_input + i)->ClearVal();
1722 }
1723 MaybeMarkCompleted(input_frame, input_iter, id);
1724 // Continue to process the nodes in 'inline_ready'.
1725 completed = NodeDone(s, item.node, ready, stats, &inline_ready);
1726 continue;
1727 }
1728
1729 // Set up compute params.
1730 OpKernel* op_kernel = item.kernel;
1731 params.op_kernel = op_kernel;
1732 params.frame_iter = FrameAndIter(input_frame->frame_id, input_iter);
1733 params.is_input_dead = is_input_dead;
1734 params.output_attr_array = item.output_attrs();
1735 params.forward_from_array = item.forward_from();
1736
1737 if (item.kernel_is_async) {
1738 // Asynchronous computes.
1739 AsyncOpKernel* async = item.kernel->AsAsync();
1740 DCHECK(async != nullptr);
1741 launched_asynchronously = true;
1742 AsyncState* state =
1743 new AsyncState(params, tagged_node, &item, first_input, stats);
1744
1745 auto done = [this, state]() {
1746 Device* device = impl_->params_.device;
1747 NodeExecStatsInterface* stats = state->stats; // Shorthand
1748 Entry* first_input = state->first_input; // Shorthand
1749
1750 nodestats::SetOpEnd(stats);
1751 EntryVector outputs;
1752 Status s = ProcessOutputs(*state->item, &state->ctx, &outputs, stats);
1753 nodestats::SetMemory(stats, &state->ctx);
1754 if (vlog_) {
1755 VLOG(2) << "Async kernel done: " << state->item->node->id()
1756 << " step " << step_id_ << " "
1757 << SummarizeNode(*state->item->node)
1758 << (state->tagged_node.is_dead ? " is dead" : "")
1759 << " device: " << device->name();
1760 }
1761
1762 // Clears inputs.
1763 const int num_inputs = state->item->num_inputs;
1764 for (int i = 0; i < num_inputs; ++i) {
1765 (first_input + i)->ClearVal();
1766 }
1767 FrameState* input_frame = state->tagged_node.input_frame;
1768 const int64 input_iter = state->tagged_node.input_iter;
1769 const int id = state->tagged_node.node->id();
1770 MaybeMarkCompleted(input_frame, input_iter, id);
1771 TaggedNodeSeq ready;
1772 if (s.ok()) {
1773 PropagateOutputs(state->tagged_node, state->item, &outputs, &ready);
1774 }
1775 outputs.clear();
1776 if (s.ok() && impl_->device_record_tensor_accesses_) {
1777 // Get the list of all tensors accessed during the execution
1778 TensorReferenceVector accessed;
1779 state->ctx.retrieve_accessed_tensors(&accessed);
1780 nodestats::SetReferencedTensors(stats, accessed);
1781 // callee takes ownership of the vector
1782 device->ConsumeListOfAccessedTensors(state->ctx.op_device_context(),
1783 accessed);
1784 }
1785 const bool completed =
1786 NodeDone(s, state->item->node, ready, stats, nullptr);
1787 delete state;
1788 if (completed) ScheduleFinish();
1789 };
1790 nodestats::SetOpStart(stats);
1791 device->ComputeAsync(async, &state->ctx, done);
1792 } else {
1793 // Synchronous computes.
1794 OpKernelContext ctx(¶ms, item.num_outputs);
1795 nodestats::SetOpStart(stats);
1796
1797 if (TF_PREDICT_FALSE(
1798 MightTrace(item, event_collector_, trace_using_annotations_))) {
1799 const string& op_name = op_kernel->name();
1800 tracing::ScopedRegion region(tracing::EventCategory::kCompute,
1801 op_name);
1802 if (trace_using_annotations_) {
1803 // The OpKernel may create child activities (such as GPU kernel
1804 // launches), so use a `ScopedAnnotation` to relate these activities
1805 // in the trace.
1806 tracing::ScopedAnnotation activity(
1807 op_name, strings::StrCat(op_kernel->type_string(),
1808 "#id=", step_id_, "#"));
1809 device->Compute(op_kernel, &ctx);
1810 } else {
1811 // Use the cheaper `ScopedActivity` to trace just the OpKernel
1812 // execution.
1813 tracing::ScopedActivity activity(
1814 op_name,
1815 strings::StrCat(op_kernel->type_string(), "#id=", step_id_,
1816 "#"),
1817 item.kernel->IsExpensive());
1818 device->Compute(op_kernel, &ctx);
1819 }
1820 } else {
1821 // In the common case, avoid creating any tracing objects.
1822 if (op_kernel->IsExpensive()) {
1823 KernelTimer timer;
1824 device->Compute(op_kernel, &ctx);
1825 op_kernel->UpdateCostEstimate(timer.ElapsedCycles());
1826 } else {
1827 device->Compute(op_kernel, &ctx);
1828 }
1829 }
1830
1831 nodestats::SetOpEnd(stats);
1832 s = ProcessOutputs(item, &ctx, &outputs, stats);
1833 if (s.ok() && impl_->device_record_tensor_accesses_) {
1834 // Get the list of all tensors accessed during the execution
1835 ctx.retrieve_accessed_tensors(&accessed_tensors);
1836 device_context = ctx.op_device_context();
1837 }
1838 nodestats::SetMemory(stats, &ctx);
1839 }
1840 }
1841
1842 if (!launched_asynchronously) {
1843 if (vlog_) {
1844 VLOG(2) << "Synchronous kernel done: " << id << " step "
1845 << params.step_id << " " << SummarizeNode(*node)
1846 << (tagged_node.is_dead ? " is dead: " : "")
1847 << " device: " << device->name();
1848 }
1849
1850 // Clears inputs.
1851 const int num_inputs = item.num_inputs;
1852 for (int i = 0; i < num_inputs; ++i) {
1853 (first_input + i)->ClearVal();
1854 }
1855 MaybeMarkCompleted(input_frame, input_iter, id);
1856 // Propagates outputs.
1857 if (s.ok()) {
1858 PropagateOutputs(tagged_node, &item, &outputs, &ready);
1859 }
1860 outputs.clear();
1861 if (!accessed_tensors.empty()) {
1862 nodestats::SetReferencedTensors(stats, accessed_tensors);
1863 // device_context is set above in synchronous computes
1864 device->ConsumeListOfAccessedTensors(device_context, accessed_tensors);
1865 }
1866 if (stats) {
1867 scheduled_nsec = nodestats::NowInNsec();
1868 }
1869 // Postprocess.
1870 completed = NodeDone(s, item.node, ready, stats, &inline_ready);
1871 }
1872 } // while !inline_ready.empty()
1873
1874 // This thread of computation is done if completed = true.
1875 if (completed) ScheduleFinish();
1876 }
1877
PrepareInputs(const NodeItem & item,Entry * first_input,TensorValueVec * inputs,DeviceContextVec * input_device_contexts,AllocatorAttributeVec * input_alloc_attrs,bool * is_input_dead)1878 Status ExecutorState::PrepareInputs(const NodeItem& item, Entry* first_input,
1879 TensorValueVec* inputs,
1880 DeviceContextVec* input_device_contexts,
1881 AllocatorAttributeVec* input_alloc_attrs,
1882 bool* is_input_dead) {
1883 const Node* node = item.node;
1884
1885 inputs->clear();
1886 inputs->resize(item.num_inputs);
1887 input_device_contexts->clear();
1888 input_device_contexts->resize(item.num_inputs);
1889 input_alloc_attrs->clear();
1890 input_alloc_attrs->resize(item.num_inputs);
1891
1892 *is_input_dead = false;
1893
1894 bool is_merge = item.is_merge;
1895 for (int i = 0; i < item.num_inputs; ++i) {
1896 const bool expect_ref = IsRefType(item.input_type(i));
1897 Entry* entry = first_input + i;
1898 (*input_device_contexts)[i] = entry->device_context;
1899 (*input_alloc_attrs)[i] = entry->alloc_attr;
1900
1901 // i-th input.
1902 TensorValue* inp = &(*inputs)[i];
1903
1904 // Only merge and transfer nodes can have no-value inputs.
1905 if (!entry->has_value) {
1906 if (!is_merge) {
1907 DCHECK(IsTransferNode(node)) << node->name() << " - input " << i;
1908 DCHECK(!entry->val_field_is_set) << node->name() << " - input " << i;
1909 entry->has_value = true;
1910 entry->val_field_is_set = true;
1911 entry->val.Init(*kEmptyTensor);
1912 inp->tensor = entry->val.get();
1913 *is_input_dead = true;
1914 }
1915 continue;
1916 }
1917 if (entry->ref == nullptr) {
1918 if (expect_ref) {
1919 return AttachDef(
1920 errors::InvalidArgument(i, "-th input expects a ref type"),
1921 item.kernel->def());
1922 }
1923 inp->tensor = entry->val.get();
1924 } else {
1925 {
1926 tf_shared_lock ml(*entry->ref_mu);
1927 if (!entry->ref->IsInitialized() && !IsInitializationOp(item.node)) {
1928 return AttachDef(errors::FailedPrecondition(
1929 "Attempting to use uninitialized value ",
1930 item.kernel->requested_input(i)),
1931 item.kernel->def());
1932 }
1933 }
1934 if (expect_ref) {
1935 inp->mutex_if_ref = entry->ref_mu;
1936 inp->tensor = entry->ref;
1937 } else {
1938 // Automatically deref the tensor ref when the op expects a
1939 // tensor but is given a ref to a tensor. Need to deref it
1940 // under the mutex.
1941 {
1942 tf_shared_lock l(*(entry->ref_mu));
1943 DCHECK(!entry->val_field_is_set);
1944 entry->val.Init(*entry->ref);
1945 entry->val_field_is_set = true;
1946 }
1947 entry->ref = nullptr;
1948 entry->ref_mu = nullptr;
1949
1950 inp->tensor = entry->val.get();
1951 // The dtype of entry->ref could have been changed by another operation
1952 // that ran after the operation that "produced" it executed, so
1953 // re-validate that the type of the dereferenced tensor matches the
1954 // expected input type.
1955 if (item.input_type(i) != inp->tensor->dtype()) {
1956 return AttachDef(
1957 errors::InvalidArgument(
1958 i, "-th input expects type ",
1959 DataTypeString(item.input_type(i)),
1960 " but automatically dereferenced input tensor has type ",
1961 DataTypeString(inp->tensor->dtype())),
1962 item.kernel->def());
1963 }
1964 }
1965 }
1966 }
1967 return Status::OK();
1968 }
1969
ProcessOutputs(const NodeItem & item,OpKernelContext * ctx,EntryVector * outputs,NodeExecStatsInterface * stats)1970 Status ExecutorState::ProcessOutputs(const NodeItem& item, OpKernelContext* ctx,
1971 EntryVector* outputs,
1972 NodeExecStatsInterface* stats) {
1973 const Node* node = item.node;
1974 DCHECK_EQ(0, outputs->size());
1975 outputs->resize(item.num_outputs);
1976
1977 Status s = ctx->status();
1978 if (!s.ok()) {
1979 s = AttachDef(s, item.kernel->def());
1980 // TODO(misard) Replace with a finer-grain enabling flag once we
1981 // add better optional debugging support.
1982 if (vlog_ && VLOG_IS_ON(1)) {
1983 LOG(WARNING) << this << " Compute status: " << s;
1984 DumpState();
1985 }
1986 if (s.code() == error::RESOURCE_EXHAUSTED) {
1987 if (stats_collector_) {
1988 string err = stats_collector_->ReportAllocsOnResourceExhausted(
1989 s.error_message());
1990 s = Status(s.code(), strings::StrCat(s.error_message(), err));
1991 } else {
1992 s = Status(
1993 s.code(),
1994 strings::StrCat(
1995 s.error_message(),
1996 "\nHint: If you want to see a list of allocated tensors when "
1997 "OOM happens, add report_tensor_allocations_upon_oom "
1998 "to RunOptions for current allocation info.\n"));
1999 }
2000 }
2001 return s;
2002 }
2003
2004 // Get the device_context for this node id, if it exists.
2005 DeviceContext* device_context = nullptr;
2006 if (node->id() < device_context_map_.size()) {
2007 device_context = device_context_map_[node->id()];
2008 }
2009
2010 for (int i = 0; i < item.num_outputs; ++i) {
2011 const TensorValue val = ctx->release_output(i);
2012 if (val.tensor == nullptr) {
2013 // Unless it's a Switch or a Recv, the node must produce a
2014 // tensor value at i-th output.
2015 if (!IsSwitch(node) && !IsRecv(node)) {
2016 s.Update(errors::Internal("Missing ", i, "-th output from ",
2017 FormatNodeForError(*node)));
2018 }
2019 } else {
2020 Entry* out = &((*outputs)[i]);
2021
2022 // Set the device context of the output entry.
2023 out->device_context = device_context;
2024
2025 // Set the allocator attributes of the output entry.
2026 out->alloc_attr = ctx->output_alloc_attr(i);
2027
2028 // Sanity check of output tensor types.
2029 DataType dtype;
2030 if (val.is_ref()) {
2031 tf_shared_lock ml(*val.mutex_if_ref);
2032 dtype = MakeRefType(val->dtype());
2033 } else {
2034 dtype = val->dtype();
2035 }
2036 if (dtype == item.output_type(i)) {
2037 if (stats && val.tensor->IsInitialized()) {
2038 nodestats::SetOutput(stats, i, val.tensor);
2039 }
2040 if (val.is_ref()) {
2041 out->has_value = true;
2042 out->ref = val.tensor;
2043 out->ref_mu = val.mutex_if_ref;
2044 if (log_memory_) {
2045 Tensor to_log;
2046 {
2047 // Dereference the tensor under the lock.
2048 tf_shared_lock l(*out->ref_mu);
2049 to_log = *out->ref;
2050 }
2051 LogMemory::RecordTensorOutput(ctx->op_kernel().name(),
2052 ctx->step_id(), i, to_log);
2053 }
2054 } else {
2055 // NOTE that std::move is used here, so val.tensor goes to
2056 // uninitialized state (val.tensor->IsInitialized return false).
2057 DCHECK(!out->val_field_is_set);
2058 out->has_value = true;
2059 out->val_field_is_set = true;
2060 out->val.Init(std::move(*val.tensor));
2061 if (log_memory_) {
2062 LogMemory::RecordTensorOutput(ctx->op_kernel().name(),
2063 ctx->step_id(), i, *out->val);
2064 }
2065 }
2066 } else {
2067 s.Update(errors::Internal("Output ", i, " of type ",
2068 DataTypeString(dtype),
2069 " does not match declared output type ",
2070 DataTypeString(item.output_type(i)),
2071 " for node ", FormatNodeForError(*node)));
2072 }
2073 }
2074 if (!val.is_ref()) {
2075 // If OpKernelContext returns outputs via pass-by-value, we
2076 // don't need this trouble.
2077 delete val.tensor;
2078 }
2079 }
2080 return s;
2081 }
2082
PropagateOutputs(const TaggedNode & tagged_node,const NodeItem * item,EntryVector * outputs,TaggedNodeSeq * ready)2083 void ExecutorState::PropagateOutputs(const TaggedNode& tagged_node,
2084 const NodeItem* item, EntryVector* outputs,
2085 TaggedNodeSeq* ready) {
2086 auto activity_handle =
2087 [&]() -> std::unique_ptr<tracing::TraceCollector::Handle> {
2088 auto* trace_collector = tracing::GetTraceCollector();
2089 if (TF_PREDICT_FALSE(trace_collector != nullptr &&
2090 trace_collector->IsEnabledForActivities(
2091 false /* is_expensive */))) {
2092 const string& op_name = item->kernel->name();
2093 // Intentionally using ExecutorPropagateOutputs as the first key so that
2094 // users are aware that it's not the op invocation.
2095 return trace_collector->CreateActivityHandle(
2096 "ExecutorPropagateOutputs",
2097 strings::StrCat(op_name, "#id=", step_id_, "#"),
2098 false /* is_expensive */);
2099 } else {
2100 return nullptr;
2101 }
2102 }();
2103
2104 const Node* node = tagged_node.node;
2105 FrameState* input_frame = tagged_node.input_frame;
2106 const int64 input_iter = tagged_node.input_iter;
2107 const bool is_dead = tagged_node.is_dead;
2108
2109 // Propagates outputs along out edges, and puts newly ready nodes
2110 // into the ready queue.
2111 ready->clear();
2112 bool is_frame_done = false;
2113 FrameState* output_frame = input_frame;
2114 int64 output_iter = input_iter;
2115
2116 if (!item->is_enter_exit_or_next_iter) {
2117 // Fast path for nodes types that don't need special handling
2118 DCHECK_EQ(input_frame, output_frame);
2119 // Normal path for most nodes
2120 mutex_lock l(input_frame->mu);
2121 output_frame->ActivateNodes(item, is_dead, output_iter, outputs, ready);
2122 is_frame_done = input_frame->DecrementOutstandingOpsLocked(
2123 &impl_->gview_, input_iter, ready);
2124 } else if (item->is_enter) {
2125 FindOrCreateChildFrame(input_frame, input_iter, node, &output_frame);
2126 output_iter = 0;
2127 {
2128 const NodeItem* item = impl_->gview_.node(node->id());
2129 mutex_lock l(output_frame->mu);
2130 if (item->is_constant_enter) {
2131 // Propagate to all active iterations if this is a loop invariant.
2132 output_frame->AddLoopInv(item, (*outputs)[0], ready);
2133 } else {
2134 output_frame->ActivateNodes(item, is_dead, output_iter, outputs, ready);
2135 }
2136 output_frame->num_pending_inputs--;
2137 }
2138 is_frame_done =
2139 input_frame->DecrementOutstandingOps(&impl_->gview_, input_iter, ready);
2140 } else if (item->is_exit) {
2141 if (is_dead) {
2142 mutex_lock l(input_frame->mu);
2143 // Stop and remember this node if it is a dead exit.
2144 if (input_iter == input_frame->iteration_count) {
2145 input_frame->dead_exits.push_back(node);
2146 }
2147 is_frame_done = input_frame->DecrementOutstandingOpsLocked(
2148 &impl_->gview_, input_iter, ready);
2149 } else {
2150 output_frame = input_frame->parent_frame;
2151 output_iter = input_frame->parent_iter;
2152 {
2153 mutex_lock l(output_frame->mu);
2154 output_frame->ActivateNodes(item, is_dead, output_iter, outputs, ready);
2155 }
2156 is_frame_done = input_frame->DecrementOutstandingOps(&impl_->gview_,
2157 input_iter, ready);
2158 }
2159 } else {
2160 DCHECK(IsNextIteration(node));
2161 mutex_lock l(input_frame->mu);
2162 if (is_dead) {
2163 // Stop the deadness propagation.
2164 output_frame = nullptr;
2165 } else {
2166 if (input_iter == input_frame->iteration_count &&
2167 input_frame->num_outstanding_iterations ==
2168 input_frame->max_parallel_iterations) {
2169 // Reached the maximum for parallel iterations.
2170 input_frame->next_iter_roots.push_back({node, (*outputs)[0]});
2171 output_frame = nullptr;
2172 } else {
2173 // If this is a new iteration, start it.
2174 if (input_iter == input_frame->iteration_count) {
2175 input_frame->IncrementIteration(&impl_->gview_, ready);
2176 }
2177 output_iter = input_iter + 1;
2178 }
2179 }
2180 if (output_frame != nullptr) {
2181 // This is the case when node is not Enter, Exit, or NextIteration.
2182 DCHECK(input_frame == output_frame);
2183 output_frame->ActivateNodes(item, is_dead, output_iter, outputs, ready);
2184 }
2185 is_frame_done = input_frame->DecrementOutstandingOpsLocked(
2186 &impl_->gview_, input_iter, ready);
2187 }
2188
2189 // At this point, this node is completely done. We also know if the
2190 // completion of this node makes its frame completed.
2191 if (is_frame_done) {
2192 FrameState* parent_frame = input_frame->parent_frame;
2193 const int64 parent_iter = input_frame->parent_iter;
2194 DeleteFrame(input_frame, ready);
2195 if (parent_frame != nullptr) {
2196 // The completion of frame may cause completions in its parent frame.
2197 // So clean things up recursively.
2198 CleanupFramesIterations(parent_frame, parent_iter, ready);
2199 }
2200 }
2201 }
2202
NodeDone(const Status & s,const Node * node,const TaggedNodeSeq & ready,NodeExecStatsInterface * stats,TaggedNodeReadyQueue * inline_ready)2203 bool ExecutorState::NodeDone(const Status& s, const Node* node,
2204 const TaggedNodeSeq& ready,
2205 NodeExecStatsInterface* stats,
2206 TaggedNodeReadyQueue* inline_ready) {
2207 nodestats::SetAllEnd(stats);
2208 if (stats) {
2209 if (stats_collector_) {
2210 stats->Done(impl_->params_.device->name());
2211 } else {
2212 delete stats;
2213 }
2214 }
2215
2216 bool abort_run = false;
2217 if (!s.ok()) {
2218 // Some error happened. This thread of computation is done.
2219 mutex_lock l(mu_);
2220 if (status_.ok()) {
2221 abort_run = true;
2222 status_ = s;
2223 }
2224 }
2225 if (abort_run) {
2226 TRACEPRINTF("StartAbort: %s", s.ToString().c_str());
2227 if (rendezvous_) {
2228 rendezvous_->StartAbort(s);
2229 }
2230 if (collective_executor_) {
2231 collective_executor_->StartAbort(s);
2232 }
2233 if (cancellation_manager_) {
2234 cancellation_manager_->StartCancel();
2235 }
2236 }
2237
2238 bool completed = false;
2239 const size_t ready_size = ready.size();
2240 if (ready_size == 0 || !s.ok()) {
2241 completed = (num_outstanding_ops_.fetch_sub(1) == 1);
2242 } else if (ready_size > 1) {
2243 num_outstanding_ops_.fetch_add(ready_size - 1, std::memory_order_relaxed);
2244 }
2245
2246 // Schedule the ready nodes in 'ready'.
2247 if (s.ok()) {
2248 ScheduleReady(ready, inline_ready);
2249 }
2250 return completed;
2251 }
2252
ScheduleReady(const TaggedNodeSeq & ready,TaggedNodeReadyQueue * inline_ready)2253 void ExecutorState::ScheduleReady(const TaggedNodeSeq& ready,
2254 TaggedNodeReadyQueue* inline_ready) {
2255 if (ready.empty()) return;
2256
2257 int64 scheduled_nsec = 0;
2258 if (stats_collector_) {
2259 scheduled_nsec = nodestats::NowInNsec();
2260 }
2261
2262 if (inline_ready == nullptr) {
2263 // Schedule to run all the ready ops in thread pool.
2264 for (auto& tagged_node : ready) {
2265 runner_([=]() { Process(tagged_node, scheduled_nsec); });
2266 }
2267 return;
2268 }
2269
2270 const GraphView& gview = impl_->gview_;
2271 const TaggedNode* curr_expensive_node = nullptr;
2272 for (auto& tagged_node : ready) {
2273 const NodeItem& item = *gview.node(tagged_node.node->id());
2274 if (tagged_node.is_dead || !item.kernel->IsExpensive()) {
2275 // Inline this inexpensive node.
2276 inline_ready->push_back(tagged_node);
2277 } else {
2278 if (curr_expensive_node) {
2279 // Dispatch to another thread since there is plenty of work to
2280 // do for this thread.
2281 runner_(std::bind(&ExecutorState::Process, this, *curr_expensive_node,
2282 scheduled_nsec));
2283 }
2284 curr_expensive_node = &tagged_node;
2285 }
2286 }
2287 if (curr_expensive_node) {
2288 if (inline_ready->empty()) {
2289 // Tail recursion optimization
2290 inline_ready->push_back(*curr_expensive_node);
2291 } else {
2292 // There are inline nodes to run already. We dispatch this expensive
2293 // node to other thread.
2294 runner_(std::bind(&ExecutorState::Process, this, *curr_expensive_node,
2295 scheduled_nsec));
2296 }
2297 }
2298 }
2299
MaybeMarkCompleted(FrameState * frame,int64 iter,int64 node_id)2300 inline void ExecutorState::MaybeMarkCompleted(FrameState* frame, int64 iter,
2301 int64 node_id) {
2302 // TODO(misard) Replace with a finer-grain enabling flag once we
2303 // add better optional debugging support.
2304 if (vlog_ && VLOG_IS_ON(1)) {
2305 const NodeItem* item = impl_->gview_.node(node_id);
2306 mutex_lock l(frame->mu);
2307 frame->GetIteration(iter)->mark_completed(item->pending_id);
2308 }
2309 }
2310
GetTensorValueForDump(const Entry & input)2311 const Tensor* ExecutorState::GetTensorValueForDump(const Entry& input) {
2312 if (!input.has_value) {
2313 return kEmptyTensor;
2314 } else if (input.ref == nullptr) {
2315 return input.val.get();
2316 } else {
2317 return input.ref;
2318 }
2319 }
2320
DumpPendingNodeState(const int node_id,const Entry * input_vector,const bool show_nodes_with_no_ready_inputs)2321 void ExecutorState::DumpPendingNodeState(
2322 const int node_id, const Entry* input_vector,
2323 const bool show_nodes_with_no_ready_inputs) {
2324 const NodeItem& node_item = *impl_->gview_.node(node_id);
2325 const Node& node = *node_item.node;
2326 const int input_base = node_item.input_start;
2327 if (!show_nodes_with_no_ready_inputs) {
2328 bool has_ready_input = false;
2329 for (int i = 0; i < node.num_inputs(); ++i) {
2330 const Entry& input = input_vector[input_base + i];
2331 const Tensor* tensor = GetTensorValueForDump(input);
2332 if (tensor->IsInitialized()) {
2333 has_ready_input = true;
2334 break;
2335 }
2336 }
2337 if (!has_ready_input) {
2338 return;
2339 }
2340 }
2341 LOG(WARNING) << " Pending Node: " << node.DebugString();
2342 for (int i = 0; i < node.num_inputs(); ++i) {
2343 const Entry& input = input_vector[input_base + i];
2344 const Tensor* tensor = GetTensorValueForDump(input);
2345 if (tensor->IsInitialized()) {
2346 LOG(WARNING) << " Input " << i << ": "
2347 << strings::StrCat(
2348 "Tensor<type: ", DataTypeString(tensor->dtype()),
2349 " shape: ", tensor->shape().DebugString(), ">");
2350 } else {
2351 LOG(WARNING) << " Input " << i << ": not present";
2352 }
2353 }
2354 }
2355
DumpActiveNodeState(const int node_id,const Entry * input_vector)2356 void ExecutorState::DumpActiveNodeState(const int node_id,
2357 const Entry* input_vector) {
2358 const NodeItem& node_item = *impl_->gview_.node(node_id);
2359 const Node& node = *node_item.node;
2360 LOG(WARNING) << " Active Node: " << node.DebugString();
2361 const int input_base = node_item.input_start;
2362 for (int i = 0; i < node.num_inputs(); ++i) {
2363 const Entry& input = input_vector[input_base + i];
2364 const Tensor* tensor = GetTensorValueForDump(input);
2365 if (tensor->IsInitialized()) {
2366 LOG(WARNING) << " Input " << i << ": "
2367 << strings::StrCat(
2368 "Tensor<type: ", DataTypeString(tensor->dtype()),
2369 " shape: ", tensor->shape().DebugString(), ">");
2370 } else {
2371 LOG(WARNING) << " Input " << i << ": not present";
2372 }
2373 }
2374 }
2375
DumpIterationState(const FrameState * frame,IterationState * iteration)2376 void ExecutorState::DumpIterationState(const FrameState* frame,
2377 IterationState* iteration) {
2378 const std::vector<const Node*>* nodes = frame->nodes;
2379 // Dump any waiting nodes that are holding on to tensors.
2380 for (const Node* node : *nodes) {
2381 const int node_id = node->id();
2382 PendingCounts::Handle pending_id = impl_->gview_.node(node_id)->pending_id;
2383 if (iteration->node_state(pending_id) == PendingCounts::PENDING_NOTREADY ||
2384 iteration->node_state(pending_id) == PendingCounts::PENDING_READY) {
2385 DumpPendingNodeState(node_id, iteration->input_tensors, false);
2386 }
2387 }
2388 // Then the active nodes.
2389 for (const Node* node : *nodes) {
2390 const int node_id = node->id();
2391 PendingCounts::Handle pending_id = impl_->gview_.node(node_id)->pending_id;
2392 if (iteration->node_state(pending_id) == PendingCounts::STARTED) {
2393 DumpActiveNodeState(node_id, iteration->input_tensors);
2394 }
2395 }
2396 // Show all input tensors in use.
2397 const int total_input_tensors = frame->total_input_tensors;
2398 size_t total_bytes = 0;
2399 for (int i = 0; i < total_input_tensors; ++i) {
2400 const Entry& input = iteration->input_tensors[i];
2401 const Tensor* tensor = GetTensorValueForDump(input);
2402 if (tensor->IsInitialized()) {
2403 LOG(WARNING) << " Input " << i << ": "
2404 << strings::StrCat(
2405 "Tensor<type: ", DataTypeString(tensor->dtype()),
2406 " shape: ", tensor->shape().DebugString(),
2407 ", bytes: ", tensor->TotalBytes(), ">");
2408 total_bytes += tensor->TotalBytes();
2409 }
2410 }
2411 LOG(WARNING) << " Total bytes " << total_bytes;
2412 }
2413
DumpState()2414 void ExecutorState::DumpState() {
2415 mutex_lock l(mu_);
2416 if (!dumped_on_error_) {
2417 LOG(WARNING) << "Dumping state";
2418 for (auto& frame : outstanding_frames_) {
2419 LOG(WARNING) << frame.first;
2420 FrameState* frame_state = frame.second;
2421 mutex_lock frame_lock(frame_state->mu);
2422 for (IterationState* iteration : frame_state->iterations) {
2423 LOG(WARNING) << " Iteration:";
2424 DumpIterationState(frame_state, iteration);
2425 }
2426 }
2427 dumped_on_error_ = true;
2428 }
2429 }
2430
ScheduleFinish()2431 void ExecutorState::ScheduleFinish() {
2432 int num_deferred_ops;
2433 {
2434 mutex_lock lock(num_deferred_ops_mu_);
2435 num_deferred_ops = num_deferred_ops_;
2436 }
2437 if (num_deferred_ops > 0) {
2438 // Finish() may be blocked waiting for deferred async ops to complete. The
2439 // execution of deferred async ops may be waiting for non-enqueued ops of
2440 // other executors to complete. So running Finish() on the current thread
2441 // (inter-op threadpool thread) may lead to a deadlock due to threadpool
2442 // exhaustion. Instead, we run it on a separate thread to unblock the
2443 // threadpool thread.
2444 Env::Default()->SchedClosure([this]() { Finish(); });
2445 } else {
2446 Finish();
2447 }
2448 }
2449
Finish()2450 void ExecutorState::Finish() {
2451 mu_.lock();
2452 auto status = status_;
2453 auto done_cb = std::move(done_cb_);
2454 auto runner = std::move(runner_);
2455 mu_.unlock();
2456 CHECK(done_cb != nullptr);
2457 Device* device = impl_->params_.device;
2458
2459 // There are several potential race conditions below. To name a few:
2460 // 1. Even if the device's status is OK at the precise moment when
2461 // num_deferred_ops_ reaches 0, it could go bad before device->RefreshStatus()
2462 // is called below, caused by work enqueued onto the same device by other
2463 // concurrent ExecutorState objects.
2464 // 2. Some implementations of Device::RefreshStatus, such as
2465 // XlaDevice::RefreshStatus, may be inherently racy because it releases the
2466 // device mutex after a stream pointer is acquired and before the stream is
2467 // queried for status.
2468 // 3. It's the same for some implementations of Device::Sync, such as
2469 // XlaDevice::Sync.
2470 //
2471 // However, these race conditions are acceptable because a stream (and
2472 // therefore an XlaDevice) can only go from OK to not-OK, never the opposite,
2473 // which means we will at worst report errors when there isn't any, never the
2474 // opposite.
2475
2476 // If inc_num_deferred_ops_function has ever been called, ExecutorState must
2477 // wait for all corresponding dec_num_deferred_ops_function calls to happen
2478 // regardless of status. This ensures that dec_num_deferred_ops_function can
2479 // safely use ExecutorState's resources.
2480 {
2481 mutex_lock lock(num_deferred_ops_mu_);
2482 while (num_deferred_ops_ > 0) {
2483 no_deferred_ops_cv_.wait(lock);
2484 }
2485 }
2486
2487 // An early exit for devices don't allow sync on completion. Ops that run on
2488 // these devices should have used num_deferred_ops correctly to ensure the
2489 // device has finished all relevant work at this point.
2490 if (!device->AllowsSyncOnCompletion()) {
2491 status.Update(device->RefreshStatus());
2492 if (!status.ok()) {
2493 // In device async execution mode, it's possible for device execution to
2494 // lag behind ExecutorState scheduling so much that this is the first
2495 // place a device execution error surfaces.
2496 // If so, all ExecutorState::NodeDone calls have already happened with OK
2497 // status. This is the last defense where StartCancel must be called to
2498 // abort all computation still running on any device.
2499 // TODO(b/124523000): Always call Finish in a separate thread, so even if
2500 // StartCancel blocks the current thread's execution, we won't encounter
2501 // deadlocks caused by inter-op thread exhaustion.
2502 if (cancellation_manager_) {
2503 cancellation_manager_->StartCancel();
2504 }
2505 }
2506 delete this;
2507 runner([=]() { done_cb(status); });
2508 return;
2509 }
2510
2511 if (sync_on_finish_ && status.ok()) {
2512 // Block until the device has finished all queued operations. For
2513 // devices like GPUs that continue to execute Ops after their Compute
2514 // methods have completed, this ensures that control is not returned to
2515 // the user until the step (and its side-effects) has actually completed.
2516 device->Sync([=](Status new_status) mutable {
2517 status.Update(new_status);
2518 delete this;
2519 runner([=]() { done_cb(status); });
2520 });
2521 } else {
2522 delete this;
2523 runner([=]() { done_cb(status); });
2524 }
2525 }
2526
FindOrCreateChildFrame(FrameState * frame,int64 iter,const Node * node,FrameState ** child)2527 void ExecutorState::FindOrCreateChildFrame(FrameState* frame, int64 iter,
2528 const Node* node,
2529 FrameState** child) {
2530 // Get the child frame name.
2531 string enter_name;
2532 Status s = GetNodeAttr(node->attrs(), "frame_name", &enter_name);
2533 DCHECK(s.ok()) << s;
2534 const string child_name = MakeFrameName(frame, iter, enter_name);
2535
2536 {
2537 mutex_lock executor_lock(mu_);
2538 auto it = outstanding_frames_.find(child_name);
2539 if (it != outstanding_frames_.end()) {
2540 *child = it->second;
2541 return;
2542 }
2543 }
2544
2545 // Need to create a new frame instance.
2546 // Note that this new frame instance is created without any locks.
2547 if (vlog_) VLOG(2) << "Create frame: " << child_name;
2548
2549 int parallel_iters;
2550 s = GetNodeAttr(node->attrs(), "parallel_iterations", ¶llel_iters);
2551 DCHECK(s.ok()) << s;
2552 FrameState* temp = new FrameState(impl_, parallel_iters);
2553 temp->frame_name = child_name;
2554 temp->frame_id = Hash64(child_name);
2555 temp->parent_frame = frame;
2556 temp->parent_iter = iter;
2557 temp->InitializeFrameInfo(enter_name);
2558
2559 // 'iterations' is a fixed-length circular buffer.
2560 temp->iterations.resize(temp->max_parallel_iterations + 1);
2561 // Initialize iteration 0.
2562 temp->iterations[0] =
2563 new IterationState(temp->pending_counts, temp->total_input_tensors);
2564
2565 {
2566 mutex_lock executor_lock(mu_);
2567 auto it = outstanding_frames_.find(child_name);
2568 if (it != outstanding_frames_.end()) {
2569 *child = it->second;
2570 } else {
2571 mutex_lock frame_lock(frame->mu);
2572 frame->GetIteration(iter)->outstanding_frame_count++;
2573 outstanding_frames_[child_name] = temp;
2574 *child = temp;
2575 temp = nullptr;
2576 }
2577 }
2578 delete temp; // Not used so delete it.
2579 }
2580
DeleteFrame(FrameState * frame,TaggedNodeSeq * ready)2581 void ExecutorState::DeleteFrame(FrameState* frame, TaggedNodeSeq* ready) {
2582 // First, propagate dead_exits (if any) to the parent frame.
2583 FrameState* parent_frame = frame->parent_frame;
2584 const int64 parent_iter = frame->parent_iter;
2585 if (parent_frame != nullptr) {
2586 mutex_lock parent_frame_lock(parent_frame->mu);
2587 // Propagate all the dead exits to the parent frame.
2588 mutex_lock this_frame_lock(frame->mu);
2589 for (const Node* node : frame->dead_exits) {
2590 auto parent_iter_state = parent_frame->GetIteration(parent_iter);
2591 for (const Edge* e : node->out_edges()) {
2592 const Node* dst_node = e->dst();
2593
2594 const auto dst_pending_id =
2595 impl_->gview_.node(dst_node->id())->pending_id;
2596
2597 // TODO(yuanbyu): We don't need this if we require the subgraph
2598 // given to an executor not to contain a sink node.
2599 if (dst_node->IsSink()) continue;
2600
2601 bool dst_dead = true;
2602 bool dst_ready = false;
2603 // We know this is a dead input to dst.
2604 if (IsMerge(dst_node)) {
2605 if (e->IsControlEdge()) {
2606 parent_iter_state->decrement_pending(dst_pending_id, 2);
2607 int count = parent_iter_state->pending(dst_pending_id);
2608 int dead_cnt = parent_iter_state->dead_count(dst_pending_id);
2609 dst_dead = (dead_cnt == dst_node->num_inputs());
2610 dst_ready = (count == 0) || ((count == 1) && dst_dead);
2611 } else {
2612 parent_iter_state->increment_dead_count(dst_pending_id);
2613 const int dead_cnt = parent_iter_state->dead_count(dst_pending_id);
2614 dst_dead = (dead_cnt == dst_node->num_inputs());
2615 dst_ready =
2616 (parent_iter_state->pending(dst_pending_id) == 1) && dst_dead;
2617 }
2618 } else {
2619 parent_iter_state->increment_dead_count(dst_pending_id);
2620 dst_ready =
2621 (parent_iter_state->decrement_pending(dst_pending_id, 1) == 0);
2622 }
2623 if (dst_ready) {
2624 if (IsControlTrigger(dst_node)) dst_dead = false;
2625 ready->emplace_back(dst_node, parent_frame, parent_iter, dst_dead);
2626 parent_iter_state->outstanding_ops++;
2627 }
2628 }
2629 }
2630 }
2631
2632 // Delete the frame.
2633 const string& frame_name = frame->frame_name;
2634 if (vlog_) VLOG(2) << "Delete frame " << frame_name;
2635 {
2636 mutex_lock executor_lock(mu_);
2637 outstanding_frames_.erase(frame_name);
2638 }
2639 delete frame;
2640 }
2641
CleanupFramesIterations(FrameState * frame,int64 iter,TaggedNodeSeq * ready)2642 void ExecutorState::CleanupFramesIterations(FrameState* frame, int64 iter,
2643 TaggedNodeSeq* ready) {
2644 bool is_frame_done = false;
2645 {
2646 mutex_lock frame_lock(frame->mu);
2647 frame->GetIteration(iter)->outstanding_frame_count--;
2648 is_frame_done = frame->CleanupIterations(&impl_->gview_, iter, ready);
2649 }
2650 if (is_frame_done) {
2651 FrameState* parent_frame = frame->parent_frame;
2652 const int64 parent_iter = frame->parent_iter;
2653 DeleteFrame(frame, ready);
2654 if (parent_frame != nullptr) {
2655 // The completion of frame may cause completions in its parent frame.
2656 // So clean things up recursively.
2657 CleanupFramesIterations(parent_frame, parent_iter, ready);
2658 }
2659 }
2660 }
2661
ActivateNodes(const NodeItem * item,const bool is_dead,int64 iter,EntryVector * outputs,TaggedNodeSeq * ready)2662 void ExecutorState::FrameState::ActivateNodes(const NodeItem* item,
2663 const bool is_dead, int64 iter,
2664 EntryVector* outputs,
2665 TaggedNodeSeq* ready) {
2666 const GraphView& gview = executor->gview_;
2667 IterationState* iter_state = GetIteration(iter);
2668 const size_t num_output_edges = item->num_output_edges;
2669 const EdgeInfo* edges = item->output_edge_list();
2670 Entry* input_tensors = iter_state->input_tensors;
2671 for (size_t out_index = 0; out_index < num_output_edges; out_index++) {
2672 const EdgeInfo& e = edges[out_index];
2673 const int dst_id = e.dst_id;
2674 const NodeItem* dst_item = gview.node(dst_id);
2675 const PendingCounts::Handle dst_pending_id = dst_item->pending_id;
2676 const int src_slot = e.output_slot;
2677
2678 // TODO(yuanbyu): We don't need this if we require the subgraph
2679 // given to an executor not to contain a sink node.
2680 if (dst_item->is_sink) continue;
2681
2682 bool dst_dead = false;
2683 bool dst_ready = false;
2684 // True iff this input for dst is needed. We only set this input for
2685 // dst if this flag is true. This is needed to make the thread safety
2686 // analysis happy.
2687 const bool is_control_edge = (src_slot == Graph::kControlSlot);
2688 bool dst_need_input = !is_control_edge;
2689 if (dst_item->is_merge) {
2690 // A merge node is ready if all control inputs have arrived and either
2691 // a) a live data input becomes available or b) all data inputs are
2692 // dead. For Merge, pending's LSB is set iff a live data input has
2693 // arrived.
2694 if (is_control_edge) {
2695 iter_state->decrement_pending(dst_pending_id, 2);
2696 int count = iter_state->pending(dst_pending_id);
2697 int dead_cnt = iter_state->dead_count(dst_pending_id);
2698 dst_dead = (dead_cnt == dst_item->num_inputs);
2699 dst_ready = (count == 0) || ((count == 1) && dst_dead);
2700 } else {
2701 if ((*outputs)[src_slot].has_value) {
2702 // This is a live data input.
2703 int count = iter_state->pending(dst_pending_id);
2704 iter_state->mark_live(dst_pending_id);
2705 // Only the first live edge sets the input and (potentially)
2706 // triggers execution. The low bit of count is set if and
2707 // only if no live input has been used yet (mark_live clears
2708 // it). The node should be started if and only if this is
2709 // the first live input and there are no pending control
2710 // edges, i.e. count == 1.
2711 dst_ready = (count == 1);
2712 dst_need_input = ((count & 0x1) == 1);
2713 } else {
2714 // This is a dead data input. Note that dst_node is dead if node is
2715 // a dead enter. We need this to handle properly a while loop on
2716 // the untaken branch of a conditional.
2717 // TODO(yuanbyu): This is a bit hacky, but a good solution for
2718 // now.
2719 iter_state->increment_dead_count(dst_pending_id);
2720 const int dead_cnt = iter_state->dead_count(dst_pending_id);
2721 dst_dead = (dead_cnt == dst_item->num_inputs) || item->is_enter;
2722 dst_ready = (iter_state->pending(dst_pending_id) == 1) && dst_dead;
2723 dst_need_input = false;
2724 }
2725 }
2726 } else {
2727 const bool increment_dead =
2728 (is_dead || (!is_control_edge && !(*outputs)[src_slot].has_value));
2729 int pending, dead;
2730 iter_state->adjust_for_activation(dst_pending_id, increment_dead,
2731 &pending, &dead);
2732 dst_dead = (dead > 0);
2733 dst_ready = (pending == 0);
2734 }
2735
2736 if (dst_need_input) {
2737 const int dst_slot = e.input_slot;
2738 const int dst_loc = dst_item->input_start + dst_slot;
2739 if (e.is_last) {
2740 input_tensors[dst_loc] = std::move((*outputs)[src_slot]);
2741 } else {
2742 input_tensors[dst_loc] = (*outputs)[src_slot];
2743 }
2744 }
2745
2746 // Add dst to the ready queue if it's ready
2747 if (dst_ready) {
2748 if (dst_item->is_control_trigger) dst_dead = false;
2749 ready->emplace_back(dst_item->node, this, iter, dst_dead);
2750 iter_state->outstanding_ops++;
2751 }
2752 }
2753 }
2754
ActivateNexts(const GraphView * gview,int64 iter,TaggedNodeSeq * ready)2755 void ExecutorState::FrameState::ActivateNexts(const GraphView* gview,
2756 int64 iter,
2757 TaggedNodeSeq* ready) {
2758 // Propagate the deferred NextIteration nodes to the new iteration.
2759 for (auto& node_entry : next_iter_roots) {
2760 const Node* node = node_entry.first;
2761 const Entry& entry = node_entry.second;
2762 const bool is_dead = !entry.has_value;
2763 const NodeItem* item = gview->node(node->id());
2764 EntryVector outputs{entry};
2765 ActivateNodes(item, is_dead, iter, &outputs, ready);
2766 }
2767 next_iter_roots.clear();
2768 }
2769
ActivateLoopInvs(const GraphView * gview,int64 iter,TaggedNodeSeq * ready)2770 void ExecutorState::FrameState::ActivateLoopInvs(const GraphView* gview,
2771 int64 iter,
2772 TaggedNodeSeq* ready) {
2773 // Propagate loop invariants to the new iteration.
2774 for (auto& node_entry : inv_values) {
2775 const Node* node = node_entry.first;
2776 const Entry& entry = node_entry.second;
2777 const bool is_dead = !entry.has_value;
2778 const NodeItem* item = gview->node(node->id());
2779 EntryVector outputs{entry};
2780 ActivateNodes(item, is_dead, iter, &outputs, ready);
2781 }
2782 }
2783
AddLoopInv(const NodeItem * item,const Entry & entry,TaggedNodeSeq * ready)2784 void ExecutorState::FrameState::AddLoopInv(const NodeItem* item,
2785 const Entry& entry,
2786 TaggedNodeSeq* ready) {
2787 // Store this value.
2788 inv_values.push_back({item->node, entry});
2789
2790 // Make this value available to all iterations.
2791 const bool is_dead = !entry.has_value;
2792 for (int i = 0; i <= iteration_count; ++i) {
2793 EntryVector outputs{entry};
2794 ActivateNodes(item, is_dead, i, &outputs, ready);
2795 }
2796 }
2797
IsIterationDone(int64 iter)2798 bool ExecutorState::FrameState::IsIterationDone(int64 iter) {
2799 IterationState* iter_state = GetIteration(iter);
2800 if (iter_state->outstanding_ops == 0 &&
2801 iter_state->outstanding_frame_count == 0) {
2802 if (iter == 0) {
2803 // The enclosing frame has no pending input.
2804 return num_pending_inputs == 0;
2805 } else {
2806 // The preceding iteration is deleted (and therefore done).
2807 return (GetIteration(iter - 1) == nullptr);
2808 }
2809 }
2810 return false;
2811 }
2812
IncrementIteration(const GraphView * gview,TaggedNodeSeq * ready)2813 void ExecutorState::FrameState::IncrementIteration(const GraphView* gview,
2814 TaggedNodeSeq* ready) {
2815 iteration_count++;
2816 const int64 next_iter = iteration_count;
2817
2818 // Initialize the next iteration.
2819 IterationState* iter_state =
2820 new IterationState(pending_counts, total_input_tensors);
2821 SetIteration(next_iter, iter_state);
2822 num_outstanding_iterations++;
2823 dead_exits.clear();
2824
2825 // Activate the successors of the deferred roots in the new iteration.
2826 ActivateNexts(gview, next_iter, ready);
2827
2828 // Activate the loop invariants in the new iteration.
2829 ActivateLoopInvs(gview, next_iter, ready);
2830 }
2831
CleanupIterations(const GraphView * gview,int64 iter,TaggedNodeSeq * ready)2832 bool ExecutorState::FrameState::CleanupIterations(const GraphView* gview,
2833 int64 iter,
2834 TaggedNodeSeq* ready) {
2835 int64 curr_iter = iter;
2836 while (curr_iter <= iteration_count && IsIterationDone(curr_iter)) {
2837 // Delete the iteration curr_iter.
2838 delete GetIteration(curr_iter);
2839 SetIteration(curr_iter, nullptr);
2840 --num_outstanding_iterations;
2841 ++curr_iter;
2842
2843 // When one iteration is completed, we check for deferred iteration,
2844 // and start it if there is one.
2845 if (!next_iter_roots.empty()) {
2846 IncrementIteration(gview, ready);
2847 }
2848 }
2849 return IsFrameDone();
2850 }
2851
RunAsync(const Args & args,DoneCallback done)2852 void ExecutorImpl::RunAsync(const Args& args, DoneCallback done) {
2853 (new ExecutorState(args, this))->RunAsync(std::move(done));
2854 }
2855
2856 } // namespace
2857
NewLocalExecutor(const LocalExecutorParams & params,std::unique_ptr<const Graph> graph,Executor ** executor)2858 Status NewLocalExecutor(const LocalExecutorParams& params,
2859 std::unique_ptr<const Graph> graph,
2860 Executor** executor) {
2861 ExecutorImpl* impl = new ExecutorImpl(params, std::move(graph));
2862 const Status s = impl->Initialize();
2863 if (s.ok()) {
2864 *executor = impl;
2865 } else {
2866 delete impl;
2867 }
2868 return s;
2869 }
2870
CreateNonCachedKernel(Device * device,FunctionLibraryRuntime * flib,const NodeDef & ndef,int graph_def_version,OpKernel ** kernel)2871 Status CreateNonCachedKernel(Device* device, FunctionLibraryRuntime* flib,
2872 const NodeDef& ndef, int graph_def_version,
2873 OpKernel** kernel) {
2874 const auto device_type = DeviceType(device->attributes().device_type());
2875 auto allocator = device->GetAllocator(AllocatorAttributes());
2876 return CreateOpKernel(device_type, device, allocator, flib, ndef,
2877 graph_def_version, kernel);
2878 }
2879
DeleteNonCachedKernel(OpKernel * kernel)2880 void DeleteNonCachedKernel(OpKernel* kernel) { delete kernel; }
2881
2882 namespace {
2883
2884 class DefaultExecutorRegistrar {
2885 public:
DefaultExecutorRegistrar()2886 DefaultExecutorRegistrar() {
2887 Factory* factory = new Factory;
2888 ExecutorFactory::Register("", factory);
2889 ExecutorFactory::Register("DEFAULT", factory);
2890 }
2891
2892 private:
2893 class Factory : public ExecutorFactory {
NewExecutor(const LocalExecutorParams & params,std::unique_ptr<const Graph> graph,std::unique_ptr<Executor> * out_executor)2894 Status NewExecutor(const LocalExecutorParams& params,
2895 std::unique_ptr<const Graph> graph,
2896 std::unique_ptr<Executor>* out_executor) override {
2897 Executor* ret = nullptr;
2898 TF_RETURN_IF_ERROR(NewLocalExecutor(params, std::move(graph), &ret));
2899 out_executor->reset(ret);
2900 return Status::OK();
2901 }
2902 };
2903 };
2904 static DefaultExecutorRegistrar registrar;
2905
2906 } // namespace
2907
2908 } // namespace tensorflow
2909