[XLA] Make XLA's memory allocator return an owning smart pointer.
authorJustin Lebar <jlebar@google.com>
Wed, 9 May 2018 18:22:31 +0000 (11:22 -0700)
committerTensorFlower Gardener <gardener@tensorflow.org>
Wed, 9 May 2018 18:34:00 +0000 (11:34 -0700)
commit7baa9ffe735adfa11c987c435216943767530269
treeccb4615481b2d062a800b55f4b63d78e6c62505e
parenta01d9f7dfb58c72ea78ed560c78f99e96223ea76
[XLA] Make XLA's memory allocator return an owning smart pointer.

Previously, xla::DeviceMemoryAllocator::Allocate returned a
stream_executor::DeviceMemoryBase.  This is morally equivalent to a raw
pointer: It's on you the user to call Deallocate().

Unfortunately we ~never got this right.  Essentially all users of
Allocate() call it in a loop, and TF_RETURN_IF_ERROR within the loop.
If any of these allocations fails (mostly commonly, due to OOM), we leak
everything we've allocated up until then.

This patch changes our API so that it returns an owning pointer.  Now
things mostly Just Work.

Also worth calling out: The lambda in CpuExecutable::ExecuteOnStream
passed to ExecuteComputeFunction almost certainly had multithreaded
use-after-free bugs.  This patch fixes them.

PiperOrigin-RevId: 196000535
27 files changed:
tensorflow/compiler/jit/BUILD
tensorflow/compiler/jit/xla_launch_util.cc
tensorflow/compiler/jit/xla_launch_util.h
tensorflow/compiler/jit/xla_launch_util_test.cc
tensorflow/compiler/jit/xla_tensor.cc
tensorflow/compiler/xla/map_util.h
tensorflow/compiler/xla/service/BUILD
tensorflow/compiler/xla/service/allocation_tracker.cc
tensorflow/compiler/xla/service/allocation_tracker.h
tensorflow/compiler/xla/service/cpu/cpu_executable.cc
tensorflow/compiler/xla/service/cpu/cpu_executable.h
tensorflow/compiler/xla/service/device_memory_allocator.cc
tensorflow/compiler/xla/service/device_memory_allocator.h
tensorflow/compiler/xla/service/gpu/buffer_allocations.cc
tensorflow/compiler/xla/service/gpu/buffer_allocations.h
tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc
tensorflow/compiler/xla/service/gpu/fft_thunk.cc
tensorflow/compiler/xla/service/gpu/fft_thunk.h
tensorflow/compiler/xla/service/gpu/gpu_executable.cc
tensorflow/compiler/xla/service/owning_device_memory.cc [new file with mode: 0644]
tensorflow/compiler/xla/service/owning_device_memory.h [new file with mode: 0644]
tensorflow/compiler/xla/service/shaped_buffer.cc
tensorflow/compiler/xla/service/shaped_buffer.h
tensorflow/compiler/xla/service/transfer_manager.cc
tensorflow/compiler/xla/tests/local_client_test_base.cc
tensorflow/compiler/xla/tests/local_client_test_base.h
tensorflow/stream_executor/stream_executor_pimpl.h