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, 3 insertions, 6 deletions

diff --git a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc
index 834878426f..c43a31b167 100644
--- a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc
+++ b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc
@@ -119,8 +119,7 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsSend) {
   ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional);
 
   auto* true_computation = conditional->true_computation();
-  auto* token =
-      true_computation->AddInstruction(HloInstruction::CreateAfterAll({}));
+  auto* token = true_computation->AddInstruction(HloInstruction::CreateToken());
   auto* send = true_computation->AddInstruction(HloInstruction::CreateSend(
       true_computation->AddInstruction(
           HloInstruction::CreateConstant(LiteralUtil::CreateR0<bool>(true))),
@@ -135,8 +134,7 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsRecv) {
   ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional);
 
   auto* true_computation = conditional->true_computation();
-  auto* token =
-      true_computation->AddInstruction(HloInstruction::CreateAfterAll({}));
+  auto* token = true_computation->AddInstruction(HloInstruction::CreateToken());
   auto* recv = true_computation->AddInstruction(HloInstruction::CreateRecv(
       ShapeUtil::MakeShape(F32, {1}), token, /*channel_id=*/0));
   true_computation->AddInstruction(HloInstruction::CreateRecvDone(recv));
@@ -148,8 +146,7 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsNonRemovableInstruction) {
   auto* conditional = computation->root_instruction();
   ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional);
   auto* false_computation = conditional->false_computation();
-  auto token =
-      false_computation->AddInstruction(HloInstruction::CreateAfterAll({}));
+  auto token = false_computation->AddInstruction(HloInstruction::CreateToken());
   false_computation->AddInstruction(HloInstruction::CreateInfeed(
       ShapeUtil::MakeShape(F32, {1}), token, "config"));
   EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie());