[xla:gpu:gemm_fusion] Fix crash in HoistBitcast by copying users vector

Copy operand->users() to a local vector before iterating over it,
as the loop body calls RemoveInstructionAndUnusedOperands which
modifies the users list, leading to iterator invalidation and crash.

PiperOrigin-RevId: 933680106

Author: vwbaker

third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc

@@ -918,7 +918,10 @@ absl::StatusOr<bool> FusionSearchSpace::HoistBitcast(HloInstruction* instr) {
       CopyElementType(operand->shape(), &new_shape);
 
       *operand->mutable_shape() = new_shape;
-      for (HloInstruction* user : operand->users()) {
+      // Copy users to a new vector to avoid invalidating the iterator.
+      std::vector<HloInstruction*> users(operand->users().begin(),
+                                         operand->users().end());
+      for (HloInstruction* user : users) {
         RETURN_IF_ERROR(user->ReplaceAllUsesWith(operand));
         RETURN_IF_ERROR(
             user->parent()->RemoveInstructionAndUnusedOperands(user));

third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc

@@ -476,7 +476,7 @@ ENTRY e {
 )");
 }
 
-TEST_P(GemmFusionTestV2, DoNotHoistBitcastOverParameterWithMultipleUsers) {
+TEST_P(GemmFusionTestV2, DoNotHoistBitcastOverParameterWithNonBitcastUsers) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 ENTRY e {
   p0 = bf16[32,12] parameter(0)
@@ -514,6 +514,26 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Parameter(), m::Bitcast(m::Parameter()))));
 }
 
+TEST_P(GemmFusionTestV2, HoistBitcastOverParameterWithMultipleBitcastUsers) {
+  // Regression test for b/524943134.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f32[32,3] parameter(0)
+  p1 = f32[8,28] parameter(1)
+  b1 = f32[32,7] bitcast(p1)
+  b2 = f32[32,7] bitcast(p1)
+  add = f32[32,7] add(b1, b2)
+  ROOT d = f32[3,7] dot(p0, add),
+    lhs_contracting_dims={0}, rhs_contracting_dims={0}
+})"));
+  ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true));
+  EXPECT_THAT(
+      module->entry_computation()->root_instruction(),
+      GmockMatch(m::Fusion(m::Parameter(), m::Bitcast(m::Parameter()))));
+}
+
 TEST_P(GemmFusionTestV2, HoistBitcastAcrossTransposeWithTrivialDimension) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 ENTRY e {