//===- DialectQuant.cpp - 'quant' dialect submodule -----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include <vector>

#include "mlir-c/Dialect/Quant.h"
#include "mlir-c/IR.h"
#include "mlir/Bindings/Python/IRCore.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"

#include <mlir/Bindings/Python/IRAttributes.h>

namespace nb = nanobind;
using namespace mlir::python::nanobind_adaptors;

namespace mlir {
namespace python {
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
namespace quant {
//===-------------------------------------------------------------------===//
// QuantizedType
//===-------------------------------------------------------------------===//

struct QuantizedType : PyConcreteType<QuantizedType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAQuantizedType;
  static constexpr const char *pyClassName = "QuantizedType";
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "default_minimum_for_integer",
        [](bool isSigned, unsigned integralWidth) {
          return mlirQuantizedTypeGetDefaultMinimumForInteger(isSigned,
                                                              integralWidth);
        },
        "Default minimum value for the integer with the specified signedness "
        "and "
        "bit width.",
        nb::arg("is_signed"), nb::arg("integral_width"));
    c.def_static(
        "default_maximum_for_integer",
        [](bool isSigned, unsigned integralWidth) {
          return mlirQuantizedTypeGetDefaultMaximumForInteger(isSigned,
                                                              integralWidth);
        },
        "Default maximum value for the integer with the specified signedness "
        "and "
        "bit width.",
        nb::arg("is_signed"), nb::arg("integral_width"));
    c.def_prop_ro(
        "expressed_type",
        [](QuantizedType &type) {
          return PyType(type.getContext(),
                        mlirQuantizedTypeGetExpressedType(type))
              .maybeDownCast();
        },
        "Type expressed by this quantized type.");
    c.def_prop_ro(
        "flags",
        [](const QuantizedType &type) {
          return mlirQuantizedTypeGetFlags(type);
        },
        "Flags of this quantized type (named accessors should be preferred to "
        "this)");
    c.def_prop_ro(
        "is_signed",
        [](const QuantizedType &type) {
          return mlirQuantizedTypeIsSigned(type);
        },
        "Signedness of this quantized type.");
    c.def_prop_ro(
        "storage_type",
        [](QuantizedType &type) {
          return PyType(type.getContext(),
                        mlirQuantizedTypeGetStorageType(type))
              .maybeDownCast();
        },
        "Storage type backing this quantized type.");
    c.def_prop_ro(
        "storage_type_min",
        [](const QuantizedType &type) {
          return mlirQuantizedTypeGetStorageTypeMin(type);
        },
        "The minimum value held by the storage type of this quantized type.");
    c.def_prop_ro(
        "storage_type_max",
        [](const QuantizedType &type) {
          return mlirQuantizedTypeGetStorageTypeMax(type);
        },
        "The maximum value held by the storage type of this quantized type.");
    c.def_prop_ro(
        "storage_type_integral_width",
        [](const QuantizedType &type) {
          return mlirQuantizedTypeGetStorageTypeIntegralWidth(type);
        },
        "The bitwidth of the storage type of this quantized type.");
    c.def(
        "is_compatible_expressed_type",
        [](const QuantizedType &type, const PyType &candidate) {
          return mlirQuantizedTypeIsCompatibleExpressedType(type, candidate);
        },
        "Checks whether the candidate type can be expressed by this quantized "
        "type.",
        nb::arg("candidate"));
    c.def_prop_ro(
        "quantized_element_type",
        [](QuantizedType &type) {
          return PyType(type.getContext(),
                        mlirQuantizedTypeGetQuantizedElementType(type))
              .maybeDownCast();
        },
        "Element type of this quantized type expressed as quantized type.");
    c.def(
        "cast_from_storage_type",
        [](QuantizedType &type, const PyType &candidate) {
          MlirType castResult =
              mlirQuantizedTypeCastFromStorageType(type, candidate);
          if (!mlirTypeIsNull(castResult))
            return QuantizedType(type.getContext(), castResult);
          throw nb::type_error("Invalid cast.");
        },
        "Casts from a type based on the storage type of this quantized type to "
        "a "
        "corresponding type based on the quantized type. Raises TypeError if "
        "the "
        "cast is not valid.",
        nb::arg("candidate"));
    c.def_static(
        "cast_to_storage_type",
        [](PyType &type) {
          MlirType castResult = mlirQuantizedTypeCastToStorageType(type);
          if (!mlirTypeIsNull(castResult))
            return PyType(type.getContext(), castResult).maybeDownCast();
          throw nb::type_error("Invalid cast.");
        },
        "Casts from a type based on a quantized type to a corresponding type "
        "based on the storage type of this quantized type. Raises TypeError if "
        "the cast is not valid.",
        nb::arg("type"));
    c.def(
        "cast_from_expressed_type",
        [](QuantizedType &type, const PyType &candidate) {
          MlirType castResult =
              mlirQuantizedTypeCastFromExpressedType(type, candidate);
          if (!mlirTypeIsNull(castResult))
            return PyType(type.getContext(), castResult).maybeDownCast();
          throw nb::type_error("Invalid cast.");
        },
        "Casts from a type based on the expressed type of this quantized type "
        "to "
        "a corresponding type based on the quantized type. Raises TypeError if "
        "the cast is not valid.",
        nb::arg("candidate"));
    c.def_static(
        "cast_to_expressed_type",
        [](PyType &type) {
          MlirType castResult = mlirQuantizedTypeCastToExpressedType(type);
          if (!mlirTypeIsNull(castResult))
            return PyType(type.getContext(), castResult).maybeDownCast();
          throw nb::type_error("Invalid cast.");
        },
        "Casts from a type based on a quantized type to a corresponding type "
        "based on the expressed type of this quantized type. Raises TypeError "
        "if "
        "the cast is not valid.",
        nb::arg("type"));
    c.def(
        "cast_expressed_to_storage_type",
        [](QuantizedType &type, const PyType &candidate) {
          MlirType castResult =
              mlirQuantizedTypeCastExpressedToStorageType(type, candidate);
          if (!mlirTypeIsNull(castResult))
            return PyType(type.getContext(), castResult).maybeDownCast();
          throw nb::type_error("Invalid cast.");
        },
        "Casts from a type based on the expressed type of this quantized type "
        "to "
        "a corresponding type based on the storage type. Raises TypeError if "
        "the "
        "cast is not valid.",
        nb::arg("candidate"));
  }
};

