Change Send and Recv HLOs to take a token operand.

Send and Recv HLOs now have an additional required operand which must be token-shaped. XLA client interface for these operations is unchanged and will be updated in follow up CLs.

PiperOrigin-RevId: 202993121

1 files changed, 12 insertions, 8 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc
index 504ea3fe7a..f40cd60907 100644
--- a/tensorflow/compiler/xla/service/hlo_parser_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc
@@ -278,10 +278,11 @@ ENTRY %WhileWithScalarS32Result.v2 () -> s32[] {
 R"(HloModule TwoSendRecvBothWayRecvFist_module
 
 ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] {
-  %recv = (f32[], u32[]) recv(), channel_id=15, sharding={maximal device=1}
+  %token = token[] after-all()
+  %recv = (f32[], u32[]) recv(token[] %token), channel_id=15, sharding={maximal device=1}
   ROOT %recv-done = f32[] recv-done((f32[], u32[]) %recv), channel_id=15, sharding={maximal device=1}
   %constant = f32[] constant(2.1), sharding={maximal device=0}
-  %send = (f32[], u32[]) send(f32[] %constant), channel_id=16, sharding={maximal device=0}, control-predecessors={%recv}
+  %send = (f32[], u32[]) send(f32[] %constant, token[] %token), channel_id=16, sharding={maximal device=0}, control-predecessors={%recv}
   %send-done = () send-done((f32[], u32[]) %send), channel_id=16, sharding={maximal device=0}
 }
 
@@ -1221,10 +1222,11 @@ TEST_F(HloParserTest, UnexpectedAttribute) {
   const string original = R"(HloModule unexpected_attr_module
 
 ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] {
-  %recv = (f32[], u32[]) recv(), channel_id=15
+  %token = token[] after-all()
+  %recv = (f32[], u32[]) recv(token[] %token), channel_id=15
   %recv-done = f32[] recv-done((f32[], u32[]) %recv), channel_id=15
   ROOT %constant = f32[] constant(2.1)
-  %send = (f32[], u32[]) send(f32[] %constant), channel_id=16, calls=%recv
+  %send = (f32[], u32[]) send(f32[] %constant, token[] %token), channel_id=16, calls=%recv
   %send-done = () send-done((f32[], u32[]) %send), channel_id=16
 }
 
@@ -1237,10 +1239,11 @@ TEST_F(HloParserTest, MissingAttribute) {
   const string original = R"(HloModule missing_attr_module
 
 ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] {
-  %recv = (f32[], u32[]) recv(), channel_id=15
+  %token = token[] after-all()
+  %recv = (f32[], u32[]) recv(token[] %token), channel_id=15
   %recv-done = f32[] recv-done((f32[], u32[]) %recv), channel_id=15
   ROOT %constant = f32[] constant(-2.1)
-  %send = (f32[], u32[]) send(f32[] %constant)
+  %send = (f32[], u32[]) send(f32[] %constant, token[] %token)
   %send-done = () send-done((f32[], u32[]) %send), channel_id=16
 }
 
@@ -1253,10 +1256,11 @@ TEST_F(HloParserTest, PredecessorUndefined) {
   const string original = R"(HloModule pre_not_found_module
 
 ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] {
-  %recv = (f32[], u32[]) recv(), channel_id=15
+  %token = token[] after-all()
+  %recv = (f32[], u32[]) recv(token[] %token), channel_id=15
   %recv-done = f32[] recv-done((f32[], u32[]) %recv), channel_id=15
   ROOT %constant = f32[] constant(2.1)
-  %send = (f32[], u32[]) send(f32[] %constant), channel_id=16, control-predecessors={%done}
+  %send = (f32[], u32[]) send(f32[] %constant, token[] %token), channel_id=16, control-predecessors={%done}
   %send-done = () send-done((f32[], u32[]) %send), channel_id=16
 }