/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/lite/micro/micro_allocator.h"

#include <cstddef>
#include <cstdint>

#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/api/error_reporter.h"
#include "tensorflow/lite/core/api/flatbuffer_conversions.h"
#include "tensorflow/lite/core/api/op_resolver.h"
#include "tensorflow/lite/core/api/tensor_utils.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/micro/compatibility.h"
#include "tensorflow/lite/micro/memory_helpers.h"
#include "tensorflow/lite/micro/memory_planner/greedy_memory_planner.h"
#include "tensorflow/lite/micro/memory_planner/memory_planner.h"
#include "tensorflow/lite/micro/micro_op_resolver.h"
#include "tensorflow/lite/micro/simple_memory_allocator.h"
#include "tensorflow/lite/schema/schema_generated.h"

namespace tflite {

namespace {
// Used to hold information used during allocation calculations.
struct AllocationInfo {
  size_t bytes;
  void** output_ptr;
  int first_created;
  int last_used;
  int32_t offline_offset;
  bool needs_allocating;
};

// We align tensor buffers to 16-byte boundaries, since this is a common
// requirement for SIMD extensions.
constexpr int kBufferAlignment = 16;

constexpr char kOfflineMemAllocMetadata[] = "OfflineMemoryAllocation";

// Instance of a zero-length int to pass as tensor dims for a flatbuffer
// Tensor with no shape. Note that the second member of a TfLiteArray is a
// flexible array member, which is not strictly valid C++. However it is
// supported by both GCC and clang, as long as the flexible array element is not
// initialized, which is ok in this case as it should never be accessed.
// Declaring this as constexpr causes build errors with clang, as it requires
// the flexible array element to be initialized.
const TfLiteIntArray kZeroLengthIntArray = {0};

class MicroBuiltinDataAllocator : public BuiltinDataAllocator {
 public:
  explicit MicroBuiltinDataAllocator(SimpleMemoryAllocator* memory_allocator)
      : memory_allocator_(memory_allocator) {}

  void* Allocate(size_t size, size_t alignment_hint) override {
    return memory_allocator_->AllocateFromTail(size, alignment_hint);
  }
  void Deallocate(void* data) override {
    // Do not deallocate, builtin data needs to be available for the life time
    // of the model.
  }

 private:
  SimpleMemoryAllocator* memory_allocator_;

  TF_LITE_REMOVE_VIRTUAL_DELETE
};

#if !defined(__clang__)
// Helper function to check flatbuffer metadata correctness. This function is
// not called by default. Hence it's not linked in to the final binary code.
TfLiteStatus CheckOfflinePlannedOffsets(const Model* model,
                                        ErrorReporter* error_reporter) {
  // Suppress compile warning for unused function
  (void)CheckOfflinePlannedOffsets;

  if (model->metadata()) {
    for (size_t i = 0; i < model->metadata()->size(); ++i) {
      auto metadata = model->metadata()->Get(i);
      if (strncmp(metadata->name()->c_str(), kOfflineMemAllocMetadata,
                  strlen(kOfflineMemAllocMetadata)) == 0) {
        auto* subgraphs = model->subgraphs();
        const SubGraph* subgraph = (*subgraphs)[0];
        const flatbuffers::Vector<flatbuffers::Offset<Tensor>>* tensors =
            subgraph->tensors();
        const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* buffers =
            model->buffers();
        int nbr_tflite_tensors = tensors->size();
        auto* buffer = (*buffers)[metadata->buffer()];
        auto* array = buffer->data();
        const uint32_t* metadata_buffer = (uint32_t*)array->data();
        int version = metadata_buffer[0];
        int subgraph_idx = metadata_buffer[1];
        const int nbr_offline_offsets = metadata_buffer[2];
        int* offline_planner_offsets = (int*)&metadata_buffer[3];

        TF_LITE_REPORT_ERROR(error_reporter, "==== Model metadata info: =====");
        TF_LITE_REPORT_ERROR(error_reporter,
                             "Offline planner metadata found, version %d, "
                             "subgraph %d, nbr offline offsets %d",
                             version, subgraph_idx, nbr_offline_offsets);
        for (int j = 0; j < nbr_offline_offsets; ++j) {
          TF_LITE_REPORT_ERROR(
              error_reporter,
              "Offline planner tensor index %d, offline offset: %d", j,
              offline_planner_offsets[j]);
        }

        if (version != 1) {
          TF_LITE_REPORT_ERROR(error_reporter, "Version not supported! (%d)\n",
                               version);
          return kTfLiteError;
        }
        if (subgraph_idx != 0) {
          TF_LITE_REPORT_ERROR(error_reporter,
                               "Only 1 subgraph supported! Subgraph idx (%d)\n",
                               subgraph_idx);
          return kTfLiteError;
        }
        if (nbr_tflite_tensors != nbr_offline_offsets) {
          TF_LITE_REPORT_ERROR(error_reporter,
                               "Nbr of offline buffer offsets (%d) in metadata "
                               "not equal nbr tensors (%d)\n",
                               nbr_offline_offsets, nbr_tflite_tensors);
          return kTfLiteError;
        }
      }
    }
  }
  return kTfLiteOk;
}
#endif

// A helper class to construct AllocationInfo array. This array contains the
// lifetime of tensors / scratch_buffer and will be used to calculate the memory
// plan. Methods need to be called in order from `Init`, `Add*`, to `Finish`.
class AllocationInfoBuilder {
 public:
  AllocationInfoBuilder(ErrorReporter* reporter,
                        SimpleMemoryAllocator* allocator)
      : reporter_(reporter), allocator_(allocator) {}

