| 873 | } |
| 874 | |
| 875 | Status Service::Execute(const ExecuteRequest* arg, ExecuteResponse* result) { |
| 876 | VLOG(1) << "running execute request"; |
| 877 | if (!arg->has_handle()) { |
| 878 | return InvalidArgument("execution handle should not be empty"); |
| 879 | } |
| 880 | TF_ASSIGN_OR_RETURN(auto executable, |
| 881 | compilation_cache_.LookUp(arg->handle())); |
| 882 | |
| 883 | TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, |
| 884 | SingleComputationDeviceHandle())); |
| 885 | TF_ASSIGN_OR_RETURN( |
| 886 | std::vector<std::vector<const ShapedBuffer*>> replicated_arguments, |
| 887 | ResolveAndValidateArguments(arg->arguments(), replicas)); |
| 888 | |
| 889 | // Check that the replicated_arguments has the same shape and layout as the |
| 890 | // module config used when creating the executable. |
| 891 | const int64 num_module_args = |
| 892 | executable->module_config().entry_computation_layout().parameter_count(); |
| 893 | if (num_module_args != arg->arguments_size()) { |
| 894 | return InvalidArgument( |
| 895 | "The executable expects %lld arguments, but sees %lld.", |
| 896 | num_module_args, arg->arguments_size()); |
| 897 | } |
| 898 | for (int64 i = 0; i < num_module_args; i++) { |
| 899 | const Shape& shape_module = |
| 900 | executable->module_config().entry_computation_layout().parameter_shape( |
| 901 | i); |
| 902 | const Shape& shape_arg = replicated_arguments.front()[i]->on_host_shape(); |
| 903 | if (!ShapeUtil::Equal(shape_module, shape_arg)) { |
| 904 | return InvalidArgumentStrCat( |
| 905 | "The executable expects the ", i, "th argument in shape ", |
| 906 | ShapeUtil::HumanStringWithLayout(shape_module), " but sees ", |
| 907 | ShapeUtil::HumanStringWithLayout(shape_arg)); |
| 908 | } |
| 909 | } |
| 910 | |
| 911 | TF_ASSIGN_OR_RETURN(auto stream, |
| 912 | execute_backend_->BorrowStream( |
| 913 | execute_backend_->default_stream_executor())); |
| 914 | HloSnapshot snapshot; |
| 915 | if (executable->dumping_snapshot()) { |
| 916 | *snapshot.mutable_hlo() = *executable->hlo_proto(); |
| 917 | snapshot.set_execution_platform(execute_backend_->platform()->Name()); |
| 918 | TF_RETURN_IF_ERROR( |
| 919 | RecordArguments(replicated_arguments.front(), stream.get(), |
| 920 | execute_backend_->transfer_manager(), &snapshot)); |
| 921 | } |
| 922 | |
| 923 | TF_ASSIGN_OR_RETURN( |
| 924 | *result->mutable_output(), |
| 925 | ExecuteAndRegisterResult(executable.get(), replicated_arguments, |
| 926 | execute_backend_.get(), |
| 927 | SingleComputationDeviceHandle(), |
| 928 | "result of " + executable->module().name(), |
| 929 | result->mutable_profile())); |
| 930 | |
| 931 | if (executable->dumping_snapshot()) { |
| 932 | TF_ASSIGN_OR_RETURN( |