| 2458 | } |
| 2459 | |
| 2460 | XlaOp XlaBuilder::RecvWithToken(XlaOp token, const Shape& shape, |
| 2461 | const ChannelHandle& handle) { |
| 2462 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2463 | if (handle.type() != ChannelHandle::DEVICE_TO_DEVICE) { |
| 2464 | return InvalidArgument("Recv must use a device-to-device channel"); |
| 2465 | } |
| 2466 | |
| 2467 | // Recv instruction produces a tuple of {receive buffer, U32 context, |
| 2468 | // token}. |
| 2469 | HloInstructionProto recv_instr; |
| 2470 | *recv_instr.mutable_shape() = |
| 2471 | ShapeUtil::MakeTupleShape( |
| 2472 | {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}) |
| 2473 | .ToProto(); |
| 2474 | recv_instr.set_channel_id(handle.handle()); |
| 2475 | TF_ASSIGN_OR_RETURN(XlaOp recv, AddInstruction(std::move(recv_instr), |
| 2476 | HloOpcode::kRecv, {token})); |
| 2477 | |
| 2478 | HloInstructionProto recv_done_instr; |
| 2479 | *recv_done_instr.mutable_shape() = |
| 2480 | ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}) |
| 2481 | .ToProto(); |
| 2482 | recv_done_instr.set_channel_id(handle.handle()); |
| 2483 | return AddInstruction(std::move(recv_done_instr), HloOpcode::kRecvDone, |
| 2484 | {recv}); |
| 2485 | }); |
| 2486 | } |
| 2487 | |
| 2488 | XlaOp XlaBuilder::SendToHost(XlaOp operand, XlaOp token, |
| 2489 | const Shape& shape_with_layout, |
no test coverage detected