| 1173 | } // namespace |
| 1174 | |
| 1175 | Status XlaCompiler::CompileGraph( |
| 1176 | const XlaCompiler::CompileOptions& options, string const& name, |
| 1177 | std::unique_ptr<Graph> graph, absl::Span<const XlaCompiler::Argument> args, |
| 1178 | absl::Span<const xla::XlaBuilder::InputOutputAlias> user_aliases, |
| 1179 | CompilationResult* result) { |
| 1180 | VLOG(1) << "Executing graph symbolically to populate XlaBuilder.: " << name; |
| 1181 | |
| 1182 | TF_RETURN_IF_ERROR(PropagateConstIntoFunctionalNodes( |
| 1183 | graph.get(), options_.flib_def, local_flib_def_.get())); |
| 1184 | TF_RETURN_IF_ERROR(RearrangeFunctionArguments( |
| 1185 | [this](const NameAttrList& function, const FunctionBody** fbody) { |
| 1186 | return FindFunctionBody(function, fbody); |
| 1187 | }, |
| 1188 | graph.get(), local_flib_def_.get())); |
| 1189 | if (VLOG_IS_ON(2)) { |
| 1190 | VLOG(2) << "XlaCompiler::CompileGraph: " |
| 1191 | << DumpGraphToFile(absl::StrCat("xla_compile_graph_", name), *graph, |
| 1192 | flib_runtime_->GetFunctionLibraryDefinition()); |
| 1193 | } |
| 1194 | |
| 1195 | // Report the error here if initialization failed. |
| 1196 | TF_RETURN_IF_ERROR(initialization_status_); |
| 1197 | |
| 1198 | // Detect invalid nodes. |
| 1199 | // FunctionalizeControlFlow may remove some nodes from the graph. |
| 1200 | TF_RETURN_IF_ERROR(ValidateGraph(graph.get(), *options_.flib_def, |
| 1201 | options_.device_type, name)); |
| 1202 | |
| 1203 | xla::XlaBuilder builder(name); |
| 1204 | XlaContext* context = new XlaContext(this, &builder); |
| 1205 | core::ScopedUnref context_unref(context); |
| 1206 | |
| 1207 | std::vector<XlaCompiler::Argument> real_args(args.begin(), args.end()); |
| 1208 | int token_input_index = -1; |
| 1209 | std::unique_ptr<xla::XlaOp> token_output; |
| 1210 | if (options.add_token_input_output) { |
| 1211 | // Add extra token input. |
| 1212 | token_input_index = real_args.size(); |
| 1213 | |
| 1214 | XlaCompiler::Argument token_arg; |
| 1215 | token_arg.kind = XlaCompiler::Argument::kToken; |
| 1216 | real_args.push_back(token_arg); |
| 1217 | } |
| 1218 | |
| 1219 | std::map<int, xla::OpSharding> arg_shardings; |
| 1220 | std::map<int, xla::OpSharding> retval_shardings; |
| 1221 | TF_ASSIGN_OR_RETURN(std::tie(arg_shardings, retval_shardings), |
| 1222 | ComputeArgAndRetvalShardings(*graph)); |
| 1223 | |
| 1224 | std::vector<XlaExpression> arg_expressions; |
| 1225 | TF_RETURN_IF_ERROR(BuildArguments( |
| 1226 | *graph, real_args, options.use_tuple_arg, &builder, context, |
| 1227 | arg_shardings, &arg_expressions, &result->input_mapping, |
| 1228 | &result->xla_input_shapes, options.is_entry_computation)); |
| 1229 | context->set_args(std::move(arg_expressions)); |
| 1230 | |
| 1231 | // Propagate any aliases given to us by the user. |
| 1232 | for (const xla::XlaBuilder::InputOutputAlias& alias : user_aliases) { |