| 61 | } |
| 62 | |
| 63 | TfLiteStatus ArenaPlanner::PlanAllocations() { |
| 64 | // Invalidate any existing data. |
| 65 | TF_LITE_ENSURE_STATUS(ResetAllocations()); |
| 66 | // The alloc_queue_ is specific to the graph topology, and will be |
| 67 | // completely reconstructed from graph data here. |
| 68 | alloc_queue_.clear(); |
| 69 | |
| 70 | // Keeps track of references to each tensor. |
| 71 | std::vector<int> refcounts(graph_info_->num_tensors(), 0); |
| 72 | // `allocated` and `deallocated` are technically list of boolean values. |
| 73 | // We're saving the compiled binary size by using `vector<int>`. |
| 74 | std::vector<int> allocated(graph_info_->num_tensors(), false); |
| 75 | std::vector<int> deallocated(graph_info_->num_tensors(), false); |
| 76 | |
| 77 | auto allocate = [this, &allocated, &deallocated](int node, |
| 78 | int tensor) -> TfLiteStatus { |
| 79 | if (allocated[tensor]) { |
| 80 | return kTfLiteOk; |
| 81 | } |
| 82 | TF_LITE_ENSURE(context_, !deallocated[tensor]); |
| 83 | alloc_queue_.push_back({node, tensor, AllocationInfo::ALLOC}); |
| 84 | allocated[tensor] = true; |
| 85 | return kTfLiteOk; |
| 86 | }; |
| 87 | |
| 88 | auto deallocate = [this, &allocated, &deallocated]( |
| 89 | int node, int tensor) -> TfLiteStatus { |
| 90 | if (!allocated[tensor]) { |
| 91 | // Do not enqueue a DEALLOC if the tensor is never allocated. |
| 92 | // This happened with the constant tensors. |
| 93 | return kTfLiteOk; |
| 94 | } |
| 95 | TF_LITE_ENSURE(context_, !deallocated[tensor]); |
| 96 | alloc_queue_.push_back({node, tensor, AllocationInfo::DEALLOC}); |
| 97 | return kTfLiteOk; |
| 98 | }; |
| 99 | |
| 100 | // There will be an entry in alloc_queue_ for the allocation of each tensor |
| 101 | // and another for their deallocation. |
| 102 | alloc_queue_.reserve(2 * graph_info_->num_tensors()); |
| 103 | |
| 104 | // We must make sure the output tensors are never overwritten. We do that by |
| 105 | // artificially adding one to their ref-counts so they are never selected |
| 106 | // for deallocation. |
| 107 | for (int tensor_index : graph_info_->outputs()) { |
| 108 | refcounts[tensor_index]++; |
| 109 | } |
| 110 | |
| 111 | // Variable tensors also should be ensured to be never overwritten and need to |
| 112 | // be alive all the time. |
| 113 | for (int tensor_index : graph_info_->variables()) { |
| 114 | refcounts[tensor_index]++; |
| 115 | } |
| 116 | |
| 117 | // Queue all graph inputs for allocation. If preserve_inputs_ is true, make |
| 118 | // sure they never be overwritten. |
| 119 | for (int tensor_index : graph_info_->inputs()) { |
| 120 | if (tensor_index != kOptionalTensor) { |