1 /* Copyright 2018 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/compiler/jit/deadness_analysis.h"
17
18 #include "absl/algorithm/container.h"
19 #include "absl/container/flat_hash_map.h"
20 #include "absl/container/flat_hash_set.h"
21 #include "absl/strings/str_join.h"
22 #include "absl/strings/string_view.h"
23 #include "tensorflow/compiler/jit/deadness_analysis_internal.h"
24 #include "tensorflow/compiler/jit/xla_cluster_util.h"
25 #include "tensorflow/compiler/xla/status_macros.h"
26 #include "tensorflow/core/framework/tensor.pb.h"
27 #include "tensorflow/core/graph/algorithm.h"
28 #include "tensorflow/core/graph/control_flow.h"
29 #include "tensorflow/core/graph/graph_node_util.h"
30 #include "tensorflow/core/graph/tensor_id.h"
31 #include "tensorflow/core/lib/hash/hash.h"
32
33 // ALGORITHM OVERVIEW
34 // ==================
35 //
36 // We map every output produced by each node in the TensorFlow graph (including
37 // control dependence) into an instance of the Predicate class. Instances of
38 // Predicate denote logical formulas and mapping a node `n` to a predicate
39 // `pred` implies that `n` is live whenever `pred` is true. Then we can deduce
40 // mismatching liveness in the inputs to node by comparing the predicate those
41 // inputs are mapped to. The core logic of this pass resides in creating the
42 // map from TensorFlow nodes to predicates.
43 //
44 //
45 // MAPPING NODES TO PREDICATES, MODULO CYCLES
46 // ------------------------------------------
47 //
48 // If we ignore cycles for a moment, computing predicates is fairly
49 // straightforward. We traverse the graph in a topological order, mapping each
50 // node to a predicate based on the predicates its inputs are mapped to. For
51 // instance a Merge(X, Y) node will be mapped to OR(PredicateFor(X),
52 // PredicateFor(Y)). Roughtly speaking, we abstractly interpret each node on
53 // the "liveness" domain, where values in the domain represent if a tensor
54 // carries a dead signal or not.
55 //
56 //
57 // DEALING WITH CYCLES
58 // -------------------
59 //
60 // We map Merge nodes that are the target of a backedge to AndRecurrence
61 // instances. An AndRecurrence with start() = S and step() = X, printed as
62 // {S,&,X}, *roughly* represents the infinite list of predicates
63 // [S,S&X,S&X&X,S&X&X, ...]. So {S,&,X} can be used to represent the predicate
64 // for Merge in a graph like:
65 //
66 // Init
67 // |
68 // v
69 // Merge <-----------+
70 // | |
71 // v |
72 // Incr |
73 // | |
74 // v |
75 // Switch <- Cond |
76 // | |
77 // v (oidx: 1) |
78 // | |
79 // +---------------+
80 //
81 // Where S is the predicate for Init and X is the predicate that asserts that
82 // Cond is true. {S,&,X} states that Merge is live on the first "iteration" iff
83 // S is true, live on the second iteration iff "S&X" is true, live on the third
84 // iteration iff "S&X&X" is true etc. There is a subtlety here, S&X&X would
85 // normally be equivalent to S&X which isn't quite what we want to represent.
86 // Instead we want {S,&,X} to denote the infinite list [S, S&X,
87 // S&X&X',S&X&X'&X'', ...] where X, X', X'' are predicates that assert Cond is
88 // true on iteration 0, 1, 2 respectively. This is made more precise in the
89 // comment on the AndRecurrence class.
90 //
91 // The general algorithm that deals with cycles does two topological-order
92 // iterations over the graph. On the first iteration it assigns a symbolic
93 // predicate to merge nodes with backedges. On the second iteration it tries
94 // to pattern match the predicates for the backedges of these merges and infer
95 // an AndRecurrence for the merge. In other words, we do a data flow analysis
96 // where the data-flow lattice has two elements, Symbolic and NonSymbolic with
97 // Symbolic > NonSymbolic. The lattice has height = 2 so two iterations are
98 // sufficient to converge.
99 //
100 // We first do an optimistic analysis and, if it does not converge, we then fall
101 // back to a pessimistic analysis. The optimistic analysis assigns the same
102 // symbolic predicate to all the merge nodes whose preceding enter nodes have
103 // the same frame name on the first iteration. On the second iteration, if all
104 // the merge nodes are pattern matched into the same AndRecurrence predicate
105 // instance, the optimistic assignment of the same symbolic predicate is correct
106 // and the analyzed result is taken.
107 //
108 // Otherwise, if the optimistic analysis fails to converge, we then obtain the
109 // result by falling back to the pessimistic analysis which assigns a unique
110 // symbolic predicate to each merge on the first iteration. We still use
111 // symbolic predicates for merges for which we can't pattern match on the
112 // backedge predicate. This is conservatively correct.
113
114 namespace tensorflow {
115
116 namespace {
117
118 using se::port::StatusOr;
119
120 // Represents a logical predicate, used as described in the algorithm overview
121 // above.
122 class Predicate {
123 public:
124 enum class Kind { kAnd, kOr, kNot, kAndRecurrence, kSymbol, kIntSymbol };
125
126 virtual string ToString() const = 0;
127
128 // An ID assigned to the Predicate at construction time. Conceptually like a
129 // pointer, except that it is stable across runs.
id() const130 int64 id() const { return id_; }
131
132 virtual absl::Span<Predicate* const> GetOperands() const = 0;
133
134 virtual Kind kind() const = 0;
~Predicate()135 virtual ~Predicate() {}
136
137 // Invokes func on p and on all of its operands recursively. Does not invoke
138 // `func` on the same Predicate instance twice. Aborts the search if `func`
139 // returns true.
140 template <typename FunctionTy>
141 static void Visit(Predicate* p, const FunctionTy& func);
142
143 protected:
Predicate(int64 id)144 explicit Predicate(int64 id) : id_(id) {}
145
146 private:
147 const int64 id_;
148
149 TF_DISALLOW_COPY_AND_ASSIGN(Predicate);
150 };
151
152 // Represents a logical conjunction of a set of predicates.
153 class AndPredicate : public Predicate {
154 public:
AndPredicate(int64 id,std::vector<Predicate * > operands)155 explicit AndPredicate(int64 id, std::vector<Predicate*> operands)
156 : Predicate(id), operands_(std::move(operands)) {}
157
ToString() const158 string ToString() const override {
159 if (operands().empty()) {
160 return "#true";
161 }
162
163 std::vector<string> operands_str;
164 std::transform(operands().begin(), operands().end(),
165 std::back_inserter(operands_str),
166 [](Predicate* pred) { return pred->ToString(); });
167
168 return absl::StrCat("(", absl::StrJoin(operands_str, " & "), ")");
169 }
170
kind() const171 Kind kind() const override { return Kind::kAnd; }
172
GetOperands() const173 absl::Span<Predicate* const> GetOperands() const override {
174 return operands_;
175 }
operands() const176 absl::Span<Predicate* const> operands() const { return operands_; }
177
178 private:
179 std::vector<Predicate*> operands_;
180 };
181
182 // Represents a logical disjunction of a set of predicates.
183 class OrPredicate : public Predicate {
184 public:
OrPredicate(int64 id,std::vector<Predicate * > operands)185 explicit OrPredicate(int64 id, std::vector<Predicate*> operands)
186 : Predicate(id), operands_(std::move(operands)) {}
187
ToString() const188 string ToString() const override {
189 if (operands().empty()) {
190 return "#false";
191 }
192
193 std::vector<string> operands_str;
194 std::transform(operands().begin(), operands().end(),
195 std::back_inserter(operands_str),
196 [](Predicate* pred) { return pred->ToString(); });
197
198 return absl::StrCat("(", absl::StrJoin(operands_str, " | "), ")");
199 }
200
kind() const201 Kind kind() const override { return Kind::kOr; }
GetOperands() const202 absl::Span<Predicate* const> GetOperands() const override {
203 return operands_;
204 }
operands() const205 absl::Span<Predicate* const> operands() const { return operands_; }
206
207 private:
208 std::vector<Predicate*> operands_;
209 };
210
211 // Represents a logical negation of a set of predicates.
212 class NotPredicate : public Predicate {
213 public:
NotPredicate(int64 id,Predicate * operand)214 explicit NotPredicate(int64 id, Predicate* operand)
215 : Predicate(id), operands_({operand}) {}
216
ToString() const217 string ToString() const override {
218 return absl::StrCat("~", operand()->ToString());
219 }
220
kind() const221 Kind kind() const override { return Kind::kNot; }
operand() const222 Predicate* operand() const { return operands_[0]; }
GetOperands() const223 absl::Span<Predicate* const> GetOperands() const override {
224 return operands_;
225 }
226
227 private:
228 std::array<Predicate*, 1> operands_;
229 };
230
231 // Represents the liveness of an induction variable. For users inside the loop
232 // this represents the "current" liveness of the induction variable. For users
233 // outside the loop it represents the "last" liveness of the induction variable.
234 //
235 // More concretely, an and recurrence {S,&,X}<loop> represents the liveness of V
236 // in the following graph:
237 //
238 // V = Merge(S', V_NextIt)
239 // V = Op(V, X')
240 // V_NextIt = NextIteration(V)
241 //
242 // where Predicate(S') = S and Predicate(X') = X.
243 //
244 // `X` may contain symbolic predicates and the operations corresponding to these
245 // symbolic predicates are either in frame `loop` or outside it. The symbols
246 // that are inside frame `loop` are loop variant (i.e. can have different
247 // liveness in each loop iteration) and the symbols that are outside frame
248 // `loop` are loop invariant (i.e. have the same liveness across all
249 // iterations).
250 class AndRecurrencePredicate : public Predicate {
251 public:
AndRecurrencePredicate(int64 id,Predicate * start,Predicate * step,std::vector<string> frame)252 explicit AndRecurrencePredicate(int64 id, Predicate* start, Predicate* step,
253 std::vector<string> frame)
254 : Predicate(id), operands_({start, step}), frame_(std::move(frame)) {}
255
start() const256 Predicate* start() const { return operands_[0]; }
step() const257 Predicate* step() const { return operands_[1]; }
frame() const258 absl::Span<const string> frame() const { return frame_; }
259
ToString() const260 string ToString() const override {
261 return absl::StrCat("{", start()->ToString(), ",&,", step()->ToString(),
262 "}<", absl::StrJoin(frame(), ";"), ">");
263 }
264
kind() const265 Kind kind() const override { return Kind::kAndRecurrence; }
266
GetOperands() const267 absl::Span<Predicate* const> GetOperands() const override {
268 return operands_;
269 }
270
271 private:
272 std::array<Predicate*, 2> operands_;
273 std::vector<string> frame_;
274 };
275
276 // Represents an uninterpreted symbol in a logical predicate.
277 //
278 // Two predicates are equivalent iff they are equivalent for all assignments to
279 // the symbols contained in them, i.e. predicates are forall qualified over
280 // symbols.
281 class SymbolPredicate : public Predicate {
282 public:
SymbolPredicate(int64 id,TensorId tensor_id,bool must_be_true)283 explicit SymbolPredicate(int64 id, TensorId tensor_id, bool must_be_true)
284 : Predicate(id),
285 tensor_id_(std::move(tensor_id)),
286 must_be_true_(must_be_true) {}
287
ToString() const288 string ToString() const override {
289 return must_be_true() ? absl::StrCat("*", tensor_id_.ToString())
290 : tensor_id_.ToString();
291 }
292
kind() const293 Kind kind() const override { return Kind::kSymbol; }
GetOperands() const294 absl::Span<Predicate* const> GetOperands() const override { return {}; }
295
296 // If `must_be_true()` is true this SymbolPredicate represents the proposition
297 // "tensor_id() is live and evaluates to true".
298 //
299 // If `must_be_true()` is false then this SymbolPredicate represents the
300 // proposition "tensor_id() is live (and may evaluate to any value)"
tensor_id() const301 TensorId tensor_id() const { return tensor_id_; }
must_be_true() const302 bool must_be_true() const { return must_be_true_; }
303
304 private:
305 TensorId tensor_id_;
306 bool must_be_true_;
307 };
308
309 // Represents an uninterpreted symbol in a logical predicate.
310 //
311 // Two predicates are equivalent iff they are equivalent for all assignments to
312 // the symbols contained in them, i.e. predicates are forall qualified over
313 // symbols.
314 class IntSymbolPredicate : public Predicate {
315 public:
IntSymbolPredicate(int64 id,TensorId tensor_id,absl::optional<int> must_have_value)316 explicit IntSymbolPredicate(int64 id, TensorId tensor_id,
317 absl::optional<int> must_have_value)
318 : Predicate(id),
319 tensor_id_(std::move(tensor_id)),
320 must_have_value_(must_have_value) {}
321
ToString() const322 string ToString() const override {
323 return must_have_value().has_value()
324 ? absl::StrCat(tensor_id_.ToString(), "=", *must_have_value_)
325 : tensor_id_.ToString();
326 }
327
kind() const328 Kind kind() const override { return Kind::kIntSymbol; }
GetOperands() const329 absl::Span<Predicate* const> GetOperands() const override { return {}; }
330
331 // If `must_have_value().has_value()` is true, then this IntSymbolPredicate
332 // represents the proposition "tensor_id() is live and evaluates to
333 // `*must_have_value()`".
334 //
335 // If `must_have_value().has_value()` is false, then this IntSymbolPredicate
336 // represents the proposition "tensor_id() is live (and may evaluate to any
337 // value)".
tensor_id() const338 TensorId tensor_id() const { return tensor_id_; }
must_have_value() const339 const absl::optional<int>& must_have_value() const {
340 return must_have_value_;
341 }
342
343 private:
344 TensorId tensor_id_;
345 absl::optional<int> must_have_value_;
346 };
347
348 template <typename FunctionTy>
Visit(Predicate * p,const FunctionTy & func)349 /*static*/ void Predicate::Visit(Predicate* p, const FunctionTy& func) {
350 absl::flat_hash_set<Predicate*> visited;
351 std::vector<Predicate*> stack;
352
353 stack.push_back(p);
354 visited.insert(p);
355
356 while (!stack.empty()) {
357 Predicate* current = stack.back();
358 stack.pop_back();
359 bool done = func(current);
360 if (done) {
361 return;
362 }
363 for (Predicate* op : current->GetOperands()) {
364 if (visited.insert(op).second) {
365 stack.push_back(op);
366 }
367 }
368 }
369 }
370
371 // Creates and owns Predicate instances. Simplifies predicates as it creates
372 // them.
373 class PredicateFactory {
374 public:
MakeAndPredicate(absl::Span<Predicate * const> operands)375 Predicate* MakeAndPredicate(absl::Span<Predicate* const> operands) {
376 return MakeAndOrImpl(operands, /*is_and=*/true);
377 }
378
MakeOrPredicate(absl::Span<Predicate * const> operands)379 Predicate* MakeOrPredicate(absl::Span<Predicate* const> operands) {
380 return MakeAndOrImpl(operands, /*is_and=*/false);
381 }
382
MakeNotPredicate(Predicate * pred)383 Predicate* MakeNotPredicate(Predicate* pred) {
384 auto it = make_not_predicate_cache_.find(pred);
385 if (it != make_not_predicate_cache_.end()) {
386 return it->second;
387 }
388
389 Predicate* result = MakeNotPredicateImpl(pred);
390
391 bool insert_successful =
392 make_not_predicate_cache_.insert({pred, result}).second;
393 (void)insert_successful;
394 DCHECK(insert_successful);
395
396 return result;
397 }
398
MakeAndRecurrencePredicate(Predicate * start,Predicate * step,std::vector<string> frame)399 Predicate* MakeAndRecurrencePredicate(Predicate* start, Predicate* step,
400 std::vector<string> frame) {
401 SignatureForAndRec signature(start, step, std::move(frame));
402 auto it = interned_and_rec_instances_.find(signature);
403 if (it != interned_and_rec_instances_.end()) {
404 return it->second.get();
405 }
406
407 std::unique_ptr<Predicate> new_pred = Make<AndRecurrencePredicate>(
408 std::get<0>(signature), std::get<1>(signature), std::get<2>(signature));
409 Predicate* new_pred_ptr = new_pred.get();
410 bool inserted =
411 interned_and_rec_instances_.emplace(signature, std::move(new_pred))
412 .second;
413 (void)inserted;
414 DCHECK(inserted);
415 return new_pred_ptr;
416 }
417
MakeSymbolPredicate(Node * node,int output_idx,bool must_be_true,Predicate ** predicate)418 Status MakeSymbolPredicate(Node* node, int output_idx, bool must_be_true,
419 Predicate** predicate) {
420 TensorId tensor_id(node->name(), output_idx);
421
422 bool is_boolean_tensor =
423 BaseType(node->output_type(tensor_id.index())) == DT_BOOL;
424 TF_RET_CHECK(!must_be_true || is_boolean_tensor);
425
426 if (node->type_string() == "Const" && must_be_true) {
427 const TensorProto* proto = nullptr;
428 TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "value", &proto));
429
430 Tensor tensor(proto->dtype());
431 TF_RET_CHECK(tensor.FromProto(*proto));
432
433 *predicate = tensor.scalar<bool>()() ? MakeTrue() : MakeFalse();
434 return Status::OK();
435 }
436
437 SignatureForSymbol signature = {tensor_id, must_be_true};
438 auto it = interned_symbol_instances_.find(signature);
439 if (it == interned_symbol_instances_.end()) {
440 std::unique_ptr<Predicate> new_pred =
441 Make<SymbolPredicate>(tensor_id, must_be_true);
442 Predicate* new_pred_ptr = new_pred.get();
443 interned_symbol_instances_.emplace(std::move(signature),
444 std::move(new_pred));
445 *predicate = new_pred_ptr;
446 } else {
447 *predicate = it->second.get();
448 }
449
450 return Status::OK();
451 }
452
MakeSymbolPredicate(Node * node,int output_idx,absl::optional<int> must_have_value,Predicate ** predicate)453 Status MakeSymbolPredicate(Node* node, int output_idx,
454 absl::optional<int> must_have_value,
455 Predicate** predicate) {
456 TensorId tensor_id(node->name(), output_idx);
457
458 TF_RET_CHECK(BaseType(node->output_type(tensor_id.index())) == DT_INT32);
459
460 if (must_have_value.has_value() && node->type_string() == "Const") {
461 const TensorProto* proto = nullptr;
462 TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "value", &proto));
463
464 Tensor tensor(proto->dtype());
465 TF_RET_CHECK(tensor.FromProto(*proto));
466
467 *predicate = tensor.scalar<int32>()() == *must_have_value ? MakeTrue()
468 : MakeFalse();
469 return Status::OK();
470 }
471 SignatureForIntSymbol signature = {tensor_id, must_have_value};
472 auto it = interned_int_symbol_instances_.find(signature);
473 if (it == interned_int_symbol_instances_.end()) {
474 std::unique_ptr<Predicate> new_pred =
475 Make<IntSymbolPredicate>(tensor_id, must_have_value);
476 Predicate* new_pred_ptr = new_pred.get();
477 interned_int_symbol_instances_.emplace(std::move(signature),
478 std::move(new_pred));
479 *predicate = new_pred_ptr;
480 } else {
481 *predicate = it->second.get();
482 }
483
484 return Status::OK();
485 }
486
MakeTrue()487 Predicate* MakeTrue() { return MakeAndPredicate({}); }
MakeFalse()488 Predicate* MakeFalse() { return MakeOrPredicate({}); }
489
~PredicateFactory()490 ~PredicateFactory() {
491 DCHECK_EQ(stack_depth_, 0) << "Unnested IncrementStackDepth?";
492 }
493
494 private:
MakeNotPredicateImpl(Predicate * pred)495 Predicate* MakeNotPredicateImpl(Predicate* pred) {
496 IncrementStackDepth stack_frame(this);
497 if (!stack_frame.HasOverflowed()) {
498 if (Predicate* simplified = SimplifyUsingDeMorgan(pred)) {
499 return simplified;
500 }
501
502 // ~~A => A
503 if (auto* not_pred = dynamic_cast<NotPredicate*>(pred)) {
504 return not_pred->operand();
505 }
506 }
507
508 SignatureForNot signature = pred;
509 auto it = interned_not_instances_.find(signature);
510 if (it == interned_not_instances_.end()) {
511 std::unique_ptr<Predicate> new_pred = Make<NotPredicate>(pred);
512 Predicate* new_pred_ptr = new_pred.get();
513 interned_not_instances_.emplace(signature, std::move(new_pred));
514 return new_pred_ptr;
515 } else {
516 return it->second.get();
517 }
518 }
519
SimplifyUsingDeMorgan(Predicate * pred)520 Predicate* SimplifyUsingDeMorgan(Predicate* pred) {
521 // ~(A & B & C & ...) => ~A | ~B | ~C | ~...
522 // ~(A | B | C | ...) -> ~A & ~B & ~C & ~...
523 Predicate::Kind kind = pred->kind();
524
525 if (kind == Predicate::Kind::kAnd || kind == Predicate::Kind::kOr) {
526 std::vector<Predicate*> new_operands;
527 absl::c_transform(pred->GetOperands(), std::back_inserter(new_operands),
528 [&](Predicate* p) { return MakeNotPredicate(p); });
529 return kind == Predicate::Kind::kOr ? MakeAndPredicate(new_operands)
530 : MakeOrPredicate(new_operands);
531 }
532
533 return nullptr;
534 }
535
536 template <typename PredicateT, typename... Args>
Make(Args &&...args)537 std::unique_ptr<Predicate> Make(Args&&... args) {
538 // If we ever expose the Predicate class outside this .cc file then we may
539 // want to make this hard to misuse (by accidentally passing in an arbitrary
540 // integer to the Predicate constructor for instance).
541 return std::unique_ptr<PredicateT>(
542 new PredicateT(id_counter_++, std::forward<Args>(args)...));
543 }
544
545 Predicate* MakeAndOrImpl(absl::Span<Predicate* const> operands, bool is_and);
546 Predicate* MakeInternedAndOr(std::vector<Predicate*> simplified_ops,
547 Predicate::Kind pred_kind);
548
549 // Predicate instances are interned, meaning that there is only a single
550 // instance of a Predicate object with a given content. This makes checking
551 // for structural equality super-cheap -- we can just compare pointers.
552 //
553 // We intern predicates by maintaining a map from the content of a Predicate
554 // to the only instance of said predicate we allow to exist in the
555 // interned_and_or_instances_, interned_not_instances_ and
556 // interned_symbol_instances_ fields. These maps also double up as storage
557 // for the owning pointers to predicate instances.
558
559 using SignatureForAndOr =
560 std::pair<Predicate::Kind, absl::Span<Predicate* const>>;
561 using SignatureForNot = Predicate*;
562 using SignatureForAndRec =
563 std::tuple<Predicate*, Predicate*, std::vector<string>>;
564 using SignatureForSymbol = std::pair<SafeTensorId, bool>;
565 using SignatureForIntSymbol = std::pair<SafeTensorId, absl::optional<int32>>;
566
567 struct HashSignatureForAndOr {
operator ()tensorflow::__anonaeebdbfa0111::PredicateFactory::HashSignatureForAndOr568 size_t operator()(const SignatureForAndOr& signature) const {
569 size_t hash = ::tensorflow::hash<Predicate::Kind>()(signature.first);
570 for (Predicate* p : signature.second) {
571 hash = Hash64Combine(hash, ::tensorflow::hash<Predicate*>()(p));
572 }
573 return hash;
574 }
575 };
576
577 struct HashSignatureForSymbol {
operator ()tensorflow::__anonaeebdbfa0111::PredicateFactory::HashSignatureForSymbol578 size_t operator()(const SignatureForSymbol& signature) const {
579 return Hash64Combine(SafeTensorId::Hasher()(signature.first),
580 ::tensorflow::hash<bool>()(signature.second));
581 }
582 };
583
584 struct HashSignatureForIntSymbol {
operator ()tensorflow::__anonaeebdbfa0111::PredicateFactory::HashSignatureForIntSymbol585 size_t operator()(const SignatureForIntSymbol& signature) const {
586 return Hash64Combine(
587 SafeTensorId::Hasher()(signature.first),
588 Hash64Combine(
589 ::tensorflow::hash<bool>()(signature.second.has_value()),
590 ::tensorflow::hash<int32>()(
591 signature.second.has_value() ? *signature.second : 0)));
592 }
593 };
594
595 // Used to limit recursion to avoid blowing up the stack and cap compile time.
596 class IncrementStackDepth {
597 public:
IncrementStackDepth(PredicateFactory * parent)598 explicit IncrementStackDepth(PredicateFactory* parent) : parent_(parent) {
599 parent_->stack_depth_++;
600 }
601
HasOverflowed() const602 bool HasOverflowed() const {
603 const int kMaxStackDepth = 8;
604 return parent_->stack_depth_ >= kMaxStackDepth;
605 }
606
~IncrementStackDepth()607 ~IncrementStackDepth() { parent_->stack_depth_--; }
608
609 private:
610 PredicateFactory* parent_;
611 };
612
613 // A cache for the MakeNotPredicate function.
614 //
615 // NB! This is *not* the same as `interned_not_instances_`.
616 // `interned_not_instances_` maps ensures pointer identity for `NotPredicate`
617 // instances, i.e., it ensures there at most one instance of Not(predicate)
618 // for any given predicate whereas `make_not_predicate_cache_` simply caches
619 // the result of the `MakeNotPredicate` function. The values in
620 // `interned_not_instances_` are always instance of `NotPredicate` whereas the
621 // values in `make_not_predicate_cache_` may not be (for instance it will map
622 // Not(Not(A)) to A).
623 absl::flat_hash_map<Predicate*, Predicate*> make_not_predicate_cache_;
624
625 absl::flat_hash_map<SignatureForAndOr, std::unique_ptr<Predicate>,
626 HashSignatureForAndOr>
627 interned_and_or_instances_;
628 absl::flat_hash_map<SignatureForNot, std::unique_ptr<Predicate>>
629 interned_not_instances_;
630 absl::flat_hash_map<SignatureForAndRec, std::unique_ptr<Predicate>>
631 interned_and_rec_instances_;
632 absl::flat_hash_map<SignatureForSymbol, std::unique_ptr<Predicate>,
633 HashSignatureForSymbol>
634 interned_symbol_instances_;
635 absl::flat_hash_map<SignatureForIntSymbol, std::unique_ptr<Predicate>,
636 HashSignatureForIntSymbol>
637 interned_int_symbol_instances_;
638 int64 id_counter_ = 0;
639 int stack_depth_ = 0;
640 };
641
MakeInternedAndOr(std::vector<Predicate * > simplified_ops,Predicate::Kind pred_kind)642 Predicate* PredicateFactory::MakeInternedAndOr(
643 std::vector<Predicate*> simplified_ops, Predicate::Kind pred_kind) {
644 std::stable_sort(
645 simplified_ops.begin(), simplified_ops.end(),
646 [](Predicate* a, Predicate* b) { return a->id() < b->id(); });
647
648 auto it = interned_and_or_instances_.find({pred_kind, simplified_ops});
649 if (it != interned_and_or_instances_.end()) {
650 return it->second.get();
651 }
652
653 simplified_ops.shrink_to_fit();
654 // NB! Because we'll use a non-owning reference to simplified_ops in the
655 // key for interned_and_or_instances_ we need to be careful to std::move()
656 // it all the way through.
657 absl::Span<Predicate* const> operands_slice = simplified_ops;
658 std::unique_ptr<Predicate> new_pred =
659 pred_kind == Predicate::Kind::kAnd
660 ? Make<AndPredicate>(std::move(simplified_ops))
661 : Make<OrPredicate>(std::move(simplified_ops));
662
663 Predicate* new_pred_ptr = new_pred.get();
664 interned_and_or_instances_.emplace(
665 SignatureForAndOr(pred_kind, operands_slice), std::move(new_pred));
666 return new_pred_ptr;
667 }
668
669 // Common code to create AndPredicate or OrPredicate instances.
MakeAndOrImpl(absl::Span<Predicate * const> operands,bool is_and)670 Predicate* PredicateFactory::MakeAndOrImpl(
671 absl::Span<Predicate* const> operands, bool is_and) {
672 Predicate::Kind pred_kind =
673 is_and ? Predicate::Kind::kAnd : Predicate::Kind::kOr;
674
675 IncrementStackDepth stack_frame(this);
676 if (stack_frame.HasOverflowed()) {
677 return MakeInternedAndOr(
678 std::vector<Predicate*>(operands.begin(), operands.end()), pred_kind);
679 }
680
681 Predicate::Kind other_pred_kind =
682 is_and ? Predicate::Kind::kOr : Predicate::Kind::kAnd;
683 absl::flat_hash_set<Predicate*> simplified_ops_set;
684 std::vector<Predicate*> simplified_ops;
685 for (Predicate* op : operands) {
686 // Simplify A&A => A and A|A => A.
687 if (!simplified_ops_set.insert(op).second) {
688 continue;
689 }
690
691 if (op->kind() == pred_kind) {
692 // "Inline" the operands of an inner And/Or into the parent And/Or.
693 for (Predicate* subop : op->GetOperands()) {
694 if (simplified_ops_set.insert(subop).second) {
695 simplified_ops.push_back(subop);
696 }
697 }
698 } else {
699 simplified_ops.push_back(op);
700 }
701 }
702
703 if (simplified_ops.size() == 1) {
704 return simplified_ops[0];
705 }
706
707 // Simplify "A&~A=>False" and "A|~A=>True".
708 absl::flat_hash_set<Predicate*> negated_ops;
709 for (Predicate* op : simplified_ops) {
710 if (negated_ops.count(op)) {
711 // Simple case:
712 //
713 // A & ~A & ... == False
714 // A | ~A | ... == True
715 return is_and ? MakeFalse() : MakeTrue();
716 }
717
718 Predicate* negated_op = MakeNotPredicate(op);
719 if (negated_op->kind() == pred_kind) {
720 // Slightly more complicated case:
721 //
722 // (~A | ~B | ~C) & A & B & C & ... ==
723 // ~(A & B & C) & (A & B & C) & ... == False
724 //
725 // (~A & ~B & ~C) | A | B | C | ... ==
726 // ~(A | B | C) | (A | B | C) | ... == True
727 if (absl::c_all_of(negated_op->GetOperands(), [&](Predicate* p) {
728 return simplified_ops_set.contains(p);
729 })) {
730 return is_and ? MakeFalse() : MakeTrue();
731 }
732 }
733 negated_ops.insert(negated_op);
734 }
735
736 // Simplify {S,&,X} & ~X & ... => S & ...
737 if (is_and) {
738 absl::flat_hash_set<Predicate*> to_remove;
739 std::vector<Predicate*> to_add;
740 for (Predicate* op : simplified_ops) {
741 if (op->kind() == Predicate::Kind::kAndRecurrence) {
742 auto* and_rec = static_cast<AndRecurrencePredicate*>(op);
743 if (negated_ops.contains(and_rec->step())) {
744 // Remove and_rec and ~X and insert S. Note that checking the
745 // existence of ~X through negated_ops is sufficient since it makes
746 // sure the predicate is in the input operands. It does not need to
747 // be in simplified_ops if it was already cancelled out.
748 to_remove.insert(and_rec);
749 to_remove.insert(MakeNotPredicate(and_rec->step()));
750 to_add.push_back(and_rec->start());
751 }
752 }
753 }
754 auto it = simplified_ops.begin();
755 while (it != simplified_ops.end()) {
756 if (to_remove.contains(*it)) {
757 it = simplified_ops.erase(it);
758 } else {
759 ++it;
760 }
761 }
762 simplified_ops.insert(simplified_ops.end(), to_add.begin(), to_add.end());
763 }
764
765 // If all ops contain the same subop, then factor it out thanks to the
766 // distributive property. Such as:
767 // - (A & B) | (A & C) | (A & D) => A & (B | C | D)
768 // - (A | B) & (A | C) & (A | D) => A | (B & C & D)
769 //
770 // First find any predicates contained in all subops.
771 std::vector<Predicate*> common_inner_operands;
772 absl::flat_hash_set<Predicate*> common_inner_operands_set;
773 for (Predicate* op : simplified_ops) {
774 if (op->kind() != other_pred_kind) {
775 common_inner_operands.clear();
776 break;
777 }
778
779 if (common_inner_operands.empty()) {
780 common_inner_operands.insert(common_inner_operands.end(),
781 op->GetOperands().begin(),
782 op->GetOperands().end());
783 } else {
784 common_inner_operands.clear();
785 absl::c_copy_if(op->GetOperands(),
786 std::back_inserter(common_inner_operands),
787 [&](Predicate* sub_op) {
788 return common_inner_operands_set.count(sub_op) == 1;
789 });
790 }
791 if (common_inner_operands.empty()) break;
792 common_inner_operands_set.clear();
793 common_inner_operands_set.insert(common_inner_operands.begin(),
794 common_inner_operands.end());
795 }
796
797 if (common_inner_operands.empty()) {
798 return MakeInternedAndOr(std::move(simplified_ops), pred_kind);
799 }
800
801 // For all predicates that can be factored out, remove them and recreate the
802 // subops.
803 std::vector<Predicate*> factored_ops;
804 for (Predicate* op : simplified_ops) {
805 std::vector<Predicate*> new_sub_op_ops;
806 absl::c_copy_if(op->GetOperands(), std::back_inserter(new_sub_op_ops),
807 [&](Predicate* sub_op) {
808 return std::find(common_inner_operands.begin(),
809 common_inner_operands.end(),
810 sub_op) == common_inner_operands.end();
811 });
812 factored_ops.push_back(MakeAndOrImpl(new_sub_op_ops, !is_and));
813 }
814
815 Predicate* new_inner_op = MakeAndOrImpl(factored_ops, is_and);
816 std::vector<Predicate*> outer_ops;
817 outer_ops.push_back(new_inner_op);
818 outer_ops.insert(outer_ops.end(), common_inner_operands.begin(),
819 common_inner_operands.end());
820 return MakeAndOrImpl(outer_ops, !is_and);
821 }
822
823 class DeadnessAnalysisImpl : public DeadnessAnalysis {
824 public:
DeadnessAnalysisImpl(const Graph * graph)825 explicit DeadnessAnalysisImpl(const Graph* graph)
826 : graph_(*graph), vlog_(VLOG_IS_ON(2)) {}
827
828 Status Populate(bool enable_optimistic);
829 Status PopulateFrame(absl::Span<Node* const> topo, bool use_optimistic_mode,
830 bool* success);
831 StatusOr<DeadnessAnalysis::DeadnessPredicate> GetPredicateFor(
832 Node* n, int oidx) const override;
833 void Print() const override;
834 absl::flat_hash_map<TensorId, string, TensorId::Hasher> PredicateMapAsString()
835 const;
836
837 private:
838 enum class EdgeKind { kDataAndControl, kDataOnly, kControlOnly };
839
840 Status GetInputPreds(Node* n, EdgeKind edge_kind,
841 std::vector<Predicate*>* result);
842
843 // Sets the predicate for output `output_idx` of `n` to `pred`. Sets the i'th
844 // bit of `should_revisit` if `pred` is different from the current predicate
845 // for the `output_idx` output of `n`.
SetPredicate(Node * n,int output_idx,Predicate * pred,std::vector<bool> * should_revisit)846 void SetPredicate(Node* n, int output_idx, Predicate* pred,
847 std::vector<bool>* should_revisit) {
848 auto insert_result =
849 predicate_map_.insert({TensorId(n->name(), output_idx), pred});
850 if (!insert_result.second && insert_result.first->second != pred) {
851 VLOG(4) << "For " << n->name() << ":" << output_idx << " from "
852 << insert_result.first->second->ToString() << " "
853 << insert_result.first->second << " to " << pred->ToString()
854 << " " << pred;
855 insert_result.first->second = pred;
856 if (should_revisit != nullptr) {
857 for (const Edge* e : n->out_edges()) {
858 (*should_revisit)[e->dst()->id()] = true;
859 }
860 }
861 }
862 }
863
SetPredicate(Node * n,absl::Span<const int> output_idxs,Predicate * pred,std::vector<bool> * should_revisit)864 void SetPredicate(Node* n, absl::Span<const int> output_idxs, Predicate* pred,
865 std::vector<bool>* should_revisit) {
866 for (int output_idx : output_idxs) {
867 SetPredicate(n, output_idx, pred, should_revisit);
868 }
869 }
870
871 Status HandleSwitch(Node* n, std::vector<bool>* should_revisit);
872 Status HandleMerge(Node* n, std::vector<bool>* should_revisit,
873 bool use_optimistic_mode);
874 Status HandleRecv(Node* n, std::vector<bool>* should_revisit);
875 Status HandleGeneric(Node* n, std::vector<bool>* should_revisit);
876 Status HandleNode(Node* n, std::vector<bool>* should_revisit,
877 bool use_optimistic_mode = false);
878
879 Status GetFrameBasedTopologicalOrder(std::vector<Node*>* order);
880
IsRootEnter(const Node * n) const881 bool IsRootEnter(const Node* n) const {
882 return IsEnter(n) && control_flow_info_[n->id()].parent_frame->IsSource();
883 }
884
IsRootExit(const Node * n) const885 bool IsRootExit(const Node* n) const {
886 return IsExit(n) && control_flow_info_[n->id()].parent_frame->IsSource();
887 }
888
889 const Graph& graph_;
890 absl::flat_hash_map<TensorId, Predicate*, TensorId::Hasher> predicate_map_;
891 PredicateFactory predicate_factory_;
892 std::vector<ControlFlowInfo> control_flow_info_;
893 bool vlog_;
894 absl::flat_hash_map<absl::string_view, Node*> frame_to_merge_node_;
895 };
896
InputEdgeToTensorId(const Edge * e)897 TensorId InputEdgeToTensorId(const Edge* e) {
898 return TensorId(e->src()->name(), e->src_output());
899 }
900
GetInputPreds(Node * n,DeadnessAnalysisImpl::EdgeKind edge_kind,std::vector<Predicate * > * result)901 Status DeadnessAnalysisImpl::GetInputPreds(
902 Node* n, DeadnessAnalysisImpl::EdgeKind edge_kind,
903 std::vector<Predicate*>* result) {
904 result->clear();
905 for (const Edge* in_edge : n->in_edges()) {
906 bool should_process =
907 edge_kind == EdgeKind::kDataAndControl ||
908 (in_edge->IsControlEdge() && edge_kind == EdgeKind::kControlOnly) ||
909 (!in_edge->IsControlEdge() && edge_kind == EdgeKind::kDataOnly);
910
911 if (should_process) {
912 auto it = predicate_map_.find(InputEdgeToTensorId(in_edge));
913 if (it == predicate_map_.end()) {
914 GraphCycles graph_cycles;
915 TF_RETURN_IF_ERROR(
916 CreateCycleDetectionGraph(&graph_, &graph_cycles).status());
917
918 // If we didn't return with an error above then the graph is probably
919 // fine and we have a bug in deadness analysis.
920 return errors::Internal("Could not find input ", in_edge->DebugString(),
921 " to ", n->name(),
922 " when visiting the graph in post-order. Most "
923 "likely indicates a bug in deadness analysis.");
924 }
925 result->push_back(it->second);
926 }
927 }
928 return Status::OK();
929 }
930
HandleSwitch(Node * n,std::vector<bool> * should_revisit)931 Status DeadnessAnalysisImpl::HandleSwitch(Node* n,
932 std::vector<bool>* should_revisit) {
933 std::vector<Predicate*> input_preds;
934 TF_RETURN_IF_ERROR(GetInputPreds(n, EdgeKind::kDataAndControl, &input_preds));
935 const Edge* pred_edge;
936 TF_RETURN_IF_ERROR(n->input_edge(1, &pred_edge));
937
938 if (n->type_string() != "_SwitchN") { // bool pred branch selector.
939 Predicate* true_switch;
940 TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate(
941 pred_edge->src(), pred_edge->src_output(),
942 /*must_be_true=*/true, &true_switch));
943
944 Predicate* false_switch = predicate_factory_.MakeNotPredicate(true_switch);
945
946 // Output 0 is alive iff all inputs are alive and the condition is false.
947 input_preds.push_back(false_switch);
948 SetPredicate(n, 0, predicate_factory_.MakeAndPredicate(input_preds),
949 should_revisit);
950 input_preds.pop_back();
951
952 // Output 1 is alive iff all inputs are alive and the condition is true.
953 input_preds.push_back(true_switch);
954 SetPredicate(n, 1, predicate_factory_.MakeAndPredicate(input_preds),
955 should_revisit);
956 input_preds.pop_back();
957 } else { // N-way switch case. Exactly one of N branches is alive.
958 Predicate* branch_pred;
959 for (int i = 0; i < n->num_outputs() - 1; i++) {
960 TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate(
961 pred_edge->src(), pred_edge->src_output(),
962 /*must_have_value=*/absl::optional<int32>(i), &branch_pred));
963 input_preds.push_back(branch_pred);
964 SetPredicate(n, i, predicate_factory_.MakeAndPredicate(input_preds),
965 should_revisit);
966 input_preds.pop_back();
967 input_preds.push_back(predicate_factory_.MakeNotPredicate(branch_pred));
968 }
969 // The default (last) branch does not need its own symbol, is simply the
970 // nor of all other branches.
971 SetPredicate(n, n->num_outputs() - 1,
972 predicate_factory_.MakeAndPredicate(input_preds),
973 should_revisit);
974 }
975
976 // Control is alive iff all inputs are alive.
977 SetPredicate(n, Graph::kControlSlot,
978 predicate_factory_.MakeAndPredicate(input_preds),
979 should_revisit);
980
981 return Status::OK();
982 }
983
984 namespace {
CreateMultipleNextIterationInputsError(Node * merge)985 Status CreateMultipleNextIterationInputsError(Node* merge) {
986 std::vector<string> backedges;
987 for (const Edge* backedge : merge->in_edges()) {
988 if (backedge->src()->IsNextIteration()) {
989 backedges.push_back(absl::StrCat(" ", SummarizeNode(*backedge->src())));
990 }
991 }
992 return errors::InvalidArgument(
993 "Multiple NextIteration inputs to merge node ",
994 FormatNodeForError(*merge), ": \n", absl::StrJoin(backedges, "\n"),
995 "\nMerge nodes can have at most one incoming NextIteration edge.");
996 }
997
FindUniqueBackedge(Node * merge,const Edge ** result)998 Status FindUniqueBackedge(Node* merge, const Edge** result) {
999 *result = nullptr;
1000 CHECK(merge->IsMerge());
1001 for (const Edge* e : merge->in_edges()) {
1002 if (e->src()->IsNextIteration()) {
1003 if (*result != nullptr) {
1004 return CreateMultipleNextIterationInputsError(merge);
1005 }
1006 *result = e;
1007 }
1008 }
1009 return Status::OK();
1010 }
1011
1012 // If `backedge_predicate` is equal to `symbolic_predicate` & Step where Step
1013 // does not contain `symbolic_predicate` as an inner (not top-level) operand
1014 // then returns `Step`. Otherwise returns nullptr.
DeduceStepPredicate(PredicateFactory * predicate_factory,Predicate * symbolic_predicate,Predicate * backedge_predicate)1015 Predicate* DeduceStepPredicate(PredicateFactory* predicate_factory,
1016 Predicate* symbolic_predicate,
1017 Predicate* backedge_predicate) {
1018 CHECK(dynamic_cast<SymbolPredicate*>(symbolic_predicate));
1019 if (backedge_predicate->kind() != Predicate::Kind::kAnd) {
1020 return nullptr;
1021 }
1022
1023 std::vector<Predicate*> and_ops;
1024 absl::Span<Predicate* const> recurrent_pred_ops =
1025 backedge_predicate->GetOperands();
1026
1027 bool found_sym = false;
1028 for (Predicate* and_op : recurrent_pred_ops) {
1029 // We want the `symbol_predicate` to be the one of the operands of
1030 // `backedge_predicate`,
1031 if (and_op == symbolic_predicate) {
1032 found_sym = true;
1033 continue;
1034 }
1035
1036 // but we don't want it to be present anywhere else in the formula. E.g. we
1037 // don't want the recurrent predicate to be
1038 // symbol_predicate&(X|symbol_predicate).
1039 bool found_sym_as_inner_operand = false;
1040 auto has_self_as_inner_operand = [&](Predicate* p) {
1041 if (p == symbolic_predicate) {
1042 found_sym_as_inner_operand = true;
1043 return true; // Stop searching, we're done.
1044 }
1045
1046 // Continue searching.
1047 return false;
1048 };
1049
1050 Predicate::Visit(and_op, has_self_as_inner_operand);
1051 if (found_sym_as_inner_operand) {
1052 return nullptr;
1053 }
1054 and_ops.push_back(and_op);
1055 }
1056
1057 return found_sym ? predicate_factory->MakeAndPredicate(and_ops) : nullptr;
1058 }
1059
GetFullFrame(const Node * n,absl::Span<const ControlFlowInfo> cfi_infos,std::vector<string> * frame)1060 Status GetFullFrame(const Node* n, absl::Span<const ControlFlowInfo> cfi_infos,
1061 std::vector<string>* frame) {
1062 int depth = 0;
1063 for (const ControlFlowInfo* cfi_iter = &cfi_infos[n->id()]; !n->IsSource();
1064 n = cfi_iter->parent_frame, cfi_iter = &cfi_infos[n->id()]) {
1065 frame->push_back(cfi_iter->frame_name);
1066
1067 if (depth++ > 5000) {
1068 return errors::Internal(
1069 "Frame of depth > 5000: Probably malformed graph or a bug in "
1070 "BuildControlFlowInfo");
1071 }
1072 }
1073
1074 return Status::OK();
1075 }
1076
1077 // If the node is inside some frames, get the name of the outermost non-empty
1078 // frame. Otherwise, get an empty frame name.
GetRootFrame(const Node * n,absl::Span<const ControlFlowInfo> cfi_infos,absl::string_view * frame)1079 Status GetRootFrame(const Node* n, absl::Span<const ControlFlowInfo> cfi_infos,
1080 absl::string_view* frame) {
1081 int depth = 0;
1082 const ControlFlowInfo* cfi_iter = &cfi_infos[n->id()];
1083 while (!cfi_iter->parent_frame->IsSource()) {
1084 n = cfi_iter->parent_frame;
1085 cfi_iter = &cfi_infos[n->id()];
1086
1087 if (depth++ > 5000) {
1088 return errors::Internal(
1089 "Frame of depth > 5000: Probably malformed graph or a bug in "
1090 "BuildControlFlowInfo");
1091 }
1092 }
1093
1094 *frame = cfi_iter->frame_name;
1095 return Status::OK();
1096 }
1097 } // namespace
1098
HandleMerge(Node * n,std::vector<bool> * should_revisit,bool use_optimistic_mode)1099 Status DeadnessAnalysisImpl::HandleMerge(Node* n,
1100 std::vector<bool>* should_revisit,
1101 bool use_optimistic_mode) {
1102 // Merge ignores deadness of its control inputs. A merge that isn't the
1103 // target of a backedge has is alive iff any of its data inputs are. The
1104 // liveness of a merge that is the target of a backedge can sometimes be
1105 // represented using a AndRecurrencePredicate. If neither apply, we represent
1106 // the liveness of the merge symbolically.
1107
1108 bool has_unvisited_backedge = false;
1109 for (const Edge* e : n->in_edges()) {
1110 if (!e->IsControlEdge() && e->src()->IsNextIteration()) {
1111 has_unvisited_backedge |= !predicate_map_.count(InputEdgeToTensorId(e));
1112 }
1113 }
1114
1115 auto it = predicate_map_.find(TensorId(n->name(), 0));
1116 if (it == predicate_map_.end()) {
1117 if (has_unvisited_backedge) {
1118 // We're visiting this merge for the first time and it has an unvisited
1119 // backedge.
1120 Predicate* input_data_pred;
1121 if (use_optimistic_mode) {
1122 // In the optimistic mode, we use the first-seen Merge node per
1123 // frame as the representative Merge node. It is just convenient and
1124 // does not affect the result after pattern-matching into the
1125 // AndRecurrence form.
1126 absl::string_view frame_name = control_flow_info_[n->id()].frame_name;
1127 auto insert_result = frame_to_merge_node_.insert({frame_name, n});
1128 Node* representative = insert_result.first->second;
1129 TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate(
1130 representative, /*output_idx=*/0, /*must_be_true=*/false,
1131 &input_data_pred));
1132 } else {
1133 TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate(
1134 n, /*output_idx=*/0, /*must_be_true=*/false, &input_data_pred));
1135 }
1136
1137 SetPredicate(n, {0, 1, Graph::kControlSlot}, input_data_pred,
1138 should_revisit);
1139 return Status::OK();
1140 }
1141
1142 std::vector<Predicate*> input_preds;
1143 TF_RETURN_IF_ERROR(GetInputPreds(n, EdgeKind::kDataOnly, &input_preds));
1144
1145 // We're visiting this merge for the first time and it is an acyclic merge.
1146 Predicate* input_data_pred =
1147 predicate_factory_.MakeOrPredicate(input_preds);
1148 SetPredicate(n, {0, 1, Graph::kControlSlot}, input_data_pred,
1149 should_revisit);
1150 return Status::OK();
1151 }
1152
1153 if (it->second->kind() == Predicate::Kind::kSymbol) {
1154 // Last time we visited this merge we only got a symbolic predicate because
1155 // of an unvisited backedge. Try to pattern match the predicate expression
1156 // for that backedge (which should be visited now) into an and recurrence
1157 // for the merge node.
1158 const Edge* unique_backedge;
1159 TF_RETURN_IF_ERROR(FindUniqueBackedge(n, &unique_backedge));
1160 if (unique_backedge) {
1161 if (Predicate* step = DeduceStepPredicate(
1162 &predicate_factory_, it->second,
1163 predicate_map_[InputEdgeToTensorId(unique_backedge)])) {
1164 // If the predicate for the backedge is "Sym&X" where "Sym" is the
1165 // predicate for the merge then the merge has predicate {S,&,X} where S
1166 // is the predicate for the merge ignoring the backedge.
1167 std::vector<Predicate*> non_recurrent_inputs;
1168 for (const Edge* e : n->in_edges()) {
1169 if (e != unique_backedge) {
1170 non_recurrent_inputs.push_back(
1171 predicate_map_[InputEdgeToTensorId(e)]);
1172 }
1173 }
1174
1175 Predicate* start =
1176 predicate_factory_.MakeOrPredicate(non_recurrent_inputs);
1177 std::vector<string> frame;
1178 TF_RETURN_IF_ERROR(GetFullFrame(n, control_flow_info_, &frame));
1179 Predicate* and_rec = predicate_factory_.MakeAndRecurrencePredicate(
1180 start, step, std::move(frame));
1181 SetPredicate(n, {0, 1, Graph::kControlSlot}, and_rec, should_revisit);
1182 return Status::OK();
1183 }
1184 }
1185 }
1186 return Status::OK();
1187 }
1188
HandleRecv(Node * n,std::vector<bool> * should_revisit)1189 Status DeadnessAnalysisImpl::HandleRecv(Node* n,
1190 std::vector<bool>* should_revisit) {
1191 // In addition to being alive or dead based on the inputs, a _Recv can also
1192 // acquire a dead signal from a _Send.
1193 std::vector<Predicate*> input_preds;
1194 TF_RETURN_IF_ERROR(GetInputPreds(n, EdgeKind::kDataAndControl, &input_preds));
1195 Predicate* signal_is_alive;
1196 TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate(
1197 n, /*output_idx=*/0, /*must_be_true=*/false, &signal_is_alive));
1198 input_preds.push_back(signal_is_alive);
1199 SetPredicate(n, {0, Graph::kControlSlot},
1200 predicate_factory_.MakeAndPredicate(input_preds),
1201 should_revisit);
1202 return Status::OK();
1203 }
1204
HandleGeneric(Node * n,std::vector<bool> * should_revisit)1205 Status DeadnessAnalysisImpl::HandleGeneric(Node* n,
1206 std::vector<bool>* should_revisit) {
1207 // Generally nodes are alive iff all their inputs are alive.
1208 std::vector<Predicate*> input_preds;
1209 TF_RETURN_IF_ERROR(GetInputPreds(n, EdgeKind::kDataAndControl, &input_preds));
1210 Predicate* pred = predicate_factory_.MakeAndPredicate(input_preds);
1211 for (int output_idx = 0; output_idx < n->num_outputs(); output_idx++) {
1212 SetPredicate(n, output_idx, pred, should_revisit);
1213 }
1214 SetPredicate(n, Graph::kControlSlot, pred, should_revisit);
1215 return Status::OK();
1216 }
1217
HandleNode(Node * n,std::vector<bool> * should_revisit,bool use_optimistic_mode)1218 Status DeadnessAnalysisImpl::HandleNode(Node* n,
1219 std::vector<bool>* should_revisit,
1220 bool use_optimistic_mode) {
1221 if (n->IsSwitch()) {
1222 TF_RETURN_IF_ERROR(HandleSwitch(n, should_revisit));
1223 } else if (n->IsMerge()) {
1224 TF_RETURN_IF_ERROR(HandleMerge(n, should_revisit, use_optimistic_mode));
1225 } else if (n->IsControlTrigger()) {
1226 SetPredicate(n, Graph::kControlSlot, predicate_factory_.MakeTrue(),
1227 nullptr);
1228 } else if (n->IsRecv() || n->IsHostRecv()) {
1229 TF_RETURN_IF_ERROR(HandleRecv(n, should_revisit));
1230 } else if (n->IsNextIteration()) {
1231 TF_RETURN_IF_ERROR(HandleGeneric(n, should_revisit));
1232 } else {
1233 TF_RETURN_IF_ERROR(HandleGeneric(n, should_revisit));
1234 }
1235 return Status::OK();
1236 }
1237
1238 // Compute a special topological order for the Graph, where nodes having the
1239 // same root frame are placed adjacent to each other. The traversal uses a
1240 // variant of Kahn's algorithm. num_ready_inputs is used to keep track of how
1241 // many inputs of each node are ready; a node is ready to be scheduled if all
1242 // of its inputs are ready.
1243 // Ref. to https://en.wikipedia.org/wiki/Topological_sorting for details.
GetFrameBasedTopologicalOrder(std::vector<Node * > * order)1244 Status DeadnessAnalysisImpl::GetFrameBasedTopologicalOrder(
1245 std::vector<Node*>* order) {
1246 absl::flat_hash_map<absl::string_view, size_t> num_enters_for_frame;
1247 absl::flat_hash_map<absl::string_view, size_t> num_exits_for_frame;
1248 std::vector<size_t> num_ready_inputs(graph_.num_node_ids(), 0);
1249 Node* src_node = graph_.source_node();
1250 for (const auto* node : graph_.op_nodes()) {
1251 const ControlFlowInfo& cf = control_flow_info_[node->id()];
1252 if (IsRootEnter(node)) {
1253 // Since we care only the root-level frame, full frame names are the same
1254 // as frame names.
1255 ++num_enters_for_frame[cf.frame_name];
1256 } else if (IsRootExit(node)) {
1257 ++num_exits_for_frame[cf.frame_name];
1258 }
1259 // Edge NextIteration->Merge is counted before starting the traversal to
1260 // break the backedges.
1261 if (IsMerge(node)) {
1262 for (const Edge* e : node->in_edges()) {
1263 if (IsNextIteration(e->src())) {
1264 ++num_ready_inputs[node->id()];
1265 }
1266 }
1267 }
1268 }
1269
1270 // dequeue is used to ensure that the nodes are first-in-first-out. This
1271 // order guarantees that the exits in the ready queue are visited before
1272 // nodes that will become ready in the future.
1273 std::deque<Node*> ready;
1274 ready.push_back(src_node);
1275 // ready_enters_per_frame and ready_exits serve as a staging area to buffer
1276 // the ready enters/exits before they are moved to the `ready` queue for
1277 // controlling the start and end of a processing frame.
1278 absl::flat_hash_map<absl::string_view, std::vector<Node*>>
1279 ready_enters_per_frame;
1280 // Exit nodes shall all be from the same frame, as we process a frame at a
1281 // time. So, one vector is enough.
1282 std::vector<Node*> ready_exits;
1283 while (!ready.empty()) {
1284 Node* curr_node = ready.front();
1285 ready.pop_front();
1286
1287 VLOG(4) << "Visiting " << curr_node->name();
1288 order->push_back(curr_node);
1289
1290 for (const Edge* out_edge : curr_node->out_edges()) {
1291 Node* out = out_edge->dst();
1292 int out_id = out->id();
1293 if (IsNextIteration(curr_node) && IsMerge(out)) {
1294 // Edge NextIteration->Merge has been counted.
1295 continue;
1296 }
1297 ++num_ready_inputs[out->id()];
1298 if (!out->IsOp()) continue; // Skip Sink/Source nodes.
1299 if (num_ready_inputs[out->id()] != out->in_edges().size()) continue;
1300
1301 absl::string_view frame_name = control_flow_info_[out_id].frame_name;
1302 if (IsRootEnter(out)) {
1303 ready_enters_per_frame[frame_name].push_back(out);
1304 } else if (IsRootExit(out)) {
1305 ready_exits.push_back(out);
1306 } else {
1307 ready.push_back(out);
1308 }
1309 }
1310
1311 if (ready.empty()) {
1312 // Try moving nodes from ready_enters_per_frame and ready_exits to
1313 // `ready`.
1314 if (!ready_exits.empty()) {
1315 // If there are nodes in ready_exits we must process them before
1316 // processing ready_enters_per_frame to make sure all nodes in the
1317 // currently processing frame are visited before starting processing
1318 // other frames.
1319 absl::string_view frame_name =
1320 control_flow_info_[ready_exits.front()->id()].frame_name;
1321 CHECK_EQ(ready_exits.size(), num_exits_for_frame[frame_name]);
1322 ready.insert(ready.end(), ready_exits.begin(), ready_exits.end());
1323 ready_exits.clear();
1324 } else {
1325 // Otherwise, try moving nodes from ready_enters to `ready`.
1326 for (auto iter = ready_enters_per_frame.begin();
1327 iter != ready_enters_per_frame.end(); ++iter) {
1328 absl::string_view frame_name = iter->first;
1329 const std::vector<Node*>& ready_enters = iter->second;
1330 if (ready_enters.size() == num_enters_for_frame[frame_name]) {
1331 ready.insert(ready.end(), ready_enters.begin(), ready_enters.end());
1332 ready_enters_per_frame.erase(iter);
1333 break;
1334 }
1335 }
1336 }
1337 }
1338 }
1339
1340 if (!ready_enters_per_frame.empty() || !ready_exits.empty()) {
1341 return errors::InvalidArgument(
1342 "Some enters/exits have never been visited in the traversal."
1343 " Most probably the input graph is malformed.");
1344 }
1345 return Status::OK();
1346 }
1347
1348 // We populate the nodes along a special topological order where nodes having
1349 // the same root frame are placed adjacent to each other. This grouping enables
1350 // processing the graph per root frame at a time and guarantees that when a root
1351 // frame is being processed, nodes in the downstream frames have not yet been
1352 // processed. This property is important because we need to process an entire
1353 // frame to know whether the optimistic mode converges or not. In other words,
1354 // nodes in the downstream frames shall not be populated until all of its
1355 // upstream frames are populated. In effect, this order enables processing each
1356 // (nested) tf.while one-by-one, as each (nested) tf.while creates a unique
1357 // (root) frame. Note that we don't separate while loops belonging to the same
1358 // nested while, as there is no clean cut for separating them in the topological
1359 // order.
Populate(bool enable_optimistic)1360 Status DeadnessAnalysisImpl::Populate(bool enable_optimistic) {
1361 std::vector<string> unreachable_nodes;
1362 // Compute the loop structure of the graph.
1363 TF_RETURN_IF_ERROR(
1364 BuildControlFlowInfo(&graph_, &control_flow_info_, &unreachable_nodes));
1365
1366 // Do some opportunistic error checking:
1367 if (!unreachable_nodes.empty()) {
1368 if (unreachable_nodes.size() > 5) {
1369 unreachable_nodes.erase(unreachable_nodes.begin() + 5,
1370 unreachable_nodes.end());
1371 }
1372
1373 return errors::InvalidArgument(
1374 "Found unreachable nodes, most likely source and sink nodes not "
1375 "connected: ",
1376 absl::StrJoin(unreachable_nodes, ", "));
1377 }
1378
1379 std::vector<Node*> topo;
1380 TF_RETURN_IF_ERROR(GetFrameBasedTopologicalOrder(&topo));
1381
1382 size_t frame_start = 0;
1383 while (frame_start < topo.size()) {
1384 // Batching nodes who have the same root frame.
1385 absl::string_view cur_frame_name;
1386 TF_RETURN_IF_ERROR(
1387 GetRootFrame(topo[frame_start], control_flow_info_, &cur_frame_name));
1388 size_t frame_end = frame_start;
1389 for (size_t i = frame_start + 1; i < topo.size(); ++i) {
1390 absl::string_view i_frame_name;
1391 TF_RETURN_IF_ERROR(
1392 GetRootFrame(topo[i], control_flow_info_, &i_frame_name));
1393 if (i_frame_name == cur_frame_name) {
1394 frame_end = i;
1395 } else {
1396 break;
1397 }
1398 }
1399 absl::Span<Node*> sub_topo(topo.data() + frame_start,
1400 /*length=*/frame_end - frame_start + 1);
1401 frame_start = frame_end + 1;
1402
1403 // First, try the optimistic mode.
1404 bool success = false;
1405 if (enable_optimistic && !cur_frame_name.empty()) {
1406 TF_RETURN_IF_ERROR(
1407 PopulateFrame(sub_topo, /*use_optimistic_mode=*/true, &success));
1408 }
1409 if (!success) {
1410 // The optimistic mode does not converge. Let's fall back to the
1411 // pessimistic mode.
1412 TF_RETURN_IF_ERROR(
1413 PopulateFrame(sub_topo, /*use_optimistic_mode=*/false, nullptr));
1414 }
1415 VLOG(2) << "Done populating frame " << cur_frame_name << " using the "
1416 << (success ? "optimistic" : "pessimistic") << " mode.";
1417 }
1418
1419 return Status::OK();
1420 }
1421
PopulateFrame(absl::Span<Node * const> topo,bool use_optimistic_mode,bool * success)1422 Status DeadnessAnalysisImpl::PopulateFrame(absl::Span<Node* const> topo,
1423 bool use_optimistic_mode,
1424 bool* success) {
1425 CHECK(use_optimistic_mode && success != nullptr ||
1426 !use_optimistic_mode && success == nullptr);
1427
1428 // This an abstract interpretation over the deadness propagation semantics of
1429 // the graph executor.
1430 //
1431 // We iterate over the graph twice, each time in a topological order. On the
1432 // first iteration merge nodes with backedges are mapped to symbolic
1433 // predicates. On the second iteration we use the predicates assigned to the
1434 // backedges in the previous iteration to infer a more precise predicate for
1435 // the backedge merge nodes and all the nodes that transitively use it.
1436 //
1437 // We don't track the output indices for should_revisit. Instead, putting a
1438 // node in `should_revisit` denotes that the deadness flowing out from any
1439 // output from said node may have changed. This is fine; only switches
1440 // propagate different deadness along different output edges, and since the
1441 // delta is solely due to the input *values* (and not input deadness), the
1442 // delta should not change in the second iteration.
1443 std::vector<bool> should_revisit;
1444 should_revisit.resize(graph_.num_node_ids());
1445 for (Node* n : topo) {
1446 VLOG(4) << "Visiting " << n->name();
1447 TF_RETURN_IF_ERROR(
1448 HandleNode(n, /*should_revisit=*/nullptr, use_optimistic_mode));
1449 if (n->IsNextIteration()) {
1450 // If this is a backedge for a merge node then remember to reprocess the
1451 // merge the next time we run.
1452 for (const Edge* e : n->out_edges()) {
1453 if (e->dst()->IsMerge()) {
1454 should_revisit[e->dst()->id()] = true;
1455 }
1456 }
1457 }
1458 }
1459
1460 for (Node* n : topo) {
1461 // The nodes added to should_revisit in the previous loop need to be
1462 // revisited now. Reprocessing these initial nodes may add *their*
1463 // consumers to should_revisit, and these newly added nodes will also be
1464 // processed by this very same loop. Since we're traversing the graph in
1465 // topological order (producers before consumers) and HandleNode(n) can only
1466 // ever add n's consumers to should_revisit, we won't "miss" an addition to
1467 // should_revisit.
1468 if (should_revisit[n->id()]) {
1469 VLOG(4) << "Revisiting " << n->name();
1470 TF_RETURN_IF_ERROR(HandleNode(n, &should_revisit));
1471 }
1472 }
1473
1474 // Check if the optimistic analysis converges. Specifically, check whether
1475 // all the predicates of the merge nodes in the same frame are the same. If
1476 // yes, report success. If not, report failure and clear the assigned
1477 // predicates.
1478 if (use_optimistic_mode) {
1479 bool is_converged = true;
1480 absl::flat_hash_map<absl::string_view, Predicate*> frame_to_pred;
1481 for (Node* n : topo) {
1482 if (!n->IsMerge()) {
1483 continue;
1484 }
1485 const Edge* e;
1486 TF_RETURN_IF_ERROR(FindUniqueBackedge(n, &e));
1487 if (e == nullptr) {
1488 // Skip acyclic merge nodes.
1489 continue;
1490 }
1491 Node* merge = n;
1492 // Note that here uses frame names instead of root frame names. In the
1493 // case of a nested while loop, each level of while loops can have merges
1494 // with different predicate instances, while the merge nodes on the same
1495 // level must have the same predicate instances.
1496 absl::string_view frame_name = control_flow_info_[merge->id()].frame_name;
1497 auto it = predicate_map_.find(TensorId(merge->name(), 0));
1498 Predicate* merge_pred = it->second;
1499 if (merge_pred->kind() != Predicate::Kind::kAndRecurrence) {
1500 is_converged = false;
1501 VLOG(2) << "Running the optimistic mode on frame " << frame_name
1502 << " does not converge because node " << merge->name()
1503 << " cannot be mapped into the AndRecurrence form.";
1504 break;
1505 }
1506
1507 auto insert_result = frame_to_pred.insert({frame_name, merge_pred});
1508 if (!insert_result.second) {
1509 // If we have already seen this frame name, verify the predicate is the
1510 // same as the previously seen one's.
1511 Predicate* curr_andrec = merge_pred;
1512 Predicate* prev_andrec = insert_result.first->second;
1513 if (curr_andrec != prev_andrec) {
1514 is_converged = false;
1515 VLOG(2) << "Running the optimistic mode on frame " << frame_name
1516 << " does not converge. Seeing different Merge predicates: \n"
1517 << curr_andrec->ToString() << " and \n"
1518 << prev_andrec->ToString();
1519 break;
1520 }
1521 }
1522 }
1523
1524 // Clear the assigned predicates if the optimistic mode does not converge.
1525 if (!is_converged) {
1526 for (Node* n : topo) {
1527 for (int oid = 0; oid < n->num_outputs(); ++oid) {
1528 predicate_map_.erase(TensorId(n->name(), oid));
1529 }
1530 predicate_map_.erase(TensorId(n->name(), Graph::kControlSlot));
1531 }
1532 }
1533
1534 if (success != nullptr) {
1535 *success = is_converged;
1536 }
1537 }
1538
1539 return Status::OK();
1540 }
1541
1542 StatusOr<DeadnessAnalysis::DeadnessPredicate>
GetPredicateFor(Node * n,int oidx) const1543 DeadnessAnalysisImpl::GetPredicateFor(Node* n, int oidx) const {
1544 auto it = predicate_map_.find(TensorId(n->name(), oidx));
1545 TF_RET_CHECK(it != predicate_map_.end())
1546 << "could not find " << TensorId(n->name(), oidx).ToString()
1547 << " in predicate map";
1548 return MakeDeadnessPredicate(it->second);
1549 }
1550
Print() const1551 void DeadnessAnalysisImpl::Print() const {
1552 std::vector<TensorId> tensor_ids;
1553 for (const auto& kv_pair : predicate_map_) {
1554 tensor_ids.push_back(kv_pair.first);
1555 }
1556
1557 std::sort(tensor_ids.begin(), tensor_ids.end());
1558
1559 for (TensorId tensor_id : tensor_ids) {
1560 auto it = predicate_map_.find(tensor_id);
1561 CHECK(it != predicate_map_.end()) << tensor_id.ToString();
1562 VLOG(2) << tensor_id.ToString() << " -> " << it->second->ToString();
1563 }
1564 }
1565
1566 } // namespace
1567
~DeadnessAnalysis()1568 DeadnessAnalysis::~DeadnessAnalysis() {}
1569
Run(const Graph & graph,std::unique_ptr<DeadnessAnalysis> * result)1570 /*static*/ Status DeadnessAnalysis::Run(
1571 const Graph& graph, std::unique_ptr<DeadnessAnalysis>* result) {
1572 std::unique_ptr<DeadnessAnalysisImpl> analysis(
1573 new DeadnessAnalysisImpl(&graph));
1574 TF_RETURN_IF_ERROR(analysis->Populate(/*enable_optimistic=*/true));
1575
1576 if (VLOG_IS_ON(2)) {
1577 analysis->Print();
1578 }
1579
1580 *result = std::move(analysis);
1581 return Status::OK();
1582 }
1583
1584 absl::flat_hash_map<TensorId, string, TensorId::Hasher>
PredicateMapAsString() const1585 DeadnessAnalysisImpl::PredicateMapAsString() const {
1586 absl::flat_hash_map<TensorId, string, TensorId::Hasher> result;
1587 for (const auto& kv_pair : predicate_map_) {
1588 CHECK(result.insert({kv_pair.first, kv_pair.second->ToString()}).second);
1589 }
1590 return result;
1591 }
1592
1593 namespace deadness_analysis_internal {
ComputePredicates(const Graph & graph,PredicateMapTy * out_predicate_map,bool enable_optimistic)1594 Status ComputePredicates(const Graph& graph, PredicateMapTy* out_predicate_map,
1595 bool enable_optimistic) {
1596 DeadnessAnalysisImpl impl(&graph);
1597 TF_RETURN_IF_ERROR(impl.Populate(enable_optimistic));
1598 *out_predicate_map = impl.PredicateMapAsString();
1599 return Status::OK();
1600 }
1601
1602 } // namespace deadness_analysis_internal
1603
DebugString(DeadnessPredicate predicate) const1604 string DeadnessAnalysis::DebugString(DeadnessPredicate predicate) const {
1605 return static_cast<Predicate*>(predicate.pred_)->ToString();
1606 }
1607
1608 } // namespace tensorflow
1609