TypeDetail.h

Go to the documentation of this file.

//===- TypeDetail.h - QuantOps Type detail ----------------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#ifndef TYPE_DETAIL_H_
#define TYPE_DETAIL_H_
 
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/TypeSupport.h"
#include "mlir/IR/Types.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/bit.h"
 
namespace mlir {
namespace quant {
namespace detail {
 
struct QuantizedTypeStorage : public mlir::TypeStorage {
  QuantizedTypeStorage(unsigned flags, Type storageType, Type expressedType,
                       int64_t storageTypeMin, int64_t storageTypeMax)
      : flags(flags), storageType(storageType), expressedType(expressedType),
        storageTypeMin(storageTypeMin), storageTypeMax(storageTypeMax) {}
 
  /// Flags corresponding to the bitmapped enum QuantizationFlags::FlagValue.
  unsigned flags;
 
  // Integral type for the storage point representation.
  Type storageType;
 
  // Floating point type that the quantized type approximates.
  Type expressedType;
 
  // The minimum value storageType can take.
  int64_t storageTypeMin;
 
  // The maximum value storageType can take.
  int64_t storageTypeMax;
};
 
struct AnyQuantizedTypeStorage : public QuantizedTypeStorage {
  struct KeyTy {
    KeyTy(unsigned flags, Type storageType, Type expressedType,
          int64_t storageTypeMin, int64_t storageTypeMax)
        : flags(flags), storageType(storageType), expressedType(expressedType),
          storageTypeMin(storageTypeMin), storageTypeMax(storageTypeMax) {}
    unsigned flags;
    Type storageType;
    Type expressedType;
    int64_t storageTypeMin;
    int64_t storageTypeMax;
 
    // Check for equality of two structures that share KeyTy data members
    // (by name).
    template <typename T, typename U>
    static bool genericIsEqual(const T &lhs, const U &rhs) {
      return lhs.flags == rhs.flags && lhs.storageType == rhs.storageType &&
             lhs.expressedType == rhs.expressedType &&
             lhs.storageTypeMin == rhs.storageTypeMin &&
             lhs.storageTypeMax == rhs.storageTypeMax;
    }
 
    bool operator==(const KeyTy &other) const {
      return genericIsEqual(*this, other);
    }
 
    unsigned getHashValue() const {
      return llvm::hash_combine(flags, storageType, expressedType,
                                storageTypeMin, storageTypeMax);
    }
  };
 
  AnyQuantizedTypeStorage(const KeyTy &key)
      : QuantizedTypeStorage(key.flags, key.storageType, key.expressedType,
                             key.storageTypeMin, key.storageTypeMax) {}
 
  bool operator==(const KeyTy &key) const {
    return KeyTy::genericIsEqual(*this, key);
  }
 
  /// Construction.
  static AnyQuantizedTypeStorage *construct(TypeStorageAllocator &allocator,
                                            const KeyTy &key) {
    return new (allocator.allocate<AnyQuantizedTypeStorage>())
        AnyQuantizedTypeStorage(key);
  }
 
  static unsigned hashKey(const KeyTy &key) { return key.getHashValue(); }
};
 
struct UniformQuantizedTypeStorage : public QuantizedTypeStorage {
  struct KeyTy {
    KeyTy(unsigned flags, Type storageType, Type expressedType, double scale,
          int64_t zeroPoint, int64_t storageTypeMin, int64_t storageTypeMax)
        : flags(flags), storageType(storageType), expressedType(expressedType),
          scale(scale), zeroPoint(zeroPoint), storageTypeMin(storageTypeMin),
          storageTypeMax(storageTypeMax) {}
    /// Flags corresponding to the bitmapped enum QuantizationFlags::FlagValue.
    unsigned flags;
 
    // Integral type for the storage point representation.
    Type storageType;
 
    // Floating point type that the quantized type approximates.
    Type expressedType;
 
    double scale;
    int64_t zeroPoint;
    int64_t storageTypeMin;
    int64_t storageTypeMax;
 
