1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
|
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/xla/service/logical_buffer_analysis.h"
#include <utility>
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/logging.h"
namespace xla {
namespace {
// Gather fusion instructions from 'instruction' into 'fusion_instructions'.
void GatherFusionInstructions(
HloInstruction* instruction,
std::vector<HloInstruction*>* fusion_instructions) {
CHECK_EQ(HloOpcode::kFusion, instruction->opcode());
for (auto* fused : instruction->fused_instructions()) {
if (fused->opcode() == HloOpcode::kFusion) {
GatherFusionInstructions(fused, fusion_instructions);
}
}
fusion_instructions->push_back(instruction);
}
} // namespace
/* static */ StatusOr<std::unique_ptr<LogicalBufferAnalysis>>
LogicalBufferAnalysis::Run(const HloModule* module) {
std::unique_ptr<LogicalBufferAnalysis> analysis(
new LogicalBufferAnalysis(module));
TF_RETURN_IF_ERROR(analysis->Analyze());
return std::move(analysis);
}
Status LogicalBufferAnalysis::Analyze() {
// Empirically we usually have a few more logical buffers than instructions,
// so reserve 10% more than the number of instructions to avoid frequent
// resizes.
logical_buffers_.clear();
logical_buffers_.reserve((module_->instruction_count() * 11) / 10);
// We filter out fusion computations, and get to them through fusion
// instructions. This is because it's possible to have orphaned (unreachable)
// fusion computations, and we don't want to try to assign buffers to those.
std::vector<HloInstruction*> fusion_instructions;
for (auto* computation : module_->MakeNonfusionComputations()) {
TF_RETURN_IF_ERROR(computation->Accept(this));
for (auto* instruction : computation->instructions()) {
if (instruction->opcode() != HloOpcode::kFusion) {
continue;
}
GatherFusionInstructions(instruction, &fusion_instructions);
}
}
for (auto* instruction : fusion_instructions) {
TF_RETURN_IF_ERROR(instruction->fused_expression_root()->Accept(this));
}
return Status::OK();
}
LogicalBuffer& LogicalBufferAnalysis::GetBuffer(LogicalBuffer::Id id) const {
CHECK_GE(id, 0);
CHECK_LT(id, logical_buffers_.size());
return *logical_buffers_[id];
}
LogicalBuffer& LogicalBufferAnalysis::GetBuffer(HloInstruction* instruction,
const ShapeIndex& index) const {
return *output_buffers_.at(std::make_pair(instruction, index));
}
void LogicalBufferAnalysis::NewLogicalBuffer(HloInstruction* instruction,
const ShapeIndex& index) {
CHECK_EQ(logical_buffers_.size(), next_buffer_id_);
logical_buffers_.emplace_back(
absl::make_unique<LogicalBuffer>(instruction, index, next_buffer_id_));
output_buffers_[std::make_pair(instruction, index)] =
logical_buffers_.back().get();
++next_buffer_id_;
}
Status LogicalBufferAnalysis::DefaultAction(HloInstruction* hlo_instruction) {
// Create a logical buffer for each output of the instruction.
ShapeUtil::ForEachSubshape(
hlo_instruction->shape(),
[this, hlo_instruction](const Shape& shape, const ShapeIndex& index) {
NewLogicalBuffer(hlo_instruction, index);
});
return Status::OK();
}
Status LogicalBufferAnalysis::HandleGetTupleElement(HloInstruction*) {
// GetTupleElement does not create buffers.
return Status::OK();
}
Status LogicalBufferAnalysis::HandleCopy(HloInstruction* copy) {
// The top-level buffer (index={}) for kCopy is newly created, but all other
// buffers (in the case of a tuple shape) come from the operand
NewLogicalBuffer(copy, /*index=*/{});
return Status::OK();
}
Status LogicalBufferAnalysis::HandleBitcast(HloInstruction*) {
// A kBitcast instruction aliases its operand. That is, the buffer of its
// result *is* the buffer of its operand.
return Status::OK();
}
Status LogicalBufferAnalysis::HandleDomain(HloInstruction*) {
// A kDomain instruction aliases its operand. That is, the buffer of its
// result *is* the buffer of its operand.
return Status::OK();
}
Status LogicalBufferAnalysis::HandleRecvDone(HloInstruction* recv_done) {
// RecvDone produces a two-element tuple containing the data value (which
// aliases part of its operand) and a token. Only the tuple index table and
// the token are defined by the RecvDone.
NewLogicalBuffer(recv_done, /*index=*/{});
NewLogicalBuffer(recv_done, /*index=*/{1});
return Status::OK();
}
Status LogicalBufferAnalysis::HandleSend(HloInstruction* send) {
// Send creates new buffers for the top-level tuple, the context (tuple
// element at {1}), and the token (tuple element at {2}). Tuple element at {0}
// is an alias of the Send operand, so we don't need to create a new Logical
// Buffer for that.
NewLogicalBuffer(send, /*index=*/{});
NewLogicalBuffer(send, /*index=*/{1});
NewLogicalBuffer(send, /*index=*/{2});
return Status::OK();
}
Status LogicalBufferAnalysis::HandleTuple(HloInstruction* tuple) {
// A Tuple instruction only creates the top-level buffer.
NewLogicalBuffer(tuple, /*index=*/{});
return Status::OK();
}
Status LogicalBufferAnalysis::HandleTupleSelect(HloInstruction* tuple_select) {
// Select allocates a new buffer and then shallow copies the on_true or
// on_false buffer into this new buffer.
NewLogicalBuffer(tuple_select, /*index=*/{});
return Status::OK();
}
} // namespace xla
|
