// Copyright (c) 2020-2022, ARM Limited.
//
//    Licensed under the Apache License, Version 2.0 with LLVM Exceptions
//    (the "License"); you may not use this file except in compliance with
//    the License. You may obtain a copy of the License at
//
//         https://llvm.org/LICENSE.txt
//
//    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.

// TOSA flatbuffer generation

#include "mlir/Dialect/Quant/QuantTypes.h"   // from @llvm-project
#include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project
#include "mlir/Dialect/Tosa/IR/TosaOps.h"    // from @llvm-project
#include "mlir/IR/Matchers.h"                // from @llvm-project
#include "mlir/Pass/Pass.h"                  // from @llvm-project
#include "mlir/Support/LogicalResult.h"      // from @llvm-project
#include "tosa_serialization_handler.h"
#include <functional>
#include <map>
#include <unordered_map>

// The namespace might be confusing here. We have mlir::tosa:: defined in MLIR
// and tosa:: defined in serialization library
// TODO: align the namespace
using namespace tosa;

namespace cl = llvm::cl;

llvm::cl::opt<std::string> tosa_flatbuffer_filename(
    "tosa-flatbuffer-filename", llvm::cl::desc("<tosa flatbuffer filename>"),
    llvm::cl::init("tosa_dump.tosa"), llvm::cl::value_desc("filename"));

llvm::cl::opt<std::string> tosa_flatbuffer_schema(
    "tosa-flatbuffer-schema", llvm::cl::desc("<tosa flatbuffer schema file>"),
    llvm::cl::init(""), llvm::cl::value_desc("filename"));

// Specialize mlir::Value for std::hash<T> and std::equal_to<T> to be able to
// build std::unordered_map<mlir::Value, ...>
namespace std {

template <> struct hash<mlir::Value> {
  std::size_t operator()(const mlir::Value &val) const {
    return static_cast<std::size_t>(mlir::hash_value(val));
  }
};

template <> struct equal_to<mlir::Value> {
  bool operator()(const mlir::Value &lhs, const mlir::Value &rhs) const {
    return (lhs == rhs);
  }
};

} // namespace std

ResizeMode ResizeModeStr2Enum(const std::string &mode_str) {
  if (mode_str == "NEAREST_NEIGHBOR")
    return ResizeMode_NEAREST;
  else if (mode_str == "BILINEAR")
    return ResizeMode_BILINEAR;
  else
    return ResizeMode_UNKNOWN;
}

DType Type2DType(mlir::Type element_type) {
  if (element_type.isF64() || element_type.isF32() || element_type.isF16() ||
      element_type.isBF16()) {
    return DType_FLOAT;
  } else if (element_type.isUnsignedInteger(8)) {
    return DType_UINT8;
  } else if (element_type.isInteger(4)) {
    return DType_INT8;
  } else if (element_type.isInteger(8)) {
    return DType_INT8;
  } else if (element_type.isInteger(16)) {
    return DType_INT16;
  } else if (element_type.isInteger(32)) {
    return DType_INT32;
  } else if (element_type.isInteger(48)) {
    return DType_INT48;
  }
  // boolean in MLIR treated as integer with bitwidth 1
  else if (element_type.isInteger(1)) {
    return DType_BOOL;
  }
  return DType_UNKNOWN;
}

int GetQuantizedParameter(mlir::Type type, std::vector<float> &scale,
                          std::vector<int32_t> &zeropoint,
                          int32_t &quantized_dimension, int64_t &quant_min,
                          int64_t &quant_max) {
  if (auto qtype = type.dyn_cast<mlir::quant::UniformQuantizedType>()) {
    scale.push_back(qtype.getScale());
    zeropoint.push_back(qtype.getZeroPoint());
    quantized_dimension = 0;

    quant_min = qtype.getStorageTypeMin();
    quant_max = qtype.getStorageTypeMax();
  } else if (auto qtype =
                 type.dyn_cast<mlir::quant::UniformQuantizedPerAxisType>()) {
    scale.assign(qtype.getScales().begin(), qtype.getScales().end());
    zeropoint.assign(qtype.getZeroPoints().begin(),
                     qtype.getZeroPoints().end());
    quantized_dimension = qtype.getQuantizedDimension();

    quant_min = qtype.getStorageTypeMin();
    quant_max = qtype.getStorageTypeMax();
  } else {
    return 1;
  }

  return 0;
}

TosaQuantInfoBase *
GetUnaryQuantInfo(mlir::tosa::UnaryOpQuantizationAttr quant_info) {
  int32_t input_zp = quant_info.input_zp().getInt();
  int32_t output_zp = quant_info.output_zp().getInt();

  TosaQuantInfoBase *qinfo = new TosaUnaryQuantInfo(input_zp, output_zp);

  return qinfo;
}

TosaQuantInfoBase *
GetConvQuantInfo(mlir::tosa::ConvOpQuantizationAttr quant_info) {
  int32_t input_zp = quant_info.input_zp().getInt();
  int32_t weight_zp = quant_info.weight_zp().getInt();

  TosaQuantInfoBase *qinfo = new TosaConvQuantInfo(input_zp, weight_zp);

  return qinfo;
}

TosaQuantInfoBase *
GetPadQuantInfo(mlir::tosa::PadOpQuantizationAttr quant_info) {
  int32_t input_zp = quant_info.input_zp().getInt();

  TosaQuantInfoBase *qinfo = new TosaPadQuantInfo(input_zp);

  return qinfo;
}

TosaQuantInfoBase *
GetMatMulQuantInfo(mlir::tosa::MatMulOpQuantizationAttr quant_info) {
  int32_t a_zp = quant_info.a_zp().getInt();
  int32_t b_zp = quant_info.b_zp().getInt();

  TosaQuantInfoBase *qinfo = new TosaMatMulQuantInfo(a_zp, b_zp);

  return qinfo;
}

class TosaSerializationBlockBuilder;

class TosaSerializationOperatorBuilder {
public:
  TosaSerializationOperatorBuilder(
      TosaSerializationBlockBuilder *_block_builder)
      : block_builder(_block_builder) {}
  template <typename T>
  TosaSerializationOperator *build(mlir::Operation &op) const;

private:
  std::string GetTensorName(mlir::Value val) const;
  TosaSerializationOperator *BuildPoolOpFromMlirOp(mlir::Operation &op,
                                                   Op opcode) const;
  TosaSerializationOperator *BuildEwiseBinaryOpFromMlirOp(mlir::Operation &op,
                                                          Op opcode) const;
  TosaSerializationOperator *BuildEwiseUnaryOpFromMlirOp(mlir::Operation &op,
                                                         Op opcode) const;
  TosaSerializationOperator *BuildReductionOpFromMlirOp(mlir::Operation &op,
                                                        Op opcode) const;
  TosaSerializationBlockBuilder *block_builder;
};