    // Check for equality of two structures that share KeyTy data members
    // (by name).
    template <typename T, typename U>
    static bool genericIsEqual(const T &lhs, const U &rhs) {
      return lhs.flags == rhs.flags && lhs.storageType == rhs.storageType &&
             lhs.expressedType == rhs.expressedType && lhs.scale == rhs.scale &&
             lhs.zeroPoint == rhs.zeroPoint &&
             lhs.storageTypeMin == rhs.storageTypeMin &&
             lhs.storageTypeMax == rhs.storageTypeMax;
    }
 
    bool operator==(const KeyTy &other) const {
      return genericIsEqual(*this, other);
    }
 
    unsigned getHashValue() const {
      int64_t scaleBits = llvm::bit_cast<int64_t>(scale);
      return llvm::hash_combine(flags, storageType, expressedType, scaleBits,
                                zeroPoint, storageTypeMin, storageTypeMax);
    }
  };
 
  UniformQuantizedTypeStorage(const KeyTy &key)
      : QuantizedTypeStorage(key.flags, key.storageType, key.expressedType,
                             key.storageTypeMin, key.storageTypeMax),
        scale(key.scale), zeroPoint(key.zeroPoint) {}
 
  bool operator==(const KeyTy &key) const {
    return KeyTy::genericIsEqual(*this, key);
  }
 
  /// Construction.
  static UniformQuantizedTypeStorage *construct(TypeStorageAllocator &allocator,
                                                const KeyTy &key) {
    return new (allocator.allocate<UniformQuantizedTypeStorage>())
        UniformQuantizedTypeStorage(key);
  }
 
  static unsigned hashKey(const KeyTy &key) { return key.getHashValue(); }
 
  double scale;
  int64_t zeroPoint;
};
 
struct UniformQuantizedPerAxisTypeStorage : public QuantizedTypeStorage {
  struct KeyTy {
    KeyTy(unsigned flags, Type storageType, Type expressedType,
          ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
          int32_t quantizedDimension, int64_t storageTypeMin,
          int64_t storageTypeMax)
        : flags(flags), storageType(storageType), expressedType(expressedType),
          scales(scales), zeroPoints(zeroPoints),
          quantizedDimension(quantizedDimension),
          storageTypeMin(storageTypeMin), storageTypeMax(storageTypeMax) {}
    /// Flags corresponding to the bitmapped enum QuantizationFlags::FlagValue.
    unsigned flags;
 
    // Integral type for the storage point representation.
    Type storageType;
 
    // Floating point type that the quantized type approximates.
    Type expressedType;
 
    ArrayRef<double> scales;
    ArrayRef<int64_t> zeroPoints;
    int32_t quantizedDimension;
    int64_t storageTypeMin;
    int64_t storageTypeMax;
 
    ArrayRef<double> getScales() const { return scales; }
 
    ArrayRef<int64_t> getZeroPoints() const { return zeroPoints; }
 
    // Check for equality of two structures that share KeyTy data members
    // (by name).
    template <typename T, typename U>
    static bool genericIsEqual(const T &lhs, const U &rhs) {
      return lhs.flags == rhs.flags && lhs.storageType == rhs.storageType &&
             lhs.expressedType == rhs.expressedType &&
             lhs.getScales() == rhs.getScales() &&
             lhs.getZeroPoints() == rhs.getZeroPoints() &&
             lhs.quantizedDimension == rhs.quantizedDimension &&
             lhs.storageTypeMin == rhs.storageTypeMin &&
             lhs.storageTypeMax == rhs.storageTypeMax;
    }
 
    bool operator==(const KeyTy &other) const {
      return genericIsEqual(*this, other);
    }
 
    unsigned getHashValue() const {
      int64_t *scalesCast = llvm::bit_cast<int64_t *>(scales.data());
      ArrayRef<int64_t> scalesBits(scalesCast, scales.size());
      return llvm::hash_combine(flags, storageType, expressedType,
                                llvm::hash_combine_range(scalesBits),
                                llvm::hash_combine_range(zeroPoints),
                                storageTypeMin, storageTypeMax);
    }
  };
 
  // We pass scales and zeroPoints in directly rather than relying on KeyTy
  // because we have to create new reallocated versions in `construct` below.
  UniformQuantizedPerAxisTypeStorage(const KeyTy &key, ArrayRef<double> scales,
                                     ArrayRef<int64_t> zeroPoints)
      : QuantizedTypeStorage(key.flags, key.storageType, key.expressedType,
                             key.storageTypeMin, key.storageTypeMax),
        scaleElements(scales.data()), zeroPointElements(zeroPoints.data()),
        quantParamsSize(scales.size()),
        quantizedDimension(key.quantizedDimension) {}
 
