diff options
| author |  Mark Heffernan <meheff@google.com> | 2018-07-02 12:22:15 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-02 12:27:03 -0700 |
| commit | 20e27ad56b95e19ebeb23e34db1aff22e0bd473e (patch) | |
| tree | 0ff0a458b28a527acbc82b6ae1f5a36aeb96d1c8 /tensorflow/compiler/xla/service/hlo_instruction.h | |
| parent | 0967cbb9a34b69ec14238802460971abbec9cbb4 (diff) | |
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
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_instruction.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_instruction.h | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 0459072127..34e7dcb43d 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -477,7 +477,7 @@ class HloInstruction { const Shape& outfeed_shape, HloInstruction* operand, HloInstruction* token_operand, tensorflow::StringPiece outfeed_config); // Overload which does not require a token. - // TODO(b/80000000): Remove this overload when all uses of infeed are + // TODO(b/80000000): Remove this overload when all uses of outfeed are // converted to take tokens. static std::unique_ptr<HloInstruction> CreateOutfeed( const Shape& outfeed_shape, HloInstruction* operand, @@ -487,6 +487,7 @@ class HloInstruction { // initiates sending the operand data to a unique receive instruction in // another computation that has the same channel id. static std::unique_ptr<HloInstruction> CreateSend(HloInstruction* operand, + HloInstruction* token, int64 channel_id); // Blocks until data transfer for the Send instruction (operand) is complete. @@ -498,6 +499,7 @@ class HloInstruction { // which allocates resources to receive data of the given shape from a unique // send instruction in another computation that has the same channel id. static std::unique_ptr<HloInstruction> CreateRecv(const Shape& shape, + HloInstruction* token, int64 channel_id); // Blocks until data transfer for the Recv instruction (operand) is complete |