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

diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index 17cc6d35cc..dac5f1cf86 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -687,11 +687,19 @@ class HloInstruction {
       const Shape& shape, HloInstruction* operand,
       tensorflow::gtl::ArraySlice<int64> dimensions);
 
-  // Creates a token instruction used for joining or creating new values of
-  // token type which thread through side-effecting operations.
+  // Creates a Afterall instruction used for joining or creating new values of
+  // token type which thread through side-effecting operations. Operands must
+  // all be tokens, and there must be at least one operand.
   static std::unique_ptr<HloInstruction> CreateAfterAll(
       tensorflow::gtl::ArraySlice<HloInstruction*> operands);
 
+  // Creates an AfterAll instruction which creates a token type out of thin air
+  // (no operands). This is a separate method from CreateAfterAll to facility
+  // the removal of operand-less AfterAll instructions.
+  // TODO(b/110532604): Remove this capability of creating a token from nothing
+  // when we plumb a primordial token from the entry computation.
+  static std::unique_ptr<HloInstruction> CreateToken();
+
   // Creates an instance of GatherDimensionNumbers.
   static GatherDimensionNumbers MakeGatherDimNumbers(
       tensorflow::gtl::ArraySlice<int64> output_window_dims,