/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file src/relay/qnn/op/concatenate.cc
 * \brief QNN concatenate operator. It concatenates quantized input tensors along a given axis.
 */

#include <tvm/tir/expr.h>
#include <tvm/relay/analysis.h>
#include <tvm/relay/op_attr_types.h>
#include <tvm/relay/qnn/attrs.h>
#include "../../op/tensor/transform.h"
#include "../../transforms/pattern_util.h"
#include "../../transforms/infer_layout_util.h"
#include "../util.h"

namespace tvm {
namespace relay {
namespace qnn {

bool QnnConcatenateRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                       const TypeReporter& reporter) {
  CHECK_EQ(types.size(), 6);

  // Check the scale and zero point types
  const auto* input_scales_tuple = types[1].as<TupleTypeNode>();
  if (input_scales_tuple == nullptr) {
    throw Error(
        ErrorBuilder()
        << "qnn concatenate requires a tuple of scales as the second argument, found "
        << PrettyPrint(types[1]));
  }
  for (const auto& input_scale : input_scales_tuple->fields) {
    CHECK(IsScalarType(input_scale, DataType::Float(32)));  // input_scales[idx]
  }

  const auto* input_zero_points_tuple = types[2].as<TupleTypeNode>();
  if (input_zero_points_tuple == nullptr) {
    throw Error(
        ErrorBuilder()
        << "qnn concatenate requires a tuple of zero_points as the third argument, found "
        << PrettyPrint(types[2]));
  }
  for (const auto& input_zero_point : input_zero_points_tuple->fields) {
    CHECK(IsScalarType(input_zero_point, DataType::Int(32)));  // input_zero_points[idx]
  }

  CHECK(IsScalarType(types[3], DataType::Float(32)));  // output_scale
  CHECK(IsScalarType(types[4], DataType::Int(32)));    // output_zero_point

  // Collect the input tensor and output tensor devoid of scale and zero points to reuse Relay
  // Concatenate infer type function.
  Array<Type> tensor_types = {types[0], types[5]};
  return ConcatenateRel<ConcatenateAttrs>(tensor_types, 2, attrs, reporter);
}

Array<Array<Layout>> QnnConcatenateLayout(const Attrs& attrs, const Array<Layout>& new_in_layouts,
                                          const Array<Layout>& old_in_layouts,
                                          const Array<tvm::relay::Type>& old_in_types) {
  // Collect the layouts and types to reuse Relay Concatenate Infer Correct Layout.
  CHECK_EQ(old_in_types.size(), 5);
  auto input_tuple_type = old_in_types[0].as<TupleTypeNode>();
  CHECK(input_tuple_type);
  auto num_input_tensors = input_tuple_type->fields.size();

  Array<Layout> relay_new_in_layouts(nullptr);
  if (new_in_layouts.defined()) {
    relay_new_in_layouts =
        Array<Layout>(new_in_layouts.begin(), new_in_layouts.begin() + num_input_tensors);
  }
  Array<Layout> relay_old_in_layouts(nullptr);
  if (old_in_layouts.defined()) {
    relay_old_in_layouts =
        Array<Layout>(old_in_layouts.begin(), old_in_layouts.begin() + num_input_tensors);
  }

  // Use Relay Concatenate Infer Correct layout to infer the layouts for data tensors.
  auto layouts =
      ConcatenateLayout(attrs, relay_new_in_layouts, relay_old_in_layouts, {old_in_types[0]});

  // Fill the layouts of remaining input tensors - scales and zero points. The layouts of these
  // tensors can be treated as channel layout. Total number of these tensors are 2 * num of data
  // tensors (scale and zero point for each input data tensor) + 2 for the output data tensor.
  Layout channel_layout = Layout("C");
  Array<Layout> input_layouts = layouts[0];

  for (size_t i = 0; i < 2 * num_input_tensors + 2; i++) {
    input_layouts.push_back(channel_layout);
  }
  Array<Layout> output_layouts = layouts[1];
  return {input_layouts, output_layouts};
}

Expr MakeQnnConcatenate(Expr data, Expr input_scales, Expr input_zero_points, Expr output_scale,
                        Expr output_zero_point, int axis) {
  auto attrs = make_object<ConcatenateAttrs>();
  attrs->axis = axis;
  static const Op& op = Op::Get("qnn.concatenate");
  return Call(op,
                        {data, input_scales, input_zero_points, output_scale, output_zero_point},
                        Attrs(attrs), {});
}

/*
 * \brief Canonicalizes the QNN concatenate op.
 * \param attrs The QNN concatenate attrs.
 * \param new_args The new mutated args to the call node.
 * \param arg_types The types of input and output.
 * \return The sequence of Relay ops for concatenate op.
 */
Expr ConcatenateQnnCanonicalize(const Attrs& attrs, const Array<Expr>& new_args,
                                const Array<tvm::relay::Type>& arg_types) {
  // Get the attrs.
  CHECK_EQ(new_args.size(), 5);
  auto& data = new_args[0];
  auto& input_scales = new_args[1];
  auto& input_zero_points = new_args[2];
  auto& output_scale = new_args[3];
  auto& output_zero_point = new_args[4];
  const auto* concatenate_attrs = attrs.as<ConcatenateAttrs>();
  CHECK(concatenate_attrs != nullptr);

  // Get the input dtype and shape.
  CHECK_GE(arg_types.size(), 1);
  auto tuple_type = arg_types[0].as<TupleTypeNode>();
  CHECK(tuple_type != nullptr);

  // FIXME (anijain2305) - The lowering can be further optimized. Instead of inserting requantize in
  // the start, we can insert requantize at the end if and only if all the input tensors have same
  // qnn params. This can be done in future.

  // If the output qnn params do not match the input qnn params, we can call requantize on the input
  // expr first, followed by a concatenate on the requantized input exprs.

  auto tuple_data = data.as<TupleNode>();
  CHECK(tuple_data != nullptr);

  auto tuple_input_scales = input_scales.as<TupleNode>();
  CHECK(tuple_input_scales != nullptr);

  auto tuple_input_zero_points = input_zero_points.as<TupleNode>();
  CHECK(tuple_input_zero_points != nullptr);

  int idx = 0;
  Array<Expr> requantized_exprs;
  for (auto quantized_expr : tuple_data->fields) {
    // Get the input scale for the idx quantized input tensor.
    auto input_scale = tuple_input_scales->fields[idx];

    // Get the zero point for the idx quantized input tensor.
    auto input_zero_point = tuple_input_zero_points->fields[idx];

    // Check if output and input qnn params are same. If not, requantize.
    if (!IsEqualScalar(input_scale, output_scale) ||
        !IsEqualScalar(input_zero_point, output_zero_point)) {
      // Get the input shape and dtype.
      auto tensor_type = tuple_type->fields[idx].as<TensorTypeNode>();
      auto input_dtype = tensor_type->dtype;
      auto input_shape = tensor_type->shape;

      // Requantize the input.
      auto requantized_expr = Requantize(quantized_expr, input_shape, input_scale, input_zero_point,
                                         output_scale, output_zero_point, input_dtype);
      requantized_exprs.push_back(requantized_expr);
    } else {
      requantized_exprs.push_back(quantized_expr);
    }
    idx++;
  }
  return MakeConcatenate(Tuple(requantized_exprs), concatenate_attrs->axis);
}

RELAY_REGISTER_OP("qnn.concatenate")
.describe(R"code(Concatenate the quantized input tensors along the given axis.
)code" TVM_ADD_FILELINE)
.set_attrs_type<ConcatenateAttrs>()
.set_num_inputs(5)
.add_argument("data", "Tensor", "The tensor to concatenate.")
.add_argument("input_scales", "Tensor", "The quantization scales of the input tensors.")
.add_argument("input_zero_points", "Tensor", "The quantization zero_points of the input tensors.")
.add_argument("output_scale", "Tensor", "The quantization scale of the output tensor.")
.add_argument("output_zero_point", "Tensor", "The quantization zero_point of the output tensor.")
.set_support_level(11)
.add_type_rel("QnnConcatenate", QnnConcatenateRel)
.set_attr<FTVMLegalize>("FTVMQnnCanonicalize", ConcatenateQnnCanonicalize)
.set_attr<FInferCorrectLayout>("FInferCorrectLayout", QnnConcatenateLayout);

TVM_REGISTER_GLOBAL("relay.qnn.op._make.concatenate")
.set_body_typed(MakeQnnConcatenate);

}  // namespace qnn
}  // namespace relay
}  // namespace tvm
