| 409 | } |
| 410 | |
| 411 | Status CapturedFunction::Instantiate( |
| 412 | IteratorContext* ctx, std::unique_ptr<InstantiatedCapturedFunction>* |
| 413 | instantiated_captured_function) { |
| 414 | // The context's runtime will be used for all subsequent calls. |
| 415 | FunctionLibraryRuntime* lib = ctx->flr(); |
| 416 | FunctionLibraryRuntime::InstantiateOptions inst_opts; |
| 417 | // TODO(b/141576771): Propagate ctx ConfigProto into inst_opts.config_proto. |
| 418 | inst_opts.lib_def = metadata_->lib_def(); |
| 419 | inst_opts.create_kernels_eagerly = true; |
| 420 | if (!metadata_->use_inter_op_parallelism()) { |
| 421 | inst_opts.executor_type = "SINGLE_THREADED_EXECUTOR"; |
| 422 | } |
| 423 | inst_opts.is_multi_device_function = metadata_->is_multi_device_function(); |
| 424 | |
| 425 | // We infer the target device from the function library runtime. |
| 426 | DCHECK(lib->device() != nullptr); |
| 427 | inst_opts.target = lib->device()->name(); |
| 428 | |
| 429 | if (metadata_->is_multi_device_function()) { |
| 430 | // Compute devices of non-captured inputs. |
| 431 | // |
| 432 | // We infer the number of non-captured inputs by subtracting the number |
| 433 | // of captured inputs from the number of input arguments and we infer the |
| 434 | // input devices from the function library runtime. |
| 435 | const FunctionDef* fdef = |
| 436 | metadata_->lib_def()->Find(metadata_->func().name()); |
| 437 | if (fdef == nullptr) { |
| 438 | return errors::InvalidArgument( |
| 439 | "Failed to find function ", metadata_->func().name(), |
| 440 | " in function library: ", lib->GetFunctionLibraryDefinition()); |
| 441 | } |
| 442 | size_t num_non_captured_inputs = |
| 443 | fdef->signature().input_arg_size() - captured_inputs_.size(); |
| 444 | for (size_t i = 0; i < num_non_captured_inputs; ++i) { |
| 445 | inst_opts.input_devices.push_back(inst_opts.target); |
| 446 | } |
| 447 | // Compute devices of captured inputs. |
| 448 | // TODO(jsimsa): Correctly handle tensors on devices other than CPU:0. |
| 449 | Device* cpu_device; |
| 450 | TF_RETURN_IF_ERROR(lib->device_mgr()->LookupDevice("CPU:0", &cpu_device)); |
| 451 | std::unordered_map<int, DtypeAndPartialTensorShape>& |
| 452 | input_resource_variable_dtypes_and_shapes = |
| 453 | inst_opts.input_resource_dtypes_and_shapes; |
| 454 | for (size_t i = 0; i < captured_inputs_.size(); ++i) { |
| 455 | const auto& input = captured_inputs_[i]; |
| 456 | DataType dtype = input.dtype(); |
| 457 | if (dtype == DT_RESOURCE) { |
| 458 | const ResourceHandle& handle = input.flat<ResourceHandle>()(0); |
| 459 | inst_opts.input_devices.push_back(handle.device()); |
| 460 | const auto& dtypes_and_shapes = handle.dtypes_and_shapes(); |
| 461 | // Set dtypes and shapes for resource variable inputs. |
| 462 | if (!dtypes_and_shapes.empty()) { |
| 463 | input_resource_variable_dtypes_and_shapes[num_non_captured_inputs + |
| 464 | i] = |
| 465 | dtypes_and_shapes.at(0); |
| 466 | } |
| 467 | } else if (MTypeFromDType(dtype) == HOST_MEMORY) { |
| 468 | inst_opts.input_devices.push_back(cpu_device->name()); |