/* Copyright 2018 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/tools/optimize/subgraph_quantizer.h"

#include <algorithm>
#include <limits>

#include "flatbuffers/flexbuffers.h"
#include "absl/memory/memory.h"
#include "tensorflow/lite/context.h"
#include "tensorflow/lite/core/api/error_reporter.h"
#include "tensorflow/lite/kernels/internal/round.h"
#include "tensorflow/lite/kernels/internal/tensor_utils.h"
#include "tensorflow/lite/kernels/internal/types.h"
#include "tensorflow/lite/model.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/tools/optimize/quantization_utils.h"

namespace tflite {
namespace optimize {
namespace internal {

namespace {
TfLiteStatus AddQuantizationParams(const std::vector<float>& scales,
                                   const std::vector<int64_t>& zero_point,
                                   int quantized_dimension,
                                   const uint8_t* buffer_data,
                                   size_t buffer_size, TensorType output_type,
                                   ModelT* model, TensorT* tensor) {
  tensor->quantization = absl::make_unique<QuantizationParametersT>();
  tensor->quantization->scale.assign(scales.begin(), scales.end());
  if (zero_point.size() != scales.size()) {
    return kTfLiteError;
  }
  tensor->quantization->zero_point.assign(zero_point.begin(), zero_point.end());
  tensor->quantization->quantized_dimension = quantized_dimension;
  model->buffers[tensor->buffer]->data.assign(buffer_data,
                                              buffer_data + buffer_size);
  // Update the tensor type.
  tensor->type = output_type;
  return kTfLiteOk;
}

bool OpHasOptionalBiasTensor(BuiltinOperator op_code) {
  return op_code == BuiltinOperator_CONV_2D ||
         op_code == BuiltinOperator_DEPTHWISE_CONV_2D;
}

struct OpWithBiasTensors {
  int activation_input_index;
  int weights_input_index;
  int bias_input_index;
  int index_for_channel_in_weights;
};

const OpWithBiasTensors* GetInfoForOpWithBiasTensor(BuiltinOperator op_code) {
  if (op_code == BuiltinOperator_CONV_2D) {
    static OpWithBiasTensors op_info = {/* activation_input_index */ 0,
                                        /* weights_input_index */ 1,
                                        /* bias_input_index */ 2,
                                        /* index_for_channel_in_weights */ 0};
    return &op_info;
  }
  if (op_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
    static OpWithBiasTensors op_info = {/* bias_input_index */ 0,
                                        /* bias_input_index */ 1,
                                        /* bias_input_index */ 2,
                                        /* index_for_channel_in_weights */ 3};
    return &op_info;
  }

  return nullptr;
}

// Symmetrically Quantizes the given tensor as int8 values.
TfLiteStatus SymmetricPerChannelQuantizeTensor(ModelT* model, TensorT* tensor,
                                               int32_t channel_dim_index,
                                               ErrorReporter* error_reporter) {
  if (tensor->shape.size() != 4) {
    error_reporter->Report("Only dims=4 is supported, tensor dims: %d",
                           tensor->shape.size());
    return kTfLiteError;
  }

  // Get dimensions.
  uint64_t num_elements;
  TF_LITE_ENSURE_STATUS(utils::NumElements(*tensor, &num_elements));
  const int32_t channel_dim_size = tensor->shape[channel_dim_index];

  // Get input float data.
  BufferT* buffer = model->buffers[tensor->buffer].get();
  float* float_input_data = reinterpret_cast<float*>(buffer->data.data());

  // Create container for output scale and output data.
  std::vector<float> scales(channel_dim_size);
  std::vector<int8_t> final_buffer(num_elements);

  // Quantize the input data with respect to channel_dim_index.
  const std::vector<int> tensor_dims = {tensor->shape[0], tensor->shape[1],
                                        tensor->shape[2], tensor->shape[3]};
  utils::SymmetricPerChannelQuantization(
      float_input_data, tensor_dims, channel_dim_index, &scales, &final_buffer);

  // Set the buffers and output type.
  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
  const size_t buffer_size = num_elements * sizeof(int8_t);
  std::vector<int64_t> zero_point(scales.size(), 0);
  return AddQuantizationParams(scales, zero_point, channel_dim_index,
                               uint8_buffer, buffer_size, TensorType_INT8,
                               model, tensor);
}