  // Initializes the builder by allocating AllocationInfo array from the
  // simple memory allocator.
  TfLiteStatus Init(size_t tensor_count, size_t scratch_buffer_count) {
    tensor_count_ = tensor_count;
    buffer_count_ = scratch_buffer_count;
    return Allocate();
  }

  // Check if model contains offline planned buffer offsets.
  //  - If there's no metadata available, offline_planner_offsets is not set
  //  - If there's metadata available, offline_planner_offsets will point to the
  //    first offset in the metadata buffer list.
  TfLiteStatus GetOfflinePlannedOffsets(
      const Model* model, const int32_t** offline_planner_offsets);

  // Add allocaiton information for the tensors.
  TfLiteStatus AddTensors(const SubGraph* subgraph,
                          const int32_t* offline_offsets,
                          TfLiteTensor* runtime_tensors);

  // Add allocation information for the scratch buffers.
  TfLiteStatus AddScratchBuffers(internal::ScratchBufferHandle* buffer_handles);

  // Returns a pointer to the built AllocationInfo array.
  const AllocationInfo* Finish() const { return info_; }
  size_t Size() const { return tensor_count_ + buffer_count_; }

 private:
  // Allocate the output AllocationInfo array from the allocator_;
  TfLiteStatus Allocate();

  ErrorReporter* reporter_ = nullptr;
  SimpleMemoryAllocator* allocator_ = nullptr;
  size_t tensor_count_ = 0;
  size_t buffer_count_ = 0;
  AllocationInfo* info_ = nullptr;
};

TfLiteStatus AllocationInfoBuilder::Allocate() {
  size_t bytes = sizeof(AllocationInfo) * Size();
  info_ = reinterpret_cast<AllocationInfo*>(
      allocator_->AllocateFromTail(bytes, alignof(AllocationInfo)));
  if (info_ == nullptr) {
    TF_LITE_REPORT_ERROR(
        reporter_,
        "Failed to allocate memory for allocation_info, %d bytes required",
        bytes);
    return kTfLiteError;
  }
  return kTfLiteOk;
}

TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph,
                                               const int32_t* offline_offsets,
                                               TfLiteTensor* runtime_tensors) {
  // Set up allocation info for all tensors.
  for (size_t i = 0; i < tensor_count_; ++i) {
    AllocationInfo* current = &info_[i];
    // TfLiteTensor.uint8 field is deprecated so use .data field instead.
    current->output_ptr = &(runtime_tensors[i].data.data);
    current->bytes = runtime_tensors[i].bytes;
    current->first_created = -1;
    current->last_used = -1;
    current->needs_allocating = (runtime_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;
    }
  }

  for (size_t i = 0; 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; i < subgraph->outputs()->size(); ++i) {
    const int tensor_index = subgraph->outputs()->Get(i);
    AllocationInfo* current = &info_[tensor_index];
    current->last_used = subgraph->operators()->size() - 1;
  }

  // Figure out when the first and last use of each tensor is.
  for (int i = (subgraph->operators()->size() - 1); i >= 0; --i) {
    const auto* op = subgraph->operators()->Get(i);
    for (size_t n = 0; 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; 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;
      }
    }
  }

  // Work out which tensors need to be allocated.
  for (size_t i = 0; i < tensor_count_; ++i) {
    AllocationInfo* current = &info_[i];
    const bool is_read_only =
        (current->first_created == -1) && (current->last_used != -1);
    if (is_read_only) {
      current->needs_allocating = false;
    }
    const bool has_partial_lifetime =
        !is_read_only &&
        ((current->first_created == -1) || (current->last_used == -1));
    if (has_partial_lifetime && current->needs_allocating) {
      TF_LITE_REPORT_ERROR(
          reporter_,
          "Logic error in memory planner, tensor %d has an invalid lifetime: "
          "first_created: %d, last_used: %d",
          i, current->first_created, current->last_used);
      return kTfLiteError;
    }
  }
  return kTfLiteOk;
}

