PR #21104: [NVIDIA GPU] Preserve backend config when folding transpose

Imported from GitHub PR #21104

Transpose folding pass doesn't preserve backend config when creating the new dot with transpose folded. Changing the behavior to copy the old dot's config to the new dot.
Copybara import of the project:

--
d2d6b62 by TJ Xu <[email protected]>:

Preserve backend config when folding transpose

--
6b5fa3a by TJ Xu <[email protected]>:

use SetupDerivedInstruction instead of just copying the backend config

Merging this change closes #21104

COPYBARA_INTEGRATE_REVIEW=#21104 from Tixxx:tixxx/transpose_folding 6b5fa3a
PiperOrigin-RevId: 715204523

Author: Tixxx

xla/service/transpose_folding.cc

@@ -101,10 +101,11 @@ absl::Status FoldTransposeIntoDot(InstructionOperandsPair& pair) {
       rhs = rhs->mutable_operand(0);
     }
   }
-
-  return dot->parent()->ReplaceWithNewInstruction(
-      dot, HloInstruction::CreateDot(dot->shape(), lhs, rhs, new_dot_dims,
-                                     dot->precision_config()));
+  HloInstruction* new_dot =
+      dot->parent()->AddInstruction(HloInstruction::CreateDot(
+          dot->shape(), lhs, rhs, new_dot_dims, dot->precision_config()));
+  dot->SetupDerivedInstruction(new_dot);
+  return dot->parent()->ReplaceInstruction(dot, new_dot);
 }
 
 // Folds the operands of `convolution` that are foldable transposes.

xla/service/transpose_folding_test.cc

@@ -559,5 +559,25 @@ ENTRY entry_computation {
   EXPECT_THAT(TransposeFolding().Run(module.get()), IsOkAndHolds(false));
 }
 
+TEST_F(TransposeFoldingTest, FoldTransposeWithBackendConfig) {
+  constexpr absl::string_view kHloString = R"(
+HloModule FoldTransposeWithBackendConfig
+
+ENTRY entry_computation {
+  x = f32[7,2,7,3]{3,2,1,0} parameter(0)
+  y = f32[7,2,7,3]{3,2,1,0} parameter(1)
+  transpose = f32[7,3,7,2]{3,2,1,0} transpose(y), dimensions={0,3,2,1}
+  ROOT dot = f32[7,7,2,2]{3,2,1,0} dot(x, transpose), lhs_contracting_dims={3},
+            rhs_contracting_dims={1}, lhs_batch_dims={0,2}, rhs_batch_dims={0,2}, backend_config={"force_earliest_schedule":false}
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(kHloString));
+
+  EXPECT_THAT(TransposeFolding().Run(module.get()), IsOkAndHolds(true));
+  EXPECT_TRUE(
+      module->entry_computation()->root_instruction()->has_backend_config());
+}
+
 }  // namespace
 }  // namespace xla