| 202 | } |
| 203 | |
| 204 | Status CreateExecuteOutput(OpKernelContext* context, |
| 205 | XRTMemoryManager* memory_manager, |
| 206 | RefPtr<XRTTupleAllocation> output_tuple, |
| 207 | bool return_exploded_tuple) { |
| 208 | if (return_exploded_tuple && output_tuple->on_host_shape().IsTuple()) { |
| 209 | int64 tuple_element_count = |
| 210 | xla::ShapeUtil::TupleElementCount(output_tuple->on_device_shape()); |
| 211 | Tensor* output_tensor; |
| 212 | TF_RETURN_IF_ERROR(context->allocate_output( |
| 213 | 0, TensorShape({tuple_element_count}), &output_tensor)); |
| 214 | |
| 215 | for (int64 i = 0; i < tuple_element_count; ++i) { |
| 216 | XRTTupleAllocation* suballocation; |
| 217 | TF_RETURN_IF_ERROR(XRTTupleAllocation::MakeSubBuffer( |
| 218 | output_tuple.get(), {i}, &suballocation, |
| 219 | /*alias_parent_allocation=*/false)); |
| 220 | output_tensor->vec<int64>()(i) = memory_manager->Register(suballocation); |
| 221 | } |
| 222 | } else { |
| 223 | Tensor* output_tensor; |
| 224 | TF_RETURN_IF_ERROR( |
| 225 | context->allocate_output(0, TensorShape({}), &output_tensor)); |
| 226 | output_tensor->scalar<int64>()() = |
| 227 | memory_manager->Register(std::move(output_tuple)); |
| 228 | } |
| 229 | return Status::OK(); |
| 230 | } |
| 231 | |
| 232 | Status ExecuteChained(OpKernelContext* context, |
| 233 | const RefPtr<XRTMemoryManager>& memory_manager, |
no test coverage detected