experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Make HLO snapshot dumping work on the LocalClient::RunAsync path.

Browse Source 
PiperOrigin-RevId: 259956061
This commit is contained in:
Peter Hawkins 2019-07-25 09:00:53 -07:00 committed by TensorFlower Gardener
parent 6710b20acc
commit 18f1467496
1 changed files with 43 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
tensorflow/compiler/xla/client/local_client.cc
View File

@ -189,8 +189,49 @@ StatusOr<ScopedShapedBuffer> LocalExecutable::RunAsync(
ExecutableRunOptions run_options) {
TF_ASSIGN_OR_RETURN(auto options_and_stream,
RunHelper(arguments, run_options));
return executable_->ExecuteAsyncOnStream(&options_and_stream.first,
arguments);
se::Stream* stream = run_options.stream();
std::shared_ptr<HloSnapshot> snapshot;
if (executable_->dumping_snapshot()) {
snapshot = std::make_shared<HloSnapshot>();
snapshot->set_execution_platform(backend_->platform()->Name());
*snapshot->mutable_hlo() = *executable_->hlo_proto();
for (const ShapedBuffer* arg : arguments) {
auto literal = std::make_shared<Literal>(arg->on_host_shape());
backend_->transfer_manager()->TransferLiteralFromDevice(
stream, *arg, literal.get(), [snapshot, literal](Status status) {
if (!status.ok()) {
LOG(ERROR) << "TransferLiteralFromDevice for HLO snapshot inputs "
"failed: "
<< status;
return;
}
*snapshot->add_arguments() = literal->ToProto();
});
}
}
TF_ASSIGN_OR_RETURN(
ScopedShapedBuffer outputs,
executable_->ExecuteAsyncOnStream(&options_and_stream.first, arguments));
// Transfer the outputs and save the snapshot to disk.
if (snapshot) {
auto literal = std::make_shared<Literal>(outputs.on_host_shape());
backend_->transfer_manager()->TransferLiteralFromDevice(
stream, outputs, literal.get(), [snapshot, literal](Status status) {
if (status.ok()) {
*snapshot->mutable_result() = literal->ToProto();
} else {
LOG(ERROR)
<< "TransferLiteralFromDevice for HLO snapshot outputs failed: "
<< status;
}
DumpHloSnapshotIfEnabled(*snapshot, GetDebugOptionsFromFlags());
});
}
return std::move(outputs);
}
StatusOr<ScopedShapedBuffer> LocalExecutable::ExecuteAndDump(