[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:
parent
edfb0bba7f
commit
2cb0880812
@ -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(),
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user