// Copyright 2020 Google LLC
//
// 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
//
//      https://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 "bindings/tflite/interpreter.h"

#include "bindings/tflite/model.h"
#include "bindings/tflite/shim.h"
#include "bindings/tflite/tensor.h"
#include "iree/base/threading.h"
#include "iree/base/tracing.h"
#include "iree/hal/drivers/init.h"
#include "iree/modules/hal/hal_module.h"

//===----------------------------------------------------------------------===//
// HAL / driver support
//===----------------------------------------------------------------------===//

static iree_once_flag _TfLiteInterpreterRegisterDriverFlag =
    IREE_ONCE_FLAG_INIT;
static void _TfLiteInterpreterRegisterDrivers(void) {
  IREE_IGNORE_ERROR(iree_hal_register_all_available_drivers(
      iree_hal_driver_registry_default()));
}

// TODO(#3977): if already provided a HAL device in the options use that.
static iree_status_t _TfLiteInterpreterPrepareHAL(
    TfLiteInterpreter* interpreter) {
  iree_call_once(&_TfLiteInterpreterRegisterDriverFlag,
                 _TfLiteInterpreterRegisterDrivers);

  iree_hal_driver_registry_t* driver_registry =
      iree_hal_driver_registry_default();

  iree_hal_driver_info_t* driver_infos = NULL;
  iree_host_size_t driver_info_count = 0;
  IREE_RETURN_IF_ERROR(iree_hal_driver_registry_enumerate(
      driver_registry, interpreter->allocator, &driver_infos,
      &driver_info_count));

  // TODO(benvanik): figure out how we want to emulate device selection; may
  // just say "whatever is first" on a query.
  // iree_string_view_t driver_name = driver_infos[0].driver_name;
  // NOTE: currently the sample file is compiled only with vmla.
  iree_string_view_t driver_name = iree_make_cstring_view("vmla");

  // TODO(benvanik): switch to iree_hal_driver_registry_try_create when
  // implemented.
  iree_status_t status = iree_hal_driver_registry_try_create_by_name(
      driver_registry, driver_name, interpreter->allocator,
      &interpreter->driver);
  iree_allocator_free(interpreter->allocator, driver_infos);
  IREE_RETURN_IF_ERROR(status, "failed to create driver '%.*s'",
                       (int)driver_name.size, driver_name.data);

  IREE_RETURN_IF_ERROR(
      iree_hal_driver_create_default_device(
          interpreter->driver, interpreter->allocator, &interpreter->device),
      "failed creating the default device for driver '%.*s'",
      (int)driver_name.size, driver_name.data);

  IREE_RETURN_IF_ERROR(iree_hal_module_create(
      interpreter->device, interpreter->allocator, &interpreter->hal_module));

  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Model shape function query/mutation utilities
//===----------------------------------------------------------------------===//

// On-stack storage for shape function invocations.
// Avoids all allocations and allows for reuse when running down lists of
// inputs and outputs calling shape functions.
typedef struct {
  // Inlined list for the !vm.list in the shape function arguments.
  uint8_t
      shape_list_storage[128 + sizeof(int32_t) * IREE_BINDINGS_TFLITE_MAX_RANK];
  iree_vm_list_t* shape_list;

  // Inlined list for the shape function arguments.
  uint8_t arg_list_storage[128 + sizeof(uintptr_t) * 2];
  iree_vm_list_t* arg_list;
} _TfLiteInterpreterShapeFrame;

// Initializes an on-stack shape frame. Existing contents are discarded.
static iree_status_t _TfLiteInterpreterShapeFrameInitialize(
    _TfLiteInterpreterShapeFrame* frame) {
  // [int32...] storage for the shape dimension inputs/outputs.
  iree_vm_type_def_t dim_type =
      iree_vm_type_def_make_value_type(IREE_VM_VALUE_TYPE_I32);
  IREE_RETURN_IF_ERROR(iree_vm_list_initialize(
      iree_make_byte_span(frame->shape_list_storage,
                          IREE_ARRAYSIZE(frame->shape_list_storage)),
      &dim_type, IREE_BINDINGS_TFLITE_MAX_RANK, &frame->shape_list));

  // (%index : i32, %shape : !vm.list<i32>)
  IREE_RETURN_IF_ERROR(iree_vm_list_initialize(
      iree_make_byte_span(frame->arg_list_storage,
                          IREE_ARRAYSIZE(frame->arg_list_storage)),
      /*element_type=*/NULL, /*index*/ 1 + /*shape*/ 1, &frame->arg_list));
  IREE_RETURN_IF_ERROR(iree_vm_list_resize(frame->arg_list, 2));

  // Arg 1 is always the shape list for all I/O, so do that once here.
  iree_vm_ref_t shape_list_ref = {0};
  IREE_RETURN_IF_ERROR(iree_vm_ref_wrap_assign(
      frame->shape_list, iree_vm_list_type_id(), &shape_list_ref));
  IREE_RETURN_IF_ERROR(
      iree_vm_list_set_ref_retain(frame->arg_list, 1, &shape_list_ref));

  return iree_ok_status();
}

// Deinitializes an on-stack shape frame.
// Though this does not free the frame memory (it's on the stack, afterall) it
// will release any resources that may be retained and is required.
static void _TfLiteInterpreterShapeFrameDeinitialize(
    _TfLiteInterpreterShapeFrame* frame) {
  iree_vm_list_deinitialize(frame->shape_list);
  iree_vm_list_deinitialize(frame->arg_list);
}

// Reads the shape value in the frame storage from the prior application.
static iree_status_t _TfLiteInterpreterShapeFrameReadValue(
    _TfLiteInterpreterShapeFrame* frame, int32_t* out_shape_rank,
    int32_t* out_shape_dims) {
  *out_shape_rank = (int32_t)iree_vm_list_size(frame->shape_list);
  for (int32_t i = 0; i < *out_shape_rank; ++i) {
    iree_vm_value_t dim;
    IREE_RETURN_IF_ERROR(iree_vm_list_get_value_as(
        frame->shape_list, i, IREE_VM_VALUE_TYPE_I32, &dim));
    out_shape_dims[i] = dim.i32;
  }
  return iree_ok_status();
}

// Writes the shape value to the current frame storage for future applications.
static iree_status_t _TfLiteInterpreterShapeFrameWriteValue(
    _TfLiteInterpreterShapeFrame* frame, int32_t shape_rank,
    const int32_t* shape_dims) {
  IREE_RETURN_IF_ERROR(iree_vm_list_resize(frame->shape_list, shape_rank));
  for (int32_t i = 0; i < shape_rank; ++i) {
    iree_vm_value_t dim = iree_vm_value_make_i32(shape_dims[i]);
    IREE_RETURN_IF_ERROR(iree_vm_list_set_value(frame->shape_list, i, &dim));
  }
  return iree_ok_status();
}

// Calls the |apply_fn| with the current shape frame state.
static iree_status_t _TfLiteInterpreterShapeFrameApply(
    _TfLiteInterpreterShapeFrame* frame, TfLiteInterpreter* interpreter,
    iree_vm_function_t apply_fn, int32_t index) {
  // Populate shape_list with the shape dimensions for this particular output.
  iree_vm_value_t index_value = iree_vm_value_make_i32(index);
  IREE_IGNORE_ERROR(iree_vm_list_set_value(frame->arg_list, 0, &index_value));
  return iree_vm_invoke(interpreter->context, apply_fn,
                        /*policy=*/NULL, frame->arg_list, /*outputs=*/NULL,
                        interpreter->allocator);
}

//===----------------------------------------------------------------------===//
// Shape I/O queries
//===----------------------------------------------------------------------===//

// Queries all input shapes from the module; some may still be dynamic (-1).
static iree_status_t _TfLiteInterpreterRefreshInputShapes(
    TfLiteInterpreter* interpreter, _TfLiteInterpreterShapeFrame* frame) {
  // NOTE: we could optimize this more by using iree_vm_invoke_within, but that
  // shouldn't be needed (it's just stack pointer manipulation).
  for (int32_t i = 0; i < interpreter->model->input_count; ++i) {
    TfLiteTensor* tensor = &interpreter->input_tensors[i];
    IREE_RETURN_IF_ERROR(_TfLiteInterpreterShapeFrameApply(
        frame, interpreter, interpreter->model->exports._query_input_shape, i));
    IREE_RETURN_IF_ERROR(_TfLiteInterpreterShapeFrameReadValue(
        frame, &tensor->shape_rank, tensor->shape_dims));
  }
  return iree_ok_status();
}

// Queries all output shapes from the module allowing it use the current input
// shapes to compute the possibly dynamic values.
static iree_status_t _TfLiteInterpreterRefreshOutputShapes(
    TfLiteInterpreter* interpreter, _TfLiteInterpreterShapeFrame* frame) {
  // NOTE: we could optimize this more by using iree_vm_invoke_within, but that
  // shouldn't be needed (it's just stack pointer manipulation).
  for (int32_t i = 0; i < interpreter->model->output_count; ++i) {
    TfLiteTensor* tensor = &interpreter->output_tensors[i];
    IREE_RETURN_IF_ERROR(_TfLiteInterpreterShapeFrameApply(
        frame, interpreter, interpreter->model->exports._query_output_shape,
        i));
    IREE_RETURN_IF_ERROR(_TfLiteInterpreterShapeFrameReadValue(
        frame, &tensor->shape_rank, tensor->shape_dims));
  }
  return iree_ok_status();
}

// Refreshes both input and output tensor shapes by querying the module.
// This should be called after each shape change so that we can let the module
// run "shape propagation" and compute the new output shapes.
static iree_status_t _TfLiteInterpreterRefreshIOShapes(
    TfLiteInterpreter* interpreter) {
  IREE_TRACE_ZONE_BEGIN(z0);
  _TfLiteInterpreterShapeFrame frame;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, _TfLiteInterpreterShapeFrameInitialize(&frame));

  // Query all shapes.
  iree_status_t status = iree_ok_status();
  if (iree_status_is_ok(status)) {
    status = _TfLiteInterpreterRefreshInputShapes(interpreter, &frame);
  }
  if (iree_status_is_ok(status)) {
    status = _TfLiteInterpreterRefreshOutputShapes(interpreter, &frame);
  }

  _TfLiteInterpreterShapeFrameDeinitialize(&frame);
  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// Creation and static initialization
//===----------------------------------------------------------------------===//

// Computes the storage requirement for the TfLiteInterpreter struct.
static iree_host_size_t _TfLiteInterpreterCalculateSize(
    const TfLiteModel* model) {
  iree_host_size_t total_size =
      iree_math_align(sizeof(TfLiteInterpreter), iree_max_align_t);

  iree_vm_type_def_t buffer_view_type_def =
      iree_vm_type_def_make_ref_type(iree_hal_buffer_type_id());
  total_size +=
      iree_vm_list_storage_size(&buffer_view_type_def, model->input_count);
  total_size +=
      iree_vm_list_storage_size(&buffer_view_type_def, model->output_count);
  total_size += sizeof(TfLiteTensor) * model->input_count;
  total_size += sizeof(TfLiteTensor) * model->output_count;

  return total_size;
}

// Allocates the interpreter slab and populates all internal pointers to the
// appropriate offsets.
static iree_status_t _TfLiteInterpreterAllocate(
    const TfLiteModel* model, TfLiteInterpreter** out_interpreter) {
  iree_host_size_t interpreter_size = _TfLiteInterpreterCalculateSize(model);
  TfLiteInterpreter* interpreter = NULL;
  IREE_RETURN_IF_ERROR(iree_allocator_malloc(model->allocator, interpreter_size,
                                             (void**)&interpreter));
  memset(interpreter, 0, interpreter_size);
  interpreter->allocator = model->allocator;
  _TfLiteInterpreterOptionsSetDefaults(&interpreter->options);
  *out_interpreter = interpreter;

  interpreter->model = (TfLiteModel*)model;
  _TfLiteModelRetain(interpreter->model);

  uint8_t* p = (uint8_t*)interpreter +
               iree_math_align(sizeof(*interpreter), iree_max_align_t);

  iree_vm_type_def_t buffer_view_type_def =
      iree_vm_type_def_make_ref_type(iree_hal_buffer_type_id());

  iree_byte_span_t input_list_storage = iree_make_byte_span(
      p, iree_vm_list_storage_size(&buffer_view_type_def, model->input_count));
  IREE_RETURN_IF_ERROR(
      iree_vm_list_initialize(input_list_storage, &buffer_view_type_def,
                              model->input_count, &interpreter->input_list));
  p += input_list_storage.data_length;

  iree_byte_span_t output_list_storage = iree_make_byte_span(
      p, iree_vm_list_storage_size(&buffer_view_type_def, model->output_count));
  IREE_RETURN_IF_ERROR(
      iree_vm_list_initialize(output_list_storage, &buffer_view_type_def,
                              model->output_count, &interpreter->output_list));
  p += output_list_storage.data_length;

  interpreter->input_tensors = (TfLiteTensor*)p;
  p += sizeof(TfLiteTensor) * model->input_count;
  interpreter->output_tensors = (TfLiteTensor*)p;
  // p += sizeof(TfLiteTensor) * model->output_count;

  return iree_ok_status();
}

// Populates the input and output tensor lists with static metadata from the
// model and prepares for allocation/invocation.
static iree_status_t _TfLiteInterpreterPopulateIO(
    TfLiteInterpreter* interpreter) {
  iree_vm_function_t main_fn = interpreter->model->exports._main;
  iree_string_view_t io_names_attr = iree_vm_function_reflection_attr(
      &main_fn, iree_make_cstring_view("tfl.io.names"));
  iree_string_view_t io_types_attr = iree_vm_function_reflection_attr(
      &main_fn, iree_make_cstring_view("tfl.io.types"));
  iree_string_view_t io_quant_attr = iree_vm_function_reflection_attr(
      &main_fn, iree_make_cstring_view("tfl.io.quant"));

  // Setup static tensor metadata.
  for (iree_host_size_t i = 0; i < interpreter->model->input_count; ++i) {
    TfLiteTensor* tensor = &interpreter->input_tensors[i];
    memset(tensor, 0, sizeof(*tensor));
    iree_string_view_t io_name_part = iree_string_view_empty();
    iree_string_view_split(io_names_attr, ';', &io_name_part, &io_names_attr);
    iree_string_view_t io_type_part = iree_string_view_empty();
    iree_string_view_split(io_types_attr, ';', &io_type_part, &io_types_attr);
    iree_string_view_t io_quant_part = iree_string_view_empty();
    iree_string_view_split(io_quant_attr, ';', &io_quant_part, &io_quant_attr);
    IREE_RETURN_IF_ERROR(_TfLiteTensorParseNameAttr(tensor, io_name_part,
                                                    interpreter->allocator));
    IREE_RETURN_IF_ERROR(_TfLiteTensorParseTypeAttr(tensor, io_type_part));
    IREE_RETURN_IF_ERROR(_TfLiteTensorParseQuantAttr(tensor, io_quant_part));
  }
  for (iree_host_size_t i = 0; i < interpreter->model->output_count; ++i) {
    TfLiteTensor* tensor = &interpreter->output_tensors[i];
    memset(tensor, 0, sizeof(*tensor));
    iree_string_view_t io_name_part = iree_string_view_empty();
    iree_string_view_split(io_names_attr, ';', &io_name_part, &io_names_attr);
    iree_string_view_t io_type_part = iree_string_view_empty();
    iree_string_view_split(io_types_attr, ';', &io_type_part, &io_types_attr);
    iree_string_view_t io_quant_part = iree_string_view_empty();
    iree_string_view_split(io_quant_attr, ';', &io_quant_part, &io_quant_attr);
    IREE_RETURN_IF_ERROR(_TfLiteTensorParseNameAttr(tensor, io_name_part,
                                                    interpreter->allocator));
    IREE_RETURN_IF_ERROR(_TfLiteTensorParseTypeAttr(tensor, io_type_part));
    IREE_RETURN_IF_ERROR(_TfLiteTensorParseQuantAttr(tensor, io_quant_part));
  }

  // Prepare the IO lists we use when calling into the model.
  // The actual contents of these cannot be set until
  // TfLiteInterpreterAllocateTensors has been called.
  IREE_RETURN_IF_ERROR(iree_vm_list_reserve(interpreter->input_list,
                                            interpreter->model->input_count));
  IREE_RETURN_IF_ERROR(iree_vm_list_reserve(interpreter->output_list,
                                            interpreter->model->output_count));

  return iree_ok_status();
}

static iree_status_t _TfLiteInterpreterCreate(
    const TfLiteModel* model, const TfLiteInterpreterOptions* optional_options,
    TfLiteInterpreter** out_interpreter) {
  *out_interpreter = NULL;

  // We allocate a large majority of the interpreter structures as a single
  // slab. There's still some allocations that we could prevent (like internal
  // VM stuff) but this at least covers half of it.
  IREE_RETURN_IF_ERROR(_TfLiteInterpreterAllocate(model, out_interpreter));
  TfLiteInterpreter* interpreter = *out_interpreter;

  if (optional_options) {
    memcpy(&interpreter->options, optional_options,
           sizeof(interpreter->options));
  }

  interpreter->user_module = model->module;
  iree_vm_module_retain(interpreter->user_module);

  // External contexts could possibly used to emulate sharing this, but really
  // if a user is running with multiple models the tflite API is insufficient.
  IREE_RETURN_IF_ERROR(
      iree_vm_instance_create(interpreter->allocator, &interpreter->instance));
  IREE_RETURN_IF_ERROR(_TfLiteInterpreterPrepareHAL(interpreter));

  // Context will contain both the user-provided bytecode and the HAL module.
  // If we were to support custom ops we would also have a
  // tflite_resolver_module that we would register to resolve tflite ops into
  // IREE functions that will call custom ops through TfLiteRegistrations.
  IREE_RETURN_IF_ERROR(iree_vm_context_create_with_modules(
      interpreter->instance, interpreter->all_modules,
      IREE_ARRAYSIZE(interpreter->all_modules), interpreter->allocator,
      &interpreter->context));

  // Setup all I/O tensors and buffer views.
  IREE_RETURN_IF_ERROR(_TfLiteInterpreterPopulateIO(interpreter));

  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Core API
//===----------------------------------------------------------------------===//

TFL_CAPI_EXPORT extern TfLiteInterpreter* TfLiteInterpreterCreate(
    const TfLiteModel* model,
    const TfLiteInterpreterOptions* optional_options) {
  IREE_TRACE_ZONE_BEGIN(z0);
  TfLiteInterpreter* interpreter = NULL;
  iree_status_t status =
      _TfLiteInterpreterCreate(model, optional_options, &interpreter);
  if (iree_status_is_ok(iree_status_consume_code(status))) {
    IREE_TRACE_ZONE_APPEND_TEXT(z0, "num_threads=", strlen("num_threads="));
    IREE_TRACE_ZONE_APPEND_VALUE(z0, interpreter->options.num_threads);
  } else {
    IREE_TRACE_MESSAGE(ERROR, "failed interpreter creation");
    TfLiteInterpreterDelete(interpreter);
    interpreter = NULL;
  }
  IREE_TRACE_ZONE_END(z0);
  return interpreter;
}

TFL_CAPI_EXPORT extern TfLiteInterpreter*
TfLiteInterpreterCreateWithSelectedOps(
    const TfLiteModel* model, const TfLiteInterpreterOptions* options) {
  // No different from TfLiteInterpreterCreate: we don't have "ops" :)
  return TfLiteInterpreterCreate(model, options);
}

TFL_CAPI_EXPORT extern void TfLiteInterpreterDelete(
    TfLiteInterpreter* interpreter) {
  IREE_TRACE_ZONE_BEGIN(z0);

  for (iree_host_size_t i = 0; i < interpreter->model->input_count; ++i) {
    _TfLiteTensorReset(&interpreter->input_tensors[i], interpreter->allocator);
  }
  for (iree_host_size_t i = 0; i < interpreter->model->output_count; ++i) {
    _TfLiteTensorReset(&interpreter->output_tensors[i], interpreter->allocator);
  }
  iree_vm_list_deinitialize(interpreter->input_list);
  iree_vm_list_deinitialize(interpreter->output_list);

  iree_vm_context_release(interpreter->context);
  iree_vm_module_release(interpreter->hal_module);
  iree_vm_module_release(interpreter->user_module);
  iree_vm_instance_release(interpreter->instance);

  _TfLiteModelRelease(interpreter->model);
  iree_allocator_free(interpreter->allocator, interpreter);

  IREE_TRACE_ZONE_END(z0);
}

TFL_CAPI_EXPORT extern TfLiteStatus TfLiteInterpreterResetVariableTensors(
    TfLiteInterpreter* interpreter) {
  IREE_TRACE_ZONE_BEGIN(z0);

  // The compiler emits a special function we can use to reset just variables.
  // NOTE: the function is optional if the model had no variables.
  iree_status_t status = iree_ok_status();
  iree_vm_function_t reset_variables_fn =
      interpreter->model->exports._reset_variables;
  if (!iree_vm_function_is_null(reset_variables_fn)) {
    status = iree_vm_invoke(interpreter->context, reset_variables_fn,
                            /*policy=*/NULL, /*inputs=*/NULL, /*outputs=*/NULL,
                            interpreter->allocator);
  }

  IREE_TRACE_ZONE_END(z0);
  return _TfLiteStatusFromIREEStatus(status);
}

TFL_CAPI_EXPORT extern int32_t TfLiteInterpreterGetInputTensorCount(
    const TfLiteInterpreter* interpreter) {
  return interpreter->model->input_count;
}

TFL_CAPI_EXPORT extern TfLiteTensor* TfLiteInterpreterGetInputTensor(
    const TfLiteInterpreter* interpreter, int32_t input_index) {
  if (input_index < 0 || input_index >= interpreter->model->input_count) {
    return NULL;
  }
  return &interpreter->input_tensors[input_index];
}

static iree_status_t _TfLiteInterpreterResizeInputTensor(
    TfLiteInterpreter* interpreter, int32_t input_index, const int* input_dims,
    int32_t input_dims_size) {
  if (input_index < 0 || input_index >= interpreter->model->input_count) {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "input_index out of range (0 <= %d < %d)",
                            input_index, interpreter->model->input_count);
  }
  if (iree_vm_function_is_null(
          interpreter->model->exports._resize_input_shape)) {
    // TODO(#3975): check if this is a no-op success in tflite.
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "model has no dynamic shapes");
  }

  _TfLiteInterpreterShapeFrame frame;
  IREE_RETURN_IF_ERROR(_TfLiteInterpreterShapeFrameInitialize(&frame));

  // Poke the model and let it update its internal shape.
  // TODO(#3975): return bool to allow model to say it failed.
  TfLiteTensor* tensor = &interpreter->input_tensors[input_index];
  iree_status_t status = _TfLiteInterpreterShapeFrameWriteValue(
      &frame, tensor->shape_rank, tensor->shape_dims);
  if (iree_status_is_ok(status)) {
    status = _TfLiteInterpreterShapeFrameApply(
        &frame, interpreter, interpreter->model->exports._resize_input_shape,
        input_index);
  }

  // NOTE: the allocation may now not match the requested shape. This is just
  // how the tflite API works unfortunately; until
  // TfLiteInterpreterAllocateTensors it will remain in an indeterminate state.

  _TfLiteInterpreterShapeFrameDeinitialize(&frame);
  return status;
}

