[XLA] Replace HloComputation::ReplaceUsesOfInstruction with HloInstruction::ReplaceAllUsesWith.

RAUW used to be *almost* synonymous with RUOI, except RAUW didn't update
the computation's root.  This was a dangerous footgun -- if you
accidentally called RAUW when you wanted RUOI (which you almost always
did), your code would work perfectly, except when the relevant node
happened to be the root of a computation.

This change simplifies our APIs so there's just one Right Way To Do It,
by making RAUW update the computation.

PiperOrigin-RevId: 170290230

1 files changed, 1 insertions, 11 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 2d07784619..e880900320 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -245,15 +245,6 @@ Status HloComputation::RemoveInstruction(HloInstruction* instruction) {
   return Status::OK();
 }
 
-Status HloComputation::ReplaceUsesOfInstruction(
-    HloInstruction* instruction_to_replace, HloInstruction* instruction) {
-  TF_RETURN_IF_ERROR(instruction_to_replace->ReplaceAllUsesWith(instruction));
-  if (instruction_to_replace == root_instruction()) {
-    set_root_instruction(instruction);
-  }
-  return Status::OK();
-}
-
 void HloComputation::set_root_instruction(
     HloInstruction* new_root_instruction) {
   // The shape of the root (ignoring layout) is an invariant of the computation
@@ -569,8 +560,7 @@ Status HloComputation::ReplaceInstruction(HloInstruction* old_instruction,
   if (new_instruction->metadata().op_name().empty()) {
     new_instruction->set_metadata(old_instruction->metadata());
   }
-  TF_RETURN_IF_ERROR(
-      ReplaceUsesOfInstruction(old_instruction, new_instruction));
+  TF_RETURN_IF_ERROR(old_instruction->ReplaceAllUsesWith(new_instruction));
   return RemoveInstructionAndUnusedOperands(old_instruction);
 }