Unify ScopedDeviceMemory and OwningDeviceMemory.

Previously, ScopedDeviceMemory was parameterized by a type it holds,
and was only constructible by a stream executor, while OwningDeviceMemory was
constructible from an arbitrary allocator, but was not parameterized.

This refactoring leaves a single class ScopedDeviceMemory which is parameterized,
can be constructed from an arbitrary allocator, and has helper constructors
for constructing from a stream executor.
OwningDeviceMemory is left as a typedef for ScopedDeviceMemory parameterized on
uint8 (holds arbitrary bytes, no alignment guaranteed).

PiperOrigin-RevId: 248234830

Author: cheshire

tensorflow/compiler/jit/xla_launch_util.h

@@ -30,7 +30,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
 #include "tensorflow/stream_executor/device_memory_allocator.h"
-#include "tensorflow/stream_executor/owning_device_memory.h"
 
 namespace tensorflow {
 class XlaAllocator;

tensorflow/compiler/jit/xla_tensor.cc

@@ -63,7 +63,7 @@ Status XlaTensor::AllocateShapedBuffer(DataType dtype,
                         client->backend().memory_allocator()->Allocate(
                             device_ordinal, size, /*retry_on_failure=*/false));
     // Move our buffer into shaped_buffer, which takes ownership of it.
-    index_to_buffer.second = buffer.Forget();
+    index_to_buffer.second = buffer.Release();
   }
 
   VLOG(4) << shaped_buffer.ToString();

tensorflow/compiler/xla/python/BUILD

@@ -214,7 +214,6 @@ tf_pybind_extension(
         "//tensorflow/compiler/xla/client/lib:self_adjoint_eig",
         "//tensorflow/compiler/xla/client/lib:svd",
         "//tensorflow/compiler/xla/service:computation_placer",
-        "//tensorflow/stream_executor:device_memory_allocator",
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service:hlo_graph_dumper",
         "//tensorflow/compiler/xla/service:name_uniquer",
@@ -229,6 +228,7 @@ tf_pybind_extension(
         # without any TF dependencies as "jaxlib" on Pypi, and "jaxlib" does
         # not require Tensorflow.
         "//tensorflow/core:lib_internal_impl",  # buildcleaner: keep
+        "//tensorflow/stream_executor:device_memory_allocator",
     ] + xla_python_default_plugins(),
 )
 

tensorflow/compiler/xla/python/shared_device_buffer.cc