  bool operator==(const KeyTy &key) const {
    return KeyTy::genericIsEqual(*this, key);
  }
 
  /// Construction.
  static UniformQuantizedPerAxisTypeStorage *
  construct(TypeStorageAllocator &allocator, const KeyTy &key) {
    ArrayRef<double> scales = allocator.copyInto(key.scales);
    ArrayRef<int64_t> zeroPoints = allocator.copyInto(key.zeroPoints);
    return new (allocator.allocate<UniformQuantizedPerAxisTypeStorage>())
        UniformQuantizedPerAxisTypeStorage(key, scales, zeroPoints);
  }
 
  static unsigned hashKey(const KeyTy &key) { return key.getHashValue(); }
 
  ArrayRef<double> getScales() const {
    return ArrayRef<double>(scaleElements, quantParamsSize);
  }
 
  ArrayRef<int64_t> getZeroPoints() const {
    return ArrayRef<int64_t>(zeroPointElements, quantParamsSize);
  }
 
  const double *scaleElements;
  const int64_t *zeroPointElements;
  unsigned quantParamsSize;
  int32_t quantizedDimension;
};
 
struct UniformQuantizedSubChannelTypeStorage : public QuantizedTypeStorage {
  struct KeyTy {
    KeyTy(unsigned flags, Type storageType, Type expressedType,
          DenseElementsAttr scales, DenseElementsAttr zeroPoints,
          ArrayRef<int32_t> quantizedDimensions, ArrayRef<int64_t> blockSizes,
          int64_t storageTypeMin, int64_t storageTypeMax)
        : flags(flags), storageType(storageType), expressedType(expressedType),
          scales(scales), zeroPoints(zeroPoints),
          quantizedDimensions(quantizedDimensions), blockSizes(blockSizes),
          storageTypeMin(storageTypeMin), storageTypeMax(storageTypeMax) {}
    /// Flags corresponding to the bitmapped enum QuantizationFlags::FlagValue.
    unsigned flags;
 
    // Integral type for the storage point representation.
    Type storageType;
 
    // Floating point type that the quantized type approximates.
    Type expressedType;
 
    DenseElementsAttr scales;
    DenseElementsAttr zeroPoints;
    ArrayRef<int32_t> quantizedDimensions;
    ArrayRef<int64_t> blockSizes;
    int64_t storageTypeMin;
    int64_t storageTypeMax;
 
    DenseElementsAttr getScales() const { return scales; }
 
    DenseElementsAttr getZeroPoints() const { return zeroPoints; }
 
    // Check for equality of two structures that share KeyTy data members
    // (by name).
    template <typename T, typename U>
    static bool genericIsEqual(const T &lhs, const U &rhs) {
      return lhs.flags == rhs.flags && lhs.storageType == rhs.storageType &&
             lhs.expressedType == rhs.expressedType &&
             lhs.scales == rhs.scales && lhs.zeroPoints == rhs.zeroPoints &&
             lhs.quantizedDimensions == rhs.quantizedDimensions &&
             lhs.blockSizes == rhs.blockSizes &&
             lhs.storageTypeMin == rhs.storageTypeMin &&
             lhs.storageTypeMax == rhs.storageTypeMax;
    }
 
    bool operator==(const KeyTy &other) const {
      return genericIsEqual(*this, other);
    }
 
    unsigned getHashValue() const {
      // Hash the scalar attributes.
      unsigned hash = llvm::hash_combine(flags, storageType, expressedType,
                                         storageTypeMin, storageTypeMax);
 
      // Hash the scales.
      for (auto scaleAttr : scales.getValues<APFloat>()) {
        hash = llvm::hash_combine(
            hash, llvm::bit_cast<int64_t>(scaleAttr.convertToDouble()));
      }
 
      // Hash the zero points.  (Assumed to be integers, adjust if needed).
      for (auto zeroPointAttr : zeroPoints.getValues<APInt>()) {
        hash = llvm::hash_combine(hash, zeroPointAttr.getSExtValue());
      }
 
      // Hash the quantized dimensions and block sizes.
      hash = llvm::hash_combine(hash,
                                llvm::hash_combine_range(quantizedDimensions),
                                llvm::hash_combine_range(blockSizes));
 
      return hash;
    }
  };
 
