| 2533 | } |
| 2534 | |
| 2535 | XlaOp XlaBuilder::RecvFromHost(XlaOp token, const Shape& shape, |
| 2536 | const ChannelHandle& handle) { |
| 2537 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2538 | if (!LayoutUtil::HasLayout(shape)) { |
| 2539 | return InvalidArgument("Shape passed to RecvFromHost must have a layout"); |
| 2540 | } |
| 2541 | |
| 2542 | // TODO(b/111544877): Support tuple shapes. |
| 2543 | if (!shape.IsArray()) { |
| 2544 | return InvalidArgument( |
| 2545 | "RecvFromHost only supports array shapes, shape: %s", |
| 2546 | ShapeUtil::HumanString(shape)); |
| 2547 | } |
| 2548 | |
| 2549 | if (handle.type() != ChannelHandle::HOST_TO_DEVICE) { |
| 2550 | return InvalidArgument("RecvFromHost must use a host-to-device channel"); |
| 2551 | } |
| 2552 | |
| 2553 | // Recv instruction produces a tuple of {receive buffer, U32 context, |
| 2554 | // token}. |
| 2555 | HloInstructionProto recv_instr; |
| 2556 | *recv_instr.mutable_shape() = |
| 2557 | ShapeUtil::MakeTupleShape( |
| 2558 | {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}) |
| 2559 | .ToProto(); |
| 2560 | recv_instr.set_channel_id(handle.handle()); |
| 2561 | recv_instr.set_is_host_transfer(true); |
| 2562 | TF_ASSIGN_OR_RETURN(XlaOp recv, AddInstruction(std::move(recv_instr), |
| 2563 | HloOpcode::kRecv, {token})); |
| 2564 | |
| 2565 | HloInstructionProto recv_done_instr; |
| 2566 | *recv_done_instr.mutable_shape() = |
| 2567 | ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}) |
| 2568 | .ToProto(); |
| 2569 | recv_done_instr.set_channel_id(handle.handle()); |
| 2570 | recv_done_instr.set_is_host_transfer(true); |
| 2571 | return AddInstruction(std::move(recv_done_instr), HloOpcode::kRecvDone, |
| 2572 | {recv}); |
| 2573 | }); |
| 2574 | } |
| 2575 | |
| 2576 | XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64 dimension) { |
| 2577 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
no test coverage detected