[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

1 files changed, 3 insertions, 4 deletions

diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
index 980cc89fa0..04b4f7aef1 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
@@ -286,8 +286,8 @@ StatusOr<ScopedShapedBuffer> GpuExecutable::ExecuteOnStream(
   se::StreamExecutor* executor = run_options->stream()->parent();
   TF_ASSIGN_OR_RETURN(
       auto buffer_allocations,
-      buffer_allocations_builder.Build(*assignment_, executor->device_ordinal(),
-                                       memory_allocator));
+      buffer_allocations_builder.Build(
+          assignment_.get(), executor->device_ordinal(), memory_allocator));
 
   bool block_host_until_done =
       !memory_allocator->AllowsAsynchronousDeallocation();
@@ -329,8 +329,7 @@ StatusOr<ScopedShapedBuffer> GpuExecutable::ExecuteOnStream(
         buffers_in_result.insert(src_base);
         return Status::OK();
       }));
-  TF_RETURN_IF_ERROR(
-      buffer_allocations->TearDown(buffers_in_result, *assignment_));
+  TF_RETURN_IF_ERROR(buffer_allocations->TearDown(buffers_in_result));
 
   return std::move(shaped_buffer);
 }