| 2433 | } |
| 2434 | |
| 2435 | XlaOp XlaBuilder::Recv(const Shape& shape, const ChannelHandle& handle) { |
| 2436 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2437 | // Recv HLO takes a single token operand. Generate the token to pass into |
| 2438 | // the Recv and RecvDone instructions. |
| 2439 | // TODO(b/80000000): Remove this when clients have been updated to handle |
| 2440 | // tokens. |
| 2441 | HloInstructionProto token_instr; |
| 2442 | *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); |
| 2443 | TF_ASSIGN_OR_RETURN(XlaOp token, AddInstruction(std::move(token_instr), |
| 2444 | HloOpcode::kAfterAll, {})); |
| 2445 | |
| 2446 | XlaOp recv = RecvWithToken(token, shape, handle); |
| 2447 | |
| 2448 | // The RecvDone instruction produces a tuple of the data and a token |
| 2449 | // type. Return XLA op containing the data. |
| 2450 | // TODO(b/80000000): Remove this when clients have been updated to handle |
| 2451 | // tokens. |
| 2452 | HloInstructionProto recv_data; |
| 2453 | *recv_data.mutable_shape() = shape.ToProto(); |
| 2454 | recv_data.set_tuple_index(0); |
| 2455 | return AddInstruction(std::move(recv_data), HloOpcode::kGetTupleElement, |
| 2456 | {recv}); |
| 2457 | }); |
| 2458 | } |
| 2459 | |
| 2460 | XlaOp XlaBuilder::RecvWithToken(XlaOp token, const Shape& shape, |
| 2461 | const ChannelHandle& handle) { |
no test coverage detected