// The tensor offsets will be encoded in the metadata:[Metadata] field of the
// Model. The following encoding applies:
//
// | Metadata component |                 Value                                |
// |    name:string     | “OfflineMemoryAllocation”                            |
// |    buffer:unit     | Index of buffer containing memory allocation data    |
//
// The buffer contents for the memory allocation is a list of 32-bit integers.
// The number of tensors, n, must be equal to the number of tensors defined in
// the model. The following encoding applies:
//
// |  Offset |                            Value                                |
// |    0    | Offline allocation format version – set to 0                    |
// |    1    | Subgraph index to which this allocation applies                 |
// |    2    | Number offsets following: n                                     |
// |    3    | Arena byte offset of tensor #0 or -1 to allocate at runtime     |
// |    4    | Arena byte offset of tensor #1 or -1 to allocate at runtime     |
// | 3+(n-1) | Arena byte offset of tensor #(n-1) or -1 to allocate at runtime |
TfLiteStatus AllocationInfoBuilder::GetOfflinePlannedOffsets(
    const Model* model, const int32_t** offline_planner_offsets) {
  if (model->metadata()) {
    for (size_t i = 0; i < model->metadata()->size(); ++i) {
      auto metadata = model->metadata()->Get(i);
      if (strncmp(metadata->name()->c_str(), kOfflineMemAllocMetadata,
                  strlen(kOfflineMemAllocMetadata)) == 0) {
        const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* buffers =
            model->buffers();
        auto* buffer = (*buffers)[metadata->buffer()];
        auto* array = buffer->data();
        const uint32_t* metadata_buffer =
            reinterpret_cast<const uint32_t*>(array->data());
        const size_t nbr_tensors = static_cast<size_t>(metadata_buffer[2]);
        *offline_planner_offsets =
            reinterpret_cast<const int32_t*>(&metadata_buffer[3]);

        if (tensor_count_ != nbr_tensors) {
          TF_LITE_REPORT_ERROR(reporter_,
                               "Nbr of offline buffer offsets (%d) in metadata "
                               "not equal nbr tensors (%d)\n",
                               nbr_tensors, tensor_count_);
          return kTfLiteError;
        }
      }
    }
  }
  return kTfLiteOk;
}

TfLiteStatus AllocationInfoBuilder::AddScratchBuffers(
    internal::ScratchBufferHandle* buffer_handles) {
  // Set up allocation info for buffers.
  for (size_t i = tensor_count_; i < tensor_count_ + buffer_count_; ++i) {
    AllocationInfo* current = &info_[i];
    internal::ScratchBufferHandle* handle =
        &(buffer_handles[i - tensor_count_]);
    current->output_ptr = reinterpret_cast<void**>(&handle->data);
    current->bytes = handle->bytes;
    current->first_created = handle->node_idx;
    current->last_used = handle->node_idx;
    current->needs_allocating = true;
    current->offline_offset = kOnlinePlannedBuffer;
  }
  return kTfLiteOk;
}

TfLiteStatus CreatePlan(ErrorReporter* error_reporter,
                        GreedyMemoryPlanner* planner,
                        const AllocationInfo* allocation_info,
                        size_t allocation_info_size) {
  // Add the tensors to our allocation plan.
  for (size_t i = 0; i < allocation_info_size; ++i) {
    const AllocationInfo* current = &allocation_info[i];
    if (current->needs_allocating) {
      size_t aligned_bytes_required =
          AlignSizeUp(current->bytes, kBufferAlignment);
      if (current->offline_offset == kOnlinePlannedBuffer) {
        TF_LITE_ENSURE_STATUS(
            planner->AddBuffer(error_reporter, aligned_bytes_required,
                               current->first_created, current->last_used));
      } else {
        TF_LITE_ENSURE_STATUS(planner->AddBuffer(
            error_reporter, aligned_bytes_required, current->first_created,
            current->last_used, current->offline_offset));
      }
    }
  }
  return kTfLiteOk;
}

TfLiteStatus CommitPlan(ErrorReporter* error_reporter, MemoryPlanner* planner,
                        uint8_t* starting_point,
                        const AllocationInfo* allocation_info,
                        size_t allocation_info_size) {
  // Figure out the actual memory addresses for each buffer, based on the plan.
  int planner_index = 0;
  for (size_t i = 0; i < allocation_info_size; ++i) {
    const AllocationInfo* current = &allocation_info[i];
    if (current->needs_allocating) {
      int offset = -1;
      TF_LITE_ENSURE_STATUS(
          planner->GetOffsetForBuffer(error_reporter, planner_index, &offset));
      *current->output_ptr = reinterpret_cast<void*>(starting_point + offset);
      ++planner_index;
    }
  }
  return kTfLiteOk;
}
}  // namespace

