Add WithToken variants to Send/Recv/Infeed/Outfeed ops in XLA builder.

In HLO side-effecting ops now take and produce a token-shaped value used for ordering. This CL exposes the tokens in the XLA builder interface by adding builder variants suffixed with "WithToken". These variants have an extra token-shaped operand and produce a token, sometimes as a element in the tuple-shaped output of the op.

Also, at the HLO level, add new method CreateToken for creating an instruction which produces token-shaped value out of thin air using an AfterAll instruction. This is explicitly broken out from CreateAfterAll to facilitate remove of zero-operand AfterAll instructions (b/110532604).

PiperOrigin-RevId: 203834271

1 files changed, 4 insertions, 5 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_dce_test.cc b/tensorflow/compiler/xla/service/hlo_dce_test.cc
index 4fa13c975a..26e3736e01 100644
--- a/tensorflow/compiler/xla/service/hlo_dce_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_dce_test.cc
@@ -75,7 +75,7 @@ TEST_F(HloDceTest, InstructionsWithSideEffect) {
   auto builder = HloComputation::Builder(TestName());
   auto constant = builder.AddInstruction(
       HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(42.0f)));
-  auto token = builder.AddInstruction(HloInstruction::CreateAfterAll({}));
+  auto token = builder.AddInstruction(HloInstruction::CreateToken());
   builder.AddInstruction(
       HloInstruction::CreateSend(constant, token, /*channel_id=*/0));
   builder.AddInstruction(HloInstruction::CreateTuple({}));
@@ -235,8 +235,7 @@ TEST_F(HloDceTest, CalledComputationWithSideEffect) {
   {
     auto param = body_builder.AddInstruction(
         HloInstruction::CreateParameter(0, shape, "param"));
-    auto token =
-        body_builder.AddInstruction(HloInstruction::CreateAfterAll({}));
+    auto token = body_builder.AddInstruction(HloInstruction::CreateToken());
     auto infeed = body_builder.AddInstruction(
         HloInstruction::CreateInfeed(shape, token, ""));
     body_builder.AddInstruction(
@@ -280,8 +279,8 @@ TEST_F(HloDceTest, CalledComputationWithNestedSideEffect) {
   {
     auto param = nested_callee_builder.AddInstruction(
         HloInstruction::CreateParameter(0, shape, "param"));
-    auto token = nested_callee_builder.AddInstruction(
-        HloInstruction::CreateAfterAll({}));
+    auto token =
+        nested_callee_builder.AddInstruction(HloInstruction::CreateToken());
     nested_callee_builder.AddInstruction(
         HloInstruction::CreateOutfeed(shape, param, token, ""));
   }