| 498 | } // namespace |
| 499 | |
| 500 | Status GraphExecutionState::PruneGraph( |
| 501 | const BuildGraphOptions& options, Graph* graph, |
| 502 | subgraph::RewriteGraphMetadata* out_rewrite_metadata) { |
| 503 | std::vector<std::unique_ptr<subgraph::PruneRewrite>> feed_rewrites; |
| 504 | feed_rewrites.reserve(options.callable_options.feed_size()); |
| 505 | std::vector<std::unique_ptr<subgraph::PruneRewrite>> fetch_rewrites; |
| 506 | fetch_rewrites.reserve(options.callable_options.fetch_size()); |
| 507 | if (options.use_function_convention) { |
| 508 | std::vector<TensorAndDevice> tensors_and_devices; |
| 509 | for (int i = 0; i < options.callable_options.feed_size(); ++i) { |
| 510 | // WARNING: feed MUST be a reference, since ArgFeedRewrite and |
| 511 | // tensors_and_devices holds on to its address. |
| 512 | const string& feed = options.callable_options.feed(i); |
| 513 | const DeviceAttributes* device_info; |
| 514 | TF_RETURN_IF_ERROR(LookupDevice(*device_set_, feed, |
| 515 | options.callable_options.feed_devices(), |
| 516 | &device_info)); |
| 517 | feed_rewrites.emplace_back( |
| 518 | new subgraph::ArgFeedRewrite(&feed, device_info, i)); |
| 519 | tensors_and_devices.push_back({ParseTensorName(feed), device_info}); |
| 520 | } |
| 521 | if (!options.callable_options.fetch_devices().empty() && |
| 522 | !options.callable_options.fetch_skip_sync()) { |
| 523 | return errors::Unimplemented( |
| 524 | "CallableOptions.fetch_skip_sync = false is not yet implemented. You " |
| 525 | "can set it to true instead, but MUST ensure that Device::Sync() is " |
| 526 | "invoked on the Device corresponding to the fetched tensor before " |
| 527 | "dereferencing the Tensor's memory."); |
| 528 | } |
| 529 | for (int i = 0; i < options.callable_options.fetch_size(); ++i) { |
| 530 | // WARNING: fetch MUST be a reference, since RetvalFetchRewrite and |
| 531 | // tensors_and_devices holds on to its address. |
| 532 | const string& fetch = options.callable_options.fetch(i); |
| 533 | const DeviceAttributes* device_info; |
| 534 | TF_RETURN_IF_ERROR(LookupDevice(*device_set_, fetch, |
| 535 | options.callable_options.fetch_devices(), |
| 536 | &device_info)); |
| 537 | fetch_rewrites.emplace_back( |
| 538 | new subgraph::RetvalFetchRewrite(&fetch, device_info, i)); |
| 539 | tensors_and_devices.push_back({ParseTensorName(fetch), device_info}); |
| 540 | } |
| 541 | TF_RETURN_IF_ERROR( |
| 542 | ValidateFeedAndFetchDevices(*graph, tensors_and_devices)); |
| 543 | } else { |
| 544 | if (!options.callable_options.feed_devices().empty() || |
| 545 | !options.callable_options.fetch_devices().empty()) { |
| 546 | return errors::Unimplemented( |
| 547 | "CallableOptions::feed_devices and CallableOptions::fetch_devices " |
| 548 | "to configure feeding/fetching tensors to/from device memory is not " |
| 549 | "yet supported when using a remote session."); |
| 550 | } |
| 551 | const DeviceAttributes* device_info = |
| 552 | &device_set_->client_device()->attributes(); |
| 553 | for (const string& feed : options.callable_options.feed()) { |
| 554 | feed_rewrites.emplace_back( |
| 555 | new subgraph::RecvFeedRewrite(&feed, device_info)); |
| 556 | } |
| 557 | for (const string& fetch : options.callable_options.fetch()) { |
no test coverage detected