| 616 | } |
| 617 | |
| 618 | Status XlaCompiler::CompileFunction( |
| 619 | const XlaCompiler::CompileOptions& options, |
| 620 | const NameAttrList& fn_name_attrs, |
| 621 | absl::Span<const XlaCompiler::Argument> args, |
| 622 | XlaCompiler::CompilationResult* result) { |
| 623 | const string function_id = |
| 624 | Canonicalize(fn_name_attrs.name(), AttrSlice(&fn_name_attrs.attr())); |
| 625 | VLOG(1) << "XlaCompiler::CompileFunction " << function_id; |
| 626 | |
| 627 | const std::vector<XlaCompiler::Argument> arg_vector(args.begin(), args.end()); |
| 628 | auto it = cache_.find({function_id, arg_vector}); |
| 629 | if (it != cache_.end()) { |
| 630 | *result = it->second; |
| 631 | return Status::OK(); |
| 632 | } |
| 633 | |
| 634 | const FunctionBody* fbody; |
| 635 | TF_RETURN_IF_ERROR(FindFunctionBody(fn_name_attrs, &fbody)); |
| 636 | |
| 637 | TF_RETURN_WITH_CONTEXT_IF_ERROR( |
| 638 | CheckSignature(fbody->arg_types, args), |
| 639 | "Signature check failure while compiling: ", fn_name_attrs.name()); |
| 640 | |
| 641 | // Set shapes for _Arg nodes. They are useful for constant folding (e.g. an |
| 642 | // Xla op requires a compile-time constant input, and that input is shape of |
| 643 | // an _Arg node. |
| 644 | for (int i = 0; i < args.size(); i++) { |
| 645 | // Skip resource variables and tensor lists. |
| 646 | DataType dtype; |
| 647 | TF_RETURN_IF_ERROR(GetNodeAttr(fbody->arg_nodes[i]->def(), "T", &dtype)); |
| 648 | if (dtype == DT_RESOURCE || dtype == DT_VARIANT) { |
| 649 | continue; |
| 650 | } |
| 651 | |
| 652 | if (absl::holds_alternative<xla::Shape>(args[i].shape)) { |
| 653 | xla::Shape xla_shape = absl::get<xla::Shape>(args[i].shape); |
| 654 | TensorShape tensor_shape; |
| 655 | if (XLAShapeToTensorShape(xla_shape, &tensor_shape).ok()) { |
| 656 | fbody->arg_nodes[i]->ClearAttr("_output_shapes"); |
| 657 | fbody->arg_nodes[i]->AddAttr("_output_shapes", |
| 658 | std::vector<TensorShape>{tensor_shape}); |
| 659 | } |
| 660 | } else { |
| 661 | TensorShape tensor_shape = absl::get<TensorShape>(args[i].shape); |
| 662 | fbody->arg_nodes[i]->ClearAttr("_output_shapes"); |
| 663 | fbody->arg_nodes[i]->AddAttr("_output_shapes", |
| 664 | std::vector<TensorShape>{tensor_shape}); |
| 665 | } |
| 666 | } |
| 667 | |
| 668 | std::unique_ptr<Graph> graph = GetGraph(fbody); |
| 669 | |
| 670 | // _Arg and _Retval nodes don't exist in the stored subgraph for the function; |
| 671 | // they are added by the function body looked up. Therefore, they don't have |
| 672 | // core assignments here. |
| 673 | // Attempt to assign a core to each _Retval and _Arg. Chooses the |
| 674 | // lowest-numbered core that consumes the argument. We choose the |
| 675 | // lowest-numbered core so the assignment is deterministic. |