diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-06-27 05:19:14 -0700 |
|---|---|---|
| committer |  Gunhan Gulsoy <gunan@google.com> | 2018-06-28 21:37:43 -0700 |
| commit | c2d369373d7e0cbdb01be9f556a5a36ff3ce6cf6 (patch) | |
| tree | cbf006b41429f239d72c42e1638cdc728a86426a /tensorflow/compiler/xla/service/hlo_runner.cc | |
| parent | 8f1061f846c76c882b75de42d9bda395822cf666 (diff) | |
[XLA] Support asynchronous execution through XLA
PiperOrigin-RevId: 202292422
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_runner.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_runner.cc | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index 4f0569f405..b2725e2918 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -180,8 +180,12 @@ StatusOr<ScopedShapedBuffer> HloRunner::ExecuteWithDeviceBuffers( TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable, CreateExecutable(std::move(module), run_hlo_passes)); - return executable->ExecuteOnStreamWrapper(&service_run_options, - /*profile=*/profile, arguments); + TF_ASSIGN_OR_RETURN( + ScopedShapedBuffer retval, + executable->ExecuteOnStreamWrapper(&service_run_options, + /*profile=*/profile, arguments)); + TF_RETURN_IF_ERROR(stream.BlockHostUntilDone()); + return std::move(retval); } StatusOr<ScopedShapedBuffer> HloRunner::ExecuteWithDeviceBuffers( @@ -309,6 +313,7 @@ StatusOr<std::vector<std::unique_ptr<Literal>>> HloRunner::ExecuteReplicated( std::vector<std::unique_ptr<Literal>> exec_results; for (int64 i = 0; i < options.num_replicas; ++i) { + TF_RETURN_IF_ERROR(streams[i]->BlockHostUntilDone()); TF_ASSIGN_OR_RETURN(std::unique_ptr<Literal> literal, backend().transfer_manager()->TransferLiteralFromDevice( streams[i].get(), results[i])); |