[XLA:GPU] Don't allow to fuse DUS with shared operands.

DynamicUpdateSlice is an in-place operation. Therefore we cannot fuse two such
ops together if they share the same operand.
Also adjust the check whether the in-place emitter can be used to handle such a
case.

PiperOrigin-RevId: 702288316

Author: akuegel

xla/service/gpu/ir_emission_utils.cc

@@ -347,6 +347,9 @@ absl::StatusOr<bool> CanEmitFusedDynamicUpdateSliceInPlaceForGpu(
               user_start_indices != dus_start_indices) {
             return false;
           }
+        } else if (user != dus &&
+                   user.opcode() == HloOpcode::kDynamicUpdateSlice) {
+          return false;
         } else if (user != dus && !user.instruction().IsElementwise() &&
                    user.opcode() != HloOpcode::kBitcast &&
                    user.opcode() != HloOpcode::kTuple) {

xla/service/gpu/ir_emission_utils_test.cc

@@ -870,6 +870,49 @@ ENTRY main {
               IsOkAndHolds(true));
 }
 
+TEST_F(
+    IrEmissionUtilsTest,
+    CanEmitFusedDynamicUpdateSliceInPlaceForGpu_HandlesMultiOutputFusionSharedParameter) {  // NOLINT
+  const char* hlo = R"(
+HloModule MultipleInplaceDus, is_scheduled=true, input_output_alias={ {0}: (0, {}), {1}: (2, {}) }
+
+fused_computation {
+  p0 = bf16[10,11,12] parameter(0)
+  p1 = bf16[1,11,12] parameter(1)
+  p2 = bf16[1,11,12] parameter(2)
+  p3 = s32[] parameter(3)
+  c0 = s32[] constant(0)
+  cmp = pred[] compare(p3, c0), direction=EQ
+  broadcast = pred[1,11,12] broadcast(cmp), dimensions={}
+  select = bf16[1,11,12] select(broadcast, p1, p2)
+  dus0 = bf16[10,11,12] dynamic-update-slice(p0, select, c0, c0, c0)
+  dus1 = bf16[10,11,12] dynamic-update-slice(p0, select, c0, c0, c0)
+  ROOT tuple = (bf16[10,11,12], bf16[10,11,12]) tuple(dus0, dus1)
+}
+
+ENTRY main {
+  p0 = bf16[10,11,12] parameter(0)
+  p1 = bf16[1,11,12] parameter(1)
+  p2 = bf16[1,11,12] parameter(2)
+  p3 = s32[] parameter(3)
+  ROOT fusion_root_multiple = (bf16[10,11,12], bf16[10,11,12]) fusion(p0, p1, p2, p3), kind=kLoop, calls=fused_computation
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(hlo));
+  auto fusion = module->entry_computation()->root_instruction();
+  BufferAllocation alloc(/*index=*/0, /*size=*/1024, /*color=*/0);
+  BufferAllocation::Slice slice0(&alloc, 0, 10);
+  auto adaptor = HloFusionAdaptor::ForInstruction(fusion);
+  EXPECT_THAT(CanEmitFusedDynamicUpdateSliceInPlaceForGpu(
+                  *adaptor,
+                  [&slice0](const HloInstruction*, const ShapeIndex&) {
+                    return slice0;
+                  },
+                  fusion),
+              IsOkAndHolds(false));
+}
+
 TEST_F(
     IrEmissionUtilsTest,
     CanEmitFusedDynamicUpdateSliceInPlaceForGpu_HandlesMultiOutputFusionWithTransposeBitcasts) {  // NOLINT

xla/service/gpu/transforms/horizontal_loop_fusion.cc

@@ -155,6 +155,19 @@ bool IsConcatenationInputFusion(const HloInstruction& instr) {
          instr.fused_expression_root()->opcode() == HloOpcode::kConcatenate;
 }
 
+bool IsDynamicUpdateSliceFusion(const HloInstruction* instr) {
+  if (instr->opcode() != HloOpcode::kFusion) {
+    return false;
+  }
+  auto root = instr->fused_expression_root();
+  if (root->opcode() == HloOpcode::kTuple) {
+    return absl::c_any_of(root->operands(), [&](const HloInstruction* operand) {
+      return operand->opcode() == HloOpcode::kDynamicUpdateSlice;
+    });
+  }
+  return root->opcode() == HloOpcode::kDynamicUpdateSlice;
+}
+
 bool IsFusibleCandidate(const HloInstruction& instr,
                         const se::DeviceDescription& device_description) {
   // For now, we do not support fusing instruction with control flow.
@@ -247,16 +260,15 @@ bool IsProfitableFusionCandidate(const HloInstruction& instr,
   return true;
 }
 
-// Returns whether any operand of `instr` is a parameter instruction that
-// is shared with `fusion_instrs`.
-bool AnyOpndIsParamSharedAmongFusions(
+// Returns whether any operand of `instr` is an instruction that is shared with
+// `fusion_instrs`.
+bool AnyOperandIsSharedAmongFusions(
     const HloInstruction* instr,
     const absl::flat_hash_set<HloInstruction*>& fusion_instrs) {
   return absl::c_any_of(instr->operands(), [&](const HloInstruction* opnd) {
-    return opnd->opcode() == HloOpcode::kParameter &&
-           absl::c_any_of(opnd->users(), [&](const HloInstruction* user) {
-             return user != instr && fusion_instrs.contains(user);
-           });
+    return absl::c_any_of(opnd->users(), [&](const HloInstruction* user) {
+      return user != instr && fusion_instrs.contains(user);
+    });
   });
 }
 
@@ -298,11 +310,11 @@ void HorizontalLoopFusionImpl::FusionCandidates::Initialize(
               << " rejects may-not-be profitable fusion instr"
               << instr->ToString();
       continue;
-    } else if (sliced_input_fusion_ &&
-               AnyOpndIsParamSharedAmongFusions(instr, fusible_candidates)) {
-      // Don't fuse fusions whose operands are parameter instructions that are
-      // shared among fusions because we cannot i/o alias the produced
-      // horizontal fusion due to the concat insertion.
+    } else if ((sliced_input_fusion_ || IsDynamicUpdateSliceFusion(instr)) &&
+               AnyOperandIsSharedAmongFusions(instr, fusible_candidates)) {
+      // Don't fuse fusions with at least one shared operand because we cannot
+      // i/o alias the produced horizontal fusion due to the concat insertion
+      // (or run into aliasing problems with DynamicUpdateSlice fusions).
       VLOG(2) << "sliced_input_fusion=" << sliced_input_fusion_
               << " rejects the fusion instr because it shares parameter with"
               << " other fusion candidates, instr: " << instr->ToString();

xla/service/gpu/transforms/horizontal_loop_fusion_test.cc

@@ -531,7 +531,7 @@ TEST_F(HorizontalLoopFusionTest, RMSPropLike) {
 
 TEST_F(HorizontalLoopFusionTest, DynamicUpdateSlice) {
   auto module = ParseAndReturnVerifiedModule(R"(
-  HloModule NegativeTestForDynamicUpdateSlice
+  HloModule DynamicUpdateSlice
 
   fusion.1 {
     p.0 = f16[5,9,10]{2,1,0} parameter(0)
@@ -576,6 +576,39 @@ TEST_F(HorizontalLoopFusionTest, DynamicUpdateSlice) {
   EXPECT_TRUE(RunAndCompareNoHloPasses(std::move(module), ErrorSpec{0, 0}));
 }
 
+TEST_F(HorizontalLoopFusionTest, DontFuseDynamicUpdateSlice) {
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule DynamicUpdateSliceFusionsShareParameter
+
+fused_dynamic_update_slice {
+  p0 = s32[3,3]{1,0} parameter(0)
+  p1 = pred[3,2]{1,0} parameter(1)
+  convert = s32[3,2]{1,0} convert(p1)
+  zero = s32[] constant(0)
+  ROOT dynamic-update-slice = s32[3,3]{1,0} dynamic-update-slice(p0, convert, zero, zero)
+}
+
+fused_dynamic_update_slice.1 {
+  p0 = s32[3,3]{1,0} parameter(0)
+  p1 = pred[2,3]{1,0} parameter(1)
+  convert = s32[2,3]{1,0} convert(p1)
+  zero = s32[] constant(0)
+  ROOT dynamic-update-slice = s32[3,3]{1,0} dynamic-update-slice(p0, convert, zero, zero)
+}
+
+ENTRY main {
+  param_0 = s32[3,3]{1,0} parameter(0)
+  param_1 = pred[2,3]{1,0} parameter(1)
+  param_2 = pred[3,2]{1,0} parameter(2)
+  loop_dynamic_update_slice_fusion = s32[3,3]{1,0} fusion(param_0, param_2), kind=kLoop, calls=fused_dynamic_update_slice
+  loop_dynamic_update_slice_fusion.1 = s32[3,3]{1,0} fusion(param_0, param_1), kind=kLoop, calls=fused_dynamic_update_slice.1
+  ROOT tuple.11.0 = (s32[3,3]{1,0}, s32[3,3]{1,0}) tuple(loop_dynamic_update_slice_fusion.1, loop_dynamic_update_slice_fusion)
+}
+)"));
+  EXPECT_FALSE(
+      HorizontalLoopFusion{device_description_}.Run(module.get()).value());
+}
+
 TEST_F(HorizontalLoopFusionTest,
        AllowSharedParametersWhenNotUsingConcatenation) {
   auto module = ParseAndReturnVerifiedModule(R"(