| 545 | } |
| 546 | |
| 547 | std::unique_ptr<Graph> XlaCompiler::GetGraph(const FunctionBody* fbody) { |
| 548 | std::unique_ptr<Graph> graph(new Graph(options_.flib_def)); |
| 549 | CopyGraph(*fbody->graph, graph.get()); |
| 550 | |
| 551 | // Performs a first function inlining pass before shape inference, since |
| 552 | // otherwise shape inference can't see inside functions and a comprehensive |
| 553 | // shape_map, including function ops, is needed to constant-propagate Shape |
| 554 | // Ops below. |
| 555 | auto flags = GetBuildXlaOpsPassFlags(); |
| 556 | OptimizerOptions opts; |
| 557 | opts.set_opt_level(OptimizerOptions::L0); |
| 558 | opts.set_do_common_subexpression_elimination(false); |
| 559 | opts.set_do_function_inlining(true); |
| 560 | opts.set_do_constant_folding(!flags->tf_xla_disable_constant_folding); |
| 561 | GraphOptimizer optimizer(opts); |
| 562 | // Do not constant fold nodes that output DT_VARIANT type tensors. |
| 563 | // XLA does not support Const nodes of Variant type since it needs |
| 564 | // to know the original ops to be able to compile them to the relevant |
| 565 | // XLA form. |
| 566 | // TODO(srbs): This filter is a little conservative. E.g. a subgraph of |
| 567 | // the form: |
| 568 | // Const |
| 569 | // | |
| 570 | // EmptyTensorList -> TensorListPushBack -> TensorListPopBack -> Op |
| 571 | // | |
| 572 | // (Discard popped list) |
| 573 | // |
| 574 | // Would have been reduced to "Const -> Op" without this filter. |
| 575 | // However since we are only allowed to specify the filter at the "Node" |
| 576 | // level there is no good way to allow the above behavior. So we |
| 577 | // disallow any sort of constant folding on Variant nodes for now. |
| 578 | auto cf_consider_fn = [](const Node* n) { |
| 579 | for (const auto& output_arg : n->op_def().output_arg()) { |
| 580 | if (output_arg.type() == DT_VARIANT) { |
| 581 | return false; |
| 582 | } |
| 583 | } |
| 584 | return true; |
| 585 | }; |
| 586 | GraphOptimizer::Options graph_optimizer_options; |
| 587 | graph_optimizer_options.cf_consider_fn = cf_consider_fn; |
| 588 | graph_optimizer_options.inline_multi_device_functions = true; |
| 589 | graph_optimizer_options.inline_impl_selection_group_functions = true; |
| 590 | optimizer.Optimize(flib_runtime_, flib_runtime_->env(), |
| 591 | /*device=*/nullptr, &graph, graph_optimizer_options); |
| 592 | |
| 593 | // Run shape inference on the graph and optimize the graph again. |
| 594 | GraphShapeInfo shape_info; |
| 595 | InferShapes(graph.get(), /*arg_shapes=*/{}, |
| 596 | flib_runtime_->GetFunctionLibraryDefinition(), &shape_info) |
| 597 | .IgnoreError(); |
| 598 | auto node_name_index = graph->BuildNodeNameIndex(); |
| 599 | std::unordered_map<string, std::vector<PartialTensorShape>> shape_map; |
| 600 | for (const auto& node_shape_info : shape_info) { |
| 601 | const string& node_name = node_shape_info.first; |
| 602 | const std::vector<InferredShape>& output_shapes = node_shape_info.second; |
| 603 | const auto& node_iter = node_name_index.find(node_name); |
| 604 | if (node_iter != node_name_index.end()) { |
no test coverage detected