namespace internal {

// Handles architecture safe mapping of flatbuffer vectors to a TfLite*Array
// struct. Matching types are required (e.g. float and TfLiteFloatArray).
template <typename kFlatBufferVectorType, typename kTfLiteArrayType>
TfLiteStatus FlatBufferVectorToTfLiteTypeArray(
    SimpleMemoryAllocator* allocator, ErrorReporter* error_reporter,
    const flatbuffers::Vector<kFlatBufferVectorType>* flatbuffer_array,
    kTfLiteArrayType** result) {
  TFLITE_DCHECK(error_reporter != nullptr);
  TFLITE_DCHECK(flatbuffer_array != nullptr);
  // Only two conversions are supported - float and int32 - ensure that these
  // match at compile time instead of duplicating functions here:
  static_assert((std::is_same<kFlatBufferVectorType, int32_t>() &&
                 std::is_same<kTfLiteArrayType, TfLiteIntArray>()) ||
                (std::is_same<kFlatBufferVectorType, float>() &&
                 std::is_same<kTfLiteArrayType, TfLiteFloatArray>()));
  if (FLATBUFFERS_LITTLEENDIAN) {
    // On little-endian machines, TfLite*Array happens to have the same memory
    // layout as flatbuffers:Vector<kFlatBufferVectorType>, so we can
    // reinterpret_cast the flatbuffer vector and avoid a copy and malloc.
    *result = const_cast<kTfLiteArrayType*>(
        reinterpret_cast<const kTfLiteArrayType*>(flatbuffer_array));
  } else {
    // Big-endian architecture can not use the same memory layout as
    // flatbuffers::Vector<kFlatBufferVectorType>. Allocate from the tail and
    // copy values from the flatbuffer into the newly allocated chunk.
    kTfLiteArrayType* array =
        reinterpret_cast<kTfLiteArrayType*>(allocator->AllocateFromTail(
            TfLiteIntArrayGetSizeInBytes(flatbuffer_array->Length()),
            alignof(kTfLiteArrayType)));
    if (array == nullptr) {
      TF_LITE_REPORT_ERROR(
          error_reporter,
          "Failed to allocate %d bytes of memory to copy an array.",
          TfLiteIntArrayGetSizeInBytes(flatbuffer_array->Length()));
      return kTfLiteError;
    }
    array->size = flatbuffer_array->Length();
    for (int i = 0; i < array->size; ++i) {
      array->data[i] = flatbuffer_array->Get(i);
    }
    *result = array;
  }
  return kTfLiteOk;
}

TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
    SimpleMemoryAllocator* allocator, const tflite::Tensor& flatbuffer_tensor,
    const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* buffers,
    ErrorReporter* error_reporter, TfLiteTensor* result) {
  *result = {};
  // Make sure the serialized type is one we know how to deal with, and convert
  // it from a flatbuffer enum into a constant used by the kernel C API.
  TF_LITE_ENSURE_STATUS(ConvertTensorType(flatbuffer_tensor.type(),
                                          &result->type, error_reporter));
  // Make sure we remember if the serialized tensor is designated as a variable.
  result->is_variable = flatbuffer_tensor.is_variable();

  // We need to figure out where the actual contents of this tensor are stored
  // in memory. We'll check to see if there's a serialized buffer (pretty much
  // the same as a constant op in TensorFlow) associated with this tensor first,
  // and if there is update the runtime structure to point to its location in
  // memory.
  // First see if there's any buffer information in the serialized tensor.
  if (auto* buffer = (*buffers)[flatbuffer_tensor.buffer()]) {
    // If we've found a buffer, does it have any data?
    if (auto* array = buffer->data()) {
      // If it has any data, is the data size larger than zero?
      if (array->size()) {
        // We've found a buffer with valid data, so update the runtime tensor
        // data structure to point to it.
        result->data.data =
            const_cast<void*>(static_cast<const void*>(array->data()));
        // We set the data from a serialized buffer, so record tha.
        result->allocation_type = kTfLiteMmapRo;
      }
    }
    // TODO(petewarden): It's not clear in what circumstances we could have a
    // buffer in the serialized tensor, but it doesn't have any data in it. Is
    // that a validly-generated file, and if so what does it mean, or is it an
    // error condition? It would be good to tighten up the specification to make
    // it less ambiguous.
  }

  // TODO(petewarden): Some of these paths aren't getting enough testing
  // coverage, so we should figure out some tests that exercise them.
  if (result->data.data == nullptr) {
    // The tensor contents haven't been set from a serialized buffer, so
    // make a note that they will be allocated from memory. The actual
    // allocation won't happen until later.
    result->allocation_type = kTfLiteArenaRw;
  }

  // Figure out what the size in bytes of the buffer is and store it.
  size_t type_size;
  TF_LITE_ENSURE_STATUS(BytesRequiredForTensor(
      flatbuffer_tensor, &result->bytes, &type_size, error_reporter));

  if (flatbuffer_tensor.shape() == nullptr) {
    // flatbuffer_tensor.shape() can return a nullptr in the case of a scalar
    // tensor.
    result->dims = const_cast<TfLiteIntArray*>(&kZeroLengthIntArray);
  } else {
    // TFLM doesn't allow reshaping the tensor which requires dynamic memory
    // allocation so it is safe to drop the const qualifier. In the future, if
    // we really want to update the tensor shape, we can always pass in a new
    // TfLiteIntArray - especially we have to do so if the dimension is
    TF_LITE_ENSURE_STATUS(FlatBufferVectorToTfLiteTypeArray(
        allocator, error_reporter, flatbuffer_tensor.shape(), &(result->dims)));
  }

  // Copy the quantization information from the serialized data.
  const auto* src_quantization = flatbuffer_tensor.quantization();
  if (src_quantization && src_quantization->scale() &&
      (src_quantization->scale()->size() > 0) &&
      src_quantization->zero_point() &&
      (src_quantization->zero_point()->size() > 0)) {
    // Always populate the TfLiteTensor.params field, even if there are
    // per-channel quantization parameters.
    result->params.scale = src_quantization->scale()->Get(0);
    // Note that the zero_point field in the FlatBuffers schema is a 64-bit
    // integer, but the zero_point field in the TfLiteQuantizationParams struct
    // is a 32-bit integer.
    result->params.zero_point =
        static_cast<int32_t>(src_quantization->zero_point()->Get(0));

    // Populate per-channel quantization params.
    int channels = src_quantization->scale()->size();
    TfLiteAffineQuantization* quantization =
        reinterpret_cast<TfLiteAffineQuantization*>(
            allocator->AllocateFromTail(sizeof(TfLiteAffineQuantization),
                                        alignof(TfLiteAffineQuantization)));
    if (quantization == nullptr) {
      TF_LITE_REPORT_ERROR(error_reporter,
                           "Unable to allocate TfLiteAffineQuantization.\n");
      return kTfLiteError;
    }

    // TODO(b/153688719): Reduce tail allocation by using a global zero-point
    // buffer. This value can not be reused from the flatbuffer since the
    // zero_point is stored as a int64_t.
    quantization->zero_point =
        reinterpret_cast<TfLiteIntArray*>(allocator->AllocateFromTail(
            TfLiteIntArrayGetSizeInBytes(channels), alignof(TfLiteIntArray)));
    if (quantization->zero_point == nullptr) {
      TF_LITE_REPORT_ERROR(error_reporter,
                           "Unable to allocate quantization->zero_point.\n");
      return kTfLiteError;
    }

    TF_LITE_ENSURE_STATUS(FlatBufferVectorToTfLiteTypeArray(
        allocator, error_reporter, src_quantization->scale(),
        &quantization->scale));

    quantization->zero_point->size = channels;
    int* zero_point_data = quantization->zero_point->data;
    for (int i = 0; i < channels; i++) {
      zero_point_data[i] = src_quantization->zero_point()->Get(i);
    }
    // TODO(rocky): Need to add a micro_allocator test case that fails when
    // this is not copied:
    quantization->quantized_dimension = src_quantization->quantized_dimension();

    result->quantization = {kTfLiteAffineQuantization, quantization};
  }
  return kTfLiteOk;
}

}  // namespace internal

