[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
This commit is contained in:
Bixia Zheng 2019-01-30 09:52:02 -08:00 committed by TensorFlower Gardener
parent edfb0bba7f
commit 2cb0880812
3 changed files with 18 additions and 3 deletions

View File

@ -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(),

View File

@ -658,8 +658,14 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder,
tensorflow::int64 parameter_number;
if (!ParseToken(TokKind::kLparen,
"expects '(' before parameter number") ||
!ParseInt64(&parameter_number) ||
!ParseToken(TokKind::kRparen, "expects ')' after parameter number") ||
!ParseInt64(&parameter_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;
}

View File

@ -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