TFL_CAPI_EXPORT extern TfLiteStatus TfLiteInterpreterResizeInputTensor(
    TfLiteInterpreter* interpreter, int32_t input_index, const int* input_dims,
    int32_t input_dims_size) {
  IREE_TRACE_ZONE_BEGIN(z0);
  iree_status_t status = _TfLiteInterpreterResizeInputTensor(
      interpreter, input_index, input_dims, input_dims_size);
  IREE_TRACE_ZONE_END(z0);
  return _TfLiteStatusFromIREEStatus(status);
}

static iree_status_t _TfLiteInterpreterAllocateTensors(
    TfLiteInterpreter* interpreter) {
  // NOTE: we could slab allocate like tflite does, but then if any single
  // tensor has any single dimension that is resized the whole thing gets
  // reallocated upon resize. That's no good. Instead, we realloc each tensor
  // if their size has changed.

  // Refresh all shapes from the model. It should have all of the
  // non-data-dependent output shapes.
  IREE_RETURN_IF_ERROR(_TfLiteInterpreterRefreshIOShapes(interpreter));

  // Drop all input tensors we hang on to in the input list. This way we aren't
  // double-allocating during the resize.
  IREE_RETURN_IF_ERROR(iree_vm_list_resize(interpreter->input_list, 0));

  // Reallocate input tensors (if needed).
  for (iree_host_size_t i = 0; i < interpreter->model->input_count; ++i) {
    TfLiteTensor* tensor = &interpreter->input_tensors[i];
    IREE_RETURN_IF_ERROR(_TfLiteTensorReallocateIfNeeded(
        tensor, iree_hal_device_allocator(interpreter->device),
        interpreter->allocator));
    iree_vm_ref_t buffer_ref = iree_hal_buffer_retain_ref(tensor->buffer);
    IREE_RETURN_IF_ERROR(
        iree_vm_list_push_ref_move(interpreter->input_list, &buffer_ref));
  }

  // TODO(benvanik): preallocate outputs when we support using them.
  // We could stash the buffer views in interpreter->output_list.
  // For now we just drop them all.
  for (iree_host_size_t i = 0; i < interpreter->model->output_count; ++i) {
    _TfLiteTensorDiscardBuffer(&interpreter->output_tensors[i]);
  }

  return iree_ok_status();
}

