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_BUFFER_LIVENESS_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_LIVENESS_H_ 18 19 #include <memory> 20 #include <string> 21 #include <utility> 22 23 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 24 #include "tensorflow/compiler/xla/service/hlo_module.h" 25 #include "tensorflow/compiler/xla/service/hlo_ordering.h" 26 #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" 27 #include "tensorflow/compiler/xla/statusor.h" 28 #include "tensorflow/compiler/xla/types.h" 29 #include "tensorflow/core/lib/core/status.h" 30 #include "tensorflow/core/lib/gtl/flatmap.h" 31 #include "tensorflow/core/lib/gtl/flatset.h" 32 33 namespace xla { 34 35 // Class which computes liveness of the output buffers of HLOs and their 36 // interference. 37 class BufferLiveness { 38 public: 39 using Colorer = std::function<Status(const BufferLiveness& buffer_liveness)>; 40 41 // Constructs a buffer liveness object for the given module assuming the given 42 // HLO instruction ordering. 43 static StatusOr<std::unique_ptr<BufferLiveness>> Run( 44 const HloModule* module, std::unique_ptr<HloOrdering> hlo_ordering); 45 46 // Returns true if the live range of the buffer containing the output of 'a' 47 // may overlap with the live range of the buffer of 'b'. If instruction 'a' 48 // interferes with instruction 'b' then they cannot share the same buffer. 49 bool MayInterfere(const LogicalBuffer& a, const LogicalBuffer& b) const; 50 51 // Returns true if the buffer for the given instruction may be live out of the 52 // module. That is, the instruction's buffer may be included in the output of 53 // the entry computation. 54 bool MaybeLiveOut(const LogicalBuffer& buffer) const; 55 56 // Returns the complete set of buffers that may be live out of the module. 57 const PointsToSet::BufferSet& maybe_live_out_buffers() const { 58 return maybe_live_out_buffers_; 59 } 60 61 // Returns the underlying points-to analysis used for this liveness analysis. 62 const TuplePointsToAnalysis& points_to_analysis() const { 63 return *points_to_analysis_; 64 } 65 66 // Returns the underlying hlo ordering used for this liveness analysis. 67 const HloOrdering& hlo_ordering() const { return *hlo_ordering_; } 68 69 const HloModule& module() const { return *module_; } 70 71 string ToString() const; 72 73 static Colorer DefaultColorer() { 74 return [](const BufferLiveness& buffer_liveness) { 75 for (LogicalBuffer::Id id = 0; 76 id < buffer_liveness.points_to_analysis().num_logical_buffers(); 77 id++) { 78 auto& buffer = buffer_liveness.points_to_analysis().logical_buffer(id); 79 buffer.set_color(LogicalBuffer::Color(0)); 80 } 81 return Status::OK(); 82 }; 83 } 84 85 private: 86 explicit BufferLiveness(const HloModule* module, 87 std::unique_ptr<HloOrdering> hlo_ordering) 88 : module_(module), hlo_ordering_(std::move(hlo_ordering)) {} 89 90 // Perform buffer liveness analysis. This method must be called prior to 91 // MayInterfere or MaybeLiveOut. 92 tensorflow::Status Analyze(); 93 94 // Returns true if the live range of the buffer of 'a' is strictly before the 95 // live range of the buffer of 'b' (they do not overlap). 96 bool live_range_strictly_before(const LogicalBuffer& a, 97 const LogicalBuffer& b) const; 98 99 const HloModule* module_; 100 std::unique_ptr<HloOrdering> hlo_ordering_; 101 102 // Set of LogicalBuffers which are aliased in the output of other 103 // instructions. For example, a LogicalBuffer which is inserted into a tuple 104 // is considered to be aliased and will be in this set. 105 tensorflow::gtl::FlatSet<const LogicalBuffer*> aliased_buffers_; 106 107 // LogicalBuffers that may be live out of the entry computation. 108 PointsToSet::BufferSet maybe_live_out_buffers_; 109 110 std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_; 111 }; 112 113 } // namespace xla 114 115 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_LIVENESS_H_ 116
