[XLA] Keep frontend attributes on their owners in ConvertMover

When `ConvertMover` moves a `kConvert` across a commuting operation, preserve the original `kConvert` frontend attributes on the new `kConvert` instructions and keep the commuting operation's attributes on the cloned instruction instead of copying the outer instruction's attributes onto the replacement root.

PiperOrigin-RevId: 1004414597

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 <deque>
+#include <memory>
+#include <utility>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -160,8 +162,16 @@ absl::StatusOr<bool> 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<HloInstruction> 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<bool> MoveConvertPrecisionOps(HloComputation* comp) {
     absl::InlinedVector<HloInstruction*, 8> 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;
   }
 

third_party/xla/xla/hlo/transforms/simplifiers/convert_mover_test.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #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<HloModule> 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<HloModule> 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