Update HloBufferDonorConfig::Verify and HloVerifier. The overlap …

…between `buffer_donor_config` and `input_output_alias_config` is not allowed. PiperOrigin-RevId: 565803412
openxla · Sep 15, 2023 · 49314f4 · 49314f4
1 parent de6e187
commit 49314f4
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 1 deletion.
diff --git a/xla/hlo/ir/hlo_input_output_alias_config.cc b/xla/hlo/ir/hlo_input_output_alias_config.cc
@@ -320,6 +320,7 @@ bool HloBufferDonorConfig::ParameterIsBufferDonor(
 
 Status HloBufferDonorConfig::Verify(const HloModule& module) const {
  const HloComputation* entry = module.entry_computation();
+ const auto& alias_config = module.input_output_alias_config();
  for (const auto& donor : buffer_donor_) {
  TF_RET_CHECK(donor.param_number >= 0);
  TF_RET_CHECK(donor.param_number < entry->num_parameters());
@@ -331,6 +332,13 @@ Status HloBufferDonorConfig::Verify(const HloModule& module) const {
  const Shape& param_subshape =
  ShapeUtil::GetSubshape(param_shape, donor.param_index);
  TF_RET_CHECK(LayoutUtil::IsDenseArray(param_subshape));
+
+ if (alias_config.ParameterHasAlias(donor.param_number, donor.param_index)) {
+ return InternalError(
+ "Input %lld at index %s is registered as a buffer donor. However, it "
+ "is also in the input output alias config.",
+ donor.param_number, donor.param_index.ToString());
+ }
  }
 
  // Since buffer_donor_ is a set, we do not need to check if one input has

diff --git a/xla/hlo/ir/hlo_input_output_alias_config.h b/xla/hlo/ir/hlo_input_output_alias_config.h
@@ -193,7 +193,8 @@ class HloBufferDonorConfig {
  const HloBufferDonorProto& proto);
 
  // Verifies that the given config is valid for the given module.
- // Specifically, the config's input should be in-bound.
+ // The config's input should be in-bound and this config cannot overlap with
+ // the given module's input_output_alias_config.
  Status Verify(const HloModule& module) const;
 
  // Returns the registered buffer donors

diff --git a/xla/service/hlo_input_output_alias_config_test.cc b/xla/service/hlo_input_output_alias_config_test.cc
@@ -285,5 +285,28 @@ ENTRY main {
  ASSERT_IS_NOT_OK(config.Verify(*module));
 }
 
+TEST_F(HloBufferDonorConfigTest, BufferDonorInputOutputAliasOverlap) {
+ const std::string module_str = R"(
+HloModule TEST
+
+ENTRY main {
+ param = (f32[], f32[]) parameter(0)
+ gte1 = f32[] get-tuple-element(%param), index=0
+ gte2 = f32[] get-tuple-element(%param), index=1
+ ROOT root = (f32[], f32[]) tuple(%gte1, %gte2)
+}
+)";
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+
+ HloBufferDonorConfig config;
+
+ TF_ASSERT_OK(config.AddBufferDonor(0, {0}));
+ TF_ASSERT_OK(config.Verify(*module));
+
+ TF_ASSERT_OK(module->input_output_alias_config().SetUpAlias({0}, 0, {0}));
+ ASSERT_IS_NOT_OK(config.Verify(*module));
+}
+
 } // namespace
 } // namespace xla
diff --git a/xla/service/hlo_verifier.cc b/xla/service/hlo_verifier.cc
@@ -2868,6 +2868,7 @@ StatusOr<bool> HloVerifier::Run(
  }
  }));
 
+ TF_RETURN_IF_ERROR(module->buffer_donor_config().Verify(*module));
  TF_RETURN_IF_ERROR(VerifyLayoutConstrainedAllReduce(*module));
  return false;
  }();