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 // Defines the data returned by the XLA buffer assignment packages. 17 18 #include "tensorflow/compiler/xla/service/buffer_liveness.h" 19 20 #include <utility> 21 #include <vector> 22 23 #include "absl/strings/str_format.h" 24 #include "absl/strings/str_join.h" 25 #include "tensorflow/compiler/xla/service/hlo_computation.h" 26 #include "tensorflow/compiler/xla/service/logical_buffer.h" 27 #include "tensorflow/compiler/xla/shape_util.h" 28 #include "tensorflow/compiler/xla/status_macros.h" 29 #include "tensorflow/compiler/xla/statusor.h" 30 #include "tensorflow/compiler/xla/types.h" 31 #include "tensorflow/compiler/xla/util.h" 32 #include "tensorflow/core/lib/core/errors.h" 33 #include "tensorflow/core/platform/logging.h" 34 35 namespace xla { 36 37 /* static */ 38 StatusOr<std::unique_ptr<BufferLiveness>> BufferLiveness::Run( 39 const HloModule* module, std::unique_ptr<HloOrdering> hlo_ordering) { 40 std::unique_ptr<BufferLiveness> liveness( 41 new BufferLiveness(module, std::move(hlo_ordering))); 42 TF_RETURN_IF_ERROR(liveness->Analyze()); 43 return std::move(liveness); 44 } 45 46 Status BufferLiveness::Analyze() { 47 TF_ASSIGN_OR_RETURN(points_to_analysis_, TuplePointsToAnalysis::Run(module_)); 48 for (auto* computation : module_->computations()) { 49 if (computation->IsFusionComputation()) { 50 continue; 51 } 52 // Gather all instructions whose buffers might alias other instructions into 53 // the set aliased_buffers_. This includes those contained as a tuple 54 // element in other instruction's output. 55 for (const auto& instruction : computation->instructions()) { 56 for (const LogicalBuffer* aliased_buffer : 57 points_to_analysis_->GetPointsToSet(instruction) 58 .CreateFlattenedSet()) { 59 if (aliased_buffer->instruction() != instruction) { 60 aliased_buffers_.insert(aliased_buffer); 61 } 62 } 63 } 64 65 if (computation == module_->entry_computation()) { 66 const HloInstruction* root = computation->root_instruction(); 67 maybe_live_out_buffers_ = 68 points_to_analysis_->GetPointsToSet(root).CreateFlattenedSet(); 69 } 70 } 71 72 XLA_VLOG_LINES(3, ToString()); 73 return Status::OK(); 74 } 75 76 string BufferLiveness::ToString() const { 77 std::vector<string> pieces; 78 pieces.push_back( 79 absl::StrFormat("BufferLiveness(module=%s):", module_->name())); 80 pieces.push_back("HloOrdering:"); 81 pieces.push_back(hlo_ordering_->ToString()); 82 pieces.push_back("Aliased buffers:"); 83 for (const LogicalBuffer* buffer : aliased_buffers_) { 84 pieces.push_back(absl::StrFormat(" %s", buffer->ToString())); 85 } 86 pieces.push_back("Live out buffers:"); 87 for (const LogicalBuffer* buffer : maybe_live_out_buffers_) { 88 pieces.push_back(absl::StrFormat(" %s", buffer->ToString())); 89 } 90 return absl::StrJoin(pieces, "\n"); 91 } 92 93 bool BufferLiveness::live_range_strictly_before(const LogicalBuffer& a, 94 const LogicalBuffer& b) const { 95 TF_DCHECK_OK(points_to_analysis_->VerifyBuffer(a)); 96 TF_DCHECK_OK(points_to_analysis_->VerifyBuffer(b)); 97 98 if (!hlo_ordering_->ExecutesBefore(a.instruction(), b.instruction())) { 99 return false; 100 } 101 102 for (const BufferAlias& alias : points_to_analysis_->GetBufferAliases(a)) { 103 // Every user of 'a' must be a predecessor of 'b' or 'b' itself. 104 for (auto user : alias.instruction()->users()) { 105 if (points_to_analysis().DoesNotUseOperandBuffer(alias.instruction(), 106 alias.index(), user)) { 107 continue; 108 } 109 if (user != b.instruction() && 110 !hlo_ordering_->ExecutesBefore(user, b.instruction())) { 111 return false; 112 } 113 } 114 115 // If the root instruction aliases the buffer 'a', the live range of 'a' is 116 // until the end of the computation and can never be strictly before another 117 // buffer defined in the same computation. This is needed to prevent the 118 // root instruction's buffers from being reused by later instructions even 119 // when the root is not the last instruction in the schedule. 120 if (alias.instruction()->parent()->root_instruction() == 121 alias.instruction() && 122 alias.instruction()->parent() == b.instruction()->parent()) { 123 return false; 124 } 125 } 126 127 // If 'b' is a user of 'a' then the buffers interfere unless 'a.instruction' 128 // and 'b.instruction' emit the same shape/layout, and 'b.instruction' meets 129 // the qualifications specified in CanShareOperandBufferWithUser. 130 for (const BufferAlias& alias : points_to_analysis_->GetBufferAliases(a)) { 131 if (b.instruction()->IsUserOf(alias.instruction()) && 132 !points_to_analysis().CanShareOperandBufferWithUser( 133 alias.instruction(), alias.index(), b.instruction(), b.index())) { 134 return false; 135 } 136 } 137 return true; 138 } 139 140 namespace { 141 bool IsEntryParameter(const HloInstruction* instruction) { 142 const HloComputation* computation = instruction->parent(); 143 return instruction->opcode() == HloOpcode::kParameter && 144 computation == computation->parent()->entry_computation(); 145 } 146 } // namespace 147 148 bool BufferLiveness::MayInterfere(const LogicalBuffer& a, 149 const LogicalBuffer& b) const { 150 // Entry parameters live at the entry of the execution, thus always interfere 151 // with all other instructions executing before them in the ordering. 152 const HloInstruction* a_instruction = a.instruction(); 153 const HloInstruction* b_instruction = b.instruction(); 154 if (IsEntryParameter(a_instruction) && 155 hlo_ordering_->ExecutesBefore(b_instruction, a_instruction)) { 156 return true; 157 } 158 if (IsEntryParameter(b_instruction) && 159 hlo_ordering_->ExecutesBefore(a_instruction, b_instruction)) { 160 return true; 161 } 162 // Buffers without disjoint liveness may interfere. 163 return !live_range_strictly_before(a, b) && !live_range_strictly_before(b, a); 164 } 165 166 bool BufferLiveness::MaybeLiveOut(const LogicalBuffer& buffer) const { 167 // Verify that a buffer is actually defined at the given instruction/index 168 // (eg, its not an alias of another buffer such as occurs with a bitcast). 169 TF_CHECK_OK(points_to_analysis_->VerifyBuffer(buffer)); 170 return maybe_live_out_buffers_.count(&buffer); 171 } 172 173 } // namespace xla 174