MicroAllocator::MicroAllocator(SimpleMemoryAllocator* memory_allocator,
                               ErrorReporter* error_reporter)
    : memory_allocator_(memory_allocator),
      error_reporter_(error_reporter),
      model_is_allocating_(false) {}

MicroAllocator::~MicroAllocator() {}

MicroAllocator* MicroAllocator::Create(uint8_t* tensor_arena, size_t arena_size,
                                       ErrorReporter* error_reporter) {
  uint8_t* aligned_arena = AlignPointerUp(tensor_arena, kBufferAlignment);
  if (aligned_arena != tensor_arena) {
    TF_LITE_REPORT_ERROR(
        error_reporter,
        "%d bytes lost due to alignment. To avoid this loss, please make sure "
        "the tensor_arena is 16 bytes aligned.",
        aligned_arena - tensor_arena);
  }
  size_t aligned_arena_size = tensor_arena + arena_size - aligned_arena;
  return Create(SimpleMemoryAllocator::Create(error_reporter, aligned_arena,
                                              aligned_arena_size),
                error_reporter);
}

MicroAllocator* MicroAllocator::Create(SimpleMemoryAllocator* memory_allocator,
                                       ErrorReporter* error_reporter) {
  TFLITE_DCHECK(memory_allocator != nullptr);
  TFLITE_DCHECK(error_reporter != nullptr);

  uint8_t* allocator_buffer = memory_allocator->AllocateFromTail(
      sizeof(MicroAllocator), alignof(MicroAllocator));
  MicroAllocator* allocator =
      new (allocator_buffer) MicroAllocator(memory_allocator, error_reporter);
  return allocator;
}

