#include "tensorflow/compiler/xla/service/hlo_proto_util.h"
+#include <string>
+
+#include "tensorflow/compiler/xla/util.h"
+
namespace xla {
+namespace {
+
+// Returns the entry computation of the HLO module in the given HloProto.
+StatusOr<const HloComputationProto*> GetEntryComputation(
+ const HloProto& hlo_proto) {
+ if (!hlo_proto.has_hlo_module()) {
+ return NotFound("HloProto missing HloModuleProto.");
+ }
+
+ if (hlo_proto.hlo_module().entry_computation_name().empty()) {
+ return NotFound("HloProto has empty entry computation name.");
+ }
+
+ const string& entry_computation_name =
+ hlo_proto.hlo_module().entry_computation_name();
+ const HloComputationProto* entry_computation = nullptr;
+ for (const HloComputationProto& computation :
+ hlo_proto.hlo_module().computations()) {
+ if (computation.name() == entry_computation_name) {
+ if (entry_computation == nullptr) {
+ entry_computation = &computation;
+ } else {
+ return InvalidArgument(
+ "HloProto has multiple computations with entry computation named "
+ "%s.",
+ entry_computation_name.c_str());
+ }
+ }
+ }
+ if (entry_computation == nullptr) {
+ return InvalidArgument("HloProto has no entry computation named %s.",
+ entry_computation_name.c_str());
+ }
+ return entry_computation;
+}
+
+// Returns the root instruction of the given computation proto.
+StatusOr<const HloInstructionProto*> GetRootInstruction(
+ const HloComputationProto& computation) {
+ if (computation.root_name().empty()) {
+ return InvalidArgument("Missing root instruction name.");
+ }
+
+ const HloInstructionProto* root = nullptr;
+ for (const HloInstructionProto& instruction : computation.instructions()) {
+ if (instruction.name() == computation.root_name()) {
+ if (root == nullptr) {
+ root = &instruction;
+ } else {
+ return InvalidArgument(
+ "Computation has multiple instructions named %s.",
+ computation.root_name().c_str());
+ }
+ }
+ }
+ if (root == nullptr) {
+ return InvalidArgument("Computation has no instruction named %s.",
+ computation.root_name().c_str());
+ }
+ return root;
+}
+
+// Returns the parameters of the given computation. Parameter numbers are
+// checked for validity and contiguousness.
+StatusOr<std::vector<const HloInstructionProto*>> GetParameters(
+ const HloComputationProto& computation) {
+ std::vector<const HloInstructionProto*> parameters;
+ for (const HloInstructionProto& instruction : computation.instructions()) {
+ if (instruction.opcode() == HloOpcodeString(HloOpcode::kParameter)) {
+ parameters.push_back(&instruction);
+ }
+ }
+
+ // Verify the uniqueness and validity of the parameter numbers.
+ tensorflow::gtl::FlatSet<int64> parameter_numbers;
+ for (const HloInstructionProto* parameter : parameters) {
+ if (parameter->parameter_number() < 0 ||
+ parameter->parameter_number() >= parameters.size()) {
+ return InvalidArgument(
+ "Parameter instruction %s has invalid parameter number %lld.",
+ parameter->name().c_str(), parameter->parameter_number());
+ }
+ if (parameter_numbers.count(parameter->parameter_number()) != 0) {
+ return InvalidArgument(
+ "Multiple parameter instructions have parameter number %lld.",
+ parameter->parameter_number());
+ }
+ parameter_numbers.insert(parameter->parameter_number());
+ }
+
+ std::sort(parameters.begin(), parameters.end(),
+ [](const HloInstructionProto* a, const HloInstructionProto* b) {
+ return a->parameter_number() < b->parameter_number();
+ });
+
+ return parameters;
+}
+
+} // namespace
+
HloProto MakeHloProto(const HloModule& module,
const BufferAssignment& assignment) {
HloOrderingProto proto_ordering =
return proto;
}
+StatusOr<std::vector<const Shape*>> EntryComputationParameterShapes(
+ const HloProto& hlo_proto) {
+ TF_ASSIGN_OR_RETURN(const HloComputationProto* entry_computation,
+ GetEntryComputation(hlo_proto));
+ TF_ASSIGN_OR_RETURN(std::vector<const HloInstructionProto*> parameters,
+ GetParameters(*entry_computation));
+ std::vector<const Shape*> parameter_shapes;
+ for (const HloInstructionProto* parameter : parameters) {
+ if (!parameter->has_shape()) {
+ return InvalidArgument("Parameter instruction %s is missing shape.",
+ parameter->name().c_str());
+ }
+ parameter_shapes.push_back(¶meter->shape());
+ }
+ return parameter_shapes;
+}
+
+StatusOr<const Shape*> EntryComputationOutputShape(const HloProto& hlo_proto) {
+ TF_ASSIGN_OR_RETURN(const HloComputationProto* entry_computation,
+ GetEntryComputation(hlo_proto));
+
+ TF_ASSIGN_OR_RETURN(const HloInstructionProto* root,
+ GetRootInstruction(*entry_computation));
+ if (!root->has_shape()) {
+ return InvalidArgument("Instruction %s is missing shape.",
+ root->name().c_str());
+ }
+
+ return &root->shape();
+}
+
} // namespace xla
--- /dev/null
+/* Copyright 2018 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/compiler/xla/service/hlo_proto_util.h"
+
+#include "tensorflow/compiler/xla/service/hlo.pb.h"
+#include "tensorflow/compiler/xla/service/hlo_opcode.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/status_macros.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+
+namespace xla {
+namespace {
+
+class HloProtoUtilTest : public ::testing::Test {};
+
+TEST_F(HloProtoUtilTest, ParamsAndOutputShape) {
+ HloProto hlo_proto;
+ HloModuleProto* module = hlo_proto.mutable_hlo_module();
+ module->set_entry_computation_name("entry");
+ HloComputationProto* computation = module->add_computations();
+ computation->set_name("entry");
+ computation->set_root_name("root");
+
+ HloInstructionProto* param0 = computation->add_instructions();
+ param0->set_opcode(HloOpcodeString(HloOpcode::kParameter));
+ param0->set_parameter_number(0);
+ *param0->mutable_shape() = ShapeUtil::MakeShape(F32, {42});
+
+ HloInstructionProto* param2 = computation->add_instructions();
+ param2->set_opcode(HloOpcodeString(HloOpcode::kParameter));
+ param2->set_parameter_number(2);
+ *param2->mutable_shape() = ShapeUtil::MakeShape(S32, {1, 2, 3});
+
+ HloInstructionProto* param1 = computation->add_instructions();
+ param1->set_opcode(HloOpcodeString(HloOpcode::kParameter));
+ param1->set_parameter_number(1);
+ *param1->mutable_shape() = ShapeUtil::MakeShape(F64, {});
+
+ HloInstructionProto* root = computation->add_instructions();
+ root->set_opcode(HloOpcodeString(HloOpcode::kAdd));
+ root->set_name("root");
+ *root->mutable_shape() = ShapeUtil::MakeShape(U8, {2});
+
+ VLOG(1) << hlo_proto.DebugString();
+
+ TF_ASSERT_OK_AND_ASSIGN(std::vector<const Shape*> parameter_shapes,
+ EntryComputationParameterShapes(hlo_proto));
+ ASSERT_EQ(parameter_shapes.size(), 3);
+ EXPECT_TRUE(
+ ShapeUtil::Equal(*parameter_shapes[0], ShapeUtil::MakeShape(F32, {42})));
+ EXPECT_TRUE(
+ ShapeUtil::Equal(*parameter_shapes[1], ShapeUtil::MakeShape(F64, {})));
+ EXPECT_TRUE(ShapeUtil::Equal(*parameter_shapes[2],
+ ShapeUtil::MakeShape(S32, {1, 2, 3})));
+
+ TF_ASSERT_OK_AND_ASSIGN(const Shape* output_shape,
+ EntryComputationOutputShape(hlo_proto));
+ EXPECT_TRUE(ShapeUtil::Equal(*output_shape, ShapeUtil::MakeShape(U8, {2})));
+}
+
+TEST_F(HloProtoUtilTest, ParamsAndOutputShapeNoParameters) {
+ HloProto hlo_proto;
+ HloModuleProto* module = hlo_proto.mutable_hlo_module();
+ module->set_entry_computation_name("entry");
+ HloComputationProto* computation = module->add_computations();
+ computation->set_name("entry");
+ computation->set_root_name("root");
+
+ HloInstructionProto* root = computation->add_instructions();
+ root->set_opcode(HloOpcodeString(HloOpcode::kAdd));
+ root->set_name("root");
+ *root->mutable_shape() = ShapeUtil::MakeShape(U8, {2});
+
+ TF_ASSERT_OK_AND_ASSIGN(std::vector<const Shape*> parameter_shapes,
+ EntryComputationParameterShapes(hlo_proto));
+ ASSERT_EQ(parameter_shapes.size(), 0);
+}
+
+TEST_F(HloProtoUtilTest, ParamsAndOutputShapeMissingModule) {
+ HloProto hlo_proto;
+
+ auto status = EntryComputationParameterShapes(hlo_proto).status();
+ ASSERT_FALSE(status.ok());
+ ASSERT_THAT(status.error_message(),
+ ::testing::HasSubstr("missing HloModuleProto"));
+}
+
+TEST_F(HloProtoUtilTest, ParamsAndOutputShapeMissingEntryComputation) {
+ HloProto hlo_proto;
+ HloModuleProto* module = hlo_proto.mutable_hlo_module();
+ module->set_entry_computation_name("entry");
+ HloComputationProto* computation = module->add_computations();
+ computation->set_name("not_entry");
+
+ auto status = EntryComputationParameterShapes(hlo_proto).status();
+ ASSERT_FALSE(status.ok());
+ ASSERT_THAT(status.error_message(),
+ ::testing::HasSubstr("has no entry computation named"));
+}
+
+TEST_F(HloProtoUtilTest, OutputShapeMissingEntryRoot) {
+ HloProto hlo_proto;
+ HloModuleProto* module = hlo_proto.mutable_hlo_module();
+ module->set_entry_computation_name("entry");
+ HloComputationProto* computation = module->add_computations();
+ computation->set_name("entry");
+ computation->set_root_name("root");
+
+ auto status = EntryComputationOutputShape(hlo_proto).status();
+ ASSERT_FALSE(status.ok());
+ ASSERT_THAT(status.error_message(),
+ ::testing::HasSubstr("has no instruction named"));
+}
+
+TEST_F(HloProtoUtilTest, ParamsShapesMissingParameterNumbers) {
+ HloProto hlo_proto;
+ HloModuleProto* module = hlo_proto.mutable_hlo_module();
+ module->set_entry_computation_name("entry");
+ HloComputationProto* computation = module->add_computations();
+ computation->set_name("entry");
+ computation->set_root_name("root");
+
+ HloInstructionProto* param0 = computation->add_instructions();
+ param0->set_opcode(HloOpcodeString(HloOpcode::kParameter));
+ param0->set_parameter_number(0);
+ *param0->mutable_shape() = ShapeUtil::MakeShape(F32, {42});
+
+ HloInstructionProto* param2 = computation->add_instructions();
+ param2->set_opcode(HloOpcodeString(HloOpcode::kParameter));
+ param2->set_parameter_number(2);
+ *param2->mutable_shape() = ShapeUtil::MakeShape(S32, {1, 2, 3});
+
+ HloInstructionProto* root = computation->add_instructions();
+ root->set_opcode(HloOpcodeString(HloOpcode::kAdd));
+ root->set_name("root");
+ *root->mutable_shape() = ShapeUtil::MakeShape(U8, {2});
+
+ auto status = EntryComputationParameterShapes(hlo_proto).status();
+ ASSERT_FALSE(status.ok());
+ ASSERT_THAT(status.error_message(),
+ ::testing::HasSubstr("invalid parameter number"));
+}
+
+} // namespace
+} // namespace xla
projects / platform / upstream / tensorflow.git / commitdiff