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

[XLA:Python] Fix potential issue in CopyToDevice where source device buffer might get deallocated before transfer completes.

Browse Source 
PiperOrigin-RevId: 257216762
This commit is contained in:
A. Unique TensorFlower 2019-07-09 10:03:22 -07:00 committed by TensorFlower Gardener
parent 4bb0f6e87f
commit af8772fa69
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/xla/python/local_client.cc
View File

@ -652,6 +652,10 @@ StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::CopyToDevice(
output_buffer)); output_buffer));
} }
// We hold on to the `src_device_buffer` until the transfer is finished.
src_device.ThenRelease(src_device_to_device_stream,
std::move(src_device_buffer));
// Write new tuple buffers. The destination buffers have different addresses, // Write new tuple buffers. The destination buffers have different addresses,
// so we must construct tuple buffers from scratch instead of copying them. // so we must construct tuple buffers from scratch instead of copying them.
if (dst_buffer.on_device_shape().IsTuple()) { if (dst_buffer.on_device_shape().IsTuple()) {