// This builder assumes each region only has only one block
class TosaSerializationBlockBuilder {
public:
  friend class TosaSerializationOperatorBuilder;
  TosaSerializationBlockBuilder(TosaSerializationBasicBlock *_block,
                                TosaSerializationHandler *_tsh,
                                mlir::Region *_region)
      : block(_block), tsh(_tsh), region(_region) {}

  mlir::LogicalResult
  BuildAllOpsInRegion(std::vector<mlir::Value> &return_values);
  TosaSerializationBasicBlock *GetBlock() { return block; }
  TosaSerializationHandler *GetTsh() { return tsh; }

private:
  TosaSerializationOperator *BuildTosaSerializationOperator(
      const TosaSerializationOperatorBuilder &op_builder, mlir::Operation &op);
  TosaSerializationTensor *
  BuildTosaSerializationTensor(mlir::Value val, const std::string &name);

  TosaSerializationBasicBlock *block;
  TosaSerializationHandler *tsh;
  mlir::Region *region;
  std::unordered_map<mlir::Value, std::string> tensor_map;
  std::unordered_map<mlir::Value, std::string> input_tensor_map;
};

std::string
TosaSerializationOperatorBuilder::GetTensorName(mlir::Value val) const {
  if (block_builder->tensor_map.find(val) == block_builder->tensor_map.end()) {
    llvm::errs() << "ERROR: Failed to get mlir::Value from tensor_map";
    assert(0);
  }
  return block_builder->tensor_map[val];
}

// Main template to catch unimplemented translation.
template <typename T>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build(mlir::Operation &op) const {
  llvm::errs() << "Translation of operator " << op.getName().getStringRef()
               << " is not implemented yet\n";
  return nullptr;
}

/* Start translating TOSA operator */

#define ASSERT_VECTOR_LENGTH(VECTOR, LENGTH)                                   \
  if (VECTOR.size() != LENGTH) {                                               \
    std::string msg;                                                           \
    msg = std::string(#VECTOR) + " is [" + std::to_string(VECTOR.size()) +     \
          "], but expected to be [" + std::to_string(LENGTH) + "]\n";          \
    op.emitOpError(msg.c_str());                                               \
    return nullptr;                                                            \
  }

TosaSerializationOperator *
TosaSerializationOperatorBuilder::BuildPoolOpFromMlirOp(mlir::Operation &op,
                                                        Op opcode) const {
  std::vector<int> pad, stride, kernel;

  auto pad_attr = op.getAttr("pad").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : pad_attr) {
    pad.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(pad, 4);

  auto stride_attr =
      op.getAttr("stride").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : stride_attr) {
    stride.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(stride, 2);

  auto kernel_attr =
      op.getAttr("kernel").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : kernel_attr) {
    kernel.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(kernel, 2);

  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaPoolAttribute attribute(pad, kernel, stride);
  auto quant_info = op.getAttrOfType<mlir::tosa::UnaryOpQuantizationAttr>(
      "quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_UnaryQuantInfo;
    qinfo = GetUnaryQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      opcode, Attribute_PoolAttribute, &attribute, qinfo_type, qinfo,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  delete qinfo;

  return tyop;
}

TosaSerializationOperator *
TosaSerializationOperatorBuilder::BuildEwiseBinaryOpFromMlirOp(
    mlir::Operation &op, Op opcode) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      opcode, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input0_name, input1_name},
      std::vector<std::string>{output_name});

  return tyop;
}

TosaSerializationOperator *
TosaSerializationOperatorBuilder::BuildEwiseUnaryOpFromMlirOp(
    mlir::Operation &op, Op opcode) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      opcode, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

TosaSerializationOperator *
TosaSerializationOperatorBuilder::BuildReductionOpFromMlirOp(
    mlir::Operation &op, Op opcode) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  int32_t axis = op.getAttr("axis").dyn_cast<mlir::IntegerAttr>().getInt();
  TosaAxisAttribute attribute(axis);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      opcode, Attribute_AxisAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

