| 111 | } |
| 112 | |
| 113 | xla::StatusOr<InputBuffers> GetChainedOpInputs( |
| 114 | const xrt::XRTChainedExecuteOp& op, |
| 115 | absl::Span<const RefPtr<XRTTupleAllocation>> op_inputs) { |
| 116 | InputBuffers input_buffers; |
| 117 | input_buffers.input_tuples.reserve(op.inputs_size()); |
| 118 | input_buffers.input_allocations.reserve(op.inputs_size()); |
| 119 | input_buffers.input_pointers.reserve(op.inputs_size()); |
| 120 | for (int i = 0; i < op.inputs_size(); ++i) { |
| 121 | auto& input = op.inputs(i); |
| 122 | input_buffers.input_tuples.emplace_back(op_inputs[i]); |
| 123 | // Thanks to the greatness of proto3, there is no way to query for |
| 124 | // explicitly set fields, so the default for output_index (zero) means no |
| 125 | // sub-index. As consequence, the real index is output_index - 1. |
| 126 | if (input.output_index() == 0) { |
| 127 | TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, |
| 128 | input_buffers.input_tuples.back()->ToShapedBuffer()); |
| 129 | input_buffers.input_allocations.emplace_back(std::move(shaped_buffer)); |
| 130 | } else { |
| 131 | TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, |
| 132 | input_buffers.input_tuples.back()->ToShapedBuffer()); |
| 133 | TF_ASSIGN_OR_RETURN( |
| 134 | xla::ShapedBuffer sub_shaped_buffer, |
| 135 | shaped_buffer.SubShapedBuffer({input.output_index() - 1})); |
| 136 | input_buffers.input_allocations.emplace_back( |
| 137 | std::move(sub_shaped_buffer)); |
| 138 | } |
| 139 | } |
| 140 | for (size_t i = 0; i < input_buffers.input_allocations.size(); ++i) { |
| 141 | input_buffers.input_pointers.push_back(&input_buffers.input_allocations[i]); |
| 142 | } |
| 143 | return std::move(input_buffers); |
| 144 | } |
| 145 | |
| 146 | xla::StatusOr<RefPtr<XRTTupleAllocation>> RunExecutable( |
| 147 | OpKernelContext* context, XRTGenericDeviceAccessor::ScopedRef* device_ref, |
nothing calls this directly
no test coverage detected