Add CalibratedQuantizedType to quant dialect

[teijeong]

Summary

This type supports a calibrated type with min, max provided.

This will be used for importing calibration values of intermediate tensors (e.g. LSTM) which can't be imported with QuantStats op.

This type was initially suggested in the following RFC: https://llvm.discourse.group/t/rfc-a-proposal-for-implementing-quantization-transformations-in-mlir/655

[stellaraccident]

Wow, that has been a minute since that was discussed. I'm not 100% sure these days that I would encode this in the type system if I were doing it, but having the representation to do it is fine with me.

[teijeong]

Wow what a coincidence, thanks for quick review!

I don't have commit permission yet - can someone help me submitting this?

[liufengdb]

mlir/include/mlir/Dialect/Quant/QuantTypes.h

378

remove this storage type.

mlir/lib/Dialect/Quant/IR/QuantOps.cpp

29

alphabeta order?

[teijeong]

Updating D91584: Add CalibratedQuantizedType to quant dialect

mlir/include/mlir/Dialect/Quant/QuantTypes.h

 class QuantizedIntegerType;
 
 namespace detail {
 
 struct QuantizedTypeStorage;
 struct AnyQuantizedTypeStorage;
 struct UniformQuantizedTypeStorage;
 struct UniformQuantizedPerAxisTypeStorage;
+struct CalibratedQuantizedTypeStorage;
 
 } // namespace detail
 
 /// Enumeration of bit-mapped flags related to quantized types.
 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.
public:
   bool isFixedPoint() const {
     if (!isSigned())
       return false;
     return llvm::all_of(getZeroPoints(),
                         [](int64_t zeroPoint) { return zeroPoint != 0; });
   }
 };
 
+/// A quantized type that infers its range from given min/max values.
+///
+/// Typical syntax:
+///   quant.calibrated<f32<-0.922,0.981>>
+class CalibratedQuantizedType
+    : public Type::TypeBase<CalibratedQuantizedType, QuantizedType,
+                            detail::CalibratedQuantizedTypeStorage> {
+public:
+  using Base::Base;
+
+  /// Gets an instance of the type with all parameters specified but not
+  /// checked.
+  static CalibratedQuantizedType get(Type expressedType, double min,
+                                     double max);
+
+  /// Gets an instance of the type with all specified parameters checked.
+  /// Returns a nullptr convertible type on failure.
+  static CalibratedQuantizedType getChecked(Type expressedType, double min,
+                                            double max, Location location);
+
+  /// Verifies construction invariants and issues errors/warnings.
+  static LogicalResult verifyConstructionInvariants(Location loc,
+                                                    Type expressedType,
+                                                    double min, double max);
+  double getMin() const;
+  double getMax() const;
+};
+
 } // namespace quant
 } // namespace mlir
 
 #endif // MLIR_DIALECT_QUANT_QUANT_TYPES_H_

mlir/lib/Dialect/Quant/IR/QuantOps.cpp

 #include "llvm/Support/MathExtras.h"
 #include <numeric>
 
 using namespace mlir;
 using namespace mlir::quant;
 using namespace mlir::quant::detail;
 
 void QuantizationDialect::initialize() {
-  addTypes<AnyQuantizedType, UniformQuantizedType,
+  addTypes<AnyQuantizedType, CalibratedQuantizedType, UniformQuantizedType,
            UniformQuantizedPerAxisType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Quant/QuantOps.cpp.inc"
       >();
 }
 
 OpFoldResult StorageCastOp::fold(ArrayRef<Attribute> operands) {
   // Matches x -> [scast -> scast] -> y, replacing the second scast with the
   // value of x if the casts invert each other.

mlir/lib/Dialect/Quant/IR/QuantTypes.cpp

 
 ArrayRef<int64_t> UniformQuantizedPerAxisType::getZeroPoints() const {
   return getImpl()->getZeroPoints();
 }
 
 int32_t UniformQuantizedPerAxisType::getQuantizedDimension() const {
   return getImpl()->quantizedDimension;
 }
+
+CalibratedQuantizedType CalibratedQuantizedType::get(Type expressedType,
+                                                     double min, double max) {
+  return Base::get(expressedType.getContext(), expressedType, min, max);
+}
+
+CalibratedQuantizedType CalibratedQuantizedType::getChecked(Type expressedType,
+                                                            double min,
+                                                            double max,
+                                                            Location location) {
+  return Base::getChecked(location, expressedType, min, max);
+}
+
+LogicalResult CalibratedQuantizedType::verifyConstructionInvariants(
+    Location loc, Type expressedType, double min, double max) {
+  // Verify that the expressed type is floating point.
+  // If this restriction is ever eliminated, the parser/printer must be
+  // extended.
+  if (!expressedType.isa<FloatType>())
+    return emitError(loc, "expressed type must be floating point");
+  if (max <= min)
+    return emitError(loc, "illegal min and max: (") << min << ":" << max << ")";
+
+  return success();
+}
+
+double CalibratedQuantizedType::getMin() const { return getImpl()->min; }
+
+double CalibratedQuantizedType::getMax() const { return getImpl()->max; }

mlir/lib/Dialect/Quant/IR/TypeDetail.h

struct UniformQuantizedPerAxisTypeStorage : public QuantizedTypeStorage {
   }
 
   const double *scaleElements;
   const int64_t *zeroPointElements;
   unsigned quantParamsSize;
   int32_t quantizedDimension;
 };
 