// Symmetrically quantizes the bias for ops like Conv and DepthwiseConv.
// The scale of bias if weight_per_channel_scale[channel] * input_scale
TfLiteStatus SymmetricPerChannelBiasQuantize(const TensorT* input_tensor,
                                             const TensorT* weight_tensor,
                                             int channel_dim_index,
                                             ModelT* model, TensorT* tensor,
                                             ErrorReporter* error_reporter) {
  if (tensor->shape.size() != 1) {
    error_reporter->Report("Expected bias tensor shape to be 1.");
    return kTfLiteError;
  }

  if (tensor->type != TensorType_FLOAT32) {
    return kTfLiteOk;
  }

  // TODO(shashishekhar): Make this support scalar biases.
  if (tensor->shape[0] != weight_tensor->shape[channel_dim_index]) {
    error_reporter->Report(
        "Channel mismatch between bias and weight tensors %d vs %d",
        tensor->shape[0], weight_tensor->shape[channel_dim_index]);
    return kTfLiteError;
  }
  int32_t channel_dim_size = tensor->shape[0];
  if (!input_tensor->quantization ||
      input_tensor->quantization->scale.size() != 1) {
    error_reporter->Report("Input tensor missing quantization information");
    return kTfLiteError;
  }
  TF_LITE_ENSURE(error_reporter, weight_tensor->quantization);
  const std::vector<float>& weight_scales = weight_tensor->quantization->scale;

  if (weight_scales.size() != channel_dim_size) {
    error_reporter->Report("Mismatch weight scale dimension: %d",
                           weight_scales.size());
    return kTfLiteError;
  }

  // Compute scales.
  std::vector<float> scales(channel_dim_size);
  for (size_t i = 0; i < channel_dim_size; i++) {
    scales[i] = input_tensor->quantization->scale[0] * weight_scales[i];
  }

  BufferT* buffer = model->buffers[tensor->buffer].get();
  float* float_data = reinterpret_cast<float*>(buffer->data.data());
  uint64_t num_elements;
  TF_LITE_ENSURE_STATUS(utils::NumElements(*tensor, &num_elements));

  std::vector<int32_t> final_buffer(num_elements);
  const int32_t kScale = std::numeric_limits<int32_t>::max();

  for (int32_t channel_idx = 0; channel_idx < channel_dim_size; channel_idx++) {
    float scaling_factor = scales[channel_idx];
    float scaling_factor_inv = (scaling_factor == 0) ? 0 : 1.0 / scaling_factor;
    const int32_t quantized_value = static_cast<int32_t>(
        TfLiteRound(float_data[channel_idx] * scaling_factor_inv));
    final_buffer[channel_idx] =
        std::min(kScale, std::max(-kScale, quantized_value));
  }

  // Set the buffers and output type.
  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
  size_t buffer_size = num_elements * sizeof(int32_t);
  std::vector<int64_t> zero_point(scales.size(), 0);
  return AddQuantizationParams(scales, zero_point, channel_dim_index,
                               uint8_buffer, buffer_size, TensorType_INT32,
                               model, tensor);
}
}  // namespace

TfLiteStatus SubgraphQuantizer::AsymmetricQuantizeTensor(
    BuiltinOperator op_code, int32_t tensor_idx) {
  TensorT* tensor = subgraph_->tensors[tensor_idx].get();
  if (tensor->type != TensorType_FLOAT32) {
    return kTfLiteOk;
  }

  if (model_->buffers[tensor->buffer]->data.data() != nullptr) {
    return kTfLiteError;
  }
  if (!tensor->quantization || tensor->quantization->min.empty() ||
      tensor->quantization->max.empty()) {
    error_reporter_->Report(
        "Missing required min/max information for tensor_idx %d of operation: "
        "%s",
        tensor_idx, EnumNameBuiltinOperator(op_code));
    return kTfLiteError;
  }
  utils::GetAsymmetricQuantizationParams(
      tensor->quantization->min[0], tensor->quantization->max[0],
      std::numeric_limits<int8_t>::min(), std::numeric_limits<int8_t>::max(),
      tensor->quantization.get());
  tensor->type = TensorType_INT8;
  return kTfLiteOk;
}

