wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/dtensor/tests/BUILD
+104
View File
@@ -0,0 +1,104 @@
# Tests for DTensor interfaces.
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
tf_cc_test(
name = "dtensor_operation_test",
srcs = ["dtensor_operation_test.cc"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/dtensor/cc:dtensor_operation",
"@com_google_googletest//:gtest_main",
],
)
tf_cc_test(
name = "slice_util_test",
srcs = ["slice_util_test.cc"],
deps = [
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/dtensor/cc:slice_util",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/status:status_matchers",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:status_matchers",
],
)
tf_cc_test(
name = "tensor_layout_test",
srcs = ["tensor_layout_test.cc"],
deps = [
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/proto:layout_proto_cc",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "layout_to_xla_sharding_test",
srcs = ["layout_to_xla_sharding_test.cc"],
deps = [
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:layout_to_xla_sharding",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/strings",
"@com_google_benchmark//:benchmark",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/hlo/ir:hlo",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "executable_manager_test",
srcs = ["executable_manager_test.cc"],
deps = [
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:refcount",
"//tensorflow/dtensor/cc:dtensor_device_util",
"//tensorflow/dtensor/cc:dtensor_operation",
"//tensorflow/dtensor/cc:tensor_layout",
"@com_google_absl//absl/status:status_matchers",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tf_cc_test(
name = "spmd_expander_test",
srcs = ["spmd_expander_test.cc"],
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:errors",
"//tensorflow/dtensor/cc:dstatus",
"//tensorflow/dtensor/cc:tensor_layout",
"//tensorflow/dtensor/mlir:spmd_expander",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
],
)
tensorflow/dtensor/tests/dtensor_operation_test.cc
+63
View File
@@ -0,0 +1,63 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/dtensor/cc/dtensor_operation.h"
#include <gtest/gtest.h>
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/op.h"
namespace tensorflow {
namespace dtensor {
namespace {
// Register a few dummy ops with resource and stateful traits.
REGISTER_OP("OutputResource").Output("resource: resource");
REGISTER_OP("InputResource").Input("resource: resource");
REGISTER_OP("Stateful").SetIsStateful();
REGISTER_OP("Pure");
TEST(DTensorOperationTest, TestEagerIsNotPure) {
DTensorOperation output{"OutputResource", nullptr, {}, {}};
DTensorOperation input{"InputResource", nullptr, {}, {}};
DTensorOperation stateful{"Stateful", nullptr, {}, {}};
DTensorOperation pure{"Pure", nullptr, {}, {}};
EXPECT_FALSE(output.is_pure());
EXPECT_FALSE(input.is_pure());
EXPECT_FALSE(stateful.is_pure());
EXPECT_TRUE(pure.is_pure());
}
TEST(DTensorOperationTest, TestFunctionIsNotPure) {
FunctionDef fdef;
DTensorOperation op{"func", &fdef, {}, {}};
EXPECT_FALSE(op.is_pure());
}
TEST(DTensorOperationTest, TestIsFunc) {
FunctionDef fdef;
DTensorOperation func_op{"func", &fdef, {}, {}};
DTensorOperation eager_op{"Pure", nullptr, {}, {}};
EXPECT_TRUE(func_op.is_func());
EXPECT_FALSE(eager_op.is_func());
}
} // namespace
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/tests/executable_manager_test.cc
+69
View File
@@ -0,0 +1,69 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status_matchers.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/dtensor/cc/dtensor_device_util.h"
#include "tensorflow/dtensor/cc/dtensor_operation.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace {
using ::testing::HasSubstr;
using ::tsl::error::UNAVAILABLE;
class ExecutableManagerTest : public ::testing::Test {
protected:
DTensorOperation CreateTestDTensorOperation() {
return DTensorOperation{"test_fn", nullptr, empty_mesh_, {}};
}
Mesh empty_mesh_ = Mesh::Empty();
core::RefCountPtr<ExecutableManager<ExecutionFunctions>> function_manager_{
new ExecutableManager<ExecutionFunctions>()};
};
TEST_F(ExecutableManagerTest, ShouldFoldInputUnavailable) {
auto result =
function_manager_->ShouldFoldInput(CreateTestDTensorOperation(), {}, 0);
EXPECT_THAT(result,
absl_testing::StatusIs(
UNAVAILABLE, HasSubstr("ExecutionFunctions manager can not "
"check if the input is foldable")));
}
TEST_F(ExecutableManagerTest, GetCachedExecutableUnavailable) {
DTensorOperation doperation = CreateTestDTensorOperation();
NameAttrList func_attr;
func_attr.set_name(doperation.name);
auto result = function_manager_->GetCachedExecutable(
doperation, func_attr,
{nullptr}, // Dummy input to trigger ShouldFoldInput check.
{});
EXPECT_THAT(result, absl_testing::StatusIs(UNAVAILABLE));
}
} // namespace
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/tests/layout_to_xla_sharding_test.cc
+314
View File
@@ -0,0 +1,314 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/dtensor/cc/xla_spmd/layout_to_xla_sharding.h"
#include <cstdint>
#include <string>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "benchmark/benchmark.h" // from @com_google_benchmark
#include "xla/hlo/ir/hlo_sharding.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace {
StatusOr<std::string> ConvertLayoutStrToHloShardingStr(std::string layout_str) {
TF_ASSIGN_OR_RETURN(const Layout layout, Layout::FromString(layout_str));
TF_ASSIGN_OR_RETURN(const xla::OpSharding op_sharding,
ConvertLayoutToXlaOpSharding(layout));
TF_ASSIGN_OR_RETURN(const auto hlo_sharding,
xla::HloSharding::FromProto(op_sharding));
return hlo_sharding.ToString();
}
TEST(LayoutToXLAShardingTest, ReplicatedLayout1D) {
std::string layout_str =
"sharding_specs:unsharded, "
"mesh:|x=2|0,1|0,1|/job:localhost/task:0/device:CPU:0,/job:localhost/"
"task:0/device:CPU:1";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{replicated}", sharding);
}
TEST(LayoutToXLAShardingTest, ReplicatedLayout2D) {
std::string layout_str =
"sharding_specs:unsharded,unsharded "
"mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{replicated}", sharding);
}
TEST(LayoutToXLAShardingTest, ReplicatedLayout3D) {
std::string layout_str =
"sharding_specs:unsharded,unsharded,unsharded, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{replicated}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedLayout1D) {
std::string layout_str =
"sharding_specs:x, "
"mesh:|x=3|0,1,2|0,1,2|/job:localhost/task:0/device:CPU:0,/job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{devices=[3]0,1,2}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedLayout2D) {
std::string layout_str =
"sharding_specs:x,y, "
"mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{devices=[2,2]0,1,2,3}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedLayout2DAsymmetricMesh) {
std::string layout_str =
"sharding_specs:y,x, "
"mesh:|x=2,y=4|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/job:localhost/task:0/device:CPU:1,/job:localhost/task:0/"
"device:CPU:2,/job:localhost/task:0/device:CPU:3,/job:localhost/task:0/"
"device:CPU:4,/job:localhost/task:0/device:CPU:5,/job:localhost/task:0/"
"device:CPU:6,/job:localhost/task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{devices=[4,2]0,4,1,5,2,6,3,7}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedPermutedLayout2D) {
std::string layout_str =
"sharding_specs:y,x, "
"mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Devices should now be ordered 'y' axis first.
EXPECT_EQ("{devices=[2,2]0,2,1,3}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedLayout3D) {
std::string layout_str =
"sharding_specs:x,y,z, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{devices=[2,2,2]0,1,2,3,4,5,6,7}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedPermutedLayout3D_1) {
std::string layout_str =
"sharding_specs:z,x,y, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Devices are permuted in z axis first and then x and y. It helps to manually
// draw this to confirm it.
EXPECT_EQ("{devices=[2,2,2]0,2,4,6,1,3,5,7}", sharding);
}
TEST(LayoutToXLAShardingTest, FullyShardedPermutedLayout3D_2) {
std::string layout_str =
"sharding_specs:z,y,x, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Devices are permuted in reverse order, it helps to draw this out manually
// to understand this is correct.
EXPECT_EQ("{devices=[2,2,2]0,4,2,6,1,5,3,7}", sharding);
}
TEST(LayoutToXLAShardingTest, PartiallyShardedLayout2D) {
std::string layout_str =
"sharding_specs:x,unsharded, "
"mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
EXPECT_EQ("{devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}", sharding);
}
TEST(LayoutToXLAShardingTest, PartiallyShardedPermutedLayout2D) {
std::string layout_str =
"sharding_specs:y,unsharded, "
"mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Permuted on the Y dimension.
EXPECT_EQ("{devices=[2,1,2]0,2,1,3 last_tile_dim_replicate}", sharding);
}
TEST(LayoutToXLAShardingTest, PartiallyShardedLayout3D_1) {
std::string layout_str =
"sharding_specs:x,y,unsharded, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Last dim is two since every replication group is size 2.
EXPECT_EQ("{devices=[2,2,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}",
sharding);
}
TEST(LayoutToXLAShardingTest, PartiallyShardedLayout3D_2) {
std::string layout_str =
"sharding_specs:x,unsharded,unsharded, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Last dim is four since every replication group is size 4.
EXPECT_EQ("{devices=[2,1,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}",
sharding);
}
TEST(LayoutToXLAShardingTest, PartiallyShardedPermutedLayout3D_1) {
std::string layout_str =
"sharding_specs:z,y,unsharded, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Last dim is two since every replication group is size 2.
// Same permutation as 'z', 'y', 'x'.
EXPECT_EQ("{devices=[2,2,1,2]0,4,2,6,1,5,3,7 last_tile_dim_replicate}",
sharding);
}
TEST(LayoutToXLAShardingTest, PartiallyShardedPermutedLayout3D_2) {
std::string layout_str =
"sharding_specs:y,unsharded,z, "
"mesh:|x=2,y=2,z=2|0,1,2,3,4,5,6,7|0,1,2,3,4,5,6,7|/job:localhost/task:0/"
"device:CPU:0,/"
"job:localhost/"
"task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/job:localhost/"
"task:0/device:CPU:3,/job:localhost/task:0/device:CPU:4,/job:localhost/"
"task:0/device:CPU:5,/job:localhost/task:0/device:CPU:6,/job:localhost/"
"task:0/device:CPU:7";
TF_ASSERT_OK_AND_ASSIGN(std::string sharding,
ConvertLayoutStrToHloShardingStr(layout_str));
// Last dim is two since every replication group is size 2.
// Same permutation as 'y', 'z'.
EXPECT_EQ("{devices=[2,1,2,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}",
sharding);
}
void BM_65536Devices(benchmark::State& state) {
std::vector<int64_t> device_ids(65536);
absl::c_iota(device_ids, 0);
std::vector<std::string> devices_str(65536);
absl::c_generate(devices_str, [n = 0]() mutable {
return absl::StrCat("/job:localhost/task:0/device:CPU:", n++);
});
auto mesh = Mesh::CreateMesh(/*mesh_name=*/"", /*dim_names=*/{"x", "y", "z"},
/*mesh_shape=*/{8, 128, 64},
/*global_device_ids=*/device_ids,
/*global_devices_str=*/{},
/*local_device_ids=*/device_ids,
/*local_devices_str=*/devices_str);
TF_ASSERT_OK_AND_ASSIGN(auto layout,
Layout::GetLayout({"x", "y", "z"}, mesh));
for (auto s : state) {
TF_EXPECT_OK(ConvertLayoutToXlaOpSharding(layout).status());
}
}
BENCHMARK(BM_65536Devices);
} // namespace
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/tests/slice_util_test.cc
+637
View File
@@ -0,0 +1,637 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/dtensor/cc/slice_util.h"
#include <cstdint>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include "absl/status/status_matchers.h"
#include "xla/tsl/platform/status_matchers.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace tensorflow {
namespace dtensor {
namespace slice_util {
namespace {
using ::testing::SizeIs;
TEST(TokenTest, NormalizeDynamic) {
auto spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/0, /*stride=*/1,
/*dynamic_mask=*/true,
/*begin_mask=*/true, /*end_mask=*/true);
EXPECT_EQ(spec.normalize(4).begin, 0);
EXPECT_EQ(spec.normalize(4).end, 0);
EXPECT_EQ(spec.normalize(4).dynamic_mask, true);
EXPECT_EQ(spec.normalize(4).begin_mask, true);
EXPECT_EQ(spec.normalize(4).end_mask, true);
}
TEST(TokenTest, NormalizeFullPositiveStride) {
auto spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/4, /*stride=*/1);
EXPECT_EQ(spec.normalize(4).begin, 0);
EXPECT_EQ(spec.normalize(4).end, 4);
spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/4, /*stride=*/2);
EXPECT_EQ(spec.normalize(4).begin, 0);
EXPECT_EQ(spec.normalize(4).end, 4);
spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/4, /*stride=*/3);
EXPECT_EQ(spec.normalize(4).begin, 0);
EXPECT_EQ(spec.normalize(4).end, 6);
spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/4, /*stride=*/5);
EXPECT_EQ(spec.normalize(4).begin, 0);
EXPECT_EQ(spec.normalize(4).end, 5);
}
TEST(TokenTest, NormalizeFullNegativeStride) {
auto spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/-1, /*stride=*/-1);
EXPECT_EQ(spec.normalize(4).begin, 3);
EXPECT_EQ(spec.normalize(4).end, -1);
spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/-1, /*stride=*/-2);
EXPECT_EQ(spec.normalize(4).begin, 3);
EXPECT_EQ(spec.normalize(4).end, -1);
spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/-1, /*stride=*/-3);
EXPECT_EQ(spec.normalize(4).begin, 3);
EXPECT_EQ(spec.normalize(4).end, -3);
spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/-1, /*stride=*/-5);
EXPECT_EQ(spec.normalize(4).begin, 3);
EXPECT_EQ(spec.normalize(4).end, -2);
}
TEST(TokenTest, NormalizeZeroPositiveStride) {
auto spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/3, /*stride=*/1);
EXPECT_EQ(spec.normalize(7).begin, 3);
EXPECT_EQ(spec.normalize(7).end, 3);
spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/0, /*stride=*/1);
EXPECT_EQ(spec.normalize(7).begin, 0);
EXPECT_EQ(spec.normalize(7).end, 0);
}
TEST(TokenTest, NormalizeZeroNegativeStride) {
auto spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/3, /*stride=*/-1);
EXPECT_EQ(spec.normalize(7).begin, 3);
EXPECT_EQ(spec.normalize(7).end, 3);
spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/0, /*stride=*/-1);
EXPECT_EQ(spec.normalize(7).begin, 0);
EXPECT_EQ(spec.normalize(7).end, 0);
}
TEST(TokenTest, NormalizePartialPositiveStride) {
auto spec = Token(Token::REGULAR, /*begin=*/1, /*end=*/5, /*stride=*/1);
EXPECT_EQ(spec.normalize(7).begin, 1);
EXPECT_EQ(spec.normalize(7).end, 5);
spec = Token(Token::REGULAR, /*begin=*/1, /*end=*/5, /*stride=*/2);
EXPECT_EQ(spec.normalize(7).begin, 1);
EXPECT_EQ(spec.normalize(7).end, 5);
spec = Token(Token::REGULAR, /*begin=*/1, /*end=*/5, /*stride=*/3);
EXPECT_EQ(spec.normalize(7).begin, 1);
EXPECT_EQ(spec.normalize(7).end, 7);
spec = Token(Token::REGULAR, /*begin=*/1, /*end=*/5, /*stride=*/5);
EXPECT_EQ(spec.normalize(7).begin, 1);
EXPECT_EQ(spec.normalize(7).end, 6);
spec = Token(Token::REGULAR, /*begin=*/1, /*end=*/-1, /*stride=*/1);
EXPECT_EQ(spec.normalize(7).begin, 1);
EXPECT_EQ(spec.normalize(7).end, 6);
spec = Token(Token::REGULAR, /*begin=*/0, /*end=*/-1, /*stride=*/1);
EXPECT_EQ(spec.normalize(7).begin, 0);
EXPECT_EQ(spec.normalize(7).end, 6);
}
TEST(TokenTest, NormalizePartialNegativeStride) {
auto spec = Token(Token::REGULAR, /*begin=*/6, /*end=*/2, /*stride=*/-1);
EXPECT_EQ(spec.normalize(7).begin, 6);
EXPECT_EQ(spec.normalize(7).end, 2);
spec = Token(Token::REGULAR, /*begin=*/6, /*end=*/2, /*stride=*/-2);
EXPECT_EQ(spec.normalize(7).begin, 6);
EXPECT_EQ(spec.normalize(7).end, 2);
spec = Token(Token::REGULAR, /*begin=*/6, /*end=*/2, /*stride=*/-3);
EXPECT_EQ(spec.normalize(7).begin, 6);
EXPECT_EQ(spec.normalize(7).end, 0);
spec = Token(Token::REGULAR, /*begin=*/6, /*end=*/2, /*stride=*/-5);
EXPECT_EQ(spec.normalize(7).begin, 6);
EXPECT_EQ(spec.normalize(7).end, 1);
}
TEST(TokenTest, NormalizeFarFromCenter) {
auto spec = Token(Token::REGULAR, /*begin=*/100, /*end=*/102, /*stride=*/1);
EXPECT_EQ(spec.normalize(9).begin, 1);
EXPECT_EQ(spec.normalize(9).end, 3);
}
TEST(TokenTest, NormalizeBeginMask) {
auto spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/2, /*stride=*/1);
spec.begin_mask = true;
EXPECT_EQ(spec.normalize(7).begin, 0);
spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/2, /*stride=*/-1);
spec.begin_mask = true;
EXPECT_EQ(spec.normalize(7).begin, 6);
}
TEST(TokenTest, NormalizeEndMask) {
auto spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/2, /*stride=*/1);
spec.end_mask = true;
EXPECT_EQ(spec.normalize(7).end, 7);
spec = Token(Token::REGULAR, /*begin=*/3, /*end=*/2, /*stride=*/-1);
spec.end_mask = true;
EXPECT_EQ(spec.normalize(7).end, -1);
}
class InferenceTest : public ::testing::Test {
protected:
Mesh GetMesh() {
return Mesh::CreateMesh("MyMesh", /*dim_names=*/{"x", "y"},
/*mesh_shape=*/{2, 1},
/*global_device_ids=*/{0, 1},
/*global_devices_str=*/
{"/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"},
/*local_device_ids=*/{0, 1},
/*local_devices_str=*/
{"/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"},
/*use_xla_spmd=*/false);
}
};
TEST_F(InferenceTest, FullyReplicatedInputs) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, Layout::kUnshardedDim},
GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, Layout::kUnshardedDim},
GetMesh());
const auto specs = std::vector<Token>{
Token(Token::REGULAR, /*begin=*/0, /*end=*/-1, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/2, /*stride=*/2,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(
forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, -1);
EXPECT_EQ(forward->local_tokens()[1].end, 2);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(
backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(
backward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->local_tokens()[0].end, -1);
EXPECT_EQ(backward->local_tokens()[1].end, 2);
}
TEST_F(InferenceTest, NewAxisMask) {
const Layout input_layout =
*Layout::GetLayout(std::vector<std::string>{"x", "y"}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, Layout::kUnshardedDim,
"x", "y"},
GetMesh());
const auto specs = std::vector<Token>{
Token(Token::NEW_AXIS, /*begin=*/0, /*end=*/0, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::NEW_AXIS, /*begin=*/0, /*end=*/0, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/2, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/4, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({"x", "y"}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(4));
EXPECT_EQ(forward->local_tokens()[0].end, 0);
EXPECT_EQ(forward->local_tokens()[1].end, 0);
EXPECT_EQ(forward->local_tokens()[2].end, 1); // dim_size(x) == 2.
EXPECT_EQ(forward->local_tokens()[3].end, 4);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>(
{Layout::kUnshardedDim, Layout::kUnshardedDim, "x", "y"}));
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_THAT(backward->local_tokens(), SizeIs(4));
EXPECT_EQ(backward->local_tokens()[0].end, 0);
EXPECT_EQ(backward->local_tokens()[1].end, 0);
EXPECT_EQ(backward->local_tokens()[2].end, 1); // dim_size(x) == 2.
EXPECT_EQ(backward->local_tokens()[3].end, 4);
}
TEST_F(InferenceTest, ShrinkAxisMask) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, Layout::kUnshardedDim},
GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim}, GetMesh());
const auto specs = std::vector<Token>{
Token(Token::REGULAR, /*begin=*/0, /*end=*/-1, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::SHRINK_AXIS, /*begin=*/0, /*end=*/2,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(
forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, -1);
EXPECT_EQ(forward->local_tokens()[1].end, 2);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_EQ(backward->local_tokens()[0].end, -1);
EXPECT_EQ(backward->local_tokens()[1].end, 2);
}
TEST_F(InferenceTest, EllipsisMask) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{"x", "y", Layout::kUnshardedDim}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{"x", "y", Layout::kUnshardedDim,
Layout::kUnshardedDim, Layout::kUnshardedDim},
GetMesh());
const auto specs =
std::vector<Token>{Token(Token::ELLIPSIS, /*begin=*/0, /*end=*/0,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::NEW_AXIS, /*begin=*/0, /*end=*/0,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::NEW_AXIS, /*begin=*/0, /*end=*/0,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4, 6});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({"x", "y", Layout::kUnshardedDim}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
// No local specs for the ellipsis axes "x" and "y".
EXPECT_THAT(forward->local_tokens(), SizeIs(3));
EXPECT_EQ(forward->local_tokens()[0].end, 0);
EXPECT_EQ(forward->local_tokens()[1].end, 0);
EXPECT_EQ(forward->local_tokens()[2].end, 0);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4, 6});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(
backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({"x", "y", Layout::kUnshardedDim,
Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(backward->expander_input_layout(), input_layout);
// No local specs for the ellipsis axes "x" and "y".
EXPECT_THAT(backward->local_tokens(), SizeIs(3));
EXPECT_EQ(backward->local_tokens()[0].end, 0);
EXPECT_EQ(backward->local_tokens()[1].end, 0);
EXPECT_EQ(backward->local_tokens()[2].end, 0);
}
TEST_F(InferenceTest, EllipsisNewAxisEndMask) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, Layout::kUnshardedDim},
GetMesh());
const auto specs = std::vector<Token>{
Token(Token::ELLIPSIS, /*begin=*/0, /*end=*/0, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::NEW_AXIS, /*begin=*/0, /*end=*/0, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/0, /*stride=*/1,
/*dynamic_mask=*/false,
/*begin_mask=*/true,
/*end_mask=*/true),
};
auto forward = CreateAndRun<ForwardLayoutInference>(specs, input_layout,
std::vector<int64_t>{2});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(3));
EXPECT_EQ(forward->local_tokens()[0].end, 0);
EXPECT_EQ(forward->local_tokens()[1].end, 0);
EXPECT_EQ(forward->local_tokens()[2].end, 2);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(
backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(backward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim}));
EXPECT_THAT(backward->local_tokens(), SizeIs(3));
EXPECT_EQ(backward->local_tokens()[0].end, 0);
EXPECT_EQ(backward->local_tokens()[1].end, 0);
EXPECT_EQ(backward->local_tokens()[2].end, 2);
}
TEST_F(InferenceTest, AdditionalAxes) {
const Layout input_layout =
*Layout::GetLayout(std::vector<std::string>{"x", "y"}, GetMesh());
const Layout output_layout =
*Layout::GetLayout(std::vector<std::string>{"x", "y"}, GetMesh());
const auto specs =
std::vector<Token>{Token(Token::REGULAR, /*begin=*/0, /*end=*/0,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/true,
/*end_mask=*/true)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({"x", "y"}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(1));
EXPECT_EQ(forward->local_tokens()[0].begin_mask, true);
EXPECT_EQ(forward->local_tokens()[0].end_mask, true);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({"x", "y"}));
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_THAT(backward->local_tokens(), SizeIs(1));
EXPECT_EQ(forward->local_tokens()[0].begin_mask, true);
EXPECT_EQ(forward->local_tokens()[0].end_mask, true);
}
TEST_F(InferenceTest, ShardingOnNonSlicedDimension) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{"x", Layout::kUnshardedDim}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{"x", Layout::kUnshardedDim}, GetMesh());
const auto specs =
std::vector<Token>{Token(Token::REGULAR, /*begin=*/0, /*end=*/2,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/2,
/*stride=*/2, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({"x", Layout::kUnshardedDim}));
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, 1); // dim_size(x) == 2
EXPECT_EQ(forward->local_tokens()[1].end, 2);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({"x", Layout::kUnshardedDim}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->local_tokens()[0].end, 1); // dim_size(x) == 2
EXPECT_EQ(backward->local_tokens()[1].end, 2);
}
TEST_F(InferenceTest, StrideOnShardedDimensionNoRelayout1) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, "x"}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, "x"}, GetMesh());
const auto specs =
std::vector<Token>{Token(Token::REGULAR, /*begin=*/0, /*end=*/2,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/4,
/*stride=*/2, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, "x"}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, 2);
EXPECT_EQ(forward->local_tokens()[1].end, 2); // dim_size(x) == 2
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, "x"}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->local_tokens()[0].end, 2); // dim_size(x) == 2
EXPECT_EQ(backward->local_tokens()[1].end, 2);
}
TEST_F(InferenceTest, StrideOnShardedDimensionNoRelayout2) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, "y"}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, "y"}, GetMesh());
const auto specs =
std::vector<Token>{Token(Token::REGULAR, /*begin=*/0, /*end=*/2,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/4,
/*stride=*/2, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, "y"}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, 2);
EXPECT_EQ(forward->local_tokens()[1].end, 4); // dim_size(x) == 1
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, "y"}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->local_tokens()[0].end, 2); // dim_size(x) == 2
EXPECT_EQ(backward->local_tokens()[1].end, 4);
}
TEST_F(InferenceTest, StrideOnShardedDimensionNoRelayout3) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, "x"}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, "x"}, GetMesh());
const auto specs =
std::vector<Token>{Token(Token::REGULAR, /*begin=*/0, /*end=*/2,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/3,
/*stride=*/2, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, "x"}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, 2);
EXPECT_EQ(forward->local_tokens()[1].end, 2); // dim_size(x) == 2
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
EXPECT_EQ(backward->expander_input_layout(), input_layout);
EXPECT_EQ(backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, "x"}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->local_tokens()[0].end, 2); // dim_size(x) == 2
EXPECT_EQ(backward->local_tokens()[1].end, 2);
}
TEST_F(InferenceTest, StrideOnShardedDimensionNeedRelayout) {
const Layout input_layout = *Layout::GetLayout(
std::vector<std::string>{"x", Layout::kUnshardedDim}, GetMesh());
const Layout output_layout = *Layout::GetLayout(
std::vector<std::string>{Layout::kUnshardedDim, Layout::kUnshardedDim},
GetMesh());
const auto specs =
std::vector<Token>{Token(Token::REGULAR, /*begin=*/0, /*end=*/-1,
/*stride=*/1, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false),
Token(Token::REGULAR, /*begin=*/0, /*end=*/4,
/*stride=*/3, /*dynamic_mask=*/false,
/*begin_mask=*/false,
/*end_mask=*/false)};
auto forward = CreateAndRun<ForwardLayoutInference>(
specs, input_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(forward, absl_testing::IsOk());
EXPECT_EQ(
forward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(forward->expander_value_layout(), output_layout);
EXPECT_THAT(forward->local_tokens(), SizeIs(2));
EXPECT_EQ(forward->local_tokens()[0].end, -1);
EXPECT_EQ(forward->local_tokens()[1].end, 4);
auto backward = CreateAndRun<BackwardLayoutInference>(
specs, output_layout, std::vector<int64_t>{2, 4});
ASSERT_THAT(backward, absl_testing::IsOk());
// The backward inferred input_layout prefers replicated layouts for this
// case.
EXPECT_EQ(
backward->expander_input_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_EQ(
backward->expander_value_layout().sharding_spec_strs(),
std::vector<std::string>({Layout::kUnshardedDim, Layout::kUnshardedDim}));
EXPECT_THAT(backward->local_tokens(), SizeIs(2));
EXPECT_EQ(backward->local_tokens()[0].end, -1);
EXPECT_EQ(backward->local_tokens()[1].end, 4);
}
} // namespace
} // namespace slice_util
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/tests/spmd_expander_test.cc
+73
View File
@@ -0,0 +1,73 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/dtensor/mlir/spmd_expander.h"
#include <memory>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "llvm/ADT/DenseMap.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
namespace tensorflow {
namespace dtensor {
namespace {
using ::testing::IsNull;
using ::testing::NotNull;
class DummyExpander : public SPMDExpanderBase {
StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override {
return absl::UnimplementedError("");
}
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& input_layouts) override {
return absl::UnimplementedError("");
}
StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
mlir::Operation* op,
const llvm::DenseMap<int, Layout>& output_layouts) override {
return absl::UnimplementedError("");
}
};
class SPMDExpanderRegistryTest : public ::testing::Test {
public:
SPMDExpanderRegistryTest() {
registry_.RegisterPropagateFn(mlir::TF::AddOp::getOperationName().str(),
std::make_unique<DummyExpander>());
}
protected:
SPMDExpanderRegistry registry_;
};
TEST_F(SPMDExpanderRegistryTest, LookupFromOpName) {
EXPECT_THAT(registry_.GetPropagateFnForFullOpName("tf.Add"), NotNull());
EXPECT_THAT(registry_.GetPropagateFnForFullOpName("Unknown"), IsNull());
}
} // namespace
} // namespace dtensor
} // namespace tensorflow
tensorflow/dtensor/tests/tensor_layout_test.cc
+779
View File
@@ -0,0 +1,779 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include <cstdint>
#include <map>
#include <ostream>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include "absl/container/inlined_vector.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/match.h"
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/dtensor/cc/dstatus.h"
#include "tensorflow/dtensor/proto/layout.pb.h"
namespace tensorflow {
namespace dtensor {
namespace {
using ::testing::ContainsRegex;
using ::testing::ElementsAre;
using ::testing::IsEmpty;
using ::testing::SizeIs;
// Simple implementation of a proto matcher comparing string representations.
// Only works as ShapeProto's textual representation is deterministic.
class ProtoStringMatcher {
public:
explicit ProtoStringMatcher(const tensorflow::protobuf::Message& expected)
: expected_(expected.SerializeAsString()) {}
template <typename Message>
bool MatchAndExplain(const Message& p,
::testing::MatchResultListener*) const {
return p.SerializeAsString() == expected_;
}
void DescribeTo(::std::ostream* os) const { *os << expected_; }
void DescribeNegationTo(::std::ostream* os) const {
*os << "not equal to expected message: " << expected_;
}
private:
const std::string expected_;
};
inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
const tensorflow::protobuf::Message& x) {
return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
}
class LayoutTest : public ::testing::Test {
protected:
Layout BatchLayout() {
return Layout::FromString("sharding_specs:x,batch, mesh:|x=4,batch=8|*TPU")
.value();
}
};
TEST(MeshTest, FromStringEmptyMesh) {
Mesh mesh = Mesh::Empty();
std::string mesh_str = mesh.ToString();
EXPECT_EQ(mesh_str, Mesh::kEmptyMeshString);
}
TEST(MeshTest, FromStringMeshWithGlobalDevices) {
StatusOr<Mesh> mesh = Mesh::FromString(
"mesh:|x=2|0,1|0|/job:localhost/task:0/device:CPU:0|/job:localhost/"
"task:0/device:CPU:0,/job:localhost/task:0/device:CPU:1");
EXPECT_THAT(mesh->global_devices(),
ElementsAre("/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"));
}
TEST(MeshTest, FromStringMeshWithXLASPMDAndGlobalDevices) {
StatusOr<Mesh> mesh = Mesh::FromString(
"mesh:|x=2|0,1|0|/job:localhost/task:0/device:CPU:0|/job:localhost/"
"task:0/device:CPU:1|use_xla_spmd");
EXPECT_TRUE(mesh->use_xla_spmd());
}
TEST(MeshTest, FromStringMeshWithXLASPMD) {
StatusOr<Mesh> mesh = Mesh::FromString(
"mesh:|x=1|0|0|/job:localhost/task:0/device:CPU:0|use_xla_spmd");
EXPECT_TRUE(mesh->use_xla_spmd());
}
TEST(MeshTest, FromStringMeshWithoutXLASPMD) {
StatusOr<Mesh> mesh =
Mesh::FromString("mesh:|x=1|0|0|/job:localhost/task:0/device:CPU:0");
EXPECT_FALSE(mesh->use_xla_spmd());
}
TEST(MeshTest, ToStringMeshWithoutXLASPMD) {
Mesh mesh = Mesh::CreateMesh("MyMesh", /*dim_names=*/{"x"},
/*mesh_shape=*/{2},
/*global_device_ids=*/{0, 1},
/*global_devices_str=*/
{"/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"},
/*local_device_ids=*/{0, 1},
/*local_devices_str=*/
{"/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"},
/*use_xla_spmd=*/false);
EXPECT_TRUE(!absl::StrContains(mesh.ToString(), Mesh::kUseXLASPMDString));
}
TEST(MeshTest, ToStringMeshWithXLASPMD) {
Mesh mesh = Mesh::CreateMesh("MyMesh", /*dim_names=*/{"x"},
/*mesh_shape=*/{2},
/*global_device_ids=*/{0, 1},
/*global_devices_str=*/
{"/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"},
/*local_device_ids=*/{0, 1},
/*local_devices_str=*/
{"/job:localhost/task:0/device:CPU:0",
"/job:localhost/task:0/device:CPU:1"},
/*use_xla_spmd=*/true);
EXPECT_THAT(mesh.ToString(), ContainsRegex(Mesh::kUseXLASPMDString));
}
TEST(MeshTest, FromStringInvalidSingleDeviceMesh) {
EXPECT_THAT(Mesh::FromString("/job:localhost/device:CPU:0"),
absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT));
}
TEST(MeshTest, FromStringSingleDeviceMesh) {
TF_ASSERT_OK_AND_ASSIGN(
Mesh mesh, Mesh::FromString("/job:localhost/task:0/device:CPU:0"));
EXPECT_EQ(mesh.ToString(), "/job:localhost/task:0/device:CPU:0");
}
TEST_F(LayoutTest, FromStringEmptyLayout) {
Layout layout = Layout::Empty();
std::string layout_str = layout.ToString();
TF_ASSERT_OK_AND_ASSIGN(Layout layout_from_str,
Layout::FromString(layout_str));
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_from_str_proto,
layout_from_str.ToProto());
EXPECT_THAT(layout.ToProto(),
absl_testing::IsOkAndHolds(EqualsProto(layout_from_str_proto)));
}
TEST_F(LayoutTest, LayoutToFromString) {
Layout layout = BatchLayout();
std::string layout_str = layout.ToString();
TF_ASSERT_OK_AND_ASSIGN(Layout layout_from_str,
Layout::FromString(layout_str));
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_from_str_proto,
layout_from_str.ToProto());
EXPECT_THAT(layout.ToProto(),
absl_testing::IsOkAndHolds(EqualsProto(layout_from_str_proto)));
}
TEST_F(LayoutTest, LayoutToFromStringNotSharded) {
std::string layout_str = "sharding_specs:x," +
std::string(Layout::kUnshardedDim) +
", mesh:|x=1|0|0|/job:localhost/task:0/device:CPU:0";
EXPECT_EQ(layout_str, Layout::FromString(layout_str)->ToString());
}
TEST_F(LayoutTest, LayoutToFromStringAny) {
std::string layout_str =
"sharding_specs:any, mesh:|x=1|0|0|/job:localhost/task:0/device:CPU:0";
EXPECT_EQ(layout_str, Layout::FromString(layout_str)->ToString());
}
TEST_F(LayoutTest, LayoutToFromStringSingleDevice) {
std::string layout_str =
"maximal:true, mesh:/job:localhost/task:0/device:CPU:0";
EXPECT_EQ(layout_str, Layout::FromString(layout_str)->ToString());
}
TEST_F(LayoutTest, AutoGenerateLayout) {
std::string layout_str = "sharding_specs:x, mesh:|x=2,y=2|*CPU";
std::string exp_layout_str =
"sharding_specs:x, "
"mesh:|x=2,y=2|0,1,2,3|0,1,2,3|/job:localhost/task:0/device:CPU:0,/"
"job:localhost/task:0/device:CPU:1,/job:localhost/task:0/device:CPU:2,/"
"job:localhost/task:0/device:CPU:3";
EXPECT_EQ(exp_layout_str, Layout::FromString(layout_str)->ToString());
}
TEST_F(LayoutTest, MeshToFromString) {
Mesh mesh = BatchLayout().mesh();
std::string mesh_str = mesh.ToString();
TF_ASSERT_OK_AND_ASSIGN(Mesh mesh_from_str, Mesh::FromString(mesh_str));
TF_ASSERT_OK_AND_ASSIGN(MeshProto mesh_from_str_proto,
mesh_from_str.ToProto());
EXPECT_THAT(mesh.ToProto(),
absl_testing::IsOkAndHolds(EqualsProto(mesh_from_str_proto)));
}
TEST_F(LayoutTest, GetType) {
Mesh mesh = BatchLayout().mesh();
EXPECT_TRUE(mesh.is_tpu_mesh());
}
TEST_F(LayoutTest, OnTPUMesh) {
Layout layout = BatchLayout();
EXPECT_TRUE(layout.mesh().is_tpu_mesh());
}
TEST_F(LayoutTest, NumShardsAsVector) {
std::vector<int32_t> shards = {4, 8};
EXPECT_EQ(BatchLayout().num_shards(), shards);
}
TEST_F(LayoutTest, IsReplicated) {
EXPECT_FALSE(BatchLayout().IsFullyReplicated());
}
TEST_F(LayoutTest, MeshDeviceLocations) {
Layout layout = BatchLayout();
absl::InlinedVector<int64_t, 4> offset = {1, 2};
EXPECT_THAT(layout.mesh().device_location(10),
absl_testing::IsOkAndHolds(offset));
offset = {2, 2};
EXPECT_THAT(layout.mesh().device_location(18),
absl_testing::IsOkAndHolds(offset));
offset = {3, 7};
EXPECT_THAT(layout.mesh().device_location(31),
absl_testing::IsOkAndHolds(offset));
EXPECT_FALSE(layout.mesh().device_location(32).ok());
EXPECT_FALSE(layout.mesh().device_location(-1).ok());
}
TEST_F(LayoutTest, ScalarLayout) {
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString("sharding_specs:scalar, mesh:|x=4,y=4|*TPU"));
EXPECT_EQ(layout.num_devices(), 16);
EXPECT_TRUE(layout.mesh().is_tpu_mesh());
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_proto, layout.ToProto());
EXPECT_EQ(layout_proto.mesh_config().mesh_dimensions(0).size(), 4);
EXPECT_EQ(layout.rank(), 0);
}
TEST_F(LayoutTest, ParseSimpleTpuMesh) {
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString("sharding_specs:x, mesh:|x=4,y=4|*TPU"));
EXPECT_EQ(layout.num_devices(), 16);
EXPECT_TRUE(layout.mesh().is_tpu_mesh());
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_proto, layout.ToProto());
EXPECT_EQ(layout_proto.mesh_config().mesh_dimensions(0).size(), 4);
}
TEST_F(LayoutTest, ParseSimpleCpuMesh) {
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString("sharding_specs:x,unsharded, mesh:|x=4,y=4|*CPU"));
EXPECT_EQ(layout.num_devices(), 16);
EXPECT_FALSE(layout.mesh().is_tpu_mesh());
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_proto, layout.ToProto());
EXPECT_EQ(layout_proto.mesh_config().mesh_dimensions(0).size(), 4);
}
TEST_F(LayoutTest, ParseFailsOnRepeatedShardingSpec) {
StatusOr<Layout> maybe_layout =
Layout::FromString("sharding_specs:x,x, mesh:|x=1,y=2|*CPU");
EXPECT_FALSE(maybe_layout.ok());
}
TEST_F(LayoutTest, ParseFailsOnInvalidScalarShardingSpec) {
StatusOr<Layout> maybe_layout =
Layout::FromString("sharding_specs:x,scalar, mesh:|x=1,y=2|*CPU");
EXPECT_FALSE(maybe_layout.ok());
}
TEST_F(LayoutTest, ParseFailsOnShardingSpecOverNonExistentMeshDim) {
StatusOr<Layout> maybe_layout =
Layout::FromString("sharding_specs:x,z, mesh:|x=1,y=2|*CPU");
EXPECT_FALSE(maybe_layout.ok());
}
TEST_F(LayoutTest, ParseFailsOnBadDeviceString) {
auto layout =
Layout::FromString("sharding_specs:x,unsharded, d:TPU mesh:x=4,y=4");
EXPECT_FALSE(layout.ok()) << layout.status();
}
TEST_F(LayoutTest, ParseReplicatedLayout) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString(
"sharding_specs:unsharded,unsharded, mesh:|x=4,y=4|*CPU"));
EXPECT_EQ(layout.num_devices(), 16);
EXPECT_FALSE(layout.mesh().is_tpu_mesh());
EXPECT_TRUE(layout.IsFullyReplicated());
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_proto, layout.ToProto());
EXPECT_EQ(layout_proto.mesh_config().mesh_dimensions(0).size(), 4);
}
TEST_F(LayoutTest, SingleHostFullyReplicatedReducedMesh) {
TF_ASSERT_OK_AND_ASSIGN(
Layout replicated_layout,
Layout::FromString(
"sharding_specs:unsharded,unsharded, mesh:|x=2,y=2|*CPU"));
Mesh reduced_mesh = replicated_layout.ReducedMesh();
EXPECT_EQ(reduced_mesh.size(), 1);
EXPECT_THAT(reduced_mesh.hosts(), SizeIs(1));
}
TEST_F(LayoutTest, SingleHostFullShardedReducedMesh) {
Layout layout = BatchLayout();
Mesh original_mesh = layout.mesh();
Mesh reduced_mesh = layout.ReducedMesh();
EXPECT_EQ(original_mesh.ToString(), reduced_mesh.ToString());
EXPECT_EQ(reduced_mesh.size(), 32);
EXPECT_THAT(reduced_mesh.hosts(), SizeIs(1));
}
TEST_F(LayoutTest, MultiHostReplicatedReducedMesh) {
StatusOr<Layout> layout = Layout::FromString(
"sharding_specs:unsharded,unsharded, "
"mesh:|x=4,y=2|0,1,2,3,4,5,6,7|4,5,6,7|"
"/job:localhost/task:1/device:CPU:0,/job:localhost/task:1/device:CPU:1,"
"/job:localhost/task:1/device:CPU:2,/job:localhost/task:1/device:CPU:3");
Mesh reduced_mesh = layout->ReducedMesh();
EXPECT_EQ(reduced_mesh.size(), 1);
EXPECT_THAT(reduced_mesh.global_device_ids(), ElementsAre(0));
EXPECT_THAT(reduced_mesh.local_device_ids(), IsEmpty());
EXPECT_THAT(reduced_mesh.local_devices(), IsEmpty());
EXPECT_THAT(reduced_mesh.hosts(), IsEmpty());
}
TEST_F(LayoutTest, MultiHostPartiallyShardedReducedMesh) {
StatusOr<Layout> layout = Layout::FromString(
"sharding_specs:x,unsharded, "
"mesh:|x=4,y=2|0,1,2,3,4,5,6,7|4,5,6,7|"
"/job:localhost/task:1/device:CPU:0,/job:localhost/task:1/device:CPU:1,"
"/job:localhost/task:1/device:CPU:2,/job:localhost/task:1/device:CPU:3");
Mesh reduced_mesh = layout->ReducedMesh();
EXPECT_EQ(reduced_mesh.size(), 4);
EXPECT_THAT(reduced_mesh.global_device_ids(), ElementsAre(0, 2, 4, 6));
EXPECT_THAT(reduced_mesh.local_device_ids(), ElementsAre(4, 6));
EXPECT_THAT(reduced_mesh.local_devices(),
ElementsAre("/job:localhost/task:1/device:CPU:0",
"/job:localhost/task:1/device:CPU:2"));
EXPECT_THAT(reduced_mesh.hosts(), SizeIs(1));
}
TEST_F(LayoutTest, MultiHostFullyShardedReducedMesh) {
StatusOr<Layout> layout = Layout::FromString(
"sharding_specs:x,y, "
"mesh:|x=4,y=2|0,1,2,3,4,5,6,7|4,5,6,7|"
"/job:localhost/task:1/device:CPU:0,/job:localhost/task:1/device:CPU:1,"
"/job:localhost/task:1/device:CPU:2,/job:localhost/task:1/device:CPU:3");
Mesh reduced_mesh = layout->ReducedMesh();
EXPECT_EQ(reduced_mesh.size(), 8);
EXPECT_THAT(reduced_mesh.global_device_ids(),
ElementsAre(0, 1, 2, 3, 4, 5, 6, 7));
EXPECT_THAT(reduced_mesh.local_device_ids(), ElementsAre(4, 5, 6, 7));
EXPECT_THAT(reduced_mesh.local_devices(),
ElementsAre("/job:localhost/task:1/device:CPU:0",
"/job:localhost/task:1/device:CPU:1",
"/job:localhost/task:1/device:CPU:2",
"/job:localhost/task:1/device:CPU:3"));
EXPECT_EQ(reduced_mesh.hosts().size(), 1);
}
// TODO(luispazos) Decide if we want this to be the case.
TEST_F(LayoutTest, FlippedShardedMultiHostMeshes) {
StatusOr<Layout> multi_host_layout_1 = Layout::FromString(
"sharding_specs:x,y, "
"mesh:|x=4,y=2|0,1,2,3,4,5,6,7|4,5,6,7|"
"/job:localhost/task:1/device:CPU:0,/job:localhost/task:1/device:CPU:1,"
"/job:localhost/task:1/device:CPU:2,/job:localhost/task:1/device:CPU:3");
StatusOr<Layout> multi_host_layout_2 = Layout::FromString(
"sharding_specs:x,y, "
"mesh:|x=4,y=2|0,1,2,3,4,5,6,7|6,7,4,5|"
"/job:localhost/task:1/device:CPU:2,/job:localhost/task:1/device:CPU:3,"
"/job:localhost/task:1/device:CPU:0,/job:localhost/task:1/device:CPU:1");
Mesh reduced_mesh_1 = multi_host_layout_1->ReducedMesh();
Mesh reduced_mesh_2 = multi_host_layout_2->ReducedMesh();
EXPECT_FALSE(reduced_mesh_1 == reduced_mesh_2);
}
TEST_F(LayoutTest, ShardEqualityOneDim) {
ShardVector shard_vec1;
Shard shard1{1};
shard_vec1.shards.push_back(shard1);
shard_vec1.num_shards_per_dim.push_back(1);
ShardVector shard_vec2;
Shard shard2{2};
Shard shard3{3};
shard_vec2.shards.push_back(shard1);
shard_vec2.shards.push_back(shard2);
shard_vec2.shards.push_back(shard3);
shard_vec2.num_shards_per_dim.push_back(3);
EXPECT_EQ(shard_vec1, shard_vec2);
}
TEST_F(LayoutTest, ShardEqualityOneDimOffset) {
ShardVector shard_vec1;
Shard shard1{3};
shard_vec1.shards.push_back(shard1);
shard_vec1.num_shards_per_dim.push_back(3);
ShardVector shard_vec2;
Shard shard2{7};
Shard shard3{8};
Shard shard4{9};
shard_vec2.shards.push_back(shard2);
shard_vec2.shards.push_back(shard3);
shard_vec2.shards.push_back(shard4);
shard_vec2.num_shards_per_dim.push_back(9);
EXPECT_EQ(shard_vec1, shard_vec2);
}
TEST_F(LayoutTest, ShardEqualityTwoDims) {
auto GenFullVector = [](std::vector<int> num_shards_per_dim) -> ShardVector {
ShardVector shard_vec;
shard_vec.num_shards_per_dim = num_shards_per_dim;
for (int i = 1; i <= num_shards_per_dim[0]; ++i)
for (int j = 1; j <= num_shards_per_dim[1]; ++j) {
Shard shard{i, j};
shard_vec.shards.push_back(shard);
}
return shard_vec;
};
std::vector<int> num_shards_per_dim_1{2, 4};
ShardVector shard_vec1 = GenFullVector(num_shards_per_dim_1);
std::vector<int> num_shards_per_dim_2{3, 3};
ShardVector shard_vec2 = GenFullVector(num_shards_per_dim_2);
EXPECT_EQ(shard_vec1, shard_vec2);
}
TEST_F(LayoutTest, Shards) {
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString("sharding_specs:x,y, mesh:|x=2,y=3|*CPU"));
ShardVector shard_vec = layout.GetShardVector();
std::string expected_shard_vec_str =
"shards:[(1,1),(1,2),(1,3),(2,1),(2,2),(2,3)] num_shards_per_dim:(2,3)";
EXPECT_EQ(shard_vec.ToString(), expected_shard_vec_str);
}
TEST_F(LayoutTest, ShardsInverted) {
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString("sharding_specs:y,x, mesh:|x=2,y=3|*CPU"));
ShardVector shards = layout.GetShardVector();
std::string expected_shards =
"shards:[(1,1),(2,1),(3,1),(1,2),(2,2),(3,2)] num_shards_per_dim:(3,2)";
EXPECT_EQ(shards.ToString(), expected_shards);
}
TEST_F(LayoutTest, HostShardMap) {
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString("sharding_specs:x,y, mesh:TPU|x=2,y=2|*TPU"));
std::string host_name = layout.mesh().hosts()[0];
auto host_map = layout.HostShardMap();
std::string expected_shards =
"shards:[(1,1),(1,2),(2,1),(2,2)] num_shards_per_dim:(2,2)";
EXPECT_EQ(host_map.find(host_name)->second.ToString(), expected_shards);
}
TEST_F(LayoutTest, MultiHostMultiDeviceShards) {
std::string host1 = "/job:localhost/task:0";
std::string host2 = "/job:localhost/task:1";
std::string device1 = "/device:TPU:0";
std::string device2 = "/device:TPU:1";
TF_ASSERT_OK_AND_ASSIGN(
Layout layout,
Layout::FromString(
"sharding_specs:x,unsharded, mesh:TPU|x=4,y=1|0,1,2,3|0,1,2,3|" +
host1 + device1 + "," + host1 + device2 + "," + host2 + device1 +
"," + host2 + device2));
std::string expected_shard_vec =
"shards:[(1,1),(2,1),(3,1),(4,1)] num_shards_per_dim:(4,1)";
EXPECT_EQ(layout.GetShardVector().ToString(), expected_shard_vec);
std::map<std::string, ShardVector> host_shard_map = layout.HostShardMap();
std::string expected_shards_host1 =
"shards:[(1,1),(2,1)] num_shards_per_dim:(4,1)";
ShardVector host1_shard_vec = host_shard_map.find(host1)->second;
EXPECT_EQ(host1_shard_vec.ToString(), expected_shards_host1);
std::string expected_shards_host2 =
"shards:[(3,1),(4,1)] num_shards_per_dim:(4,1)";
ShardVector host2_shard_vec = host_shard_map.find(host2)->second;
EXPECT_EQ(host2_shard_vec.ToString(), expected_shards_host2);
}
TEST_F(LayoutTest, MultiHostCommXYSharded) {
std::string host_0 = "/job:localhost/task:0/";
std::string host_1 = "/job:localhost/task:1/";
StatusOr<Layout> send_layout =
Layout::FromString("sharding_specs:y,x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|" +
host_0 + "device:CPU:0," + host_0 + "device:CPU:1," +
host_1 + "device:CPU:0," + host_1 + "device:CPU:1");
StatusOr<Layout> recv_layout =
Layout::FromString("sharding_specs:x,y, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|" +
host_0 + "device:TPU:0," + host_0 + "device:TPU:1," +
host_1 + "device:TPU:0," + host_1 + "device:TPU:1");
std::vector<std::string> send_hosts = send_layout->ReducedMesh().hosts();
std::vector<std::string> recv_hosts = recv_layout->ReducedMesh().hosts();
EXPECT_TRUE(send_hosts == recv_hosts);
}
TEST_F(LayoutTest, MultiHostCommXSharded) {
std::vector<std::string> hosts{"/job:localhost/task:0",
"/job:localhost/task:1"};
StatusOr<Layout> send_layout = Layout::FromString(
"sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|" + hosts[0] +
"/device:CPU:0," + hosts[0] + "/device:CPU:1," + hosts[1] +
"/device:CPU:0," + hosts[1] + "/device:CPU:1");
StatusOr<Layout> recv_layout = Layout::FromString(
"sharding_specs:x, mesh:|x=2,y=2|0,1,2,3|0,1,2,3|" + hosts[0] +
"/device:TPU:0," + hosts[0] + "/device:TPU:1," + hosts[1] +
"/device:TPU:0," + hosts[1] + "/device:TPU:1");
std::vector<std::string> send_hosts = send_layout->ReducedMesh().hosts();
std::vector<std::string> recv_hosts = recv_layout->ReducedMesh().hosts();
EXPECT_TRUE(send_hosts == recv_hosts);
std::map<std::string, ShardVector> send_host_shard_map =
send_layout->HostShardMap();
std::map<std::string, ShardVector> recv_host_shard_map =
recv_layout->HostShardMap();
// Check shards match in each host.
for (const std::string& host : hosts) {
ShardVector shard_vec_in_send_host = send_host_shard_map.find(host)->second;
ShardVector shard_vec_in_recv_host = recv_host_shard_map.find(host)->second;
EXPECT_EQ(shard_vec_in_send_host, shard_vec_in_recv_host);
}
}
TEST_F(LayoutTest, Transposed2DLayout) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString("sharding_specs:x,y, mesh:|x=2,y=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto expected_layout,
Layout::FromString("sharding_specs:y,x, mesh:|x=2,y=2|*CPU"));
EXPECT_THAT(Layout::Transposed2D(layout),
absl_testing::IsOkAndHolds(expected_layout));
}
TEST_F(LayoutTest, Transposed2DLayoutWithBatch) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString(
"sharding_specs:b1,b2,x,y, mesh:|x=2,y=2,b1=2,b2=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto expected_layout,
Layout::FromString(
"sharding_specs:b1,b2,y,x, mesh:|x=2,y=2,b1=2,b2=2|*CPU"));
EXPECT_THAT(Layout::Transposed2D(layout),
absl_testing::IsOkAndHolds(expected_layout));
}
TEST_F(LayoutTest, MeshDimensionIndex) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString("sharding_specs:x,y, mesh:|x=2,y=2|*CPU"));
EXPECT_THAT(layout.mesh().idx_for_dim("x"), absl_testing::IsOkAndHolds(0));
EXPECT_THAT(layout.mesh().idx_for_dim("y"), absl_testing::IsOkAndHolds(1));
}
TEST_F(LayoutTest, TruncateBeginning) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString("sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto expected_layout,
Layout::FromString("sharding_specs:x, mesh:CPU|x=2,y=2|*CPU"));
EXPECT_EQ(layout.Truncate(/*split_point=*/1), expected_layout);
}
TEST_F(LayoutTest, TruncateEnd) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString("sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto expected_layout,
Layout::FromString("sharding_specs:y, mesh:CPU|x=2,y=2|*CPU"));
EXPECT_EQ(layout.Truncate(/*split_point=*/1, /*end=*/true), expected_layout);
}
TEST_F(LayoutTest, Concatenate) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout_1,
Layout::FromString("sharding_specs:x, mesh:CPU|x=2,y=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto layout_2,
Layout::FromString("sharding_specs:y, mesh:CPU|x=2,y=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto expected_layout,
Layout::FromString("sharding_specs:x,y, mesh:CPU|x=2,y=2|*CPU"));
EXPECT_THAT(ConcatenateLayouts(layout_1, layout_2),
absl_testing::IsOkAndHolds(expected_layout));
}
TEST_F(LayoutTest, ConcatenateDifferentMesh) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout_1, Layout::FromString("sharding_specs:x, mesh:CPU|x=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto layout_2, Layout::FromString("sharding_specs:y, mesh:CPU|y=2|*CPU"));
auto layout = ConcatenateLayouts(layout_1, layout_2);
EXPECT_FALSE(layout.ok()) << layout.status();
}
TEST_F(LayoutTest, ConcatenateSameDimension) {
TF_ASSERT_OK_AND_ASSIGN(
auto layout_1,
Layout::FromString("sharding_specs:x, mesh:CPU|x=2,y=2|*CPU"));
TF_ASSERT_OK_AND_ASSIGN(
auto layout_2,
Layout::FromString("sharding_specs:x, mesh:CPU|x=2,y=2|*CPU"));
auto layout = ConcatenateLayouts(layout_1, layout_2);
EXPECT_FALSE(layout.ok()) << layout.status();
}
TEST_F(LayoutTest, EmptyMeshDeviceType) {
auto mesh = Mesh::Empty();
EXPECT_EQ(mesh.device_type(), std::string());
}
TEST_F(LayoutTest, ConvertMeshDeviceType) {
TF_ASSERT_OK_AND_ASSIGN(Mesh mesh,
Mesh::FromString("mesh_name|x=2,batch=1|*TPU"));
TF_ASSERT_OK_AND_ASSIGN(Mesh cpu_mesh, mesh.ToDeviceType("CPU"));
EXPECT_TRUE(cpu_mesh.is_cpu_mesh());
std::string expected_task_name = "/job:localhost/replica:0/task:0/";
TF_ASSERT_OK_AND_ASSIGN(
Mesh expected_mesh,
Mesh::FromString("|x=2,batch=1|0,1|0,1|" + expected_task_name +
"device:CPU:0," + expected_task_name + "device:CPU:1"));
EXPECT_EQ(cpu_mesh, expected_mesh);
}
TEST_F(LayoutTest, EquivalentLayout) {
TF_ASSERT_OK_AND_ASSIGN(
Layout fully_sharded,
Layout::FromString("sharding_specs:x,y, mesh:|x=2,y=1|*TPU"));
TF_ASSERT_OK_AND_ASSIGN(
Layout x_sharded,
Layout::FromString("sharding_specs:x,unsharded, mesh:|x=2,y=1|*TPU"));
TF_ASSERT_OK_AND_ASSIGN(
Layout y_sharded,
Layout::FromString("sharding_specs:unsharded,y, mesh:|x=2,y=1|*TPU"));
EXPECT_TRUE(fully_sharded.IsEquivalent(x_sharded));
EXPECT_TRUE(x_sharded.IsEquivalent(fully_sharded));
EXPECT_FALSE(fully_sharded.IsEquivalent(y_sharded));
EXPECT_FALSE(y_sharded.IsEquivalent(fully_sharded));
}
TEST_F(LayoutTest, GetSingleDeviceMeshEmptyDeviceString) {
EXPECT_THAT(Mesh::GetSingleDeviceMesh(""),
absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT));
}
TEST_F(LayoutTest, GetSingleDeviceMeshSuccess) {
TF_ASSERT_OK_AND_ASSIGN(
auto mesh, Mesh::FromString("/job:localhost/task:1/device:CPU:0"));
EXPECT_THAT(Mesh::GetSingleDeviceMesh("/job:localhost/task:1/device:CPU:0"),
absl_testing::IsOkAndHolds(mesh));
}
TEST_F(LayoutTest, GetSingleDeviceLayoutInvalidMesh) {
auto mesh = Mesh::Empty();
EXPECT_THAT(Layout::GetLayout(Layout::LayoutType::kSingleDevice, {}, mesh),
absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT));
}
TEST_F(LayoutTest, GetSingleDeviceLayoutSuccess) {
TF_ASSERT_OK_AND_ASSIGN(
auto mesh, Mesh::FromString("/job:localhost/task:1/device:CPU:0"));
TF_ASSERT_OK_AND_ASSIGN(
auto layout,
Layout::FromString(
"maximal:true, mesh:/job:localhost/task:1/device:CPU:0"));
EXPECT_THAT(Layout::GetLayout(Layout::LayoutType::kSingleDevice, {}, mesh),
absl_testing::IsOkAndHolds(layout));
}
TEST(DynamicSizeTest, IsDynamicSize) {
EXPECT_TRUE(IsDynamicSize(-1));
EXPECT_TRUE(IsDynamicSize(mlir::ShapedType::kDynamic));
EXPECT_FALSE(IsDynamicSize(10));
}
TEST_F(LayoutTest, LayoutType) {
TF_ASSERT_OK_AND_ASSIGN(
auto maximal,
Layout::FromString(
"maximal:true, mesh:/job:localhost/task:1/device:CPU:0"));
EXPECT_EQ(maximal.type(), Layout::LayoutType::kSingleDevice);
TF_ASSERT_OK_AND_ASSIGN(auto parted,
Layout::FromString("parted:x, mesh:|x=2|*TPU"));
EXPECT_EQ(parted.type(), Layout::LayoutType::kParted);
TF_ASSERT_OK_AND_ASSIGN(
auto static_layout,
Layout::FromString("sharding_specs:x, mesh:|x=2|*TPU"));
EXPECT_EQ(static_layout.type(), Layout::LayoutType::kStatic);
}
TEST_F(LayoutTest, PartedLayoutToFromString) {
TF_ASSERT_OK_AND_ASSIGN(Layout layout, BatchLayout().ToParted());
std::string layout_str = layout.ToString();
TF_ASSERT_OK_AND_ASSIGN(Layout layout_from_str,
Layout::FromString(layout_str));
TF_ASSERT_OK_AND_ASSIGN(LayoutProto layout_from_str_proto,
layout_from_str.ToProto());
EXPECT_THAT(layout.ToProto(),
absl_testing::IsOkAndHolds(EqualsProto(layout_from_str_proto)));
}
TEST_F(LayoutTest, RaggedLayoutEqual) {
TF_ASSERT_OK_AND_ASSIGN(
Layout fully_sharded,
Layout::FromString("sharding_specs:x,y, mesh:|x=2,y=1|*TPU"));
TF_ASSERT_OK_AND_ASSIGN(
Layout x_sharded,
Layout::FromString("sharding_specs:x,unsharded, mesh:|x=2,y=1|*TPU"));
TF_ASSERT_OK_AND_ASSIGN(
Layout x_parted,
Layout::FromString("parted:x,unsharded, mesh:|x=2,y=1|*TPU"));
TF_ASSERT_OK_AND_ASSIGN(Layout x_y_parted,
Layout::FromString("parted:x,y, mesh:|x=2,y=1|*TPU"));
// Test that 'IsEquivalent' and '==' take layout type into account.
EXPECT_TRUE(x_parted.IsEquivalent(x_y_parted));
EXPECT_TRUE(x_y_parted.IsEquivalent(x_parted));
EXPECT_FALSE(x_sharded.IsEquivalent(x_parted));
EXPECT_FALSE(fully_sharded.IsEquivalent(x_y_parted));
EXPECT_FALSE(x_sharded == x_parted);
EXPECT_FALSE(fully_sharded == x_y_parted);
}
} // namespace
} // namespace dtensor
} // namespace tensorflow