| 73 | } |
| 74 | |
| 75 | xla::StatusOr<InputBuffers> GetInputBuffers( |
| 76 | XRTMemoryManager::WorkingSet* working_set, xla::Backend* backend, |
| 77 | const std::vector<InputCoords>& input_coords, bool release_inputs) { |
| 78 | InputBuffers input_buffers; |
| 79 | input_buffers.input_tuples.reserve(input_coords.size()); |
| 80 | input_buffers.input_allocations.reserve(input_coords.size()); |
| 81 | input_buffers.input_pointers.reserve(input_coords.size()); |
| 82 | for (size_t i = 0; i < input_coords.size(); ++i) { |
| 83 | TF_RETURN_IF_ERROR( |
| 84 | working_set->LookupAndPin(backend, input_coords[i].handle)); |
| 85 | auto tuple = working_set->PinnedTuples().back(); |
| 86 | input_buffers.input_tuples.emplace_back(tuple); |
| 87 | if (release_inputs) { |
| 88 | // We are holding a reference to the tuple, so we can safely delete it |
| 89 | // from the resource manager here. |
| 90 | TF_RETURN_IF_ERROR( |
| 91 | working_set->MemoryManager()->Release(input_coords[i].handle)); |
| 92 | VLOG(2) << "Released allocation handle " << input_coords[i].handle; |
| 93 | } |
| 94 | if (input_coords[i].index.empty()) { |
| 95 | TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, |
| 96 | tuple->ToShapedBuffer()); |
| 97 | input_buffers.input_allocations.emplace_back(std::move(shaped_buffer)); |
| 98 | } else { |
| 99 | TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, |
| 100 | tuple->ToShapedBuffer()); |
| 101 | TF_ASSIGN_OR_RETURN(xla::ShapedBuffer sub_shaped_buffer, |
| 102 | shaped_buffer.SubShapedBuffer(input_coords[i].index)); |
| 103 | input_buffers.input_allocations.emplace_back( |
| 104 | std::move(sub_shaped_buffer)); |
| 105 | } |
| 106 | } |
| 107 | for (size_t i = 0; i < input_buffers.input_allocations.size(); ++i) { |
| 108 | input_buffers.input_pointers.push_back(&input_buffers.input_allocations[i]); |
| 109 | } |
| 110 | return std::move(input_buffers); |
| 111 | } |
| 112 | |
| 113 | xla::StatusOr<InputBuffers> GetChainedOpInputs( |
| 114 | const xrt::XRTChainedExecuteOp& op, |
nothing calls this directly
no test coverage detected