diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover.cc b/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover.cc index 61087ec1cc2..128b7485353 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover.cc @@ -16,6 +16,8 @@ limitations under the License. #include "xla/hlo/transforms/simplifiers/convert_mover.h" #include +#include +#include #include #include "absl/algorithm/container.h" @@ -160,8 +162,16 @@ absl::StatusOr MoveConvertPrecisionOps(HloComputation* comp) { new_shape.set_element_type(src_ty); HloInstruction* new_instr = comp->AddInstruction( instr->CloneWithNewOperands(new_shape, new_operands)); + // The new convert is the moved `convert_op`, so it keeps that convert's + // frontend attributes and must not inherit `instr`'s attributes, which + // still belong to the cloned `new_instr`. + std::unique_ptr new_convert = + HloInstruction::CreateConvert(instr->shape(), new_instr); + new_convert->set_frontend_attributes(convert_op->frontend_attributes()); ABSL_RETURN_IF_ERROR(comp->ReplaceWithNewInstruction( - instr, HloInstruction::CreateConvert(instr->shape(), new_instr))); + instr, std::move(new_convert), /*preserve_sharding=*/false, + /*relay_control_dependency=*/false, /*remove_unused_operands=*/true, + /*preserve_frontend_attributes=*/false)); changed = true; } @@ -197,13 +207,22 @@ absl::StatusOr MoveConvertPrecisionOps(HloComputation* comp) { absl::InlinedVector new_operands; new_operands.reserve(to_convert->operand_count()); for (HloInstruction* operand : to_convert->operands()) { - work_queue.push_front(MakeConvertToHlo(operand, dst_ty)); - new_operands.push_back(work_queue.front()); + HloInstruction* new_convert = MakeConvertToHlo(operand, dst_ty); + // The new converts are the moved `instr`, so they keep its frontend + // attributes; the cloned op below keeps `to_convert`'s own attributes. + if (new_convert != operand) { + new_convert->set_frontend_attributes(instr->frontend_attributes()); + } + work_queue.push_front(new_convert); + new_operands.push_back(new_convert); } Shape new_shape = to_convert->shape(); new_shape.set_element_type(dst_ty); ABSL_RETURN_IF_ERROR(comp->ReplaceWithNewInstruction( - instr, to_convert->CloneWithNewOperands(new_shape, new_operands))); + instr, to_convert->CloneWithNewOperands(new_shape, new_operands), + /*preserve_sharding=*/false, /*relay_control_dependency=*/false, + /*remove_unused_operands=*/true, + /*preserve_frontend_attributes=*/false)); changed = true; } diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover_test.cc index bc366e0e737..4be5a14315e 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/convert_mover_test.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include #include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/pattern_matcher_gmock.h" @@ -268,5 +269,54 @@ TEST_F(ConvertMoverTest, NoMoveDownPadBecauseConstantIsOutOfRange) { EXPECT_FALSE(changed); } +// Frontend attributes travel with the instruction they describe: moving a +// convert through a reshape must not copy the reshape's attributes onto the +// convert, and vice versa. +TEST_F(ConvertMoverTest, MoveDownKeepsFrontendAttributesOnTheirOwners) { + absl::string_view module_string = R"( + HloModule module + + ENTRY main { + p0 = f16[100] parameter(0) + c = f32[100] convert(p0), frontend_attributes={CONVERT_ATTR="c"} + ROOT r = f32[10,10] reshape(c), frontend_attributes={RESHAPE_ATTR="r"} + })"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_string)); + ConvertMover pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + SCOPED_TRACE(module->ToString()); + EXPECT_TRUE(changed); + const HloInstruction* root = module->entry_computation()->root_instruction(); + ASSERT_THAT(root, GmockMatch(MatchConvertToF32(m::Reshape(m::Parameter(0))))); + EXPECT_THAT(root->frontend_attributes().map(), + ::testing::ElementsAre(::testing::Pair("CONVERT_ATTR", "c"))); + EXPECT_THAT(root->operand(0)->frontend_attributes().map(), + ::testing::ElementsAre(::testing::Pair("RESHAPE_ATTR", "r"))); +} + +TEST_F(ConvertMoverTest, MoveUpKeepsFrontendAttributesOnTheirOwners) { + absl::string_view module_string = R"( + HloModule module + + ENTRY main { + p0 = f32[100] parameter(0) + r = f32[10,10] reshape(p0), frontend_attributes={RESHAPE_ATTR="r"} + ROOT c = f16[10,10] convert(r), frontend_attributes={CONVERT_ATTR="c"} + })"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_string)); + ConvertMover pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + SCOPED_TRACE(module->ToString()); + EXPECT_TRUE(changed); + const HloInstruction* root = module->entry_computation()->root_instruction(); + ASSERT_THAT(root, GmockMatch(m::Reshape(MatchConvertToF16(m::Parameter(0))))); + EXPECT_THAT(root->frontend_attributes().map(), + ::testing::ElementsAre(::testing::Pair("RESHAPE_ATTR", "r"))); + EXPECT_THAT(root->operand(0)->frontend_attributes().map(), + ::testing::ElementsAre(::testing::Pair("CONVERT_ATTR", "c"))); +} + } // anonymous namespace } // namespace xla