TfLiteStatus MicroAllocator::StartModelAllocation(
    const Model* model, TfLiteContext* context,
    const MicroOpResolver& op_resolver,
    NodeAndRegistration** node_and_registrations) {
  TFLITE_DCHECK(model != nullptr);
  TFLITE_DCHECK(context != nullptr);

  if (model_is_allocating_) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "MicroAllocator: Model allocation started before "
                         "finishing previously allocated model");
    return kTfLiteError;
  }

  const SubGraph* subgraph = GetSubGraphFromModel(model);
  TFLITE_DCHECK(subgraph != nullptr);
  model_is_allocating_ = true;

  TF_LITE_ENSURE_STATUS(
      InitGraphAndContextTensorData(model, context, subgraph));
  TF_LITE_ENSURE_STATUS(
      AllocateNodeAndRegistrations(subgraph, node_and_registrations));
  TF_LITE_ENSURE_STATUS(PrepareNodeAndRegistrationDataFromFlatbuffer(
      model, subgraph, op_resolver, *node_and_registrations));

  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::FinishModelAllocation(const Model* model,
                                                   TfLiteContext* context) {
  if (!model_is_allocating_) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "MicroAllocator: Model allocation finished before "
                         "starting allocating model");
    return kTfLiteError;
  }

  const SubGraph* subgraph = GetSubGraphFromModel(model);
  TFLITE_DCHECK(subgraph != nullptr);

  TF_LITE_ENSURE_STATUS(CommitStaticMemoryPlan(model, context, subgraph));
  TF_LITE_ENSURE_STATUS(AllocateVariables(context, subgraph));

  model_is_allocating_ = false;
  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::AllocatePersistentBuffer(size_t bytes,
                                                      void** ptr) {
  uint8_t* data = memory_allocator_->AllocateFromTail(bytes, kBufferAlignment);
  if (data == nullptr) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "Failed to allocate persistent buffer of size %d",
                         bytes);
    return kTfLiteError;
  }
  (*ptr) = data;
  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::RequestScratchBufferInArena(int node_id,
                                                         size_t bytes,
                                                         int* buffer_idx) {
  // A sanity check to make sure scratch_buffer_handles_ is contiguous i.e.
  // scratch_buffer_handles_ is pointing to the last allocation from memory
  // allocator.
  if (scratch_buffer_handles_ != nullptr &&
      reinterpret_cast<uint8_t*>(scratch_buffer_handles_) !=
          memory_allocator_->GetTail()) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "Internal error: AllocateFromTail can not be called "
                         "between two RequestScratchBufferInArena calls.");
    return kTfLiteError;
  }

  internal::ScratchBufferHandle* handle =
      reinterpret_cast<internal::ScratchBufferHandle*>(
          memory_allocator_->AllocateFromTail(
              sizeof(internal::ScratchBufferHandle),
              alignof(internal::ScratchBufferHandle)));
  if (handle == nullptr) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "Failed to register scratch buffer handle for node %s",
                         node_id);
    return kTfLiteError;
  }
  *handle = {};
  handle->bytes = bytes;
  handle->node_idx = node_id;
  *buffer_idx = scratch_buffer_count_;
  scratch_buffer_count_ += 1;
  // scratch_buffer_handles_ is in reverse order. The following code ensures
  // that scratch_buffers[0] is pointing to the newly allocated handle.
  scratch_buffer_handles_ = handle;
  return kTfLiteOk;
}

void* MicroAllocator::GetScratchBuffer(int buffer_idx) const {
  if (static_cast<size_t>(buffer_idx) >= scratch_buffer_count_) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "Buffer %d not found. %d buffers available.",
                         buffer_idx, scratch_buffer_count_);
    return nullptr;
  }
  // scratch_buffer_handles_ is in reverse order.
  return scratch_buffer_handles_[scratch_buffer_count_ - buffer_idx - 1].data;
}

size_t MicroAllocator::used_bytes() const {
  return memory_allocator_->GetUsedBytes();
}

