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 */
Run(const HloModule * module,std::unique_ptr<HloOrdering> hlo_ordering)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
Analyze()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
ToString() const76 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
live_range_strictly_before(const LogicalBuffer & a,const LogicalBuffer & b) const93 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 {
IsEntryParameter(const HloInstruction * instruction)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
MayInterfere(const LogicalBuffer & a,const LogicalBuffer & b) const148 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
MaybeLiveOut(const LogicalBuffer & buffer) const166 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