+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_

mlir/lib/Dialect/Quant/IR/TypeParser.cpp

static ParseResult parseStorageRange(DialectAsmParser &parser,
   }
   if (storageTypeMax > defaultIntegerMax) {
     return parser.emitError(maxLoc, "illegal storage type maximum: ")
            << storageTypeMax;
   }
   return success();
 }
 
-/// Parses a UniformQuantizedType.
+static FloatType parseExpressedTypeAndRange(DialectAsmParser &parser,
+                                            double &min, double &max) {
+  auto typeLoc = parser.getCurrentLocation();
+  FloatType type;
+
+  if (failed(parser.parseType(type))) {
+    parser.emitError(typeLoc, "expecting float expressed type");
+    return nullptr;
+  }
+
+  // Calibrated min and max values.
+  if (parser.parseLess() || parser.parseFloat(min) || parser.parseColon() ||
+      parser.parseFloat(max) || parser.parseGreater()) {
+    parser.emitError(typeLoc, "calibrated values must be present");
+    return nullptr;
+  }
+  return type;
+}
+
+/// Parses an AnyQuantizedType.
 ///
-///   uniform_per_layer ::= `any<` storage-spec (expressed-type-spec)?`>`
+///   any ::= `any<` storage-spec (expressed-type-spec)?`>`
 ///   storage-spec ::= storage-type (`<` storage-range `>`)?
 ///   storage-range ::= integer-literal `:` integer-literal
 ///   storage-type ::= (`i` | `u`) integer-literal
 ///   expressed-type-spec ::= `:` `f` integer-literal
 static Type parseAnyType(DialectAsmParser &parser, Location loc) {
   IntegerType storageType;
   FloatType expressedType;
   unsigned typeFlags = 0;
return UniformQuantizedPerAxisType::getChecked(
         quantizedDimension, storageTypeMin, storageTypeMax, loc);
   }
 
   return UniformQuantizedType::getChecked(typeFlags, storageType, expressedType,
                                           scales.front(), zeroPoints.front(),
                                           storageTypeMin, storageTypeMax, loc);
 }
 
