| 245 | } |
| 246 | |
| 247 | std::vector<XrtTensorHandle> XrtTfContext::EnqueueOp( |
| 248 | absl::string_view name, absl::Span<const XrtTensorHandle* const> inputs, |
| 249 | int output_arity, protobuf::Map<std::string, AttrValue> attrs, |
| 250 | int device_id, std::shared_ptr<XrtRecvTensorFuture> future) { |
| 251 | std::vector<XrtTensorHandle> outputs; |
| 252 | absl::MutexLock lock(&mu_); |
| 253 | Operation* op = AddOperation(); |
| 254 | |
| 255 | eager::Operation* proto = enqueue_request_->add_queue()->mutable_operation(); |
| 256 | proto->set_id(op->id); |
| 257 | proto->set_name(static_cast<std::string>(name)); |
| 258 | for (const XrtTensorHandle* input : inputs) { |
| 259 | input->Serialize(proto->add_inputs()); |
| 260 | } |
| 261 | proto->mutable_attrs()->swap(attrs); |
| 262 | proto->set_device(devices_.at(device_id).name()); |
| 263 | |
| 264 | outputs.reserve(output_arity); |
| 265 | for (int i = 0; i < output_arity; ++i) { |
| 266 | outputs.push_back( |
| 267 | XrtTensorHandle(shared_from_this(), device_id, TensorId{op->id, i})); |
| 268 | } |
| 269 | if (future) { |
| 270 | op->tensor_futures.push_back(future); |
| 271 | } |
| 272 | |
| 273 | return outputs; |
| 274 | } |
| 275 | |
| 276 | XrtTensorHandle XrtTfContext::SendTensor( |
| 277 | std::unique_ptr<TensorProto> tensor_proto, int device_id, |