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
|
/* 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/call_inliner.h"
#include <deque>
#include "tensorflow/compiler/xla/service/call_graph.h"
#include "tensorflow/compiler/xla/service/hlo_dce.h"
#include "tensorflow/core/lib/core/errors.h"
namespace xla {
namespace {
// Traverses the callee computation, inlining cloned nodes into the caller
// computation and connecting them to producers/consumers appropriately.
// When the traversal has completed, the provided call instruction is entriely
// replaced in the caller's graph.
class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
public:
// call is the call operation -- it will be replaced with the body of the
// called computation.
explicit SubcomputationInsertionVisitor(HloInstruction* call)
: call_(call), outer_(call->parent()) {
CHECK_EQ(HloOpcode::kCall, call_->opcode());
}
// Resolves the operands to the HLO instruction in the inlined (caller) graph,
// and clones the HLO instruction into that graph with the new operands.
// If the instruction is a call, it is added to the work queue.
Status DefaultAction(HloInstruction* hlo) override {
TF_RET_CHECK(hlo->opcode() != HloOpcode::kCall);
std::vector<HloInstruction*> new_operands;
for (HloInstruction* operand : hlo->operands()) {
TF_ASSIGN_OR_RETURN(HloInstruction * new_operand, Resolve(operand));
new_operands.push_back(new_operand);
}
VLOG(1) << "Cloning HLO and adding to caller: " << hlo->ToString();
auto new_hlo = hlo->CloneWithNewOperands(hlo->shape(), new_operands);
HloInstruction* new_hlo_pointer =
outer_->AddInstruction(std::move(new_hlo));
TF_RETURN_IF_ERROR(NoteMapping(hlo, new_hlo_pointer));
// Account for control edges.
for (HloInstruction* control_predecessor : hlo->control_predecessors()) {
TF_ASSIGN_OR_RETURN(HloInstruction * new_control_predecessor,
Resolve(control_predecessor));
TF_RETURN_IF_ERROR(
new_control_predecessor->AddControlDependencyTo(new_hlo_pointer));
}
return Status::OK();
}
// Does not create new nodes for the parameter; rather, notes the mapping from
// the subcomputation parameter node to the call operands in the caller
// computation.
Status HandleParameter(HloInstruction* parameter) override {
TF_RETURN_IF_ERROR(NoteMapping(
parameter, call_->mutable_operand(parameter->parameter_number())));
return Status::OK();
}
// Wires the consumers of the call to instead point at the newly created root,
// replacing the call operation in the caller computation.
Status FinishVisit(HloInstruction* root) override {
TF_ASSIGN_OR_RETURN(HloInstruction * new_root, Resolve(root));
VLOG(1) << "Replacing all uses of " << call_->ToString()
<< " with new root " << new_root->ToString();
call_->ClearCalledComputations();
return outer_->ReplaceInstruction(call_, new_root);
}
CallInliner::InlinedInstructionMap ConsumeInstructionMap() {
return std::move(subcomputation_hlo_to_new_hlo_);
}
private:
// Resolves the callee subcomputation_hlo to the new (inline) HLO in the
// caller computation, or returns a NotFound error if that subcomputation HLO
// has not been mapped.
StatusOr<HloInstruction*> Resolve(HloInstruction* subcomputation_hlo) {
auto it = subcomputation_hlo_to_new_hlo_.find(subcomputation_hlo);
if (it == subcomputation_hlo_to_new_hlo_.end()) {
return NotFound(
"Could not find mapping from subcomputation HLO %s to a cloned HLO.",
subcomputation_hlo->ToString().c_str());
}
return it->second;
}
// Notes that the given subcomputation_hlo in the callee has been mapped to
// the (inline) new_hlo in the caller computation.
//
// Returns an error status if the subcomputation_hlo is mapped more than
// once.
Status NoteMapping(HloInstruction* subcomputation_hlo,
HloInstruction* new_hlo) {
auto result = subcomputation_hlo_to_new_hlo_.insert(
std::make_pair(subcomputation_hlo, new_hlo));
TF_RET_CHECK(result.second)
<< "A mapping for the subcomputation HLO is already present.";
return Status::OK();
}
HloInstruction* call_;
HloComputation* outer_;
CallInliner::InlinedInstructionMap subcomputation_hlo_to_new_hlo_;
};
} // namespace
/* static */ StatusOr<CallInliner::InlinedInstructionMap> CallInliner::Inline(
HloInstruction* call) {
TF_RET_CHECK(call->opcode() == HloOpcode::kCall)
<< "Instruction was not a call op: " << call->opcode();
const auto& callees = call->called_computations();
TF_RET_CHECK(callees.size() == 1);
HloComputation* callee = callees[0];
// We visit the callee, cloning its body into its caller.
SubcomputationInsertionVisitor visitor(call);
TF_RETURN_IF_ERROR(callee->Accept(&visitor));
return visitor.ConsumeInstructionMap();
}
StatusOr<bool> CallInliner::Run(HloModule* module) {
std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module);
// Because call graph nodes are visited in post-order (callees before callers)
// we'll always inline kCalls into their callers in the appropriate order.
bool did_mutate = false;
TF_RETURN_IF_ERROR(
call_graph->VisitNodes([&](const CallGraphNode& node) -> Status {
for (const CallSite& callsite : node.caller_callsites()) {
VLOG(1) << "Visiting callsite: " << callsite.ToString();
if (callsite.instruction()->opcode() == HloOpcode::kCall) {
HloInstruction* call = callsite.instruction();
TF_RETURN_IF_ERROR(Inline(call).status());
did_mutate = true;
}
}
return Status::OK();
}));
if (did_mutate) {
// Run DCE to remove called computations which are now becoming unused.
// This can result then in problems if within the called computation, there
// were send/recv instructions, which the module group verifier will flag as
// error findingthe same channel ID used for multiple send/recv
// instructions.
TF_RETURN_IF_ERROR(HloDCE().Run(module).status());
}
return did_mutate;
}
} // namespace xla
|
