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

diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc
index 54fc34b862..2c3dab0a45 100644
--- a/tensorflow/compiler/xla/service/hlo_parser.cc
+++ b/tensorflow/compiler/xla/service/hlo_parser.cc
@@ -622,8 +622,12 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder,
       if (!ParseOperands(&operands) || !ParseAttributes(attrs)) {
         return false;
       }
-      instruction =
-          builder->AddInstruction(HloInstruction::CreateAfterAll(operands));
+      if (operands.empty()) {
+        instruction = builder->AddInstruction(HloInstruction::CreateToken());
+      } else {
+        instruction =
+            builder->AddInstruction(HloInstruction::CreateAfterAll(operands));
+      }
       break;
     }
     case HloOpcode::kSort: {