TfLiteStatus MicroAllocator::AllocateTfLiteTensorArray(
    TfLiteContext* context, const SubGraph* subgraph) {
  context->tensors_size = subgraph->tensors()->size();
  context->tensors =
      reinterpret_cast<TfLiteTensor*>(memory_allocator_->AllocateFromTail(
          sizeof(TfLiteTensor) * context->tensors_size, alignof(TfLiteTensor)));
  if (context->tensors == nullptr) {
    TF_LITE_REPORT_ERROR(
        error_reporter_,
        "Failed to allocate memory for context->tensors, %d bytes required",
        sizeof(TfLiteTensor) * context->tensors_size);
    return kTfLiteError;
  }
  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::PopulateTfLiteTensorArrayFromFlatbuffer(
    const Model* model, TfLiteContext* context, const SubGraph* subgraph) {
  // Initialize tensors in context_ using the flatbuffer for quantization data.
  for (size_t i = 0; i < subgraph->tensors()->size(); ++i) {
    TfLiteStatus status = internal::InitializeTfLiteTensorFromFlatbuffer(
        memory_allocator_, *subgraph->tensors()->Get(i), model->buffers(),
        error_reporter_, &context->tensors[i]);
    if (status != kTfLiteOk) {
      TF_LITE_REPORT_ERROR(error_reporter_, "Failed to initialize tensor %d",
                           i);
      return kTfLiteError;
    }
  }
  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::AllocateNodeAndRegistrations(
    const SubGraph* subgraph, NodeAndRegistration** node_and_registrations) {
  NodeAndRegistration* output = reinterpret_cast<NodeAndRegistration*>(
      memory_allocator_->AllocateFromTail(
          sizeof(NodeAndRegistration) * subgraph->operators()->size(),
          alignof(NodeAndRegistration)));
  if (output == nullptr) {
    TF_LITE_REPORT_ERROR(
        error_reporter_,
        "Failed to allocate memory for node_and_registrations.");
    return kTfLiteError;
  }
  *node_and_registrations = output;
  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
    const Model* model, const SubGraph* subgraph,
    const MicroOpResolver& op_resolver,
    NodeAndRegistration* node_and_registrations) {
  TfLiteStatus status = kTfLiteOk;
  auto* opcodes = model->operator_codes();
  MicroBuiltinDataAllocator builtin_data_allocator(memory_allocator_);
  for (size_t i = 0; i < subgraph->operators()->size(); ++i) {
    const auto* op = subgraph->operators()->Get(i);
    const size_t index = op->opcode_index();
    if (index >= opcodes->size()) {
      TF_LITE_REPORT_ERROR(error_reporter_,
                           "Missing registration for opcode_index %d\n", index);
      return kTfLiteError;
    }
    auto* opcode = (*opcodes)[index];
    status =
        GetRegistrationFromOpCode(opcode, op_resolver, error_reporter_,
                                  &(node_and_registrations[i].registration));
    if (status != kTfLiteOk) {
      TF_LITE_REPORT_ERROR(error_reporter_,
                           "Failed to get registration from op code %s\n ",
                           EnumNameBuiltinOperator(opcode->builtin_code()));
      return status;
    }
    const auto* registration = node_and_registrations[i].registration;
    if (registration == nullptr) {
      TF_LITE_REPORT_ERROR(error_reporter_, "Skipping op for opcode_index %d\n",
                           index);
      return kTfLiteError;
    }
    BuiltinOperator op_type =
        static_cast<BuiltinOperator>(registration->builtin_code);

    const char* custom_data = nullptr;
    size_t custom_data_size = 0;
    unsigned char* builtin_data = nullptr;

    if (op_type == BuiltinOperator_CUSTOM) {
      // Custom Ops may or may not have a non-null custom_options field.
      if (op->custom_options() != nullptr) {
        custom_data =
            reinterpret_cast<const char*>(op->custom_options()->data());
        custom_data_size = op->custom_options()->size();
      }
    } else {
      if (op->custom_options() != nullptr) {
        TF_LITE_REPORT_ERROR(
            error_reporter_,
            "Unsupported behavior: found builtin operator %s with custom "
            "options.\n",
            EnumNameBuiltinOperator(op_type));
        return kTfLiteError;
      }

      MicroOpResolver::BuiltinParseFunction parser =
          op_resolver.GetOpDataParser(op_type);
      if (parser == nullptr) {
        TF_LITE_REPORT_ERROR(error_reporter_, "Did not find a parser for %s",
                             EnumNameBuiltinOperator(op_type));

        return kTfLiteError;
      }
      TF_LITE_ENSURE_STATUS(parser(op, op_type, error_reporter_,
                                   &builtin_data_allocator,
                                   (void**)(&builtin_data)));
    }

    TfLiteIntArray* inputs_array;
    TF_LITE_ENSURE_STATUS(internal::FlatBufferVectorToTfLiteTypeArray(
        memory_allocator_, error_reporter_, op->inputs(), &inputs_array));

    TfLiteIntArray* outputs_array;
    TF_LITE_ENSURE_STATUS(internal::FlatBufferVectorToTfLiteTypeArray(
        memory_allocator_, error_reporter_, op->outputs(), &outputs_array));

    TfLiteNode* node = &(node_and_registrations[i].node);
    *node = {};
    node->inputs = inputs_array;
    node->outputs = outputs_array;
    node->builtin_data = reinterpret_cast<void*>(builtin_data);
    node->custom_initial_data = custom_data;
    node->custom_initial_data_size = custom_data_size;
  }

  return kTfLiteOk;
}

TfLiteStatus MicroAllocator::AllocateVariables(TfLiteContext* context,
                                               const SubGraph* subgraph) {
  for (size_t i = 0; i < context->tensors_size; ++i) {
    if (subgraph->tensors()->Get(i)->is_variable()) {
      context->tensors[i].data.data = memory_allocator_->AllocateFromTail(
          context->tensors[i].bytes, kBufferAlignment);
      // Allocation failure.
      if (context->tensors[i].data.data == nullptr) {
        TF_LITE_REPORT_ERROR(error_reporter_,
                             "Failed to allocate variable tensor of size %d",
                             context->tensors[i].bytes);
        return kTfLiteError;
      }
    }
    tflite::ResetVariableTensor(&(context->tensors[i]));
  }

  return kTfLiteOk;
}

ErrorReporter* MicroAllocator::error_reporter() const {
  return error_reporter_;
}

TfLiteStatus MicroAllocator::InitGraphAndContextTensorData(
    const Model* model, TfLiteContext* context, const SubGraph* subgraph) {
  TF_LITE_ENSURE_STATUS(AllocateTfLiteTensorArray(context, subgraph));
  TF_LITE_ENSURE_STATUS(
      PopulateTfLiteTensorArrayFromFlatbuffer(model, context, subgraph));
  return kTfLiteOk;
}

const SubGraph* MicroAllocator::GetSubGraphFromModel(const Model* model) {
  auto* subgraphs = model->subgraphs();
  if (subgraphs->size() != 1) {
    TF_LITE_REPORT_ERROR(error_reporter_,
                         "Only 1 subgraph is currently supported.\n");
    return nullptr;
  }
  return (*subgraphs)[0];
}

TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(const Model* model,
                                                    TfLiteContext* context,
                                                    const SubGraph* subgraph) {
  // Create static memory plan
  // 1. Calculate AllocationInfo to know the lifetime of each tensor/buffer.
  // 2. Add them into the planner (such as the GreedyMemoryPlanner).
  // 3. Static memory planning using the planner.
  // 4. Set tensor/buffer pointers based on the offsets from the previous step.
  // Note that AllocationInfo is only needed for creating the plan. It will be
  // thrown away when the child allocator (tmp_allocator) goes out of scope.
  {
    SimpleMemoryAllocator tmp_allocator(error_reporter_,
                                        memory_allocator_->GetHead(),
                                        memory_allocator_->GetTail());

    AllocationInfoBuilder builder(error_reporter_, &tmp_allocator);
    TF_LITE_ENSURE_STATUS(
        builder.Init(subgraph->tensors()->size(), scratch_buffer_count_));

    const int32_t* offline_planner_offsets = nullptr;
    TF_LITE_ENSURE_STATUS(
        builder.GetOfflinePlannedOffsets(model, &offline_planner_offsets));
    TF_LITE_ENSURE_STATUS(builder.AddTensors(subgraph, offline_planner_offsets,
                                             context->tensors));

    TF_LITE_ENSURE_STATUS(builder.AddScratchBuffers(scratch_buffer_handles_));
    const AllocationInfo* allocation_info = builder.Finish();

    // Remaining arena size that memory planner can use for calculating offsets.
    size_t remaining_arena_size = tmp_allocator.GetAvailableMemory();
    uint8_t* planner_arena =
        tmp_allocator.AllocateFromHead(remaining_arena_size, /*alignment=*/1);
    TF_LITE_ENSURE(error_reporter_, planner_arena != nullptr);
    GreedyMemoryPlanner planner(planner_arena, remaining_arena_size);
    TF_LITE_ENSURE_STATUS(
        CreatePlan(error_reporter_, &planner, allocation_info, builder.Size()));

    size_t actual_available_arena_size =
        memory_allocator_->GetAvailableMemory();
    // Make sure we have enough arena size.
    if (planner.GetMaximumMemorySize() > actual_available_arena_size) {
      TF_LITE_REPORT_ERROR(
          error_reporter_,
          "Arena size is too small for activation buffers. Needed %d but only "
          "%d was available.",
          planner.GetMaximumMemorySize(), actual_available_arena_size);
      return kTfLiteError;
    }

    // Commit the plan.
    TF_LITE_ENSURE_STATUS(CommitPlan(error_reporter_, &planner,
                                     memory_allocator_->GetHead(),
                                     allocation_info, builder.Size()));
    // Allocate the planned area, so the allocator knows it's used.
    uint8_t* allocated_tensor_memory =
        memory_allocator_->AllocateFromHead(planner.GetMaximumMemorySize(),
                                            /*alignment=*/1);
    TF_LITE_ENSURE(error_reporter_, allocated_tensor_memory != nullptr);
  }
  return kTfLiteOk;
}

}  // namespace tflite