TfLiteStatus AllocationInfoBuilder::AddTensors()

in src/tensorflow/lite/micro/micro_allocator.cpp [122:188]

• 56 lines of code
• 19 McCabe index (conditional complexity)

TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph,
                                               const int32_t* offline_offsets,
                                               TfLiteEvalTensor* eval_tensors) {
  TFLITE_DCHECK(eval_tensors != nullptr);

  // Set up allocation info for all tensors.
  for (size_t i = 0; i < tensor_count_; ++i) {
    AllocationInfo* current = &info_[i];
    current->output_ptr = &(eval_tensors[i].data.data);

    TF_LITE_ENSURE_STATUS(
        TfLiteEvalTensorByteLength(&eval_tensors[i], &current->bytes));

    current->first_created = -1;
    current->last_used = -1;
    current->needs_allocating = (eval_tensors[i].data.data == nullptr) &&
                                (!subgraph->tensors()->Get(i)->is_variable());
    if (offline_offsets) {
      current->offline_offset = offline_offsets[i];
    } else {
      current->offline_offset = kOnlinePlannedBuffer;
    }
  }

  uint32_t operators_size = NumSubgraphOperators(subgraph);

  for (size_t i = 0;
       subgraph->inputs() != nullptr && i < subgraph->inputs()->size(); ++i) {
    const int tensor_index = subgraph->inputs()->Get(i);
    AllocationInfo* current = &info_[tensor_index];
    current->first_created = 0;
  }

  // Mark all outputs as persistent to the end of the invocation.
  for (size_t i = 0;
       subgraph->outputs() != nullptr && i < subgraph->outputs()->size(); ++i) {
    const int tensor_index = subgraph->outputs()->Get(i);
    AllocationInfo* current = &info_[tensor_index];
    current->last_used = operators_size - 1;
  }

  // Figure out when the first and last use of each tensor is.
  for (int i = (operators_size - 1); i >= 0; --i) {
    const auto* op = subgraph->operators()->Get(i);
    for (size_t n = 0; op->inputs() != nullptr && n < op->inputs()->size();
         ++n) {
      const int tensor_index = op->inputs()->Get(n);
      AllocationInfo* current = &info_[tensor_index];
      if (((current->last_used == -1) || (current->last_used < i))) {
        current->last_used = i;
      }
    }
    for (size_t n = 0; op->outputs() != nullptr && n < op->outputs()->size();
         ++n) {
      const int tensor_index = op->outputs()->Get(n);
      AllocationInfo* current = &info_[tensor_index];
      if ((current->first_created == -1) || (current->first_created > i)) {
        current->first_created = i;
      }
      // Since operator outputs are written to, they must be marked as used.
      if ((current->last_used == -1) || (current->last_used < i)) {
        current->last_used = i;
      }
    }
  }
  return kTfLiteOk;
}