diff options
| author |  Justin Lebar <jlebar@google.com> | 2018-05-09 11:22:31 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-05-09 11:34:00 -0700 |
| commit | 7baa9ffe735adfa11c987c435216943767530269 (patch) | |
| tree | ccb4615481b2d062a800b55f4b63d78e6c62505e /tensorflow/stream_executor/stream_executor_pimpl.h | |
| parent | a01d9f7dfb58c72ea78ed560c78f99e96223ea76 (diff) | |
[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
Diffstat (limited to 'tensorflow/stream_executor/stream_executor_pimpl.h')
| -rw-r--r-- | tensorflow/stream_executor/stream_executor_pimpl.h | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/stream_executor/stream_executor_pimpl.h b/tensorflow/stream_executor/stream_executor_pimpl.h index ab6b00f660..e426cf9931 100644 --- a/tensorflow/stream_executor/stream_executor_pimpl.h +++ b/tensorflow/stream_executor/stream_executor_pimpl.h @@ -177,6 +177,9 @@ class StreamExecutor { // // Resets the internal contents of mem to be null-representative, but this // null-out effect should not be relied upon in client code. + // + // TODO(jlebar): Change this to accept a DeviceMemoryBase by value, see + // discussion in cl/195744342. void Deallocate(DeviceMemoryBase *mem); // Retrieves a mapping of active opaque device memory pointer to a string |