#define BUILD_OP_POOL2D(MLIR_OP_NAME, SCHEMA_OP_NAME)                          \
  template <>                                                                  \
  TosaSerializationOperator *                                                  \
  TosaSerializationOperatorBuilder::build<mlir::tosa::MLIR_OP_NAME##Op>(       \
      mlir::Operation & op) const {                                            \
    return BuildPoolOpFromMlirOp(op, Op_##SCHEMA_OP_NAME);                     \
  }

#define BUILD_OP_ELEMENTWISE_BINARY(MLIR_OP_NAME, SCHEMA_OP_NAME)              \
  template <>                                                                  \
  TosaSerializationOperator *                                                  \
  TosaSerializationOperatorBuilder::build<mlir::tosa::MLIR_OP_NAME##Op>(       \
      mlir::Operation & op) const {                                            \
    return BuildEwiseBinaryOpFromMlirOp(op, Op_##SCHEMA_OP_NAME);              \
  }

#define BUILD_OP_ELEMENTWISE_UNARY(MLIR_OP_NAME, SCHEMA_OP_NAME)               \
  template <>                                                                  \
  TosaSerializationOperator *                                                  \
  TosaSerializationOperatorBuilder::build<mlir::tosa::MLIR_OP_NAME##Op>(       \
      mlir::Operation & op) const {                                            \
    return BuildEwiseUnaryOpFromMlirOp(op, Op_##SCHEMA_OP_NAME);               \
  }

#define BUILD_OP_REDUCTION(MLIR_OP_NAME, SCHEMA_OP_NAME)                       \
  template <>                                                                  \
  TosaSerializationOperator *                                                  \
  TosaSerializationOperatorBuilder::build<mlir::tosa::MLIR_OP_NAME##Op>(       \
      mlir::Operation & op) const {                                            \
    return BuildReductionOpFromMlirOp(op, Op_##SCHEMA_OP_NAME);                \
  }

BUILD_OP_POOL2D(MaxPool2d, MAX_POOL2D)
BUILD_OP_POOL2D(AvgPool2d, AVG_POOL2D)

BUILD_OP_ELEMENTWISE_BINARY(Add, ADD)
BUILD_OP_ELEMENTWISE_BINARY(BitwiseAnd, BITWISE_AND)
BUILD_OP_ELEMENTWISE_BINARY(BitwiseXor, BITWISE_XOR)
BUILD_OP_ELEMENTWISE_BINARY(BitwiseOr, BITWISE_OR)
BUILD_OP_ELEMENTWISE_BINARY(Div, INTDIV)
BUILD_OP_ELEMENTWISE_BINARY(LogicalAnd, LOGICAL_AND)
BUILD_OP_ELEMENTWISE_BINARY(LogicalLeftShift, LOGICAL_LEFT_SHIFT)
BUILD_OP_ELEMENTWISE_BINARY(LogicalRightShift, LOGICAL_RIGHT_SHIFT)
BUILD_OP_ELEMENTWISE_BINARY(LogicalOr, LOGICAL_OR)
BUILD_OP_ELEMENTWISE_BINARY(LogicalXor, LOGICAL_XOR)
BUILD_OP_ELEMENTWISE_BINARY(Maximum, MAXIMUM)
BUILD_OP_ELEMENTWISE_BINARY(Minimum, MINIMUM)
BUILD_OP_ELEMENTWISE_BINARY(Pow, POW)
BUILD_OP_ELEMENTWISE_BINARY(Sub, SUB)

BUILD_OP_ELEMENTWISE_UNARY(Abs, ABS)
BUILD_OP_ELEMENTWISE_UNARY(BitwiseNot, BITWISE_NOT)
BUILD_OP_ELEMENTWISE_UNARY(Ceil, CEIL)
BUILD_OP_ELEMENTWISE_UNARY(Clz, CLZ)
BUILD_OP_ELEMENTWISE_UNARY(Exp, EXP)
BUILD_OP_ELEMENTWISE_UNARY(Floor, FLOOR)
BUILD_OP_ELEMENTWISE_UNARY(Log, LOG)
BUILD_OP_ELEMENTWISE_UNARY(LogicalNot, LOGICAL_NOT)
BUILD_OP_ELEMENTWISE_UNARY(Reciprocal, RECIPROCAL)
BUILD_OP_ELEMENTWISE_UNARY(Rsqrt, RSQRT)

BUILD_OP_REDUCTION(ReduceAny, REDUCE_ANY)
BUILD_OP_REDUCTION(ReduceAll, REDUCE_ALL)
BUILD_OP_REDUCTION(ReduceMax, REDUCE_MAX)
BUILD_OP_REDUCTION(ReduceMin, REDUCE_MIN)
BUILD_OP_REDUCTION(ReduceProd, REDUCE_PRODUCT)
BUILD_OP_REDUCTION(ReduceSum, REDUCE_SUM)

BUILD_OP_ELEMENTWISE_BINARY(Equal, EQUAL)
BUILD_OP_ELEMENTWISE_BINARY(Greater, GREATER)
BUILD_OP_ELEMENTWISE_BINARY(GreaterEqual, GREATER_EQUAL)

BUILD_OP_ELEMENTWISE_UNARY(Sigmoid, SIGMOID)
BUILD_OP_ELEMENTWISE_UNARY(Tanh, TANH)
BUILD_OP_ELEMENTWISE_UNARY(Identity, IDENTITY)
BUILD_OP_ELEMENTWISE_UNARY(Cast, CAST)

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ConstOp>(
    mlir::Operation &op) const {
  std::string output_name = GetTensorName(op.getResult(0));
  TosaSerializationTensor *ts =
      block_builder->GetBlock()->GetTensorByName(output_name);
  if (!ts) {
    op.emitOpError(
        "ERROR: serialization tensor must be built before building operator");
    return nullptr;
  }

#if 0
    // Gracefully handle constants of "constant unit" type which have no value
    // by creating a numpy value of 0.
    auto unit_val = op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::UnitAttr>();

    if (unit_val)
    {
        std::vector<float> data = { 0.0 };
        type = DType_FLOAT;
        TosaSerializationHandler::ConvertF32toU8(data, u8_data);
    }
#endif

  // Update tensor.data array with Const value attribute
  std::vector<uint8_t> u8_data;
  DType type = ts->GetDtype();
  if (type == DType_FLOAT) {
    std::vector<float> data;
    auto dense_attr = op.getAttr(llvm::StringRef("value"))
                          .dyn_cast<mlir::DenseElementsAttr>();
    auto val_attr =
        op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::FloatAttr>();

    if (dense_attr) {
      for (auto val : dense_attr.getValues<float>()) {
        data.push_back(val);
      }
    } else if (val_attr) {
      data.push_back((float)val_attr.getValueAsDouble());
    } else {
      op.emitOpError("Unknown const attribute");
      return nullptr;
    }
    TosaSerializationHandler::ConvertF32toU8(data, u8_data);
  } else if (type == DType_INT8) {
    std::vector<int8_t> data;
    auto dense_attr = op.getAttr(llvm::StringRef("value"))
                          .dyn_cast<mlir::DenseElementsAttr>();
    auto val_attr =
        op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::IntegerAttr>();

    if (dense_attr) {
      for (auto val : dense_attr.getValues<int8_t>()) {
        data.push_back(val);
      }
    } else if (val_attr) {
      data.push_back(val_attr.getInt());
    } else {
      op.emitOpError("Unknown const attribute");
      return nullptr;
    }
    TosaSerializationHandler::ConvertI8toU8(data, u8_data);
  } else if (type == DType_INT16) {
    std::vector<int16_t> data;
    auto dense_attr = op.getAttr(llvm::StringRef("value"))
                          .dyn_cast<mlir::DenseElementsAttr>();
    auto val_attr =
        op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::IntegerAttr>();

    if (dense_attr) {
      for (auto val : dense_attr.getValues<int16_t>()) {
        data.push_back(val);
      }
    } else if (val_attr) {
      data.push_back(val_attr.getInt());
    } else {
      op.emitOpError("Unknown const attribute");
      return nullptr;
    }
    TosaSerializationHandler::ConvertI16toU8(data, u8_data);
  } else if (type == DType_INT32) {
    std::vector<int32_t> data;
    auto dense_attr = op.getAttr(llvm::StringRef("value"))
                          .dyn_cast<mlir::DenseElementsAttr>();
    auto val_attr =
        op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::IntegerAttr>();

    if (dense_attr) {
      for (auto val : dense_attr.getValues<int32_t>()) {
        data.push_back(val);
      }
    } else if (val_attr) {
      data.push_back(val_attr.getInt());
    } else {
      op.emitOpError("Unknown const attribute");
      return nullptr;
    }
    TosaSerializationHandler::ConvertI32toU8(data, u8_data);
  } else if (type == DType_INT48) {
    std::vector<int64_t> data;
    auto dense_attr = op.getAttr(llvm::StringRef("value"))
                          .dyn_cast<mlir::DenseElementsAttr>();
    auto val_attr =
        op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::IntegerAttr>();

    if (dense_attr) {
      for (auto valueIt : dense_attr.getValues<mlir::APInt>()) {
        uint64_t val = valueIt.getLimitedValue();
        data.push_back(val);
      }
    } else if (val_attr) {
      data.push_back(val_attr.getInt());
    } else {
      op.emitOpError("Unknown const attribute");
      return nullptr;
    }
    TosaSerializationHandler::ConvertI48toU8(data, u8_data);
  } else if (type == DType_BOOL) {
    std::vector<bool> data;

    auto dense_attr = op.getAttr(llvm::StringRef("value"))
                          .dyn_cast<mlir::DenseElementsAttr>();
    auto val_attr =
        op.getAttr(llvm::StringRef("value")).dyn_cast<mlir::BoolAttr>();

    if (dense_attr) {
      for (auto val : dense_attr.getValues<bool>()) {
        data.push_back(val);
      }
    } else if (val_attr) {
      data.push_back(val_attr.getValue());
    } else {
      op.emitOpError("Unknown const attribute");
      return nullptr;
    }

    TosaSerializationHandler::ConvertBooltoU8(data, u8_data);
  } else {
    op.emitOpError("Unknown element type of const attribute");
    return nullptr;
  }
  ts->SetData(u8_data);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_CONST, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{}, std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::Conv2DOp>(
    mlir::Operation &op) const {
  std::vector<int> pad, stride, dilation;

  auto pad_attr = op.getAttr("pad").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : pad_attr) {
    pad.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(pad, 4);

  auto stride_attr =
      op.getAttr("stride").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : stride_attr) {
    stride.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(stride, 2);

  auto dilation_attr =
      op.getAttr("dilation").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : dilation_attr) {
    dilation.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(dilation, 2);

  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string input2_name = GetTensorName(op.getOperand(2));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaConvAttribute attribute(pad, stride, dilation);

  auto quant_info =
      op.getAttrOfType<mlir::tosa::ConvOpQuantizationAttr>("quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_ConvQuantInfo;
    qinfo = GetConvQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_CONV2D, Attribute_ConvAttribute, &attribute, qinfo_type, qinfo,
      std::vector<std::string>{input0_name, input1_name, input2_name},
      std::vector<std::string>{output_name});

  delete qinfo;

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::DepthwiseConv2DOp>(
    mlir::Operation &op) const {
  std::vector<int> pad, stride, dilation;

  auto pad_attr = op.getAttr("pad").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : pad_attr) {
    pad.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(pad, 4);

  auto stride_attr =
      op.getAttr("stride").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : stride_attr) {
    stride.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(stride, 2);

  auto dilation_attr =
      op.getAttr("dilation").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : dilation_attr) {
    dilation.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(dilation, 2);

  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string input2_name = GetTensorName(op.getOperand(2));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaConvAttribute attribute(pad, stride, dilation);

  auto quant_info =
      op.getAttrOfType<mlir::tosa::ConvOpQuantizationAttr>("quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_ConvQuantInfo;
    qinfo = GetConvQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_DEPTHWISE_CONV2D, Attribute_ConvAttribute, &attribute, qinfo_type,
      qinfo, std::vector<std::string>{input0_name, input1_name, input2_name},
      std::vector<std::string>{output_name});

  delete qinfo;

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::TransposeConv2DOp>(
    mlir::Operation &op) const {
  std::vector<int> outpad, stride, dilation, output_shape;

  auto outpad_attr =
      op.getAttr("out_pad").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : outpad_attr) {
    outpad.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(outpad, 2);

  auto stride_attr =
      op.getAttr("stride").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : stride_attr) {
    stride.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(stride, 2);

  auto dilation_attr =
      op.getAttr("dilation").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : dilation_attr) {
    dilation.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(dilation, 2);

  auto output_shape_attr =
      op.getAttr("out_shape").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : output_shape_attr) {
    output_shape.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(output_shape, 4);

  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string input2_name = GetTensorName(op.getOperand(2));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaTransposeConvAttribute attribute(outpad, stride, dilation, output_shape);

  auto quant_info =
      op.getAttrOfType<mlir::tosa::ConvOpQuantizationAttr>("quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_ConvQuantInfo;
    qinfo = GetConvQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_TRANSPOSE_CONV2D, Attribute_TransposeConvAttribute, &attribute,
      qinfo_type, qinfo,
      std::vector<std::string>{input0_name, input1_name, input2_name},
      std::vector<std::string>{output_name});

  delete qinfo;

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::FullyConnectedOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string input2_name = GetTensorName(op.getOperand(2));
  std::string output_name = GetTensorName(op.getResult(0));

  auto quant_info =
      op.getAttrOfType<mlir::tosa::ConvOpQuantizationAttr>("quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_ConvQuantInfo;
    qinfo = GetConvQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_FULLY_CONNECTED, Attribute_NONE, nullptr, qinfo_type, qinfo,
      std::vector<std::string>{input0_name, input1_name, input2_name},
      std::vector<std::string>{output_name});

  delete qinfo;

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::MatMulOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string output_name = GetTensorName(op.getResult(0));

  auto quant_info = op.getAttrOfType<mlir::tosa::MatMulOpQuantizationAttr>(
      "quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_MatMulQuantInfo;
    qinfo = GetMatMulQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_MATMUL, Attribute_NONE, nullptr, qinfo_type, qinfo,
      std::vector<std::string>{input0_name, input1_name},
      std::vector<std::string>{output_name});

  delete qinfo;

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::SelectOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string input2_name = GetTensorName(op.getOperand(2));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_SELECT, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input0_name, input1_name, input2_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ClampOp>(
    mlir::Operation &op) const {
  int32_t min_int =
      op.getAttr("min_int").dyn_cast<mlir::IntegerAttr>().getInt();
  int32_t max_int =
      op.getAttr("max_int").dyn_cast<mlir::IntegerAttr>().getInt();
  float min_fp = op.getAttr("min_fp")
                     .dyn_cast<mlir::FloatAttr>()
                     .getValue()
                     .convertToFloat();
  float max_fp = op.getAttr("max_fp")
                     .dyn_cast<mlir::FloatAttr>()
                     .getValue()
                     .convertToFloat();

  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaClampAttribute attribute(min_int, max_int, min_fp, max_fp);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_CLAMP, Attribute_ClampAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ArgMaxOp>(
    mlir::Operation &op) const {
  int32_t axis = op.getAttr("axis").dyn_cast<mlir::IntegerAttr>().getInt();

  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaAxisAttribute attribute(axis);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_ARGMAX, Attribute_AxisAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ConcatOp>(
    mlir::Operation &op) const {
  int32_t axis = op.getAttr("axis").dyn_cast<mlir::IntegerAttr>().getInt();

  std::vector<std::string> inputs;
  for (uint32_t i = 0; i < op.getNumOperands(); i++) {
    std::string input_name = GetTensorName(op.getOperand(i));
    inputs.push_back(input_name);
  }

  std::string output_name = GetTensorName(op.getResult(0));

  TosaAxisAttribute attribute(axis);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_CONCAT, Attribute_AxisAttribute, &attribute, QuantInfo_NONE, nullptr,
      inputs, std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::NegateOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  auto quant_info = op.getAttrOfType<mlir::tosa::UnaryOpQuantizationAttr>(
      "quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_UnaryQuantInfo;
    qinfo = GetUnaryQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_NEGATE, Attribute_NONE, nullptr, qinfo_type, qinfo,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});
  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ReshapeOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  std::vector<int> shape;
  auto shape_attr =
      op.getAttr("new_shape").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : shape_attr) {
    shape.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }

  TosaReshapeAttribute attribute(shape);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_RESHAPE, Attribute_ReshapeAttribute, &attribute, QuantInfo_NONE,
      nullptr, std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::PadOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  // Match padding tensor as compile-time constant attribute
  // TODO: fix when MLIR dialect changes
  mlir::ElementsAttr paddings_elems;
  if (!matchPattern(op.getOperand(1), m_Constant(&paddings_elems)))
    return nullptr;

  std::vector<int> paddings;
  for (int32_t val : paddings_elems.getValues<int32_t>()) {
    paddings.push_back(val);
  }

  TosaPadAttribute attribute(paddings, 0 /* pad_const_int */,
                             0.0f /* pad_const_fp */);

  auto quant_info =
      op.getAttrOfType<mlir::tosa::PadOpQuantizationAttr>("quantization_info");
  QuantInfo qinfo_type;
  TosaQuantInfoBase *qinfo;
  if (quant_info) {
    qinfo_type = QuantInfo_PadQuantInfo;
    qinfo = GetPadQuantInfo(quant_info);
  } else {
    qinfo_type = QuantInfo_NONE;
    qinfo = new TosaNoneQuantInfo();
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_PAD, Attribute_PadAttribute, &attribute, qinfo_type, qinfo,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});
  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::TransposeOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  // Match perm tensor as compile-time constant attribute
  // TODO: fix when MLIR dialect changes
  mlir::ElementsAttr perm_elems;
  if (!matchPattern(op.getOperand(1), m_Constant(&perm_elems)))
    return nullptr;

  std::vector<int> perm;
  auto values = perm_elems.getValues<mlir::IntegerAttr>();
  for (int32_t i = 0; i < values.size(); i++) {
    perm.push_back(values[i].getInt());
  }

  TosaTransposeAttribute attribute(perm);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_TRANSPOSE, Attribute_TransposeAttribute, &attribute, QuantInfo_NONE,
      nullptr, std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::SliceOp>(
    mlir::Operation &op) const {
  std::vector<int> start, size;
  auto begin_attr = op.getAttr("start").dyn_cast<mlir::ArrayAttr>().getValue();
  auto size_attr = op.getAttr("size").dyn_cast<mlir::ArrayAttr>().getValue();

  for (auto &int_attr : begin_attr) {
    start.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }

  for (auto &int_attr : size_attr) {
    size.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }

  TosaSliceAttribute attribute(start, size);

  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_SLICE, Attribute_SliceAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::TileOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  std::vector<int> multiples;
  auto multiples_attr =
      op.getAttr("multiples").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : multiples_attr) {
    multiples.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }

  TosaTileAttribute attribute(multiples);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_TILE, Attribute_TileAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::GatherOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_GATHER, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input0_name, input1_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ScatterOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string input2_name = GetTensorName(op.getOperand(2));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_SCATTER, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input0_name, input1_name, input2_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ResizeOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  std::vector<int> output_size;
  auto output_size_attr =
      op.getAttr("output_size").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : output_size_attr) {
    output_size.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(output_size, 2);

  std::vector<int> stride;
  auto stride_attr =
      op.getAttr("stride").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : stride_attr) {
    stride.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(stride, 2);

  std::vector<int> offset;
  auto offset_attr =
      op.getAttr("offset").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &int_attr : offset_attr) {
    offset.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }
  ASSERT_VECTOR_LENGTH(offset, 2);

  int32_t shift = op.getAttr("shift").dyn_cast<mlir::IntegerAttr>().getInt();

  std::vector<float> stride_fp;
  auto stride_fp_attr =
      op.getAttr("stride_fp").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &fp_attr : stride_fp_attr) {
    stride_fp.push_back(fp_attr.dyn_cast<mlir::FloatAttr>().getValueAsDouble());
  }
  ASSERT_VECTOR_LENGTH(stride_fp, 2);

  std::vector<float> offset_fp;
  auto offset_fp_attr =
      op.getAttr("offset_fp").dyn_cast<mlir::ArrayAttr>().getValue();
  for (auto &fp_attr : offset_fp_attr) {
    offset_fp.push_back(fp_attr.dyn_cast<mlir::FloatAttr>().getValueAsDouble());
  }
  ASSERT_VECTOR_LENGTH(offset_fp, 2);

  auto mode_str =
      op.getAttr("mode").dyn_cast<mlir::StringAttr>().getValue().str();
  ResizeMode mode = ResizeModeStr2Enum(mode_str);

  TosaResizeAttribute attribute(output_size, stride, offset, shift, stride_fp,
                                offset_fp, mode);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_RESIZE, Attribute_ResizeAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ReverseOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  int32_t axis = op.getAttr("axis").dyn_cast<mlir::IntegerAttr>().getInt();

  TosaAxisAttribute attribute(axis);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_REVERSE, Attribute_AxisAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::MulOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string output_name = GetTensorName(op.getResult(0));

  int32_t shift = op.getAttr("shift").dyn_cast<mlir::IntegerAttr>().getInt();

  TosaMulAttribute attribute(shift);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_MUL, Attribute_MulAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input0_name, input1_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::ArithmeticRightShiftOp>(
    mlir::Operation &op) const {
  std::string input0_name = GetTensorName(op.getOperand(0));
  std::string input1_name = GetTensorName(op.getOperand(1));
  std::string output_name = GetTensorName(op.getResult(0));

  bool round = op.getAttr("round").dyn_cast<mlir::BoolAttr>().getValue();

  TosaArithmeticRightShiftAttribute attribute(round);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_ARITHMETIC_RIGHT_SHIFT, Attribute_ArithmeticRightShiftAttribute,
      &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input0_name, input1_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::TableOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  // Match table tensor as compile-time constant attribute
  // TODO: fix when MLIR dialect changes
  mlir::ElementsAttr table_elems;
  if (!matchPattern(op.getOperand(1), m_Constant(&table_elems)))
    return nullptr;

  std::vector<int> table;
  auto values = table_elems.getValues<mlir::IntegerAttr>();
  for (int32_t i = 0; i < values.size(); i++) {
    table.push_back(values[i].getInt());
  }

  TosaTableAttribute attribute(table);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_TABLE, Attribute_TableAttribute, &attribute, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::RescaleOp>(
    mlir::Operation &op) const {
  int32_t input_zp =
      op.getAttr("input_zp").dyn_cast<mlir::IntegerAttr>().getInt();
  int32_t output_zp =
      op.getAttr("output_zp").dyn_cast<mlir::IntegerAttr>().getInt();
  bool scale32 = op.getAttr("scale32").dyn_cast<mlir::BoolAttr>().getValue();
  bool double_round =
      op.getAttr("double_round").dyn_cast<mlir::BoolAttr>().getValue();
  bool per_channel =
      op.getAttr("per_channel").dyn_cast<mlir::BoolAttr>().getValue();

  std::vector<int> multiplier, shift;
  auto multiplier_attr =
      op.getAttr("multiplier").dyn_cast<mlir::ArrayAttr>().getValue();
  auto shift_attr = op.getAttr("shift").dyn_cast<mlir::ArrayAttr>().getValue();

  for (auto &int_attr : multiplier_attr) {
    multiplier.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }

  for (auto &int_attr : shift_attr) {
    shift.push_back(int_attr.dyn_cast<mlir::IntegerAttr>().getInt());
  }

  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaRescaleAttribute attribute(input_zp, output_zp, multiplier, shift,
                                 scale32, double_round, per_channel);

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_RESCALE, Attribute_RescaleAttribute, &attribute, QuantInfo_NONE,
      nullptr, std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::CustomOp>(
    mlir::Operation &op) const {
  std::string input_name = GetTensorName(op.getOperand(0));
  std::string output_name = GetTensorName(op.getResult(0));

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_CUSTOM, Attribute_NONE, nullptr, QuantInfo_NONE, nullptr,
      std::vector<std::string>{input_name},
      std::vector<std::string>{output_name});

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::IfOp>(
    mlir::Operation &op) const {
  std::vector<std::string> input_names, output_names;

  mlir::Region &then_region = op.getRegion(0);
  mlir::Region &else_region = op.getRegion(1);
  std::vector<mlir::Value> then_yields, else_yields;
  TosaSerializationBasicBlock *then_block = nullptr;
  TosaSerializationBasicBlock *else_block = nullptr;

  // Building then branch block
  std::string then_block_name =
      "bb" + std::to_string(block_builder->GetTsh()->GetBlocks().size());
  then_block = new TosaSerializationBasicBlock(
      then_block_name, std::vector<TosaSerializationOperator *>(),
      std::vector<TosaSerializationTensor *>(), std::vector<std::string>(),
      std::vector<std::string>());
  assert(then_block);
  block_builder->GetTsh()->GetBlocks().push_back(then_block);

  TosaSerializationBlockBuilder then_block_builder(
      then_block, block_builder->GetTsh(), &then_region);
  if (then_block_builder.BuildAllOpsInRegion(then_yields).failed()) {
    return nullptr;
  }
  if (then_yields.size() != op.getNumResults()) {
    op.emitOpError("BuildOpCondIf: then_region yield.size() doesn't match "
                   "cond_if's output size");
    return nullptr;
  }

  // Building else branch block
  std::string else_block_name =
      "bb" + std::to_string(block_builder->GetTsh()->GetBlocks().size());
  else_block = new TosaSerializationBasicBlock(
      else_block_name, std::vector<TosaSerializationOperator *>(),
      std::vector<TosaSerializationTensor *>(), std::vector<std::string>(),
      std::vector<std::string>());
  assert(else_block);
  block_builder->GetTsh()->GetBlocks().push_back(else_block);

  TosaSerializationBlockBuilder else_block_builder(
      else_block, block_builder->GetTsh(), &else_region);
  if (else_block_builder.BuildAllOpsInRegion(else_yields).failed()) {
    return nullptr;
  }
  if (else_yields.size() != op.getNumResults()) {
    op.emitOpError("BuildOpCondIf: else_region yield.size() doesn't match "
                   "cond_if's output size");
    return nullptr;
  }

  TosaCondIfAttribute attribute(then_block->GetName(), else_block->GetName());

  for (size_t i = 0; i < op.getNumOperands(); i++) {
    std::string input_name = GetTensorName(op.getOperand(i));
    input_names.push_back(input_name);
  }

  for (size_t i = 0; i < op.getNumResults(); i++) {
    std::string output_name = GetTensorName(op.getResult(i));
    output_names.push_back(output_name);
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_COND_IF, Attribute_CondIfAttribute, &attribute, QuantInfo_NONE,
      nullptr, input_names, output_names);

  return tyop;
}

template <>
TosaSerializationOperator *
TosaSerializationOperatorBuilder::build<mlir::tosa::WhileOp>(
    mlir::Operation &op) const {
  std::vector<std::string> input_names, output_names;

  mlir::Region &cond_region = op.getRegion(0);
  mlir::Region &body_region = op.getRegion(1);
  std::vector<mlir::Value> cond_yields, body_yields;
  TosaSerializationBasicBlock *cond_block = nullptr;
  TosaSerializationBasicBlock *body_block = nullptr;

  // Building cond branch block
  std::string cond_block_name =
      "bb" + std::to_string(block_builder->GetTsh()->GetBlocks().size());
  cond_block = new TosaSerializationBasicBlock(
      cond_block_name, std::vector<TosaSerializationOperator *>(),
      std::vector<TosaSerializationTensor *>(), std::vector<std::string>(),
      std::vector<std::string>());
  assert(cond_block);
  block_builder->GetTsh()->GetBlocks().push_back(cond_block);

  TosaSerializationBlockBuilder cond_block_builder(
      cond_block, block_builder->GetTsh(), &cond_region);
  if (cond_block_builder.BuildAllOpsInRegion(cond_yields).failed()) {
    return nullptr;
  }
  if (cond_yields.size() != 1) {
    op.emitOpError("BuildOpWhileLoop: cond_region yield.size() is not 1");
    return nullptr;
  }

  // Building body branch block
  std::string body_block_name =
      "bb" + std::to_string(block_builder->GetTsh()->GetBlocks().size());
  body_block = new TosaSerializationBasicBlock(
      body_block_name, std::vector<TosaSerializationOperator *>(),
      std::vector<TosaSerializationTensor *>(), std::vector<std::string>(),
      std::vector<std::string>());
  assert(body_block);
  block_builder->GetTsh()->GetBlocks().push_back(body_block);

  TosaSerializationBlockBuilder body_block_builder(
      body_block, block_builder->GetTsh(), &body_region);
  if (body_block_builder.BuildAllOpsInRegion(body_yields).failed()) {
    return nullptr;
  }
  if (body_yields.size() != op.getNumResults()) {
    op.emitOpError("BuildOpWhileLoop: body_region yield.size() doesn't "
                   "match while_loop's output size");
    return nullptr;
  }

  TosaWhileLoopAttribute attribute(cond_block->GetName(),
                                   body_block->GetName());

  for (size_t i = 0; i < op.getNumOperands(); i++) {
    std::string input_name = GetTensorName(op.getOperand(i));
    input_names.push_back(input_name);
  }

  for (size_t i = 0; i < op.getNumResults(); i++) {
    std::string output_name = GetTensorName(op.getResult(i));
    output_names.push_back(output_name);
  }

  TosaSerializationOperator *tyop = new TosaSerializationOperator(
      Op_WHILE_LOOP, Attribute_WhileLoopAttribute, &attribute, QuantInfo_NONE,
      nullptr, input_names, output_names);

  return tyop;
}

/* End translating TOSA operator */

mlir::LogicalResult TosaSerializationBlockBuilder::BuildAllOpsInRegion(
    std::vector<mlir::Value> &return_values) {
  TosaSerializationOperator *ser_operator = nullptr;
  TosaSerializationTensor *ser_tensor = nullptr;
  size_t num_blocks_in_region = 0;
  static int input_tensor_index = 0;
  static int intermediate_tensor_index = 0;
  static int output_tensor_index = 0;
  TosaSerializationOperatorBuilder op_builder(this);

  for (auto &bb : region->getBlocks()) {
    num_blocks_in_region++;

    if (num_blocks_in_region > 1) {
      llvm::errs() << "Invalid MLIR: multiple blocks in a region\n";
      return mlir::failure();
    }

    // We always have one block for each region right now
    assert(bb.isEntryBlock());

    // Specify block input tensor name
    for (auto args : bb.getArguments()) {
      std::string block_input_name =
          "TosaInput_" + std::to_string(input_tensor_index++);
      block->GetInputs().push_back(block_input_name);
      tensor_map[args] = block_input_name;
      input_tensor_map[args] = block_input_name;
    }

    // Build tensor_map
    for (auto &op : bb) {
      if (!(llvm::isa<mlir::tosa::YieldOp>(op) ||
            llvm::isa<mlir::ReturnOp>(op) ||
            llvm::isa<mlir::tensor::CastOp>(op))) {
        for (uint32_t i = 0; i < op.getNumResults(); i++) {
          std::string intermediate_tensor_name =
              "layer_" + std::to_string(intermediate_tensor_index++);
          tensor_map[op.getResult(i)] = intermediate_tensor_name;
        }
      } else {
        if (llvm::isa<mlir::tensor::CastOp>(op))
          continue;
        // Override return tensor name
        for (auto val : op.getOperands()) {
          // Workaround to skip mlir::tensor::CastOp before return
          mlir::Operation *val_defining_op = val.getDefiningOp();
          if (val_defining_op) {
            if (llvm::isa<mlir::tensor::CastOp>(*val_defining_op))
              val = val_defining_op->getOperand(0);
          }

          // Sanity check. This mlir::Value should be built in map since graph
          // is DAG
          if (tensor_map.find(val) == tensor_map.end()) {
            llvm::errs() << "ERROR: Can't find built mlir::Value key.\n";
            return mlir::failure();
          }

          // If returned value is block input, short-circuit the tensor name
          // Otherwise, build a new output name and override the origin tensor
          // name
          if (input_tensor_map.find(val) != input_tensor_map.end()) {
            block->GetOutputs().push_back(input_tensor_map[val]);
            return_values.push_back(val);
          } else {
            std::string output_name =
                "TosaOutput_" + std::to_string(output_tensor_index++);
            tensor_map[val] = output_name;
            block->GetOutputs().push_back(output_name);
            return_values.push_back(val);
          }
        }
      }
    }

    // Build tensor

    // The tensor_map is sorted by hashed mlir::Value types.
    // For serialization, sort tensors alphabetically by name for a
    // deterministic and human-friendly ordering.
    std::map<std::string, mlir::Value> tensor_name_sort;
    for (auto pair : tensor_map)
      tensor_name_sort[pair.second] = pair.first;

    for (auto pair : tensor_name_sort) {
      ser_tensor = BuildTosaSerializationTensor(pair.second /* val */,
                                                pair.first /* name */);
      if (!ser_tensor) {
        llvm::errs() << "ERROR: Failed to build TosaSerializationTensor\n";
        return mlir::failure();
      }
      block->GetTensors().push_back(ser_tensor);
    }

    // Build operator
    for (auto &op : bb) {
      if (llvm::isa<mlir::tosa::YieldOp>(op) || llvm::isa<mlir::ReturnOp>(op) ||
          llvm::isa<mlir::tensor::CastOp>(op))
        continue;
      ser_operator = BuildTosaSerializationOperator(op_builder, op);
      if (!ser_operator) {
        llvm::errs() << "ERROR: Failed to build TosaSerializationOperator\n";
        return mlir::failure();
      }
      block->GetOperators().push_back(ser_operator);
    }
  }

  return mlir::success();
}

TosaSerializationOperator *
TosaSerializationBlockBuilder::BuildTosaSerializationOperator(
    const TosaSerializationOperatorBuilder &op_builder, mlir::Operation &op) {
  std::string full_op_name = op.getName().getStringRef().str();
  TosaSerializationOperator *target_operator = nullptr;

  if (false) {
  }
#define DEF_OPERATOR(MLIR_OP)                                                  \
  else if (llvm::isa<mlir::tosa::MLIR_OP##Op>(op)) {                           \
    target_operator = op_builder.build<mlir::tosa::MLIR_OP##Op>(op);           \
  }
#include "operator.def"
#undef DEF_OPERATOR
  else {
    llvm::errs() << "unsupported tosa operator " << op.getName().getStringRef()
                 << "\n";
  }

  if (!target_operator) {
    llvm::errs() << op.getName().getStringRef()
                 << " operator hasn't been translated to flatbuffer, skipped\n";
    return nullptr;
  }

  // Sanity check the number of inputs/outputs of TOSA dialect matches the
  // number of TOSA flatbuffer
  if (op.getNumOperands() != target_operator->GetInputTensorNames().size()) {
    llvm::errs() << "WARNING. MLIR operator has " << op.getNumOperands()
                 << " input tensors != Flatbuffer "
                    "operator has "
                 << target_operator->GetInputTensorNames().size()
                 << " input tensors\n";
  }
  if (op.getNumResults() != target_operator->GetOutputTensorNames().size()) {
    llvm::errs() << "WARNING. MLIR operator has " << op.getNumResults()
                 << " output tensors != Flatbuffer "
                    "operator has "
                 << target_operator->GetOutputTensorNames().size()
                 << " output tensors\n";
  }

  return target_operator;
}

TosaSerializationTensor *
TosaSerializationBlockBuilder::BuildTosaSerializationTensor(
    mlir::Value val, const std::string &name) {
  // If tensor already created before, use that tensor directly, create a new
  // one otherwise
  TosaSerializationTensor *ts = block->GetTensorByName(name);
  if (ts) {
    return nullptr;
  }

  mlir::RankedTensorType tensor =
      val.getType().dyn_cast<mlir::RankedTensorType>();
  std::vector<int32_t> shape(tensor.getShape().begin(),
                             tensor.getShape().end());
  DType type = Type2DType(tensor.getElementType());

  ts = new TosaSerializationTensor(name, shape, type, std::vector<uint8_t>());

  return ts;
}

mlir::LogicalResult translate2FlatBuffer(mlir::FuncOp &func,
                                         TosaSerializationHandler &tsh) {
  TosaSerializationBasicBlock *main_block;

  mlir::Region *main_region = func.getCallableRegion();
  std::vector<mlir::Value> main_returns;

  if (!main_region) {
    llvm::errs() << "Invalid MLIR: doesn't have valid \"main\" region\n";
    return mlir::failure();
  }

  if (!tsh.GetBlocks().empty()) {
    llvm::errs() << "Internal Error: TosaSerializationHandler's block list "
                    "must be empty\n";
    return mlir::failure();
  }

  main_block = new TosaSerializationBasicBlock(
      std::string("main"), std::vector<TosaSerializationOperator *>(),
      std::vector<TosaSerializationTensor *>(), std::vector<std::string>(),
      std::vector<std::string>());
  assert(main_block);
  tsh.GetBlocks().push_back(main_block);

  TosaSerializationBlockBuilder block_builder(main_block, &tsh, main_region);
  if (block_builder.BuildAllOpsInRegion(main_returns).failed()) {
    return mlir::failure();
  }

  if (main_returns.empty()) {
    llvm::errs() << "Warning: graph doesn't have return values\n";
  }

  return mlir::success();
}

mlir::LogicalResult dumpTosaFlatbuffer(mlir::FuncOp &func) {
  tosa::TosaSerializationHandler tsh;

  std::string tosa_flatbuffer_directory_fullpath;
  if (translate2FlatBuffer(func, tsh).failed()) {
    llvm::errs() << "Fail to translate TOSA MLIR to flatbuffer\n";
    return mlir::failure();
  }

  if (tsh.SaveFileTosaFlatbuffer(tosa_flatbuffer_filename.c_str())) {
    llvm::errs() << "Fail to save flatbuffer " << tosa_flatbuffer_filename
                 << "\n";
    return mlir::failure();
  }
  return mlir::success();
}

mlir::LogicalResult dumpTosaJSON(mlir::FuncOp &func) {
  tosa::TosaSerializationHandler tsh;

  const char *tosa_schema = tosa_flatbuffer_schema.c_str();

  if (!tosa_schema) {
    llvm::errs() << "Flatbuffer schema not defined\n";
    return mlir::failure();
  }

  if (tsh.LoadFileSchema(tosa_schema)) {
    llvm::errs() << "Error loading tosa schema file: " << tosa_schema << "\n";
    return mlir::failure();
  }

  std::string tosa_flatbuffer_directory_fullpath;
  if (translate2FlatBuffer(func, tsh).failed()) {
    llvm::errs() << "Fail to translate TOSA MLIR to flatbuffer\n";
    return mlir::failure();
  }

  if (tsh.SaveFileJson(tosa_flatbuffer_filename.c_str())) {
    llvm::errs() << "Fail to save flatbuffer " << tosa_flatbuffer_filename
                 << "\n";
    return mlir::failure();
  }

  return mlir::success();
}

namespace mlir {

namespace tosa {

namespace {

class TosaSerialize : public PassWrapper<TosaSerialize, FunctionPass> {
public:
  TosaSerialize() = default;

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "tosa-serialize";
  }
  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Run the TOSA serialization (flatbuffer generation) pass";
  }

  void runOnFunction() override {
    auto function = getFunction();

    if (dumpTosaFlatbuffer(function).failed()) {
      llvm::errs() << "Failed to generate TOSA flatbuffer...\n";
      return signalPassFailure();
    }
  }
};

class TosaSerializeJSON : public PassWrapper<TosaSerializeJSON, FunctionPass> {
public:
  TosaSerializeJSON() = default;

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "tosa-serialize-json";
  }
  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Run the TOSA serialization (JSON generation) pass";
  }

  void runOnFunction() override {
    auto function = getFunction();

    if (dumpTosaJSON(function).failed()) {
      llvm::errs() << "Failed to generate TOSA JSON...\n";
      return signalPassFailure();
    }
  }
};

} // anonymous namespace

// Creates an instance of the TOSA flatbuffer generation pass
std::unique_ptr<OperationPass<FuncOp>> createTosaSerializePass() {
  return std::make_unique<TosaSerialize>();
}

std::unique_ptr<OperationPass<FuncOp>> createTosaSerializeJSONPass() {
  return std::make_unique<TosaSerializeJSON>();
}

static PassRegistration<TosaSerialize> pass([] {
  return createTosaSerializePass();
});

static PassRegistration<TosaSerializeJSON> passJSON([] {
  return createTosaSerializeJSONPass();
});

} // namespace tosa
} // namespace mlir