TfLiteStatus SubgraphQuantizer::QuantizeOpWithBias(BuiltinOperator op_code,
                                                   OperatorT* op) {
  auto op_tensor_info = GetInfoForOpWithBiasTensor(op_code);
  if (!op_tensor_info) {
    error_reporter_->Report("Cannot quantize op: %s",
                            EnumNameBuiltinOperator(op_code));
    return kTfLiteError;
  }

  // Conv/Depthwise conv have 2 inputs when there is no bias, 3 otherwise.
  if (op->inputs.size() != 2 && op->inputs.size() != 3) {
    return kTfLiteError;
  }
  auto input_tensor_idx = op->inputs[op_tensor_info->activation_input_index];
  if (IsSubgraphInput(input_tensor_idx)) {
    TF_LITE_ENSURE_STATUS(AsymmetricQuantizeTensor(op_code, input_tensor_idx));
  }
  auto weights_tensor_idx = op->inputs[op_tensor_info->weights_input_index];

  TensorT* weights_tensor = subgraph_->tensors[weights_tensor_idx].get();
  int weights_channel_index = op_tensor_info->index_for_channel_in_weights;

  auto status = SymmetricPerChannelQuantizeTensor(
      model_, weights_tensor, weights_channel_index, error_reporter_);
  TF_LITE_ENSURE_STATUS(status);

  // If there is bias, quantize it.
  if (op->inputs.size() == 3) {
    auto bias_tensor_idx = op->inputs[op_tensor_info->bias_input_index];
    const TensorT* input_tensor = subgraph_->tensors[input_tensor_idx].get();
    TensorT* bias_tensor = subgraph_->tensors[bias_tensor_idx].get();
    TF_LITE_ENSURE_STATUS(SymmetricPerChannelBiasQuantize(
        input_tensor, weights_tensor, weights_channel_index, model_,
        bias_tensor, error_reporter_));
  }

  if (op->outputs.size() != 1) {
    return kTfLiteError;
  }
  auto output_tensor_idx = op->outputs[0];
  TF_LITE_ENSURE_STATUS(AsymmetricQuantizeTensor(op_code, output_tensor_idx));

  return kTfLiteOk;
}

TfLiteStatus SubgraphQuantizer::PropagateMinMaxForAvgAndMaxPool(
    BuiltinOperator op_code, OperatorT* op) {
  TF_LITE_ENSURE_EQ(this->error_reporter_, op->inputs.size(), 1);

  if (IsSubgraphInput(op->inputs[0])) {
    TF_LITE_ENSURE_STATUS(AsymmetricQuantizeTensor(op_code, op->inputs[0]));
  }

  auto output_tensor = subgraph_->tensors[op->outputs[0]].get();
  if (output_tensor->type != TensorType_FLOAT32) {
    return kTfLiteOk;
  }
  auto input_tensor = subgraph_->tensors[op->inputs[0]].get();
  if (!input_tensor->quantization) {
    error_reporter_->Report(
        "Missing required min/max information for input of operation: %s",
        EnumNameBuiltinOperator(op_code));
    return kTfLiteError;
  }
  if (input_tensor->quantization->min.size() != 1 ||
      input_tensor->quantization->max.size() != 1 ||
      input_tensor->quantization->scale.size() != 1 ||
      input_tensor->quantization->zero_point.size() != 1) {
    error_reporter_->Report(
        "Invalid quantization information for Op: %s, tensor: %s",
        EnumNameBuiltinOperator(op_code), input_tensor->name.c_str());
    return kTfLiteError;
  }
  auto quant_params = absl::make_unique<QuantizationParametersT>();
  // Nudge min, max to include the floating point zero.
  const float min = std::min(0.f, input_tensor->quantization->min[0]);
  const float max = std::max(0.f, input_tensor->quantization->max[0]);
  quant_params->min.push_back(min);
  quant_params->max.push_back(max);
  quant_params->scale.push_back(input_tensor->quantization->scale[0]);
  quant_params->zero_point.push_back(input_tensor->quantization->zero_point[0]);
  // TODO(shashishekhar): Log a warning here if overriding existing
  // min/max/scales differ from input scales.
  output_tensor->quantization = std::move(quant_params);
  output_tensor->type = TensorType_INT8;
  return kTfLiteOk;
}