//===-------------------------------------------------------------------===//
// AnyQuantizedType
//===-------------------------------------------------------------------===//

struct AnyQuantizedType : PyConcreteType<AnyQuantizedType, QuantizedType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAAnyQuantizedType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirAnyQuantizedTypeGetTypeID;
  static constexpr const char *pyClassName = "AnyQuantizedType";
  static inline const MlirStringRef name = mlirAnyQuantizedTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](unsigned flags, const PyType &storageType,
           const PyType &expressedType, int64_t storageTypeMin,
           int64_t storageTypeMax, DefaultingPyMlirContext context) {
          return AnyQuantizedType(
              context->getRef(),
              mlirAnyQuantizedTypeGet(flags, storageType, expressedType,
                                      storageTypeMin, storageTypeMax));
        },
        "Gets an instance of AnyQuantizedType in the same context as the "
        "provided storage type.",
        nb::arg("flags"), nb::arg("storage_type"), nb::arg("expressed_type"),
        nb::arg("storage_type_min"), nb::arg("storage_type_max"),
        nb::arg("context") = nb::none());
  }
};

//===-------------------------------------------------------------------===//
// UniformQuantizedType
//===-------------------------------------------------------------------===//

struct UniformQuantizedType
    : PyConcreteType<UniformQuantizedType, QuantizedType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUniformQuantizedType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirUniformQuantizedTypeGetTypeID;
  static constexpr const char *pyClassName = "UniformQuantizedType";
  static inline const MlirStringRef name = mlirUniformQuantizedTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](unsigned flags, const PyType &storageType,
           const PyType &expressedType, double scale, int64_t zeroPoint,
           int64_t storageTypeMin, int64_t storageTypeMax,
           DefaultingPyMlirContext context) {
          return UniformQuantizedType(
              context->getRef(),
              mlirUniformQuantizedTypeGet(flags, storageType, expressedType,
                                          scale, zeroPoint, storageTypeMin,
                                          storageTypeMax));
        },
        "Gets an instance of UniformQuantizedType in the same context as the "
        "provided storage type.",
        nb::arg("flags"), nb::arg("storage_type"), nb::arg("expressed_type"),
        nb::arg("scale"), nb::arg("zero_point"), nb::arg("storage_type_min"),
        nb::arg("storage_type_max"), nb::arg("context") = nb::none());
    c.def_prop_ro(
        "scale",
        [](const UniformQuantizedType &type) {
          return mlirUniformQuantizedTypeGetScale(type);
        },
        "The scale designates the difference between the real values "
        "corresponding to consecutive quantized values differing by 1.");
    c.def_prop_ro(
        "zero_point",
        [](const UniformQuantizedType &type) {
          return mlirUniformQuantizedTypeGetZeroPoint(type);
        },
        "The storage value corresponding to the real value 0 in the affine "
        "equation.");
    c.def_prop_ro(
        "is_fixed_point",
        [](const UniformQuantizedType &type) {
          return mlirUniformQuantizedTypeIsFixedPoint(type);
        },
        "Fixed point values are real numbers divided by a scale.");
  }
};