TFL_CAPI_EXPORT extern TfLiteStatus TfLiteInterpreterAllocateTensors(
    TfLiteInterpreter* interpreter) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_status_t status = _TfLiteInterpreterAllocateTensors(interpreter);

#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
  iree_device_size_t total_input_size = 0;
  for (iree_host_size_t i = 0; i < interpreter->model->input_count; ++i) {
    total_input_size +=
        iree_hal_buffer_byte_length(interpreter->input_tensors[i].buffer);
  }
  IREE_TRACE_ZONE_APPEND_VALUE(z0, total_input_size);
#endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION

  IREE_TRACE_ZONE_END(z0);
  return _TfLiteStatusFromIREEStatus(status);
}

static iree_status_t _TfLiteInterpreterInvoke(TfLiteInterpreter* interpreter) {
  // tflite models only have a single entry point and the IREE converter
  // emits it as '_main'.
  IREE_RETURN_IF_ERROR(
      iree_vm_invoke(interpreter->context, interpreter->model->exports._main,
                     /*policy=*/NULL, interpreter->input_list,
                     interpreter->output_list, interpreter->allocator));

  // Refresh output shapes.
  // TODO(#3975): just use buffer view results or at least just refresh outputs.
  IREE_RETURN_IF_ERROR(_TfLiteInterpreterRefreshIOShapes(interpreter));

  // Map the output buffers.
  // NOTE: we could defer the mapping unless requested and ensure state buffers
  // remain where they currently are for the next invocation.
  for (iree_host_size_t i = 0; i < interpreter->model->output_count; ++i) {
    iree_hal_buffer_t* buffer = (iree_hal_buffer_t*)iree_vm_list_get_ref_deref(
        interpreter->output_list, i, iree_hal_buffer_get_descriptor());
    TfLiteTensor* tensor = &interpreter->output_tensors[i];
    IREE_RETURN_IF_ERROR(_TfLiteTensorBind(tensor, buffer));
  }

  return iree_ok_status();
}

TFL_CAPI_EXPORT extern TfLiteStatus TfLiteInterpreterInvoke(
    TfLiteInterpreter* interpreter) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_status_t status = _TfLiteInterpreterInvoke(interpreter);

#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
  iree_device_size_t total_output_size = 0;
  for (iree_host_size_t i = 0; i < interpreter->model->output_count; ++i) {
    total_output_size +=
        iree_hal_buffer_byte_length(interpreter->output_tensors[i].buffer);
  }
  IREE_TRACE_ZONE_APPEND_VALUE(z0, total_output_size);
#endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION

  IREE_TRACE_ZONE_END(z0);
  return _TfLiteStatusFromIREEStatus(status);
}

TFL_CAPI_EXPORT extern int32_t TfLiteInterpreterGetOutputTensorCount(
    const TfLiteInterpreter* interpreter) {
  return interpreter->model->output_count;
}

TFL_CAPI_EXPORT extern const TfLiteTensor* TfLiteInterpreterGetOutputTensor(
    const TfLiteInterpreter* interpreter, int32_t output_index) {
  if (output_index < 0 || output_index >= interpreter->model->output_count) {
    return NULL;
  }
  return &interpreter->output_tensors[output_index];
}