Change for asynchronous Send and Recv by splitting Send into {Send, SendDone}

and Recv into {Recv, RecvDone}. See operation_semantics.md for the updated
semantics.

PiperOrigin-RevId: 175216012

1 files changed, 15 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc
index b92017c6cb..02dc49e78c 100644
--- a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc
+++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc
@@ -104,6 +104,21 @@ Status LogicalBufferAnalysis::HandleBitcast(HloInstruction*) {
   return Status::OK();
 }
 
+Status LogicalBufferAnalysis::HandleRecvDone(HloInstruction*) {
+  // RecvDone doesn't create a new buffer but rather aliases its input (Recv)
+  // tuple element at {0} to its output.
+  return Status::OK();
+}
+
+Status LogicalBufferAnalysis::HandleSend(HloInstruction* send) {
+  // Send creates new buffers for the top-level tuple and the context (tuple
+  // element at {1}). Tuple element at {0} is an alias of the Send operand, so
+  // we don't need to create a new Logical Buffer for that.
+  NewLogicalBuffer(send, /*index=*/{});
+  NewLogicalBuffer(send, /*index=*/{1});
+  return Status::OK();
+}
+
 Status LogicalBufferAnalysis::HandleTuple(HloInstruction* tuple) {
   // A Tuple instruction only creates the top-level buffer.
   NewLogicalBuffer(tuple, /*index=*/{});