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