From 2cb0880812a9a99485feb2bbbfa9b8aca430a6c6 Mon Sep 17 00:00:00 2001 From: Bixia Zheng Date: Wed, 30 Jan 2019 09:52:02 -0800 Subject: [PATCH] [XLA] Fix problems in handling kParameter HLO instruction with negative parameter number. When the input is HLO text that contains kParameter instructions with negative parameter numbers, an HLO tool such as run_hlo_module, crashes in creating the HloComptation. We fix the HLO parser to report errors instead. Add an HLO parser test case. When the input is a binary HLO proto that contains kParameter instructions with negative parameter numbers, run_hlo_module crashes in verifying the module. We fix the DynamicParameterBinding verifier to report errors instead. Add an HLO proto corpus for fuzzing. PiperOrigin-RevId: 231612816 --- .../compiler/xla/service/dynamic_parameter_binding.cc | 3 ++- tensorflow/compiler/xla/service/hlo_parser.cc | 10 ++++++++-- tensorflow/compiler/xla/service/hlo_parser_test.cc | 8 ++++++++ 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc index e9c8aa03e2a..7f0ae692f74 100644 --- a/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc @@ -112,7 +112,8 @@ Status DynamicParameterBinding::Verify(const HloModule& module) const { return ForEachBinding([&](const DynamicParameter& dynamic_parameter, const DynamicDimension& dynamic_dimension) -> Status { - TF_RET_CHECK(dynamic_parameter.parameter_num < entry->num_parameters()); + TF_RET_CHECK(dynamic_parameter.parameter_num >= 0 && + dynamic_parameter.parameter_num < entry->num_parameters()); TF_RET_CHECK(dynamic_dimension.parameter_num < entry->num_parameters()); TF_RET_CHECK(ShapeUtil::IndexIsValid( entry->parameter_instruction(dynamic_parameter.parameter_num)->shape(), diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc index 1a930370a9e..d7fa6b4a643 100644 --- a/tensorflow/compiler/xla/service/hlo_parser.cc +++ b/tensorflow/compiler/xla/service/hlo_parser.cc @@ -658,8 +658,14 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, tensorflow::int64 parameter_number; if (!ParseToken(TokKind::kLparen, "expects '(' before parameter number") || - !ParseInt64(¶meter_number) || - !ParseToken(TokKind::kRparen, "expects ')' after parameter number") || + !ParseInt64(¶meter_number)) { + return false; + } + if (parameter_number < 0) { + Error(lexer_.GetLoc(), "parameter number must be >= 0"); + return false; + } + if (!ParseToken(TokKind::kRparen, "expects ')' after parameter number") || !ParseAttributes(attrs)) { return false; } diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index ff52a89d1bd..6eee767900d 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -2558,5 +2558,13 @@ TEST_F(HloParserTest, ParseDynamicTuple) { << "actual: " << ShapeUtil::HumanString(actual); } +TEST_F(HloParserTest, NegativeParameterNumber) { + const string hlo_string = "par0 = f32[3,5] parameter(-1)"; + auto result = ParseHloString(hlo_string); + ASSERT_FALSE(result.status().ok()); + EXPECT_THAT(result.status().error_message(), + ::testing::HasSubstr("parameter number must be >= 0")); +} + } // namespace } // namespace xla