diff options
| author |  Mark Heffernan <meheff@google.com> | 2018-07-03 10:34:20 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-03 10:37:09 -0700 |
| commit | 31a5fa1ee88f8f3bb1a46f3734136b6d85e8642f (patch) | |
| tree | 8ada804c9a30d162778decc9fcef009acfa39d49 /tensorflow/compiler/xla/service/hlo_instructions.cc | |
| parent | bce9c2ef1ad72a5d962faac0d114932af6a69bf9 (diff) | |
Change Send, SendDone, Recv and RecvDone to produce tokens.
This is a follow up to cl/202069017 which added tokens as operands to Send and Recv.
PiperOrigin-RevId: 203145403
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_instructions.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_instructions.cc | 14 |
1 files changed, 9 insertions, 5 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc index dcc1e3c8af..7052e236cd 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.cc +++ b/tensorflow/compiler/xla/service/hlo_instructions.cc @@ -207,8 +207,9 @@ HloSendInstruction::HloSendInstruction(HloInstruction* operand, HloInstruction* token, int64 channel_id) : HloSendRecvInstruction( HloOpcode::kSend, - ShapeUtil::MakeTupleShape( - {CHECK_NOTNULL(operand)->shape(), ShapeUtil::MakeShape(U32, {})}), + ShapeUtil::MakeTupleShape({CHECK_NOTNULL(operand)->shape(), + ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeTokenShape()}), channel_id) { AppendOperand(operand); AppendOperand(token); @@ -224,7 +225,7 @@ std::unique_ptr<HloInstruction> HloSendInstruction::CloneWithNewOperandsImpl( } HloSendDoneInstruction::HloSendDoneInstruction(HloSendInstruction* operand) - : HloSendRecvInstruction(HloOpcode::kSendDone, ShapeUtil::MakeNil(), + : HloSendRecvInstruction(HloOpcode::kSendDone, ShapeUtil::MakeTokenShape(), CHECK_NOTNULL(operand)->channel_id()) { AppendOperand(operand); } @@ -244,7 +245,8 @@ HloRecvInstruction::HloRecvInstruction(const Shape& shape, HloInstruction* token, int64 channel_id) : HloSendRecvInstruction( HloOpcode::kRecv, - ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeShape(U32, {})}), + ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeTokenShape()}), channel_id) { AppendOperand(token); } @@ -261,7 +263,9 @@ std::unique_ptr<HloInstruction> HloRecvInstruction::CloneWithNewOperandsImpl( HloRecvDoneInstruction::HloRecvDoneInstruction(HloRecvInstruction* operand) : HloSendRecvInstruction( HloOpcode::kRecvDone, - ShapeUtil::GetTupleElementShape(operand->shape(), 0), + ShapeUtil::MakeTupleShape( + {ShapeUtil::GetTupleElementShape(operand->shape(), 0), + ShapeUtil::MakeTokenShape()}), CHECK_NOTNULL(operand)->channel_id()) { AppendOperand(operand); } |