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