QuantTypes.h

Go to the documentation of this file.

//===- QuantTypes.h - Quantization Ops and Types ----------------*- C++ -*-===//
//
// Part of the MLIR 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
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_QUANTOPS_QUANT_TYPES_H_
#define MLIR_DIALECT_QUANTOPS_QUANT_TYPES_H_

#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/StandardTypes.h"
#include "mlir/IR/Types.h"
#include "llvm/Support/MathExtras.h"

namespace mlir {
namespace quant {

class QuantizedIntegerType;

namespace detail {

struct QuantizedTypeStorage;
struct AnyQuantizedTypeStorage;
struct UniformQuantizedTypeStorage;
struct UniformQuantizedPerAxisTypeStorage;

} // namespace detail

namespace QuantizationTypes {
enum Kind {
  Any = Type::FIRST_QUANTIZATION_TYPE,
  UniformQuantized,
  UniformQuantizedPerAxis,
  LAST_USED_QUANTIZATION_TYPE = UniformQuantizedPerAxis,
};
} // namespace QuantizationTypes

namespace QuantizationFlags {
enum FlagValue {
  // Indicates that the storage type should be interpreted as a signed
  // integer. The default is to interpret it as an unsigned value.
  Signed = 1,
};
} // namespace QuantizationFlags

class QuantizedType : public Type {
public:
  using ImplType = detail::QuantizedTypeStorage;
  using Type::Type;

  static constexpr unsigned MaxStorageBits = 32;

  static LogicalResult
  verifyConstructionInvariants(Optional<Location> loc, MLIRContext *context,
                               unsigned flags, Type storageType,
                               Type expressedType, int64_t storageTypeMin,
                               int64_t storageTypeMax);

  static bool classof(Type type) {
    return type.getKind() >= Type::FIRST_QUANTIZATION_TYPE &&
           type.getKind() <= QuantizationTypes::LAST_USED_QUANTIZATION_TYPE;
  }

  static int64_t getDefaultMinimumForInteger(bool isSigned,
                                             unsigned integralWidth) {
    if (isSigned) {
      return llvm::minIntN(integralWidth);
    }
    return 0;
  }

  static int64_t getDefaultMaximumForInteger(bool isSigned,
                                             unsigned integralWidth) {
    if (isSigned) {
      return llvm::maxIntN(integralWidth);
    }
    return llvm::maxUIntN(integralWidth);
  }

  Type getExpressedType() const;

  unsigned getFlags() const;

  // Convenience helpers.
  bool isSigned() const {
    return (getFlags() & QuantizationFlags::Signed) ==
           QuantizationFlags::Signed;
  }

  Type getStorageType() const;

  int64_t getStorageTypeMin() const;

  int64_t getStorageTypeMax() const;

  unsigned getStorageTypeIntegralWidth() const;

  bool isCompatibleExpressedType(Type candidateExpressedType);

  static QuantizedType getQuantizedElementType(Type primitiveOrContainerType);

  Type castFromStorageType(Type candidateType);

  static Type castToStorageType(Type quantizedType);

  Type castFromExpressedType(Type candidateType);

  static Type castToExpressedType(Type quantizedType);

  Type castExpressedToStorageType(Type candidateType);

private:
  using Type::isBF16;
  using Type::isF16;
  using Type::isF32;
  using Type::isF64;
  using Type::isIndex;
  using Type::isInteger;
};

class AnyQuantizedType
    : public Type::TypeBase<AnyQuantizedType, QuantizedType,
                            detail::AnyQuantizedTypeStorage> {
public:
  using Base::Base;

  static bool kindof(unsigned kind) { return kind == QuantizationTypes::Any; }

  static AnyQuantizedType get(unsigned flags, Type storageType,
                              Type expressedType, int64_t storageTypeMin,
                              int64_t storageTypeMax);

  static AnyQuantizedType getChecked(unsigned flags, Type storageType,
                                     Type expressedType, int64_t storageTypeMin,
                                     int64_t storageTypeMax, Location location);

  static LogicalResult
  verifyConstructionInvariants(Optional<Location> loc, MLIRContext *context,
                               unsigned flags, Type storageType,
                               Type expressedType, int64_t storageTypeMin,
                               int64_t storageTypeMax);
};

class UniformQuantizedType
    : public Type::TypeBase<UniformQuantizedType, QuantizedType,
                            detail::UniformQuantizedTypeStorage> {
public:
  using Base::Base;

  static UniformQuantizedType get(unsigned flags, Type storageType,
                                  Type expressedType, double scale,
                                  int64_t zeroPoint, int64_t storageTypeMin,
                                  int64_t storageTypeMax);

  static UniformQuantizedType
  getChecked(unsigned flags, Type storageType, Type expressedType, double scale,
             int64_t zeroPoint, int64_t storageTypeMin, int64_t storageTypeMax,
             Location location);

  static LogicalResult verifyConstructionInvariants(
      Optional<Location> loc, MLIRContext *context, unsigned flags,
      Type storageType, Type expressedType, double scale, int64_t zeroPoint,
      int64_t storageTypeMin, int64_t storageTypeMax);

  static bool kindof(unsigned kind) {
    return kind == QuantizationTypes::UniformQuantized;
  }

  double getScale() const;

  int64_t getZeroPoint() const;

  // Fixed point values are real numbers divided by a scale.
  // Currently, only signed storage types are treated as fixed point.
  // A fixed point value can be obtained from an affine value by subtracting
  // the zeroPoint.
  // In the future, this may be explicit versus implied by type and zeroPoint.
  bool isFixedPoint() const { return isSigned() && getZeroPoint() == 0; }
};

class UniformQuantizedPerAxisType
    : public Type::TypeBase<UniformQuantizedPerAxisType, QuantizedType,
                            detail::UniformQuantizedPerAxisTypeStorage> {
public:
  using Base::Base;

  static UniformQuantizedPerAxisType
  get(unsigned flags, Type storageType, Type expressedType,
      ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
      int32_t quantizedDimension, int64_t storageTypeMin,
      int64_t storageTypeMax);

  static UniformQuantizedPerAxisType
  getChecked(unsigned flags, Type storageType, Type expressedType,
             ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
             int32_t quantizedDimension, int64_t storageTypeMin,
             int64_t storageTypeMax, Location location);

  static LogicalResult verifyConstructionInvariants(
      Optional<Location> loc, MLIRContext *context, unsigned flags,
      Type storageType, Type expressedType, ArrayRef<double> scales,
      ArrayRef<int64_t> zeroPoints, int32_t quantizedDimension,
      int64_t storageTypeMin, int64_t storageTypeMax);

  static bool kindof(unsigned kind) {
    return kind == QuantizationTypes::UniformQuantizedPerAxis;
  }

  ArrayRef<double> getScales() const;

  ArrayRef<int64_t> getZeroPoints() const;

  int32_t getQuantizedDimension() const;

  bool isFixedPoint() const {
    if (!isSigned())
      return false;
    return llvm::all_of(getZeroPoints(),
                        [](int64_t zeroPoint) { return zeroPoint != 0; });
  }
};

} // namespace quant
} // namespace mlir

#endif // MLIR_DIALECT_QUANTOPS_QUANT_TYPES_H_