":hlo",
"//tensorflow/compiler/xla:test",
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
+ "//tensorflow/compiler/xla/tools/parser:hlo_parser",
"//tensorflow/core:lib",
],
)
":hlo_matchers",
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
- "//tensorflow/compiler/xla/tools/parser:hlo_parser",
],
)
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla/tests:hlo_test_base",
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
+ "//tensorflow/compiler/xla/tools/parser:hlo_parser",
],
)
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
#include "tensorflow/core/lib/gtl/optional.h"
namespace xla {
return ::testing::MakeMatcher(
new ::xla::testing::HloShardingMatcher(sharding));
}
+// Matcher for Sharding from sharding string
+inline ::testing::Matcher<const ::xla::HloInstruction*> Sharding(
+ tensorflow::StringPiece sharding) {
+ return ::testing::MakeMatcher(new ::xla::testing::HloShardingMatcher(
+ xla::tools::ParseSharding(sharding).ValueOrDie()));
+}
// Verifies that no HloSharding is set for an HLO instruction.
inline ::testing::Matcher<const ::xla::HloInstruction*> NoSharding() {
return ::testing::MakeMatcher(
#include "tensorflow/compiler/xla/service/hlo_matchers.h"
#include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
namespace op = xla::testing::opcode_matchers;
using ::testing::_;
"param.1");
p1->set_sharding(HloSharding::AssignDevice(1));
+ auto tuple_shape = ShapeUtil::MakeTupleShape(
+ {ShapeUtil::MakeShape(F32, {7}), ShapeUtil::MakeShape(S32, {9}),
+ ShapeUtil::MakeShape(F32, {11})});
+ auto p2 = HloInstruction::CreateParameter(1, tuple_shape, "param.2");
+ Array<int64> assignment({2});
+ assignment.SetValues({0, 1});
+ auto sharding = HloSharding::Tuple(
+ tuple_shape,
+ {HloSharding::Tile(ShapeUtil::MakeShape(F32, {5}), assignment),
+ HloSharding::AssignDevice(1), HloSharding::Replicate()});
+ p2->set_sharding(sharding);
+
EXPECT_THAT(p0.get(), op::NoSharding());
EXPECT_THAT(p0.get(),
::testing::Not(op::Sharding(HloSharding::AssignDevice(1))));
::testing::Not(op::Sharding(HloSharding::AssignDevice(0))));
EXPECT_THAT(p1.get(), op::Sharding(HloSharding::AssignDevice(1)));
+ EXPECT_THAT(
+ p2.get(),
+ op::Sharding(
+ "{{f32[5] devices=[2]0,1}, {maximal device=1}, {replicated}}"));
+
EXPECT_THAT(Explain(p0.get(), op::Sharding(HloSharding::AssignDevice(1))),
"%param.0 = f32[5]{0} parameter(0) has no sharding (expected: "
"{maximal device=1})");
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/xla/service/hlo_sharding.h"
-
#include <set>
#include <unordered_map>
#include <utility>
#include "tensorflow/compiler/xla/test.h"
#include "tensorflow/compiler/xla/test_helpers.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
#include "tensorflow/compiler/xla/util.h"
namespace xla {
EXPECT_EQ(oss.str(), "{f32[3,5,7,11] devices=[1,1,2,2]0,1,2,3}");
}
+TEST_F(HloShardingTest, Parse) {
+ auto check = [](const HloSharding& sharding) {
+ TF_ASSERT_OK_AND_ASSIGN(auto parsed_sharding,
+ tools::ParseSharding(sharding.ToString()));
+ EXPECT_EQ(sharding, parsed_sharding);
+ };
+ check(HloSharding::Replicate());
+ check(HloSharding::AssignDevice(2));
+ check(HloSharding::Tile(ShapeUtil::MakeShape(F32, {3, 1, 3, 7}),
+ Array4D<int64>({{{{0}, {1}}}})));
+ // Empty tuple.
+ check(HloSharding::Tuple(ShapeUtil::MakeTupleShape({}), {}));
+ {
+ // Non-nested tuple.
+ auto tuple_shape =
+ ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {3, 1, 5, 7}),
+ ShapeUtil::MakeShape(F32, {3, 5, 7}),
+ ShapeUtil::MakeShape(F32, {3, 7})});
+ check(HloSharding::Tuple(
+ tuple_shape, {HloSharding::Tile(ShapeUtil::MakeShape(F32, {3, 1, 3, 7}),
+ Array4D<int64>({{{{0}, {1}}}})),
+ HloSharding::Replicate(), HloSharding::AssignDevice(1)}));
+ }
+ {
+ // Nested tuple.
+ auto tuple_shape = ShapeUtil::MakeTupleShape(
+ {ShapeUtil::MakeShape(F32, {3, 1, 5, 7}),
+ ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {3, 5, 7}),
+ ShapeUtil::MakeShape(F32, {3, 7})})});
+ std::vector<HloSharding> leaf_shardings = {
+ HloSharding::Tile(ShapeUtil::MakeShape(F32, {3, 1, 3, 7}),
+ Array4D<int64>({{{{0}, {1}}}})),
+ HloSharding::Replicate(), HloSharding::AssignDevice(1)};
+ ShapeTree<HloSharding> sharding_tree(tuple_shape, HloSharding::Replicate());
+ // Assign leaf_shardings to sharding_tree leaves.
+ auto it = leaf_shardings.begin();
+ for (auto& index_to_sharding : sharding_tree.leaves()) {
+ index_to_sharding.second = *it++;
+ }
+ check(HloSharding::Tuple(sharding_tree));
+ }
+}
+
} // namespace
} // namespace xla
// Returns the error information.
string GetError() const { return Join(error_, "\n"); }
+ // Stand alone parsing for sharding. The parser string is supposed to
+ // contain the body of the sharding, i.e. just the rhs of the "sharding={...}"
+ // attribute string.
+ StatusOr<HloSharding> ParseShardingOnly();
+
private:
// ParseXXX returns false if an error occurred.
bool ParseHloModule();
return true;
}
+StatusOr<HloSharding> HloParser::ParseShardingOnly() {
+ lexer_.Lex();
+ OpSharding op_sharding;
+ if (!ParseSharding(&op_sharding)) {
+ return InvalidArgument("Syntax error:\n%s", GetError().c_str());
+ }
+ if (lexer_.GetKind() != TokKind::kEof) {
+ return InvalidArgument("Syntax error:\nExtra content after sharding");
+ }
+ return HloSharding::FromProto(op_sharding);
+}
+
} // namespace
StatusOr<std::unique_ptr<HloModule>> Parse(StringPiece str,
return Parse(str, config);
}
+StatusOr<HloSharding> ParseSharding(tensorflow::StringPiece str) {
+ HloModuleConfig config;
+ HloParser parser(str, config);
+ return parser.ParseShardingOnly();
+}
+
} // namespace tools
} // namespace xla
// format, parses the string and creates a HloModule with default config.
StatusOr<std::unique_ptr<HloModule>> Parse(tensorflow::StringPiece str);
+// Parse sharding from str. str is supposed to contain the body of the
+// sharding, i.e. just the rhs of the "sharding={...}" attribute string.
+StatusOr<HloSharding> ParseSharding(tensorflow::StringPiece str);
+
} // namespace tools
} // namespace xla