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 #include "tensorflow/compiler/tf2xla/tf2xla.h"
17
18 #include <map>
19 #include <memory>
20 #include <string>
21 #include <unordered_map>
22 #include <utility>
23 #include <vector>
24
25 #include "absl/strings/str_cat.h"
26 #include "absl/strings/str_join.h"
27 #include "tensorflow/compiler/aot/aot_only_var_handle_op.h"
28 #include "tensorflow/compiler/tf2xla/graph_compiler_util.h"
29 #include "tensorflow/compiler/tf2xla/shape_util.h"
30 #include "tensorflow/compiler/tf2xla/tf2xla_util.h"
31 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
32 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
33 #include "tensorflow/compiler/xla/client/xla_computation.h"
34 #include "tensorflow/core/common_runtime/function.h"
35 #include "tensorflow/core/framework/function.h"
36 #include "tensorflow/core/framework/graph.pb.h"
37 #include "tensorflow/core/framework/graph_def_util.h"
38 #include "tensorflow/core/framework/node_def.pb.h"
39 #include "tensorflow/core/framework/op.h"
40 #include "tensorflow/core/framework/tensor_shape.h"
41 #include "tensorflow/core/framework/versions.pb.h"
42 #include "tensorflow/core/graph/algorithm.h"
43 #include "tensorflow/core/graph/graph.h"
44 #include "tensorflow/core/graph/node_builder.h"
45 #include "tensorflow/core/lib/core/errors.h"
46 #include "tensorflow/core/platform/errors.h"
47 #include "tensorflow/core/platform/logging.h"
48 #include "tensorflow/core/platform/types.h"
49 #include "tensorflow/core/util/dump_graph.h"
50
51 namespace tensorflow {
52
53 namespace {
54
55 // Converts the TensorFlow graph into an XLA computation, by executing the
56 // graph symbolically, with each op building up the XLA HLO.
ConvertGraphToXla(std::unique_ptr<Graph> graph,const tf2xla::Config & config,xla::Client * client,xla::XlaComputation * computation)57 Status ConvertGraphToXla(std::unique_ptr<Graph> graph,
58 const tf2xla::Config& config, xla::Client* client,
59 xla::XlaComputation* computation) {
60 XlaOpRegistry::RegisterCompilationKernels();
61 for (Node* node : graph->nodes()) {
62 node->set_assigned_device_name(
63 absl::StrCat("/device:", DEVICE_CPU_XLA_JIT));
64 }
65 std::vector<XlaCompiler::Argument> xla_args;
66 TF_RETURN_IF_ERROR(CreateXlaArgs(*graph, &xla_args));
67
68 PopulateXlaArgs(config, &xla_args);
69 // Compile the graph into an XLA computation.
70 XlaCompiler::Options compiler_options;
71 compiler_options.client = client;
72 compiler_options.device_type = DeviceType(DEVICE_CPU_XLA_JIT);
73 compiler_options.flib_def = &graph->flib_def();
74 compiler_options.graph_def_version = graph->versions().producer();
75 compiler_options.allow_cpu_custom_calls = true;
76
77 XlaCompiler compiler(compiler_options);
78
79 XlaCompiler::CompilationResult result;
80
81 XlaCompiler::CompileOptions options;
82 options.alias_resource_update = true;
83 TF_RETURN_IF_ERROR(compiler.CompileGraph(
84 options, "tfcompile", std::move(graph), xla_args, &result));
85 *computation = std::move(*result.computation);
86
87 int num_const_results = 0;
88 for (int i = 0, end = result.outputs.size(); i < end; ++i) {
89 // Ending up with const results (i.e. output args) is an error, since it
90 // means that one or more fetches that the user specified will be dropped
91 // from the generated function. It's most likely a configuration error,
92 // since the user shouldn't be asking for output args that end up as consts.
93 //
94 // TODO(toddw): Provide a way for the user to access const output args,
95 // e.g. perhaps hard-coded into the header, or somehow copied into the
96 // output buffers.
97 if (result.outputs[i].is_constant) {
98 ++num_const_results;
99 LOG(ERROR) << "ConstRetVal index:" << i
100 << " value:" << result.outputs[i].constant_value.DebugString();
101 }
102 }
103 if (num_const_results > 0) {
104 return errors::Unimplemented(
105 "Conversion from TensorFlow graph to XLA resulted in ",
106 num_const_results,
107 " constant results. The configuration of "
108 "the output args (i.e. fetch ids) is probably wrong.");
109 }
110 {
111 // Verify that the readonly bits on variables are set correctly by the user.
112 std::vector<bool> updated_inputs(xla_args.size());
113 for (const XlaCompiler::ResourceUpdate& update : result.resource_updates) {
114 updated_inputs[update.input_index] = true;
115 }
116 int64_t input_index = xla_args.size() - config.variable_size();
117 for (const tf2xla::Variable& variable : config.variable()) {
118 if (variable.readonly() == updated_inputs[input_index]) {
119 return errors::InvalidArgument(
120 "Variable \"", variable.node_name(), "\" is marked as ",
121 variable.readonly() ? "" : "not ", "readonly, but is ",
122 updated_inputs[input_index] ? "" : "not ",
123 "modified by the computation.");
124 }
125 ++input_index;
126 }
127 }
128 return OkStatus();
129 }
130
ConvertVarHandlesToAotVarHandles(GraphDef * graph_def)131 Status ConvertVarHandlesToAotVarHandles(GraphDef* graph_def) {
132 auto update_var_handle_op_node = [](NodeDef& node) -> Status {
133 if (node.op() == "VarHandleOp") {
134 node.set_op(tfcompile::kXlaAotOnlyVarHandleOp);
135 const auto& it = node.attr().find("allowed_devices");
136 if (it != node.attr().end()) {
137 if (!it->second.list().s().empty()) {
138 return errors::InvalidArgument(
139 "VarHandleOp with non-empty allowed devices is not supported.");
140 }
141 node.mutable_attr()->erase("allowed_devices");
142 }
143 }
144 return OkStatus();
145 };
146 for (auto& node : *graph_def->mutable_node()) {
147 TF_RETURN_IF_ERROR(update_var_handle_op_node(node));
148 }
149 for (auto& fn : *graph_def->mutable_library()->mutable_function()) {
150 for (auto& node : *fn.mutable_node_def()) {
151 TF_RETURN_IF_ERROR(update_var_handle_op_node(node));
152 }
153 }
154 return OkStatus();
155 }
156
157 } // namespace
158
ConvertGraphDefToXla(GraphDef graph_def,const tf2xla::Config & config,xla::Client * client,xla::XlaComputation * computation)159 Status ConvertGraphDefToXla(GraphDef graph_def, const tf2xla::Config& config,
160 xla::Client* client,
161 xla::XlaComputation* computation) {
162 std::unique_ptr<Graph> graph;
163 TF_RETURN_IF_ERROR(ConvertVarHandlesToAotVarHandles(&graph_def));
164 TF_RETURN_IF_ERROR(InitGraph(graph_def, config, &graph));
165 TF_RETURN_IF_ERROR(
166 ConvertGraphToXla(std::move(graph), config, client, computation));
167 return OkStatus();
168 }
169
170 } // namespace tensorflow
171