@@ -132,7 +132,7 @@ static void PopulateShapedBufferFromBuffer(
     ShapeTree<se::DeviceMemoryBase>::iterator* iterator,
     const ShapeTree<se::DeviceMemoryBase>::iterator& end) {
   CHECK(*iterator != end);
-  (*iterator)->second = buffer.device_memory().AsDeviceMemoryBase();
+  (*iterator)->second = *buffer.device_memory();
   ++*iterator;
   for (const auto& child : buffer.children()) {
     PopulateShapedBufferFromBuffer(*child, iterator, end);

tensorflow/compiler/xla/python/shared_device_buffer.h

@@ -21,7 +21,6 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/transfer_manager.h"
 #include "tensorflow/compiler/xla/shape.h"
 #include "tensorflow/stream_executor/device_memory_allocator.h"
-#include "tensorflow/stream_executor/owning_device_memory.h"
 
 namespace xla {
 

tensorflow/compiler/xla/python/shared_device_buffer_test.cc

@@ -114,11 +114,9 @@ TEST(PySharedDeviceBufferTest, AsShapedBuffer) {
             abc_tuple_buffer->on_device_shape());
 
   std::vector<se::DeviceMemoryBase> expected_buffer_sequence = {
-      abc_tuple_buffer->device_memory().AsDeviceMemoryBase(),
-      c_buffer->device_memory().AsDeviceMemoryBase(),
-      ab_tuple_buffer->device_memory().AsDeviceMemoryBase(),
-      a_buffer->device_memory().AsDeviceMemoryBase(),
-      b_buffer->device_memory().AsDeviceMemoryBase(),
+      *abc_tuple_buffer->device_memory(), *c_buffer->device_memory(),
+      *ab_tuple_buffer->device_memory(),  *a_buffer->device_memory(),
+      *b_buffer->device_memory(),
   };
   auto it = shaped_buffer.buffers().begin();
   auto expected_it = expected_buffer_sequence.begin();

tensorflow/compiler/xla/python/types.cc

@@ -17,7 +17,6 @@ limitations under the License.
 
 #include "absl/container/flat_hash_map.h"
 #include "tensorflow/compiler/xla/status_macros.h"
-#include "tensorflow/stream_executor/owning_device_memory.h"
 
 namespace xla {
 

tensorflow/compiler/xla/service/cpu/cpu_executable.cc

@@ -113,17 +113,17 @@ CpuExecutable::CreateBufferTable(
     } else {
       TF_ASSIGN_OR_RETURN(owning_buffers[i], memory_allocator->Allocate(
                                                  device_ordinal, buffer_size));
-      unowning_buffers[i] = owning_buffers[i].AsDeviceMemoryBase();
+      unowning_buffers[i] = *owning_buffers[i];
 
       VLOG(3) << "buffer #" << i << " allocated " << buffer_size << " bytes ["
-              << owning_buffers[i].opaque() << "]";
+              << owning_buffers[i]->opaque() << "]";
     }
 
     // Since the output buffer and all the temporary buffers were written into
     // by the JITed code, msan has no way of knowing their memory was
     // initialized. Mark them initialized so that msan doesn't flag loads from
     // these buffers.
-    TF_ANNOTATE_MEMORY_IS_INITIALIZED(owning_buffers[i].opaque(), buffer_size);
+    TF_ANNOTATE_MEMORY_IS_INITIALIZED(owning_buffers[i]->opaque(), buffer_size);
   }
 
   TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice result_slice,
@@ -247,7 +247,7 @@ StatusOr<ScopedShapedBuffer> CpuExecutable::CreateResultShapedBuffer(
           // ownership, and hence a buffer coming from there cannot be part of
           // the new ScopedShapedBuffer we create for the result (which assumes
           // ownership).
-          *device_memory = buffer.Forget();
+          *device_memory = buffer.Release();
         } else {
           auto output_alias = input_output_alias.GetAliasedOutput(
               slice.allocation()->parameter_number(),

tensorflow/compiler/xla/service/executable.h

@@ -39,7 +39,6 @@ limitations under the License.
 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 #include "tensorflow/stream_executor/device_memory_allocator.h"
-#include "tensorflow/stream_executor/owning_device_memory.h"
 
 namespace xla {
 

tensorflow/compiler/xla/service/gpu/buffer_allocations.cc

@@ -80,17 +80,18 @@ StatusOr<std::unique_ptr<BufferAllocations>> BufferAllocations::Builder::Build(
         se::OwningDeviceMemory buffer;
         TF_ASSIGN_OR_RETURN(
             buffer, memory_allocator->Allocate(device_ordinal, buffer_size));
-        if (reinterpret_cast<uintptr_t>(buffer.opaque()) % expected_alignment !=
+        if (reinterpret_cast<uintptr_t>(buffer->opaque()) %
+                expected_alignment !=
             0) {
           return InternalError(
               "Address returned by memory_allocator->Allocate must be a "
               "multiple of 0x%x, but was %p",
-              kXlaAllocatedBufferAlignBytes, buffer.opaque());
+              kXlaAllocatedBufferAlignBytes, buffer->opaque());
         }
         // We do manual memory management within BufferAllocations.  Be sure not
         // to do a TF_RETURN_IF_ERROR between this line and the
         // buffer_allocations->SetBuffer(buffer_address) call below!
-        buffer_address = buffer.Forget();
+        buffer_address = buffer.Release();
       }
 
       buffer_allocations->SetBuffer(i, buffer_address);