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

Merge pull request #43134 from ROCmSoftwarePlatform:google-upstream-thenexecute-deadlock

Browse Source 
PiperOrigin-RevId: 332431107
Change-Id: I84bd51a1a7079206c015a17afb130de592ac0faf
This commit is contained in:
TensorFlower Gardener 2020-09-18 05:20:22 -07:00
commit c2c2ff8032
2 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/compiler/jit/xla_device_context.cc
View File

@ -294,4 +294,12 @@ se::Stream* XlaDeviceContext::GetDeviceToDeviceStream() {
return device_to_device_stream(stream);
}
Status XlaDeviceContext::ThenExecute(Device* device,
stream_executor::Stream* stream,
std::function<void()> func) {
VLOG(2) << "XlaDeviceContext::ThenExecute";
stream->ThenDoHostCallback(std::move(func));
return Status::OK();
}
} // namespace tensorflow

3
tensorflow/compiler/jit/xla_device_context.h
View File

@ -86,6 +86,9 @@ class XlaDeviceContext : public DeviceContext {
// Returns a device-to-device stream, in round-robin fashion.
se::Stream* GetDeviceToDeviceStream();
Status ThenExecute(Device* device, stream_executor::Stream* stream,
std::function<void()> func) override;
private:
bool UseMultipleStreams() const { return stream_ != host_to_device_stream_; }