Create copy if the operands of gather/scatter instructions overlap.

A gather has two operands, input and indices. If they point to the same instruction, create a copy for indices.

A scatter has n inputs, 1 indices, and n updates (2n+1 operands in total). We allow overlap between n inputs. We also allow overlap between n updates. We need to create a copy if
* indices overlap with any input or update
* update overlap with any input

The added copy will be removed if it is redundant in the following memory related passes (e.g., CopyInsertion).

PiperOrigin-RevId: 715164959

Author: ZixuanJiang

xla/service/spmd/gather_scatter_handler.cc

@@ -957,7 +957,12 @@ absl::Status SpmdPartitioningVisitor::HandleGather(HloInstruction* hlo) {
   auto gather = Cast<HloGatherInstruction>(hlo);
   const auto& dnums = gather->gather_dimension_numbers();
   auto operand = GetPartitionedHlo(gather->operand(0));
-  auto indices = GetPartitionedHlo(gather->operand(1));
+  auto raw_indices = GetPartitionedHlo(gather->operand(1));
+  auto indices =
+      (operand.hlo() == raw_indices.hlo())
+          ? MakeACopyAndReturnItsPartitionedHlo(raw_indices, builder())
+          : raw_indices;
+
   std::vector<int64_t> batch_dims;
   for (int64_t i = 0; i < gather->shape().rank(); ++i) {
     if (!absl::c_linear_search(dnums.offset_dims(), i)) {
@@ -1822,6 +1827,13 @@ absl::Status SpmdPartitioningVisitor::HandleScatter(HloInstruction* hlo) {
   absl::c_transform(
       scatter->scatter_updates(), std::back_inserter(updates),
       [this](HloInstruction* hlo) { return GetPartitionedHlo(hlo); });
+  for (PartitionedHlo& update : updates) {
+    if (absl::c_any_of(operands, [&](const PartitionedHlo& operand) {
+          return update.hlo() == operand.hlo();
+        })) {
+      update = MakeACopyAndReturnItsPartitionedHlo(update, builder());
+    }
+  }
   if (!absl::c_all_of(updates, [&](const PartitionedHlo& update) {
         return update.sharding() == updates[0].sharding() &&
                update.base_shape() == updates[0].base_shape();
@@ -1847,6 +1859,15 @@ absl::Status SpmdPartitioningVisitor::HandleScatter(HloInstruction* hlo) {
                ? scatter_reduction_root->shape().tuple_shapes_size()
                : 1);
   auto indices = GetPartitionedHlo(scatter->scatter_indices());
+  if (absl::c_any_of(operands,
+                     [&](const PartitionedHlo& operand) {
+                       return indices.hlo() == operand.hlo();
+                     }) ||
+      absl::c_any_of(updates, [&](const PartitionedHlo& update) {
+        return indices.hlo() == update.hlo();
+      })) {
+    indices = MakeACopyAndReturnItsPartitionedHlo(indices, builder());
+  }
   auto indices_sharding = indices.sharding();
   // Reshard indices with -1 padding, which will have no effect on the result as
   // guaranteed by the scatter semantics.

xla/service/spmd/spmd_partitioner_test.cc

@@ -14871,6 +14871,33 @@ ENTRY %main.21 {
   EXPECT_THAT(updates, op::Shape("bf16[4096,64]"));
 }
 
+TEST_P(SpmdPartitioningTest, ScatterAllOperandsAreSameInstruction) {
+  const char* const hlo_string = R"(
+HloModule pjit
+
+%s32_add {
+  a = s32[] parameter(0)
+  b = s32[] parameter(1)
+  ROOT result = s32[] add(a, b)
+}
+
+ENTRY %main.21 {
+  p0 = s32[8,64] parameter(0), sharding={devices=[4,1]<=[4]}
+  ROOT scatter = s32[8,64] scatter(p0, p0, p0), update_window_dims={}, 
+    input_batching_dims={0}, scatter_indices_batching_dims={0},
+    inserted_window_dims={1}, scatter_dims_to_operand_dims={1},
+    index_vector_dim=2, to_apply=s32_add, sharding={devices=[4,1]<=[4]}
+})";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          PartitionComputation(hlo_string, /*num_devices=*/4));
+
+  auto p0 = AllOf(op::Shape("s32[2,64]"), op::Parameter(0));
+  auto p0_copy = AllOf(op::Shape("s32[2,64]"), op::Copy(p0));
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              AllOf(op::Shape("s32[2,64]"), op::Scatter(p0, p0_copy, p0_copy)));
+}
+
 TEST_P(SpmdPartitioningTest, ComplexReshardUnmerge) {
   const char* const hlo_string = R"(
 HloModule Test