| 326 | } |
| 327 | |
| 328 | StatusOr<std::vector<std::unique_ptr<GlobalData>>> Client::ExecuteParallel( |
| 329 | absl::Span<const XlaComputationInstance> computations) { |
| 330 | ExecuteGraphParallelRequest request; |
| 331 | |
| 332 | for (const XlaComputationInstance& computation : computations) { |
| 333 | ExecuteGraphRequest single_request; |
| 334 | *single_request.mutable_computation() = computation.computation.proto(); |
| 335 | for (GlobalData* argument : computation.arguments) { |
| 336 | *single_request.add_arguments() = argument->handle(); |
| 337 | } |
| 338 | *single_request.mutable_execution_options() = computation.execution_options; |
| 339 | *request.add_requests() = single_request; |
| 340 | } |
| 341 | |
| 342 | ExecuteParallelResponse response; |
| 343 | VLOG(1) << "making execute-graph-parallel request: " |
| 344 | << request.ShortDebugString(); |
| 345 | Status s = stub_->ExecuteGraphParallel(&request, &response); |
| 346 | VLOG(1) << "done with request"; |
| 347 | |
| 348 | if (!s.ok()) { |
| 349 | return s; |
| 350 | } |
| 351 | |
| 352 | std::vector<std::unique_ptr<GlobalData>> outputs; |
| 353 | for (size_t i = 0; i < response.responses_size(); ++i) { |
| 354 | outputs.push_back( |
| 355 | absl::make_unique<GlobalData>(stub_, response.responses(i).output())); |
| 356 | if (i < computations.size() && |
| 357 | computations[i].execution_profile != nullptr) { |
| 358 | *computations[i].execution_profile = response.responses(i).profile(); |
| 359 | } |
| 360 | } |
| 361 | |
| 362 | return std::move(outputs); |
| 363 | } |
| 364 | |
| 365 | StatusOr<std::vector<DeviceHandle>> Client::GetDeviceHandles( |
| 366 | int64 device_count) { |