1 /* Copyright 2018 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 #include "tensorflow/compiler/jit/encapsulate_xla_computations_pass.h"
17
18 #include <functional>
19 #include <string>
20
21 #include "absl/algorithm/container.h"
22 #include "absl/container/flat_hash_set.h"
23 #include "absl/memory/memory.h"
24 #include "absl/strings/ascii.h"
25 #include "absl/strings/str_cat.h"
26 #include "tensorflow/compiler/jit/defs.h"
27 #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
28 #include "tensorflow/compiler/jit/xla_cluster_util.h"
29 #include "tensorflow/compiler/xla/status_macros.h"
30 #include "tensorflow/core/framework/node_def.pb.h"
31 #include "tensorflow/core/framework/types.h"
32 #include "tensorflow/core/graph/graph_node_util.h"
33 #include "tensorflow/core/lib/core/stringpiece.h"
34 #include "tensorflow/core/lib/hash/hash.h"
35 #include "tensorflow/core/lib/strings/proto_serialization.h"
36 #include "tensorflow/core/lib/strings/str_util.h"
37 #include "tensorflow/core/platform/fingerprint.h"
38 #include "tensorflow/core/platform/statusor.h"
39 #include "tensorflow/core/util/dump_graph.h"
40
41 namespace tensorflow {
42
43 namespace {
44
45 const char* const kXlaClusterOutput = "XlaClusterOutput";
46
IsCpuGpuCompile(const Graph * graph)47 bool IsCpuGpuCompile(const Graph* graph) {
48 for (Node* n : graph->nodes()) {
49 string name;
50 // Only consider nodes being compiled.
51 if (!TryGetNodeAttr(n->attrs(), kXlaClusterIdAttr, &name)) continue;
52 // Early return for any node with a device that is not a CPU or GPU.
53 DeviceNameUtils::ParsedName parsed;
54 if (DeviceNameUtils::ParseFullName(n->requested_device(), &parsed)) {
55 if (parsed.type != DEVICE_CPU && parsed.type != DEVICE_GPU) {
56 return false;
57 }
58 }
59 }
60 return true;
61 }
62
63 // Checks if a graph node is marked to be a guaranteed constant.
is_guaranteed_constant(const Node & n)64 bool is_guaranteed_constant(const Node& n) {
65 bool guaranteed_constant = false;
66 if (!TryGetNodeAttr(n.attrs(), "_is_guaranteed_constant",
67 &guaranteed_constant)) {
68 return false;
69 }
70 return guaranteed_constant;
71 }
72
73 // Finds the `index` of an _Arg or _Retval node.
GetIndexAttr(const Node & n,int num_args,int * index)74 Status GetIndexAttr(const Node& n, int num_args, int* index) {
75 TF_RETURN_IF_ERROR(GetNodeAttr(n.attrs(), "index", index));
76 if (*index < 0 || *index >= num_args) {
77 return errors::InvalidArgument("Invalid ", n.type_string(), " number ",
78 *index);
79 }
80 return OkStatus();
81 }
82
83 // Returns the data type of the destination of an edge.
EdgeType(const Edge * edge)84 DataType EdgeType(const Edge* edge) {
85 return edge->dst()->input_type(edge->dst_input());
86 }
87
88 // Adds the control inputs of `node` to `*deps`.
AddControlInputs(const Node & node,absl::flat_hash_set<Node * > * deps)89 void AddControlInputs(const Node& node, absl::flat_hash_set<Node*>* deps) {
90 for (const Edge* edge : node.in_edges()) {
91 if (edge->IsControlEdge()) {
92 deps->insert(edge->src());
93 }
94 }
95 }
96
97 // Adds the control outputs of `node` to `*deps`.
AddControlOutputs(const Node & node,absl::flat_hash_set<Node * > * deps)98 void AddControlOutputs(const Node& node, absl::flat_hash_set<Node*>* deps) {
99 for (const Edge* edge : node.out_edges()) {
100 if (edge->IsControlEdge()) {
101 deps->insert(edge->dst());
102 }
103 }
104 }
105
106 // Rewrite function to be passed to EncapsulateSubgraphsInFunctions that sorts
107 // the arguments into the order expected by XlaLaunch computations:
108 // 1) arguments
109 // 2) resource variable arguments
110 // See the documentation of EncapsulateSubgraphsInFunctions for the meaning
111 // of the arguments.
112 //
113 // TODO(b/113166435): Ordering constraints on XlaLaunch op can be relaxed.
RewriteSubgraph(const std::vector<OutputTensor> & arg_source_tensors,std::unique_ptr<Graph> * graph_ptr,std::vector<int> * input_permutation,std::vector<int> * output_permutation,NodeDef * call_def)114 Status RewriteSubgraph(const std::vector<OutputTensor>& arg_source_tensors,
115 std::unique_ptr<Graph>* graph_ptr,
116 std::vector<int>* input_permutation,
117 std::vector<int>* output_permutation,
118 NodeDef* call_def) {
119 Graph* graph = graph_ptr->get();
120 const int num_args = input_permutation->size();
121 const int num_retvals = output_permutation->size();
122
123 std::vector<Node*> args;
124 std::vector<Node*> retvals;
125 args.reserve(num_args);
126 retvals.reserve(num_retvals);
127 for (Node* n : graph->nodes()) {
128 if (n->type_string() == "_Arg") {
129 // Check if this is a guaranteed constant.
130 if (is_guaranteed_constant(*n)) {
131 return errors::InvalidArgument(
132 "Guaranteed constants are not supported (", n->name(), ")");
133 }
134 args.push_back(n);
135 } else if (n->type_string() == "_Retval") {
136 retvals.push_back(n);
137 }
138 }
139
140 if (std::find(args.begin(), args.end(), nullptr) != args.end()) {
141 return errors::InvalidArgument("Missing or non-consecutive arguments");
142 }
143
144 // Reorders the arguments.
145 std::sort(args.begin(), args.end(), [&](Node* a, Node* b) {
146 // Non-resources appear before resources
147 bool a_is_resource = (a->output_type(0) == DT_RESOURCE);
148 bool b_is_resource = (b->output_type(0) == DT_RESOURCE);
149 // Uses the name as a tiebreaker so the output is deterministic.
150 StringPiece a_name(a->name());
151 StringPiece b_name(b->name());
152 return std::tie(a_is_resource, a_name) < std::tie(b_is_resource, b_name);
153 });
154
155 // Sorts the retvals by name so the order is deterministic.
156 std::sort(retvals.begin(), retvals.end(),
157 [](Node* a, Node* b) { return a->name() < b->name(); });
158
159 // Computes the permutation to produce the correct argument order, and update
160 // the argument indices.
161 int variable_start_index = num_args;
162 for (int i = 0; i < num_args; ++i) {
163 int index;
164 TF_RETURN_IF_ERROR(GetIndexAttr(*args[i], num_args, &index));
165 if (args[i]->output_type(0) == DT_RESOURCE &&
166 variable_start_index == num_args) {
167 variable_start_index = i;
168 }
169 (*input_permutation)[index] = i;
170 args[i]->AddAttr("index", i);
171 }
172 VLOG(4) << "variable_start_index: " << variable_start_index;
173
174 // Computes the permutation to produce the correct retval order, and update
175 // the argument indices.
176 for (int i = 0; i < num_retvals; ++i) {
177 int index;
178 TF_RETURN_IF_ERROR(GetIndexAttr(*retvals[i], num_retvals, &index));
179 (*output_permutation)[index] = i;
180 retvals[i]->AddAttr("index", i);
181 }
182
183 AddNodeAttr(kXlaClusterIdAttr, call_def->name(), call_def);
184 AddNodeAttr("_variable_start_index", variable_start_index, call_def);
185
186 // Uniquify the function name by computing a fingerprint of the function.
187 // Nondeterminism in serialization would not lead to incorrect results, but
188 // may cause spurious cache misses.
189 TF_ASSIGN_OR_RETURN(uint64 fingerprint, FingerprintGraph(*graph));
190 VLOG(1) << "Subgraph fingerprint:" << fingerprint;
191 call_def->set_op(absl::StrCat(call_def->op(), "_", fingerprint));
192 return OkStatus();
193 }
194
195 } // namespace
196
Encapsulate(std::unique_ptr<Graph> * graph,FunctionLibraryDefinition * flib_def)197 /*static*/ Status EncapsulateXlaComputationsPass::Encapsulate(
198 std::unique_ptr<Graph>* graph, FunctionLibraryDefinition* flib_def) {
199 // Check for undeclared outputs before Encapsulation, so we can give a better
200 // error message.
201 // TODO(phawkins): merge this with the encapsulation code to avoid the extra
202 // O(n) pass over the edges.
203 for (const Edge* e : (*graph)->edges()) {
204 if (!e->IsControlEdge() &&
205 e->src()->attrs().Find(kXlaClusterIdAttr) != nullptr &&
206 e->dst()->attrs().Find(kXlaClusterIdAttr) == nullptr &&
207 e->dst()->type_string() != kXlaClusterOutput) {
208 return errors::InvalidArgument(
209 "Undeclared output of XLA computation. Some common causes of this "
210 "error are: 1) variable initializers that depend on the XLA "
211 "computation; 2) gradient computations that depend on the XLA "
212 "computation, which can be mitigated by moving gradient computations "
213 "inside XLA computation. Offending edge: ",
214 e->src()->name(), ":", e->src_output(), " -> ", e->dst()->name(), ":",
215 e->dst_input());
216 }
217 }
218
219 auto output = std::make_unique<Graph>((*graph)->op_registry());
220 TF_RETURN_WITH_CONTEXT_IF_ERROR(
221 EncapsulateSubgraphsInFunctions(
222 kXlaClusterIdAttr, **graph, RewriteSubgraph,
223 /*reuse_existing_functions=*/true, &output, flib_def),
224 "EncapsulateXlaComputationsPass failed");
225 graph->swap(output);
226 return OkStatus();
227 }
228
BuildXlaLaunchOps(Graph * graph,const std::function<StatusOr<bool> (const Node &)> & is_xla_launch_node,const std::function<StatusOr<XlaFunctionInfo> (const Node &)> & get_xla_function_info,const bool add_edges_to_output_of_downstream_nodes)229 /*static*/ Status EncapsulateXlaComputationsPass::BuildXlaLaunchOps(
230 Graph* graph,
231 const std::function<StatusOr<bool>(const Node&)>& is_xla_launch_node,
232 const std::function<StatusOr<XlaFunctionInfo>(const Node&)>&
233 get_xla_function_info,
234 const bool add_edges_to_output_of_downstream_nodes) {
235 // Finds all of the XlaLaunch function calls, to avoid mutating the graph
236 // while iterating.
237 std::vector<Node*> launch_nodes;
238 for (Node* n : graph->nodes()) {
239 TF_ASSIGN_OR_RETURN(const bool is_xla_launch_node, is_xla_launch_node(*n));
240 if (is_xla_launch_node) launch_nodes.push_back(n);
241 }
242
243 // Replaces each launch function call together with its neighboring
244 // XlaClusterOutput nodes with a XlaLaunch node.
245 for (Node* launch : launch_nodes) {
246 TF_ASSIGN_OR_RETURN(const XlaFunctionInfo xla_function_info,
247 get_xla_function_info(*launch));
248
249 std::vector<const Edge*> in_edges;
250 TF_RETURN_IF_ERROR(launch->input_edges(&in_edges));
251
252 const int num_inputs = in_edges.size();
253 const int variable_start_index = xla_function_info.variable_start_index;
254 const int num_variables = num_inputs - variable_start_index;
255 const int num_args = variable_start_index;
256
257 VLOG(4) << "Launch node '" << launch->name() << "'"
258 << " input edges: " << in_edges.size() << " num_args: " << num_args
259 << " num_variables: " << num_variables;
260
261 std::vector<Node*> nodes_to_remove = {launch};
262
263 // Data and control inputs to the new XlaLaunch node.
264 std::vector<std::pair<Node*, int>> data_inputs(num_inputs);
265 absl::flat_hash_set<Node*> control_inputs;
266 DataTypeVector arg_types(num_args);
267
268 AddControlInputs(*launch, &control_inputs);
269
270 for (int i = 0; i < num_args; ++i) {
271 const Edge* edge = in_edges[i];
272 data_inputs[i] = {edge->src(), edge->src_output()};
273 arg_types[i] = EdgeType(edge);
274 }
275
276 // Appends the variable inputs.
277 for (int i = 0; i < num_variables; ++i) {
278 int pos = variable_start_index + i;
279 const Edge* edge = in_edges[pos];
280 data_inputs[pos] = {edge->src(), edge->src_output()};
281 }
282
283 // Outputs.
284 const int num_outputs = launch->output_types().size();
285 absl::flat_hash_set<Node*> control_outputs;
286 std::vector<std::vector<std::pair<Node*, int>>> data_outputs(num_outputs);
287 DataTypeVector output_types(num_outputs);
288
289 for (const Edge* le : launch->out_edges()) {
290 if (le->IsControlEdge()) {
291 control_outputs.insert(le->dst());
292 } else {
293 TF_RET_CHECK(le->src_output() < num_outputs);
294 Node* output_node = le->dst();
295
296 if (add_edges_to_output_of_downstream_nodes) {
297 TF_RET_CHECK(output_node->type_string() == kXlaClusterOutput)
298 << le->DebugString();
299 nodes_to_remove.push_back(output_node);
300
301 for (const Edge* oe : output_node->out_edges()) {
302 TF_RET_CHECK(!oe->IsControlEdge());
303 data_outputs[le->src_output()].push_back(
304 {oe->dst(), oe->dst_input()});
305 }
306
307 AddControlOutputs(*output_node, &control_outputs);
308 } else {
309 data_outputs[le->src_output()].push_back(
310 {le->dst(), le->dst_input()});
311 }
312 output_types[le->src_output()] = output_node->input_type(0);
313 }
314 }
315
316 NodeDef def;
317 def.set_name(launch->name());
318 MergeDebugInfo(NodeDebugInfo(launch->def()), &def);
319
320 // Target the XLA CPU/GPU backends.
321 VLOG(2) << "Replacing with XlaLaunch";
322 VLOG(2) << "Device is " << launch->requested_device();
323 def.set_op("XlaLaunch");
324 def.set_device(launch->requested_device());
325 AddNodeAttr("Tconstants", DataTypeVector{}, &def);
326 AddNodeAttr("Targs", arg_types, &def);
327 AddNodeAttr("Nresources", num_variables, &def);
328 AddNodeAttr("Tresults", output_types, &def);
329 NameAttrList function;
330 function.set_name(xla_function_info.function_name);
331 AddNodeAttr("function", function, &def);
332
333 for (Node* node : nodes_to_remove) {
334 VLOG(2) << "Deleting node " << node->DebugString();
335 // Ensure that we do not attempt to add control edges to nodes that are
336 // deleted.
337 control_inputs.erase(node);
338 control_outputs.erase(node);
339 graph->RemoveNode(node);
340 }
341
342 TF_ASSIGN_OR_RETURN(Node * xla_launch, graph->AddNode(def));
343 for (int i = 0, end = data_inputs.size(); i < end; ++i) {
344 graph->AddEdge(data_inputs[i].first, data_inputs[i].second, xla_launch,
345 i);
346 }
347 for (Node* n : control_inputs) {
348 graph->AddControlEdge(n, xla_launch);
349 }
350 for (int i = 0, end = data_outputs.size(); i < end; ++i) {
351 for (const auto& successor : data_outputs[i]) {
352 graph->AddEdge(xla_launch, i, successor.first, successor.second);
353 }
354 }
355 for (Node* n : control_outputs) {
356 graph->AddControlEdge(xla_launch, n);
357 }
358 }
359 return OkStatus();
360 }
361
BuildXlaLaunchOps(Graph * graph)362 /*static*/ Status EncapsulateXlaComputationsPass::BuildXlaLaunchOps(
363 Graph* graph) {
364 const auto is_xla_launch_node = [](const Node& node) -> StatusOr<bool> {
365 const string& name = GetNodeAttrString(node.attrs(), kXlaClusterIdAttr);
366 return !name.empty();
367 };
368
369 const auto get_xla_function_info =
370 [](const Node& node) -> StatusOr<XlaFunctionInfo> {
371 XlaFunctionInfo result;
372 TF_RETURN_IF_ERROR(GetNodeAttr(node.attrs(), "_variable_start_index",
373 &result.variable_start_index));
374 result.function_name = node.type_string();
375 return result;
376 };
377 return BuildXlaLaunchOps(graph, is_xla_launch_node, get_xla_function_info,
378 /*add_edges_to_output_of_downstream_nodes=*/true);
379 }
380
Run(const GraphOptimizationPassOptions & options)381 Status EncapsulateXlaComputationsPass::Run(
382 const GraphOptimizationPassOptions& options) {
383 VLOG(1) << "EncapsulateXlaComputations(): "
384 << DumpGraphToFile("encapsulate_xla_computations_before",
385 **options.graph, options.flib_def);
386
387 const char* additional_help =
388 IsCpuGpuCompile(options.graph->get())
389 ? xla::status_macros::kPossibleAutoJitAlternative
390 : "";
391
392 TF_RETURN_WITH_CONTEXT_IF_ERROR(Encapsulate(options.graph, options.flib_def),
393 additional_help);
394 VLOG(1) << "EncapsulateXlaComputations() half-way: "
395 << DumpGraphToFile("encapsulate_xla_computations_halfway",
396 **options.graph, options.flib_def);
397
398 TF_RETURN_WITH_CONTEXT_IF_ERROR(BuildXlaLaunchOps(options.graph->get()),
399 additional_help);
400 VLOG(1) << "EncapsulateXlaComputations() finished: "
401 << DumpGraphToFile("encapsulate_xla_computations_after",
402 **options.graph, options.flib_def);
403 return OkStatus();
404 }
405
406 } // namespace tensorflow
407