+/// Parses an CalibratedQuantizedType.
+///
+///   calibrated ::= `calibrated<` expressed-spec `>`
+///   expressed-spec ::= expressed-type `<` calibrated-range `>`
+///   expressed-type ::= `f` integer-literal
+///   calibrated-range ::= float-literal `:` float-literal
+static Type parseCalibratedType(DialectAsmParser &parser, Location loc) {
+  FloatType expressedType;
+  double min;
+  double max;
+
+  // Type specification.
+  if (parser.parseLess())
+    return nullptr;
+
+  // Expressed type.
+  expressedType = parseExpressedTypeAndRange(parser, min, max);
+  if (!expressedType) {
+    return nullptr;
+  }
+
+  if (parser.parseGreater()) {
+    return nullptr;
+  }
+
+  return CalibratedQuantizedType::getChecked(expressedType, min, max, loc);
+}
+
 /// Parse a type registered to this dialect.
 Type QuantizationDialect::parseType(DialectAsmParser &parser) const {
   Location loc = parser.getEncodedSourceLoc(parser.getNameLoc());
 
   // All types start with an identifier that we switch on.
   StringRef typeNameSpelling;
   if (failed(parser.parseKeyword(&typeNameSpelling)))
     return nullptr;
 
   if (typeNameSpelling == "uniform")
     return parseUniformType(parser, loc);
   if (typeNameSpelling == "any")
     return parseAnyType(parser, loc);
+  if (typeNameSpelling == "calibrated")
+    return parseCalibratedType(parser, loc);
 
   parser.emitError(parser.getNameLoc(),
                    "unknown quantized type " + typeNameSpelling);
   return nullptr;
 }
 
 static void printStorageType(QuantizedType type, DialectAsmPrinter &out) {
   // storage type

 static void printQuantParams(double scale, int64_t zeroPoint,
                              DialectAsmPrinter &out) {
   out << scale;
   if (zeroPoint != 0) {
     out << ":" << zeroPoint;
   }
 }
 
-/// Helper that prints a UniformQuantizedType.
+/// Helper that prints a AnyQuantizedType.
 static void printAnyQuantizedType(AnyQuantizedType type,
                                   DialectAsmPrinter &out) {
   out << "any<";
   printStorageType(type, out);
   if (Type expressedType = type.getExpressedType()) {
     out << ":" << expressedType;
   }
   out << ">";
llvm::interleave(
       llvm::seq<size_t>(0, scales.size()), out,
       [&](size_t index) {
         printQuantParams(scales[index], zeroPoints[index], out);
       },
       ",");
   out << "}>";
 }
 
+/// Helper that prints a CalibratedQuantizedType.
+static void printCalibratedQuantizedType(CalibratedQuantizedType type,
+                                         DialectAsmPrinter &out) {
+  out << "calibrated<" << type.getExpressedType();
+  out << "<" << type.getMin() << ", " << type.getMax() << ">";
+  out << ">";
+}
+
 /// Print a type registered to this dialect.
 void QuantizationDialect::printType(Type type, DialectAsmPrinter &os) const {
   if (auto anyType = type.dyn_cast<AnyQuantizedType>())
     printAnyQuantizedType(anyType, os);
   else if (auto uniformType = type.dyn_cast<UniformQuantizedType>())
     printUniformQuantizedType(uniformType, os);
   else if (auto perAxisType = type.dyn_cast<UniformQuantizedPerAxisType>())
     printUniformQuantizedPerAxisType(perAxisType, os);
+  else if (auto calibratedType = type.dyn_cast<CalibratedQuantizedType>())
+    printCalibratedQuantizedType(calibratedType, os);
   else
     llvm_unreachable("Unhandled quantized type");
 }

mlir/test/Dialect/Quant/parse-calibrated-invalid.mlir

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// -----
+// Unrecognized token: missing calibrated type maximum
+// expected-error@+2 {{calibrated values must be present}}
+// expected-error@+1 {{expected ':'}}
+!qalias = type !quant.calibrated<f32<-0.998>>
+
+// -----
+// Unrecognized token: missing closing angle bracket
+// expected-error@+1 {{expected '>'}}
+!qalias = type !quant<"calibrated<f32<-0.998:1.232>">
+
+// -----
+// Unrecognized expressed type: integer type
+// expected-error@+2 {{invalid kind of type specified}}
+// expected-error@+1 {{expecting float expressed type}}
+!qalias = type !quant.calibrated<i8<-4:3>>
+
+// -----
+// Illegal storage min/max: max - min < 0
+// expected-error@+1 {{illegal min and max: (1.000000e+00:-1.000000e+00)}}
+!qalias = type !quant.calibrated<f32<1.0:-1.0>>
+
+// -----
+// Illegal storage min/max: max - min == 0
+// expected-error@+1 {{illegal min and max: (1.000000e+00:1.000000e+00)}}
+!qalias = type !quant.calibrated<f32<1.0:1.0>>

mlir/test/Dialect/Quant/parse-calibrated.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file | FileCheck %s
+
+// -----
+// CHECK-LABEL: parseCalibrated
+// CHECK: !quant.calibrated<f32<-0.998, 1.232100e+00>
+!qalias = type !quant.calibrated<f32<-0.998:1.2321>>
+func @parseCalibrated() -> !qalias {
+  %0 = "foo"() : () -> !qalias
+  return %0 : !qalias
+}