[xla:cpu] Modernize buffer_allocations_test

PiperOrigin-RevId: 709841500
openxla · Dec 26, 2024 · 4107c21 · 4107c21
1 parent e52f275
commit 4107c21
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 32 deletions.
diff --git a/xla/backends/cpu/runtime/BUILD b/xla/backends/cpu/runtime/BUILD
@@ -55,14 +55,15 @@ xla_cc_test(
     srcs = ["buffer_allocations_test.cc"],
     deps = [
         ":buffer_allocations",
+        ":thunk_testlib",
+        "//xla:literal_util",
         "//xla/service:buffer_assignment",
-        "//xla/service:maybe_owning_device_memory",
         "//xla/stream_executor:device_memory",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:statusor",
+        "//xla/tsl/platform:test",
+        "//xla/tsl/platform:test_main",
         "@com_google_absl//absl/status",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:test",
-        "@tsl//tsl/platform:test_main",
     ],
 )
 

diff --git a/xla/backends/cpu/runtime/buffer_allocations_test.cc b/xla/backends/cpu/runtime/buffer_allocations_test.cc
@@ -15,58 +15,48 @@ limitations under the License.
 
 #include "xla/backends/cpu/runtime/buffer_allocations.h"
 
-#include <cstddef>
-#include <vector>
-
+#include "xla/backends/cpu/runtime/thunk_testlib.h"
+#include "xla/literal_util.h"
 #include "xla/service/buffer_assignment.h"
-#include "xla/service/maybe_owning_device_memory.h"
 #include "xla/stream_executor/device_memory.h"
-#include "tsl/platform/statusor.h"
-#include "tsl/platform/test.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/test.h"
 
 namespace xla::cpu {
 namespace {
 
 TEST(BufferAllocationsTest, GetDeviceAddress) {
-  std::vector<MaybeOwningDeviceMemory> buffers;
-  std::vector<float> data = {1.0, 2.0, 3.0, 4.0};
-
-  size_t size_in_bytes = data.size() * sizeof(float);
-  buffers.emplace_back(se::DeviceMemoryBase(data.data(), size_in_bytes));
+  auto data = LiteralUtil::CreateR1<float>({1.0, 2.0, 3.0, 4.0});
 
-  BufferAllocations allocations(buffers);
+  BufferAllocation alloc = CreateBufferAllocation(0, data);
+  BufferAllocation::Slice slice = CreateBufferAllocationSlice(
+      alloc, /*offset=*/2 * sizeof(float), /*size=*/sizeof(float));
 
-  BufferAllocation alloc(0, size_in_bytes, 0);
-  BufferAllocation::Slice slice(&alloc, /*offset=*/2 * sizeof(float),
-                                /*size=*/sizeof(float));
+  BufferAllocations allocations = CreateBufferAllocations(data);
 
   TF_ASSERT_OK_AND_ASSIGN(se::DeviceMemoryBase alloc_mem,
                           allocations.GetDeviceAddress(0));
-  EXPECT_EQ(alloc_mem.opaque(), &data[0]);
+  EXPECT_EQ(alloc_mem.opaque(), &data.data<float>()[0]);
 
   TF_ASSERT_OK_AND_ASSIGN(se::DeviceMemoryBase slice_mem,
                           allocations.GetDeviceAddress(slice));
-  EXPECT_EQ(slice_mem.opaque(), &data[2]);
+  EXPECT_EQ(slice_mem.opaque(), &data.data<float>()[2]);
 }
 
 TEST(BufferAllocationsTest, GetDeviceAddressUnchecked) {
-  std::vector<MaybeOwningDeviceMemory> buffers;
-  std::vector<float> data = {1.0, 2.0, 3.0, 4.0};
-
-  size_t size_in_bytes = data.size() * sizeof(float);
-  buffers.emplace_back(se::DeviceMemoryBase(data.data(), size_in_bytes));
+  auto data = LiteralUtil::CreateR1<float>({1.0, 2.0, 3.0, 4.0});
 
-  BufferAllocations allocations(buffers);
+  BufferAllocation alloc = CreateBufferAllocation(0, data);
+  BufferAllocation::Slice slice = CreateBufferAllocationSlice(
+      alloc, /*offset=*/2 * sizeof(float), /*size=*/sizeof(float));
 
-  BufferAllocation alloc(0, size_in_bytes, 0);
-  BufferAllocation::Slice slice(&alloc, /*offset=*/2 * sizeof(float),
-                                /*size=*/sizeof(float));
+  BufferAllocations allocations = CreateBufferAllocations(data);
 
   se::DeviceMemoryBase alloc_mem = allocations.GetDeviceAddressUnchecked(0);
-  EXPECT_EQ(alloc_mem.opaque(), &data[0]);
+  EXPECT_EQ(alloc_mem.opaque(), &data.data<float>()[0]);
 
   se::DeviceMemoryBase slice_mem = allocations.GetDeviceAddressUnchecked(slice);
-  EXPECT_EQ(slice_mem.opaque(), &data[2]);
+  EXPECT_EQ(slice_mem.opaque(), &data.data<float>()[2]);
 }
 
 }  // namespace