//===-------------------------------------------------------------------===//
// UniformQuantizedPerAxisType
//===-------------------------------------------------------------------===//

struct UniformQuantizedPerAxisType
    : PyConcreteType<UniformQuantizedPerAxisType, QuantizedType> {
  static constexpr IsAFunctionTy isaFunction =
      mlirTypeIsAUniformQuantizedPerAxisType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirUniformQuantizedPerAxisTypeGetTypeID;
  static constexpr const char *pyClassName = "UniformQuantizedPerAxisType";
  static inline const MlirStringRef name =
      mlirUniformQuantizedPerAxisTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](unsigned flags, const PyType &storageType,
           const PyType &expressedType, std::vector<double> scales,
           std::vector<int64_t> zeroPoints, int32_t quantizedDimension,
           int64_t storageTypeMin, int64_t storageTypeMax,
           DefaultingPyMlirContext context) {
          if (scales.size() != zeroPoints.size())
            throw nb::value_error(
                "Mismatching number of scales and zero points.");
          auto nDims = static_cast<intptr_t>(scales.size());
          return UniformQuantizedPerAxisType(
              context->getRef(),
              mlirUniformQuantizedPerAxisTypeGet(
                  flags, storageType, expressedType, nDims, scales.data(),
                  zeroPoints.data(), quantizedDimension, storageTypeMin,
                  storageTypeMax));
        },
        "Gets an instance of UniformQuantizedPerAxisType in the same context "
        "as "
        "the provided storage type.",
        nb::arg("flags"), nb::arg("storage_type"), nb::arg("expressed_type"),
        nb::arg("scales"), nb::arg("zero_points"),
        nb::arg("quantized_dimension"), nb::arg("storage_type_min"),
        nb::arg("storage_type_max"), nb::arg("context") = nb::none());
    c.def_prop_ro(
        "scales",
        [](const UniformQuantizedPerAxisType &type) {
          intptr_t nDim = mlirUniformQuantizedPerAxisTypeGetNumDims(type);
          std::vector<double> scales;
          scales.reserve(nDim);
          for (intptr_t i = 0; i < nDim; ++i) {
            double scale = mlirUniformQuantizedPerAxisTypeGetScale(type, i);
            scales.push_back(scale);
          }
          return scales;
        },
        "The scales designate the difference between the real values "
        "corresponding to consecutive quantized values differing by 1. The ith "
        "scale corresponds to the ith slice in the quantized_dimension.");
    c.def_prop_ro(
        "zero_points",
        [](const UniformQuantizedPerAxisType &type) {
          intptr_t nDim = mlirUniformQuantizedPerAxisTypeGetNumDims(type);
          std::vector<int64_t> zeroPoints;
          zeroPoints.reserve(nDim);
          for (intptr_t i = 0; i < nDim; ++i) {
            int64_t zeroPoint =
                mlirUniformQuantizedPerAxisTypeGetZeroPoint(type, i);
            zeroPoints.push_back(zeroPoint);
          }
          return zeroPoints;
        },
        "the storage values corresponding to the real value 0 in the affine "
        "equation. The ith zero point corresponds to the ith slice in the "
        "quantized_dimension.");
    c.def_prop_ro(
        "quantized_dimension",
        [](const UniformQuantizedPerAxisType &type) {
          return mlirUniformQuantizedPerAxisTypeGetQuantizedDimension(type);
        },
        "Specifies the dimension of the shape that the scales and zero points "
        "correspond to.");
    c.def_prop_ro(
        "is_fixed_point",
        [](const UniformQuantizedPerAxisType &type) {
          return mlirUniformQuantizedPerAxisTypeIsFixedPoint(type);
        },
        "Fixed point values are real numbers divided by a scale.");
  }
};

