Co-allocate call operands with callee parameters

Right now we give these parameters an additional allocation and then don't use
it.

This is interesting for the CPU backend because we use call instructions to
represent fork-join parallelism (i.e. a specially annotated kCall instruction
tells the CPU IR emitter to shard the called computation across CPU threads).

Moreover, I need this for a principled fix to b/111116907.

PiperOrigin-RevId: 204820965

1 files changed, 9 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 166a83fade..441288da1a 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -898,4 +898,13 @@ void HloComputation::UniquifyName(NameUniquer* name_uniquer) {
   name_ = name_uniquer->GetUniqueName(name_);
 }
 
+HloInstruction* HloComputation::GetInstructionWithName(
+    tensorflow::StringPiece name) {
+  auto instructions_in_computation = instructions();
+  auto it = c_find_if(instructions_in_computation, [&](HloInstruction* instr) {
+    return instr->name() == name;
+  });
+  return it == instructions_in_computation.end() ? nullptr : *it;
+}
+
 }  // namespace xla