| 2811 | } |
| 2812 | |
| 2813 | StatusOr<XlaOp> XlaBuilder::AddInstruction(HloInstructionProto&& instr, |
| 2814 | HloOpcode opcode, |
| 2815 | absl::Span<const XlaOp> operands) { |
| 2816 | TF_RETURN_IF_ERROR(first_error_); |
| 2817 | |
| 2818 | const int64 handle = GetNextId(); |
| 2819 | instr.set_id(handle); |
| 2820 | instr.set_opcode(HloOpcodeString(opcode)); |
| 2821 | if (instr.name().empty()) { |
| 2822 | instr.set_name(instr.opcode()); |
| 2823 | } |
| 2824 | for (const auto& operand : operands) { |
| 2825 | if (operand.builder_ == nullptr) { |
| 2826 | return InvalidArgument("invalid XlaOp with handle %d", operand.handle()); |
| 2827 | } |
| 2828 | if (operand.builder_ != this) { |
| 2829 | return InvalidArgument("Do not add XlaOp from builder %s to builder %s", |
| 2830 | operand.builder_->name(), this->name()); |
| 2831 | } |
| 2832 | instr.add_operand_ids(operand.handle()); |
| 2833 | } |
| 2834 | |
| 2835 | *instr.mutable_metadata() = metadata_; |
| 2836 | if (sharding_) { |
| 2837 | *instr.mutable_sharding() = *sharding_; |
| 2838 | } |
| 2839 | *instr.mutable_frontend_attributes() = frontend_attributes_; |
| 2840 | |
| 2841 | handle_to_index_[handle] = instructions_.size(); |
| 2842 | instructions_.push_back(std::move(instr)); |
| 2843 | instruction_shapes_.push_back( |
| 2844 | absl::make_unique<Shape>(instructions_.back().shape())); |
| 2845 | |
| 2846 | XlaOp op(handle, this); |
| 2847 | return op; |
| 2848 | } |
| 2849 | |
| 2850 | void XlaBuilder::AddCalledComputation(const XlaComputation& computation, |
| 2851 | HloInstructionProto* instr) { |