//===-------------------------------------------------------------------===//
// UniformQuantizedSubChannelType
//===-------------------------------------------------------------------===//

struct UniformQuantizedSubChannelType
    : PyConcreteType<UniformQuantizedSubChannelType, QuantizedType> {
  static constexpr IsAFunctionTy isaFunction =
      mlirTypeIsAUniformQuantizedSubChannelType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirUniformQuantizedSubChannelTypeGetTypeID;
  static constexpr const char *pyClassName = "UniformQuantizedSubChannelType";
  static inline const MlirStringRef name =
      mlirUniformQuantizedSubChannelTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](unsigned flags, const PyType &storageType,
           const PyType &expressedType, PyAttribute scales,
           PyAttribute zeroPoints, std::vector<int32_t> quantizedDimensions,
           std::vector<int64_t> blockSizes, int64_t storageTypeMin,
           int64_t storageTypeMax, DefaultingPyMlirContext context) {
          return UniformQuantizedSubChannelType(
              context->getRef(),
              mlirUniformQuantizedSubChannelTypeGet(
                  flags, storageType, expressedType, scales, zeroPoints,
                  static_cast<intptr_t>(blockSizes.size()),
                  quantizedDimensions.data(), blockSizes.data(), storageTypeMin,
                  storageTypeMax));
        },
        "Gets an instance of UniformQuantizedSubChannel in the same context as "
        "the provided storage type.",
        nb::arg("flags"), nb::arg("storage_type"), nb::arg("expressed_type"),
        nb::arg("scales"), nb::arg("zero_points"),
        nb::arg("quantized_dimensions"), nb::arg("block_sizes"),
        nb::arg("storage_type_min"), nb::arg("storage_type_max"),
        nb::arg("context") = nb::none());
    c.def_prop_ro(
        "quantized_dimensions",
        [](const UniformQuantizedSubChannelType &type) {
          intptr_t nDim =
              mlirUniformQuantizedSubChannelTypeGetNumBlockSizes(type);
          std::vector<int32_t> quantizedDimensions;
          quantizedDimensions.reserve(nDim);
          for (intptr_t i = 0; i < nDim; ++i) {
            quantizedDimensions.push_back(
                mlirUniformQuantizedSubChannelTypeGetQuantizedDimension(type,
                                                                        i));
          }
          return quantizedDimensions;
        },
        "Gets the quantized dimensions. Each element in the returned list "
        "represents an axis of the quantized data tensor that has a specified "
        "block size. The order of elements corresponds to the order of block "
        "sizes returned by 'block_sizes' method. It means that the data tensor "
        "is quantized along the i-th dimension in the returned list using the "
        "i-th block size from block_sizes method.");
    c.def_prop_ro(
        "block_sizes",
        [](const UniformQuantizedSubChannelType &type) {
          intptr_t nDim =
              mlirUniformQuantizedSubChannelTypeGetNumBlockSizes(type);
          std::vector<int64_t> blockSizes;
          blockSizes.reserve(nDim);
          for (intptr_t i = 0; i < nDim; ++i) {
            blockSizes.push_back(
                mlirUniformQuantizedSubChannelTypeGetBlockSize(type, i));
          }
          return blockSizes;
        },
        "Gets the block sizes for the quantized dimensions. The i-th element "
        "in "
        "the returned list corresponds to the block size for the i-th "
        "dimension "
        "in the list returned by quantized_dimensions method.");
    c.def_prop_ro(
        "scales",
        [](UniformQuantizedSubChannelType &type) {
          return PyDenseElementsAttribute(
              type.getContext(),
              mlirUniformQuantizedSubChannelTypeGetScales(type));
        },
        "The scales of the quantized type.");
    c.def_prop_ro(
        "zero_points",
        [](UniformQuantizedSubChannelType &type) {
          return PyDenseElementsAttribute(
              type.getContext(),
              mlirUniformQuantizedSubChannelTypeGetZeroPoints(type));
        },
        "The zero points of the quantized type.");
  }
};

