| 84 | } |
| 85 | |
| 86 | StatusOr<std::vector<ScopedShapedBuffer>> Executable::ExecuteOnStreams( |
| 87 | absl::Span<const ServiceExecutableRunOptions> run_options, |
| 88 | absl::Span<const absl::Span<const ShapedBuffer* const>> arguments) { |
| 89 | TF_RET_CHECK(run_options.size() == arguments.size()); |
| 90 | |
| 91 | std::vector<ScopedShapedBuffer> return_values; |
| 92 | return_values.reserve(run_options.size()); |
| 93 | |
| 94 | if (run_options.size() == 1) { |
| 95 | TF_ASSIGN_OR_RETURN(auto rv, |
| 96 | ExecuteOnStream(&run_options[0], arguments[0], |
| 97 | /*hlo_execution_profile=*/nullptr)); |
| 98 | return_values.push_back(std::move(rv)); |
| 99 | return std::move(return_values); |
| 100 | } |
| 101 | |
| 102 | for (size_t i = 0; i < run_options.size(); ++i) { |
| 103 | // We cannot BlockHostUntilDone() on the already-launched executions in case |
| 104 | // of error, since if the executions communicate, the initially launched |
| 105 | // executions may never complete if not all executions are running. |
| 106 | TF_ASSIGN_OR_RETURN( |
| 107 | auto rv, ExecuteAsyncOnStream(&run_options[i], arguments[i], |
| 108 | /*hlo_execution_profile=*/nullptr)); |
| 109 | return_values.push_back(std::move(rv)); |
| 110 | } |
| 111 | for (const auto& options : run_options) { |
| 112 | TF_RET_CHECK(options.stream() != nullptr); |
| 113 | TF_RETURN_IF_ERROR(options.stream()->BlockHostUntilDone()); |
| 114 | } |
| 115 | return std::move(return_values); |
| 116 | } |
| 117 | |
| 118 | StatusOr<ScopedShapedBuffer> Executable::ExecuteOnStreamWrapper( |
| 119 | const ServiceExecutableRunOptions* run_options, |
no test coverage detected