| 2405 | } |
| 2406 | |
| 2407 | XlaOp XlaBuilder::SendWithToken(XlaOp operand, XlaOp token, |
| 2408 | const ChannelHandle& handle) { |
| 2409 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2410 | if (handle.type() != ChannelHandle::DEVICE_TO_DEVICE) { |
| 2411 | return InvalidArgument("Send must use a device-to-device channel"); |
| 2412 | } |
| 2413 | |
| 2414 | // Send instruction produces a tuple of {aliased operand, U32 context, |
| 2415 | // token}. |
| 2416 | HloInstructionProto send_instr; |
| 2417 | TF_ASSIGN_OR_RETURN(const Shape* shape, GetShapePtr(operand)); |
| 2418 | *send_instr.mutable_shape() = |
| 2419 | ShapeUtil::MakeTupleShape({*shape, ShapeUtil::MakeShape(U32, {}), |
| 2420 | ShapeUtil::MakeTokenShape()}) |
| 2421 | .ToProto(); |
| 2422 | send_instr.set_channel_id(handle.handle()); |
| 2423 | TF_ASSIGN_OR_RETURN(XlaOp send, |
| 2424 | AddInstruction(std::move(send_instr), HloOpcode::kSend, |
| 2425 | {operand, token})); |
| 2426 | |
| 2427 | HloInstructionProto send_done_instr; |
| 2428 | *send_done_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); |
| 2429 | send_done_instr.set_channel_id(handle.handle()); |
| 2430 | return AddInstruction(std::move(send_done_instr), HloOpcode::kSendDone, |
| 2431 | {send}); |
| 2432 | }); |
| 2433 | } |
| 2434 | |
| 2435 | XlaOp XlaBuilder::Recv(const Shape& shape, const ChannelHandle& handle) { |
| 2436 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
no test coverage detected