//===-------------------------------------------------------------------===//
// CalibratedQuantizedType
//===-------------------------------------------------------------------===//

struct CalibratedQuantizedType
    : PyConcreteType<CalibratedQuantizedType, QuantizedType> {
  static constexpr IsAFunctionTy isaFunction =
      mlirTypeIsACalibratedQuantizedType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirCalibratedQuantizedTypeGetTypeID;
  static constexpr const char *pyClassName = "CalibratedQuantizedType";
  static inline const MlirStringRef name = mlirCalibratedQuantizedTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](const PyType &expressedType, double min, double max,
           DefaultingPyMlirContext context) {
          return CalibratedQuantizedType(
              context->getRef(),
              mlirCalibratedQuantizedTypeGet(expressedType, min, max));
        },
        "Gets an instance of CalibratedQuantizedType in the same context as "
        "the "
        "provided expressed type.",
        nb::arg("expressed_type"), nb::arg("min"), nb::arg("max"),
        nb::arg("context") = nb::none());
    c.def_prop_ro("min", [](const PyType &type) {
      return mlirCalibratedQuantizedTypeGetMin(type);
    });
    c.def_prop_ro("max", [](const PyType &type) {
      return mlirCalibratedQuantizedTypeGetMax(type);
    });
  }
};

static void populateDialectQuantSubmodule(nb::module_ &m) {
  QuantizedType::bind(m);

  // Set the FLAG_SIGNED class attribute after binding QuantizedType
  auto quantizedTypeClass = m.attr("QuantizedType");
  quantizedTypeClass.attr("FLAG_SIGNED") = mlirQuantizedTypeGetSignedFlag();

  AnyQuantizedType::bind(m);
  UniformQuantizedType::bind(m);
  UniformQuantizedPerAxisType::bind(m);
  UniformQuantizedSubChannelType::bind(m);
  CalibratedQuantizedType::bind(m);
}
} // namespace quant
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
} // namespace python
} // namespace mlir

NB_MODULE(_mlirDialectsQuant, m) {
  m.doc() = "MLIR Quantization dialect";

  mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::quant::
      populateDialectQuantSubmodule(m);
}