  // We pass scales and zeroPoints in directly rather than relying on KeyTy
  // because we have to create new reallocated versions in `construct` below.
  UniformQuantizedSubChannelTypeStorage(const KeyTy &key,
                                        DenseElementsAttr scales,
                                        DenseElementsAttr zeroPoints,
                                        ArrayRef<int32_t> quantizedDimensions,
                                        ArrayRef<int64_t> blockSizes)
      : QuantizedTypeStorage(key.flags, key.storageType, key.expressedType,
                             key.storageTypeMin, key.storageTypeMax),
        scales(scales), zeroPoints(zeroPoints),
        quantizedDimensions(quantizedDimensions), blockSizes(blockSizes) {}
 
  bool operator==(const KeyTy &key) const {
    return KeyTy::genericIsEqual(*this, key);
  }
 
  /// Construction.
  static UniformQuantizedSubChannelTypeStorage *
  construct(TypeStorageAllocator &allocator, const KeyTy &key) {
    DenseElementsAttr scales = key.scales;
    DenseElementsAttr zeroPoints = key.zeroPoints;
    ArrayRef<int32_t> quantizedDimensions =
        allocator.copyInto(key.quantizedDimensions);
    ArrayRef<int64_t> blockSizes = allocator.copyInto(key.blockSizes);
    return new (allocator.allocate<UniformQuantizedSubChannelTypeStorage>())
        UniformQuantizedSubChannelTypeStorage(key, scales, zeroPoints,
                                              quantizedDimensions, blockSizes);
  }
 
  static unsigned hashKey(const KeyTy &key) { return key.getHashValue(); }
 
  DenseElementsAttr getScales() const { return scales; }
 
  DenseElementsAttr getZeroPoints() const { return zeroPoints; }
 
  ArrayRef<int32_t> getQuantizedDimensions() const {
    return quantizedDimensions;
  }
 
  ArrayRef<int64_t> getBlockSizes() const { return blockSizes; }
 
  DenseElementsAttr scales;
  DenseElementsAttr zeroPoints;
  ArrayRef<int32_t> quantizedDimensions;
  ArrayRef<int64_t> blockSizes;
};
 
struct CalibratedQuantizedTypeStorage : public QuantizedTypeStorage {
  struct KeyTy {
    KeyTy(Type expressedType, double min, double max)
        : expressedType(expressedType), min(min), max(max) {}
    // Floating point type that the quantized type approximates.
    Type expressedType;
 
    double min;
    double max;
 
    // Check for equality of two structures that share KeyTy data members
    // (by name).
    template <typename T, typename U>
    static bool genericIsEqual(const T &lhs, const U &rhs) {
      return lhs.expressedType == rhs.expressedType && lhs.min == rhs.min &&
             lhs.max == rhs.max;
    }
 
    bool operator==(const KeyTy &other) const {
      return genericIsEqual(*this, other);
    }
 
    unsigned getHashValue() const {
      int64_t minBits = llvm::bit_cast<double>(min);
      int64_t maxBits = llvm::bit_cast<double>(max);
      return llvm::hash_combine(expressedType, minBits, maxBits);
    }
  };
 
  CalibratedQuantizedTypeStorage(const KeyTy &key)
      : QuantizedTypeStorage(0, NoneType(), key.expressedType, 0, 0),
        min(key.min), max(key.max) {}
 
  bool operator==(const KeyTy &key) const {
    return KeyTy::genericIsEqual(*this, key);
  }
 
  /// Construction.
  static CalibratedQuantizedTypeStorage *
  construct(TypeStorageAllocator &allocator, const KeyTy &key) {
    return new (allocator.allocate<CalibratedQuantizedTypeStorage>())
        CalibratedQuantizedTypeStorage(key);
  }
 
  static unsigned hashKey(const KeyTy &key) { return key.getHashValue(); }
 
  double min;
  double max;
};
 
} // namespace detail
} // namespace quant
} // namespace mlir
 
#endif // TYPE_DETAIL_H_