1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_
17#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_
18
19#include <ostream>
20#include <string>
21#include <vector>
22
23#include "tensorflow/compiler/xla/service/hlo_instruction.h"
24#include "tensorflow/compiler/xla/shape_tree.h"
25#include "tensorflow/compiler/xla/types.h"
26#include "tensorflow/compiler/xla/xla_data.pb.h"
27#include "tensorflow/core/lib/gtl/array_slice.h"
28#include "tensorflow/core/platform/macros.h"
29
30namespace xla {
31
32// Abstraction which identifies a specific point in the XLA graph. An
33// HloPosition specifies a ShapeIndex within the output of a specific
34// instruction.
35struct HloPosition {
36 HloInstruction* instruction;
37 ShapeIndex index;
38
39 // Returns the shape at this position.
40 const Shape& shape() const;
41
42 string ToString() const;
43
44 bool operator==(const HloPosition& other) const {
45 return instruction == other.instruction && index == other.index;
46 }
47 bool operator!=(const HloPosition& other) const { return !(*this == other); }
48
49 // Stable less-than operator using instruction id and index.
50 bool operator<(const HloPosition& other) const {
51 return instruction->unique_id() < other.instruction->unique_id() ||
52 (instruction->unique_id() == other.instruction->unique_id() &&
53 index < other.index);
54 }
55};
56
57std::ostream& operator<<(std::ostream& out, const HloPosition& position);
58
59// Defines a single use of an HLO value.
60struct HloUse {
61 // Instruction at which the value is used.
62 HloInstruction* instruction;
63
64 // The operand number in which the value is appears.
65 int64 operand_number;
66
67 // The shape index within the operand in which the value appears.
68 ShapeIndex operand_index;
69
70 string ToString() const;
71
72 bool operator==(const HloUse& other) const {
73 return instruction == other.instruction &&
74 operand_number == other.operand_number &&
75 operand_index == other.operand_index;
76 }
77
78 bool operator!=(const HloUse& other) const { return !(*this == other); }
79};
80
81std::ostream& operator<<(std::ostream& out, const HloUse& use);
82
83// Class describing a value used by the dataflow analysis. XLA arrays are
84// trivially a single HloValue. Tuples are made up of more than one HloValue: an
85// HloValue for the pointer vector, and an HloValue for each child element.
86//
87// Every HloValue is defined by a particular instruction and most instructions
88// define only a single HloValue. Instructions which define a single HloValue
89// include array-shaped instructions such as Add but also includes Tuple-shaped
90// instructions such as Tuple. The Tuple instruction defines a single HloValue
91// which is a vector of pointers to the values containing the Tuple
92// instruction's operands. Though the result of the Tuple instruction includes
93// multiple values only the top-level HloValue (the vector of pointers) is
94// defined by the Tuple instruction. The values containing the tuple elements
95// are defined by earlier instructions, usually the operands of the Tuple
96// instruction.
97//
98// Instructions which construct both the tuple *and* the tuple elements define
99// more than one HloValue. This includes (at least) tuple-shaped Constant,
100// Parameter, Infeed and While instructions. These tuple-shaped instructions do
101// not assemble a tuple from existing HloValues like the Tuple instruction does,
102// but rather define all the HloValues in the tuple.
103class HloValue {
104 public:
105 using Id = int64;
106
107 // Predicate comparing HloValues by increasing id, useful for std::sort.
108 static bool IdLessThan(const HloValue* a, const HloValue* b) {
109 return a->id() < b->id();
110 }
111
112 // Predicate comparing HloValues by equal id, useful for std::unique.
113 static bool IdEqual(const HloValue* a, const HloValue* b) {
114 return a->id() == b->id();
115 }
116
117 // Construct an HloValue defined by 'instruction' at shape index 'index'. If
118 // is_phi is true, then this value is a phi value, for example, at the
119 // parameter of a while body computation. Phi values are only used in the SSA
120 // dataflow analysis (HloDataflowAnalysis::ssa_form_ is true).
121 HloValue(Id id, HloInstruction* instruction, const ShapeIndex& index,
122 bool is_phi = false);
123
124 // Sets the positions in the module at which the HloValue appears. Updates
125 // uses. Should be called once and only once. The defining position should not
126 // be included in 'positions' as this is set at construction time.
127 void SetPositionsAndComputeUses(
128 tensorflow::gtl::ArraySlice<HloPosition> positions);
129
130 // Return a unique identifier for this HloValue. This value is used for stable
131 // sorting and iteration
132 Id id() const { return id_; }
133
134 // Returns whether this value is a phi value.
135 bool is_phi() const { return is_phi_; }
136
137 // Return the position where this value is defined.
138 const HloPosition& defining_position() const { return positions_[0]; }
139
140 // Return the instruction which defines this HloValue.
141 HloInstruction* defining_instruction() const {
142 return defining_position().instruction;
143 }
144
145 // Return the shape index at which this HloValue is defined in the output of
146 // its defining instruction.
147 const ShapeIndex& defining_index() const { return defining_position().index; }
148
149 // Return the shape of this HloValue.
150 const Shape& shape() const { return defining_position().shape(); }
151
152 // Return all positions of the HloValue in the module.
153 const std::vector<HloPosition>& positions() const { return positions_; }
154
155 // Return all uses of the HloValue.
156 const std::vector<HloUse>& uses() const { return uses_; }
157
158 // Get whether this HloValue is live out of the module.
159 bool live_out_of_module() const { return live_out_of_module_; }
160
161 bool operator==(const HloValue& other) const;
162 bool operator!=(const HloValue& other) const;
163
164 // Return a single-line string representation of the value.
165 string ToShortString() const;
166
167 string ToString(int indent = 0) const;
168
169 private:
170 // Unique identifier for this HloValue. Used for stable sorting and iteration.
171 const Id id_;
172
173 // Whether this instruction is a phi value.
174 const bool is_phi_;
175
176 // The set of positions of this HloValue. The first element is always the
177 // position of the definition.
178 std::vector<HloPosition> positions_;
179
180 // The set of uses of this HloValue.
181 std::vector<HloUse> uses_;
182
183 // Whether this value is live out of the HLO module.
184 bool live_out_of_module_ = false;
185
186 // Whether this value is live out of its computation.
187 bool live_out_of_computation_ = false;
188};
189
190std::ostream& operator<<(std::ostream& out, const HloValue& hlo_value);
191
192// A class representing the possible set of HloValues at a particular point
193// (shape index in the output of an instruction) in the XLA graph. This set
194// contains the set of reaching HloValue definitions. For a simple array-shaped
195// instruction like Add, the HloValueSet of the top-level of the instruction's
196// output trivially contains only the HloValue defined by the instruction. For
197// instructions which have non-trivial dataflow such as Tuple or Select, the
198// HloValueSets of the instruction's output contains one or more HloValues
199// defined by the instruction's operands or defined further up in the XLA graph.
200class HloValueSet {
201 public:
202 HloValueSet() = default;
203
204 explicit HloValueSet(tensorflow::gtl::ArraySlice<const HloValue*> values)
205 : values_(values.begin(), values.end()) {
206 SortAndUniquifyValues();
207 }
208
209 // Sets this value set to the union of the given value sets. Returns whether
210 // this value set changed.
211 bool AssignUnionOf(tensorflow::gtl::ArraySlice<const HloValueSet*> inputs);
212
213 // Return the vector of HloValues in the set. Values in the vector are unique
214 // and stably sorted by value id.
215 const std::vector<const HloValue*>& values() const { return values_; }
216
217 // Adds the value to the set. Returns true iff the value was added and didn't
218 // already exist in the set.
219 bool AddValue(const HloValue* value);
220
221 // Clear all values from the set.
222 void Clear() { values_.clear(); }
223
224 // Return the unique HLO value in the set. CHECKs if the set does not contain
225 // exactly one value.
226 const HloValue& GetUniqueValue() const {
227 CHECK_EQ(values_.size(), 1);
228 return *values_[0];
229 }
230
231 bool operator==(const HloValueSet& other) const {
232 if (values_.size() != other.values_.size()) return false;
233 for (size_t i = 0; i < values_.size(); ++i) {
234 if (values_[i]->id() != other.values_[i]->id()) {
235 return false;
236 }
237 }
238 return true;
239 }
240 bool operator!=(const HloValueSet& other) const { return !(*this == other); }
241
242 string ToString() const;
243
244 private:
245 // Sorts value_ and removes duplicates. This should be called after adding any
246 // elements to values_.
247 void SortAndUniquifyValues();
248
249 // HloValues sorted by HloValue::Id.
250 std::vector<const HloValue*> values_;
251};
252
253std::ostream& operator<<(std::ostream& out, const HloValueSet& hlo_value);
254
255// A class collecting the HloValues which might be contained in the output of
256// an HLO instruction. For array-shaped instructions, an InstructionValueSet
257// trivially holds a single HloValueSet. Tuple-shaped InstructionValueSets
258// hold multiple HloValueSets.
259class InstructionValueSet : public ShapeTree<HloValueSet> {
260 public:
261 InstructionValueSet(const Shape& shape) : ShapeTree<HloValueSet>(shape) {}
262
263 // Sets this value set to the union of the given value sets. Returns whether
264 // this value set changed.
265 bool AssignUnionOf(
266 tensorflow::gtl::ArraySlice<const InstructionValueSet*> inputs);
267
268 string ToString() const;
269};
270
271std::ostream& operator<<(std::ostream& out,
272 const InstructionValueSet& instruction_value_set);
273
274} // namespace xla
275
276#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_
277