TfLiteStatus SubgraphQuantizer::AsymmetricQuantizeSoftmax(
    BuiltinOperator op_code, OperatorT* op) {
  TF_LITE_ENSURE_EQ(this->error_reporter_, op->inputs.size(), 1);
  TF_LITE_ENSURE_EQ(this->error_reporter_, op->outputs.size(), 1);

  if (IsSubgraphInput(op->inputs[0])) {
    TF_LITE_ENSURE_STATUS(AsymmetricQuantizeTensor(op_code, op->inputs[0]));
  }

  auto output_tensor = subgraph_->tensors[op->outputs[0]].get();
  if (output_tensor->type != TensorType_FLOAT32) {
    return kTfLiteOk;
  }

  // Softmax output is hardcoded to have 1/256 as scale and -128 as zero point.
  output_tensor->type = TensorType_INT8;
  output_tensor->quantization->scale = {1.0f / 256.0f};
  output_tensor->quantization->zero_point = {-128};
  return kTfLiteOk;
}

TfLiteStatus SubgraphQuantizer::AsymmetricQuantizeInputsAndOutputs(
    BuiltinOperator op_code, OperatorT* op) {
  TF_LITE_ENSURE(this->error_reporter_, !op->inputs.empty());
  TF_LITE_ENSURE(this->error_reporter_, !op->outputs.empty());
  for (size_t input_idx = 0; input_idx < op->inputs.size(); ++input_idx) {
    auto input_tensor = subgraph_->tensors[op->inputs[input_idx]].get();
    if (IsSubgraphInput(op->inputs[input_idx]) &&
        input_tensor->type == TensorType_FLOAT32) {
      TF_LITE_ENSURE_STATUS(
          AsymmetricQuantizeTensor(op_code, op->inputs[input_idx]));
    }
  }

  for (size_t output_idx = 0; output_idx < op->outputs.size(); ++output_idx) {
    auto output_tensor = subgraph_->tensors[op->outputs[output_idx]].get();
    if (output_tensor->type == TensorType_FLOAT32) {
      TF_LITE_ENSURE_STATUS(
          AsymmetricQuantizeTensor(op_code, op->outputs[output_idx]));
    }
  }
  return kTfLiteOk;
}

bool SubgraphQuantizer::IsSubgraphInput(int32_t tensor_idx) const {
  return std::find(subgraph_->inputs.begin(), subgraph_->inputs.end(),
                   tensor_idx) != subgraph_->inputs.end();
}

TfLiteStatus SubgraphQuantizer::QuantizeOperator(int op_idx) {
  OperatorT* op = subgraph_->operators[op_idx].get();
  const BuiltinOperator op_code =
      model_->operator_codes[op->opcode_index]->builtin_code;
  if (OpHasOptionalBiasTensor(op_code)) {
    return QuantizeOpWithBias(op_code, op);
  }
  switch (op_code) {
    case BuiltinOperator_AVERAGE_POOL_2D:
    case BuiltinOperator_MAX_POOL_2D:
      return PropagateMinMaxForAvgAndMaxPool(op_code, op);
    case BuiltinOperator_SQUEEZE:
    case BuiltinOperator_RESHAPE:
    case BuiltinOperator_ADD:
      return AsymmetricQuantizeInputsAndOutputs(op_code, op);
    case BuiltinOperator_SOFTMAX:
      return AsymmetricQuantizeSoftmax(op_code, op);
    default:
      return kTfLiteError;
  }

  return kTfLiteError;
}

}  // namespace internal
}  // namespace optimize
}  // namespace tflite