1 /* Copyright 2017 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_LOGICAL_BUFFER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_ 18 19 #include <functional> 20 #include <iosfwd> 21 #include <string> 22 #include <vector> 23 24 #include "tensorflow/compiler/xla/service/hlo.pb.h" 25 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 26 #include "tensorflow/compiler/xla/shape_util.h" 27 #include "tensorflow/compiler/xla/types.h" 28 #include "tensorflow/compiler/xla/xla_data.pb.h" 29 #include "tensorflow/core/lib/gtl/array_slice.h" 30 #include "tensorflow/core/lib/gtl/int_type.h" 31 #include "tensorflow/core/platform/macros.h" 32 #include "tensorflow/core/platform/types.h" 33 34 namespace xla { 35 36 // Class describing a contiguous sequence of elements (ie, C array) which form 37 // the components of Shaped values in XLA. XLA arrays are trivially a 38 // single LogicalBuffer. Tuple values are made up of more than one 39 // LogicalBuffer: a LogicalBuffer for the pointers to elements, and a 40 // LogicalBuffer for each child element. 41 // 42 // Every buffer is defined by a particular instruction and most instructions 43 // define only a single buffer. Instructions which define a single buffer 44 // include array-shaped instructions such as Add but also includes Tuple-shaped 45 // instructions such as Tuple. The Tuple instruction defines a single buffer 46 // which is a vector of pointers to the buffers containing the Tuple 47 // instruction's operands. Though the result of the Tuple instruction includes 48 // multiple buffers only the top-level buffer (the vector of pointers) is 49 // defined by the Tuple instruction. The buffers containing the tuple elements 50 // are defined by earlier instructions, usually the operands of the Tuple 51 // instruction. 52 // 53 // Instructions which construct both the tuple *and* the tuple elements define 54 // more than one buffer. This includes (at least) tuple-shaped Constant, 55 // Parameter, Infeed and While instructions. The tuple-shaped instructions do 56 // not assemble a tuple from existing buffers like the Tuple instruction does, 57 // but rather define the entire tuple. 58 // 59 // Some instructions, such as Bitcast, define no buffers. These instructions 60 // simply forward buffers from their operands. 61 // 62 // The LogicalBuffer object describes which HLO instruction defines a buffer and 63 // where within that instruction's output shape the buffer is defined. The 64 // location within the output shape is indicated by LogicalBuffer::index() which 65 // is defined identically to the index used in 66 // ShapeUtil::GetSubshape(). Examples: 67 // 68 // %add = Add(%foo, %bar) 69 // %tuple_constant = Constant({1, {42, 43}}) 70 // 71 // %add defines a single array-shaped buffer LogicalBuffer(%add, {}) which holds 72 // the array result of the add operation. The nested-tuple-shaped 73 // %tuple_constant defines 5 buffers described by the following LogicalBuffer 74 // objects: 75 // 76 // LogicalBuffer(%tuple_constant, {}) // "Top-level" buffer: vector of 77 // // pointers to LogicalBuffers at 78 // // indices {0} and {1} 79 // LogicalBuffer(%tuple_constant, {0}) // Holds value "1" 80 // LogicalBuffer(%tuple_constant, {1}) // Holds nested tuple: vector of 81 // // pointers to LogicalBuffers at 82 // // indices {1, 0} and {1, 1} 83 // LogicalBuffer(%tuple_constant, {1, 0}) // Holds value "42" 84 // LogicalBuffer(%tuple_constant, {1, 1}) // Holds value "43" 85 class LogicalBuffer { 86 public: 87 TF_LIB_GTL_DEFINE_INT_TYPE(Color, int64); 88 89 // Id is a unique identifier for the LogicalBuffer to facilitate efficient 90 // collections of LogicalBuffers with stable iteration order. 91 // LogicalBuffers are typically created and accessed through 92 // TuplePointsToAnalysis, and points-to analysis assigns each LogicalBuffer a 93 // unique value. 94 using Id = int64; 95 96 // Functions which return the size and alignment of a logical buffer in bytes. 97 using SizeFunction = std::function<int64(const LogicalBuffer&)>; 98 using AlignmentFunction = std::function<int64(LogicalBuffer::Color)>; 99 100 LogicalBuffer(HloInstruction* instruction, const ShapeIndex& index, Id id); 101 id()102 Id id() const { return id_; } 103 104 // Return the instruction that defines the buffer. instruction()105 HloInstruction* instruction() const { return instruction_; } 106 107 // Return the index within the output of the instruction where the buffer is 108 // defined. Index used defined as in ShapeUtil::GetSubshape() index()109 const ShapeIndex& index() const { return index_; } 110 111 // Return the color of the logical buffer. Differently colored buffers can 112 // not be parts of the same allocation. color()113 Color color() const { 114 CHECK_NE(color_, kInvalidColor) 115 << "Should not query the color of a buffer that was never colored"; 116 return color_; 117 } 118 set_color(Color color)119 void set_color(Color color) { 120 CHECK_NE(color, kInvalidColor) 121 << "Should not set the color of a buffer to the invalid color"; 122 color_ = color; 123 } 124 has_color()125 bool has_color() const { return color_ != kInvalidColor; } 126 127 // Return the shape of the buffer. This reference points into the shape field 128 // of the instruction defining the buffer. Therefore, the returned shape will 129 // contain the layout of instruction, if any. shape()130 const Shape& shape() const { 131 return ShapeUtil::GetSubshape(instruction_->shape(), index_); 132 } 133 134 // Returns true if this buffer is the top-level output buffer of the defining 135 // HLO instruction. This is equivalent to index == {}. IsTopLevel()136 bool IsTopLevel() const { return index_.empty(); } 137 138 // Whether this buffer contains a tuple. IsTuple()139 bool IsTuple() const { return is_tuple_; } 140 141 // Whether this buffer contains an array. IsArray()142 bool IsArray() const { return is_array_; } 143 144 // operator< is required for std::set. 145 bool operator<(const LogicalBuffer& other) const { return id_ < other.id_; } 146 147 string ToString() const; 148 LogicalBufferProto ToProto(const SizeFunction& size_fn) const; 149 150 // Returns the LogicalBufferProto::Location that serializes the given 151 // instruction and index. 152 static LogicalBufferProto::Location ToLocationProto( 153 const HloInstruction& instruction, const ShapeIndex& index); 154 155 const Color kInvalidColor = Color(-1); 156 157 private: 158 HloInstruction* instruction_; 159 Id id_ : 62; 160 bool is_array_ : 1; 161 bool is_tuple_ : 1; 162 Color color_; 163 ShapeIndex index_; 164 165 // Similar to HLO constructs (HloInstruction, etc), pointers are used for 166 // comparison to equality, so disable all copying. 167 TF_DISALLOW_COPY_AND_ASSIGN(LogicalBuffer); 168 }; 169 170 std::ostream& operator<<(std::ostream& out, const LogicalBuffer& buffer); 171 172 } // namespace xla 173 174 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_ 175