• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2020 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 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PHI_GRAPH_H_
17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PHI_GRAPH_H_
18 
19 #include <iterator>
20 #include <memory>
21 #include <string>
22 #include <vector>
23 
24 #include "absl/container/flat_hash_map.h"
25 #include "absl/container/flat_hash_set.h"
26 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
27 #include "tensorflow/compiler/xla/service/hlo_module.h"
28 #include "tensorflow/compiler/xla/service/hlo_value.h"
29 
30 namespace xla {
31 // Phi graph is a graph that contains and connects phi nodes build on top of
32 // HloValues with explicit edges, as well as non-phi nodes that are direct
33 // inputs to the phi nodes. The graph can be viewed as an 'overlay' on top of
34 // HloValues, with some edges that connect them together. After optimization,
35 // some phis nodes will be simplified away and the user can then ask two useful
36 // questions:
37 //
38 // 1. Which HloValue should a phi node being replaced with?
39 // 2. TODO(yunxing): What are the set of aliased HloValues that are connecting
40 // to the same phi (Must-aliasing).
41 class PhiGraph {
42  public:
43   // Register an hlo value into the phi node.
44   void RegisterPhi(const HloValue& value,
45                    absl::Span<const HloValue* const> inputs);
46 
47   HloValue::Id GetOptimizedId(const HloValue& value);
48 
49   // Returns true if the input to a hlo value is the same as `inputs`.
50   bool InputsEqualTo(const HloValue& value,
51                      absl::Span<const HloValue* const> inputs);
52 
53   // Given `id`, returns the new id that `id` should be replaced with. If the
54   // node is not optimized, returns the same value.
55   HloValue::Id FindOptimizedValue(const HloValue::Id id);
56 
57   // Optimize the entire graph.
58   void Optimize();
59 
60   std::string ToString();
61 
62  private:
63   struct Node {
64     bool is_phi;
65     // Users of this node. Non-phi node has no operands.
66     std::vector<Node*> users;
67     // Operands of this node.
68     std::vector<Node*> operands;
69 
70     // The value that the node is originally registered with.
71     HloValue::Id value_id;
72 
73     // mark_as_dead is set to true when a phi node is simplified away
74     //
75     // Precondition: node is a phi.
76     bool mark_as_dead = false;
77   };
78 
79   Node* CreateOrReuseNode(const HloValue& value);
80 
81   // Replace `node` with `replace`. Redirect all users to the `replace` and
82   // all HloValues pointing to the `node` to `replace`. Also mark `node` as
83   // dead.
84   //
85   // Precondition: node is a phi -- It's only possible to simplify away a
86   // phi node.
87   void ReplaceNodeWith(Node* node, Node* replace);
88 
89   // A reverse mapping of a node in the phi graph and all HloValues pointing
90   // to that phi.
91   absl::flat_hash_map<Node*, std::vector<HloValue::Id>> node_to_value_id_;
92 
93   // A mapping from a HloValue to node in the phi graph.
94   absl::flat_hash_map<HloValue::Id, Node*> value_id_to_node_;
95   std::vector<std::unique_ptr<Node>> node_storage_;
96 };
97 
98 }  // namespace xla
99 
100 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PHI_GRAPH_H_
101