Skip to content

Commit

Permalink
[xla:cpu] Modernize buffer_allocations_test
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 709615096
  • Loading branch information
@ezhulenev @Google-ML-Automation
ezhulenev authored and Google-ML-Automation committed Dec 25, 2024
1 parent e268cb7 commit d2cfb4f
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 67 deletions.
9 changes: 5 additions & 4 deletions xla/backends/cpu/runtime/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
)

Expand Down
46 changes: 18 additions & 28 deletions xla/backends/cpu/runtime/buffer_allocations_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 2 additions & 7 deletions xla/backends/cpu/runtime/xnnpack/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,15 +163,11 @@ xla_cc_test(
deps = [
":xnn_dot_thunk",
"//xla:executable_run_options",
"//xla:literal",
"//xla:literal_util",
"//xla:shape_util",
"//xla/backends/cpu/runtime:buffer_allocations",
"//xla/backends/cpu/runtime:thunk",
"//xla/backends/cpu/runtime:thunk_testlib",
"//xla/service:buffer_assignment",
"//xla/service:maybe_owning_device_memory",
"//xla/stream_executor:device_memory",
"//xla/tsl/concurrency:async_value",
"//xla/tsl/platform:statusor",
"//xla/tsl/platform:test",
Expand Down Expand Up @@ -227,12 +223,11 @@ xla_cc_test(
":xnn_fusion_thunk",
":xnn_interop",
"//xla:executable_run_options",
"//xla:literal_util",
"//xla:shape_util",
"//xla/backends/cpu/runtime:buffer_allocations",
"//xla/backends/cpu/runtime:thunk",
"//xla/service:buffer_assignment",
"//xla/service:maybe_owning_device_memory",
"//xla/stream_executor:device_memory",
"//xla/backends/cpu/runtime:thunk_testlib",
"//xla/tsl/concurrency:async_value",
"//xla/tsl/platform:statusor",
"//xla/tsl/platform:test",
Expand Down
5 changes: 0 additions & 5 deletions xla/backends/cpu/runtime/xnnpack/xnn_dot_thunk_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,10 @@ limitations under the License.

#include "xla/backends/cpu/runtime/xnnpack/xnn_dot_thunk.h"

#include <vector>

#include "xla/backends/cpu/runtime/buffer_allocations.h"
#include "xla/backends/cpu/runtime/thunk.h"
#include "xla/backends/cpu/runtime/thunk_testlib.h"
#include "xla/literal_util.h"
#include "xla/service/maybe_owning_device_memory.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/tsl/concurrency/async_value_ref.h"
Expand All @@ -32,8 +29,6 @@ namespace xla::cpu {
namespace {

TEST(XnnDotThunkTest, SimpleDot) {
std::vector<MaybeOwningDeviceMemory> buffers;

auto lhs = LiteralUtil::CreateR2<float>({{1.0, 2.0}, {3.0, 4.0}});
auto rhs = LiteralUtil::CreateR2<float>({{4.0, 3.0}, {2.0, 1.0}});
auto out = LiteralUtil::CreateR2<float>({{0.0, 0.0}, {0.0, 0.0}});
Expand Down
34 changes: 11 additions & 23 deletions xla/backends/cpu/runtime/xnnpack/xnn_fusion_thunk_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,11 @@ limitations under the License.
#include "absl/types/span.h"
#include "xla/backends/cpu/runtime/buffer_allocations.h"
#include "xla/backends/cpu/runtime/thunk.h"
#include "xla/backends/cpu/runtime/thunk_testlib.h"
#include "xla/backends/cpu/runtime/xnnpack/xnn_interop.h"
#include "xla/service/buffer_assignment.h"
#include "xla/service/maybe_owning_device_memory.h"
#include "xla/literal_util.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/tsl/concurrency/async_value_ref.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
Expand Down Expand Up @@ -79,26 +78,16 @@ static absl::StatusOr<xnn_subgraph_t> CreateBinaryAdd(
}

TEST(XnnFusionThunkTest, ElementwiseAdd) {
std::vector<MaybeOwningDeviceMemory> buffers;
auto lhs = LiteralUtil::CreateR1<float>({1.0, 2.0, 3.0, 4.0});
auto rhs = LiteralUtil::CreateR1<float>({4.0, 3.0, 2.0, 1.0});
auto out = LiteralUtil::CreateR1<float>({0.0, 0.0, 0.0, 0.0});

std::vector<float> lhs = {1.0, 2.0, 3.0, 4.0};
std::vector<float> rhs = {4.0, 3.0, 2.0, 1.0};
std::vector<float> out(4, 0.0);
BufferAllocations allocations = CreateBufferAllocations(lhs, rhs, out);

size_t size_in_bytes = lhs.size() * sizeof(float);
buffers.emplace_back(se::DeviceMemoryBase(lhs.data(), size_in_bytes));
buffers.emplace_back(se::DeviceMemoryBase(rhs.data(), size_in_bytes));
buffers.emplace_back(se::DeviceMemoryBase(out.data(), size_in_bytes));

BufferAllocations allocations(buffers);

BufferAllocation lhs_alloc(0, size_in_bytes, 0);
BufferAllocation rhs_alloc(1, size_in_bytes, 0);
BufferAllocation out_alloc(2, size_in_bytes, 0);

BufferAllocation::Slice lhs_slice(&lhs_alloc, 0, size_in_bytes);
BufferAllocation::Slice rhs_slice(&rhs_alloc, 0, size_in_bytes);
BufferAllocation::Slice out_slice(&out_alloc, 0, size_in_bytes);
auto [lhs_alloc, rhs_alloc, out_alloc] =
CreateBufferAllocation(lhs, rhs, out);
auto [lhs_slice, rhs_slice, out_slice] =
CreateBufferAllocationSlice(lhs_alloc, rhs_alloc, out_alloc);

Shape shape = ShapeUtil::MakeShape(F32, {2, 2});

Expand All @@ -117,8 +106,7 @@ TEST(XnnFusionThunkTest, ElementwiseAdd) {
tsl::BlockUntilReady(execute_event);
ASSERT_FALSE(execute_event.IsError()) << execute_event.GetError();

std::vector<float> expected = {5.0, 5.0, 5.0, 5.0};
EXPECT_EQ(out, expected);
EXPECT_EQ(out, LiteralUtil::CreateR1<float>({5.0, 5.0, 5.0, 5.0}));
}

} // namespace
Expand Down

0 comments on commit d2cfb4f

Please sign in to comment.