static*/
| 109 | } |
| 110 | |
| 111 | /*static*/ xla::StatusOr<std::shared_ptr<XrtBuffer>> XrtBuffer::MakeTuple( |
| 112 | const std::shared_ptr<XrtContext>& context, |
| 113 | const std::vector<std::shared_ptr<XrtBuffer>>& elements, |
| 114 | int xrt_device_ordinal) { |
| 115 | if (elements.empty()) { |
| 116 | // XRTMakeTuple cannot construct empty tuples. Construct via a literal |
| 117 | // instead. |
| 118 | return FromLiteral(context, xrt_device_ordinal, |
| 119 | xla::LiteralUtil::MakeTuple({})); |
| 120 | } |
| 121 | |
| 122 | if (xrt_device_ordinal < 0 || |
| 123 | xrt_device_ordinal >= context->tf_device_ids().size()) { |
| 124 | return errors::InvalidArgument("Invalid XRT device ordinal ", |
| 125 | xrt_device_ordinal); |
| 126 | } |
| 127 | int tf_device_id = context->tf_device_ids().at(xrt_device_ordinal); |
| 128 | xrt::XLATupleNode tuple_description; |
| 129 | std::vector<xla::Shape> element_shapes; |
| 130 | element_shapes.reserve(elements.size()); |
| 131 | for (int index = 0; index < elements.size(); ++index) { |
| 132 | xrt::XLATupleNode* node = tuple_description.add_tuples(); |
| 133 | node->set_input_index(index); |
| 134 | element_shapes.push_back(elements[index]->shape()); |
| 135 | if (elements[index]->handle().device_id() != tf_device_id) { |
| 136 | return errors::InvalidArgument( |
| 137 | "All elements of tuple must be on the same device ( ", |
| 138 | elements[index]->handle().device_id(), " vs. ", tf_device_id, ")"); |
| 139 | } |
| 140 | } |
| 141 | auto proto = absl::make_unique<TensorProto>(); |
| 142 | proto->set_dtype(DT_STRING); |
| 143 | tuple_description.SerializeToString(proto->add_string_val()); |
| 144 | |
| 145 | XrtTensorHandle description_handle = |
| 146 | context->tf_context()->SendTensor(std::move(proto), tf_device_id, |
| 147 | /*host_memory=*/true); |
| 148 | |
| 149 | protobuf::Map<string, AttrValue> attrs; |
| 150 | attrs["Ninputs"] = MakeAttrValue(elements.size()); |
| 151 | |
| 152 | std::vector<const XrtTensorHandle*> args; |
| 153 | args.reserve(elements.size() + 1); |
| 154 | args.push_back(&description_handle); |
| 155 | for (const auto& element : elements) { |
| 156 | args.push_back(&element->handle()); |
| 157 | } |
| 158 | XrtTensorHandle buffer_handle = std::move(context->tf_context()->EnqueueOp( |
| 159 | "XRTMakeTuple", args, /*output_arity=*/1, attrs, tf_device_id)[0]); |
| 160 | return std::make_shared<XrtBuffer>( |
| 161 | std::move(buffer_handle), xrt_device_ordinal, |
| 162 | xla::ShapeUtil::MakeTupleShape(element_shapes)); |
| 163 | } |
| 164 | |
| 165 | xla::StatusOr<xla::Literal> XrtBuffer::ToLiteral() const { |
| 166 | TF_RET_CHECK(handle_.valid()); |